Repository: xgcm/xhistogram

Issue:
Support for density=True keyword
Numpy histogram has the `density` keyword, which normalized the outputs by the bin spacing.
https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html

We can easily support this in xhistogram, but it has to .be plugged in. Just opening this issue to make a note of it.

Comments:

Patch:
diff --git a/xhistogram/core.py b/xhistogram/core.py
index c14f977..9f992c2 100644
--- a/xhistogram/core.py
+++ b/xhistogram/core.py
@@ -88,7 +88,6 @@ def _histogram_2d_vectorized(*args, bins=None, weights=None, density=False,
     """Calculate the histogram independently on each row of a 2D array"""
 
     N_inputs = len(args)
-    bins = _ensure_bins_is_a_list_of_arrays(bins, N_inputs)
     a0 = args[0]
 
     # consistency checks for inputa
@@ -128,7 +127,9 @@ def _histogram_2d_vectorized(*args, bins=None, weights=None, density=False,
     # just throw out everything outside of the bins, as np.histogram does
     # TODO: make this optional?
     slices = (slice(None),) + (N_inputs * (slice(1, -1),))
-    return bin_counts[slices]
+    bin_counts = bin_counts[slices]
+
+    return bin_counts
 
 
 def histogram(*args, bins=None, axis=None, weights=None, density=False,
@@ -242,9 +243,29 @@ def histogram(*args, bins=None, axis=None, weights=None, density=False,
     else:
         weights_reshaped = None
 
-    h = _histogram_2d_vectorized(*all_args_reshaped, bins=bins,
-                                 weights=weights_reshaped,
-                                 density=density, block_size=block_size)
+    n_inputs = len(all_args_reshaped)
+    bins = _ensure_bins_is_a_list_of_arrays(bins, n_inputs)
+
+    bin_counts = _histogram_2d_vectorized(*all_args_reshaped, bins=bins,
+                                          weights=weights_reshaped,
+                                          density=density,
+                                          block_size=block_size)
+
+    if density:
+        # Normalise by dividing by bin counts and areas such that all the
+        # histogram data integrated over all dimensions = 1
+        bin_widths = [np.diff(b) for b in bins]
+        if n_inputs == 1:
+            bin_areas = bin_widths[0]
+        elif n_inputs == 2:
+            bin_areas = np.outer(*bin_widths)
+        else:
+            # Slower, but N-dimensional logic
+            bin_areas = np.prod(np.ix_(*bin_widths))
+
+        h = bin_counts / bin_areas / bin_counts.sum()
+    else:
+        h = bin_counts
 
     if h.shape[0] == 1:
         assert do_full_array
diff --git a/xhistogram/xarray.py b/xhistogram/xarray.py
index c4d41d8..8b9af1f 100644
--- a/xhistogram/xarray.py
+++ b/xhistogram/xarray.py
@@ -126,7 +126,7 @@ def histogram(*args, bins=None, dim=None, weights=None, density=False,
         axis = None
 
     h_data = _histogram(*args_data, weights=weights_data, bins=bins, axis=axis,
-                        block_size=block_size)
+                        density=density, block_size=block_size)
 
     # create output dims
     new_dims = [a.name + bin_dim_suffix for a in args[:N_args]]
@@ -155,6 +155,12 @@ def histogram(*args, bins=None, dim=None, weights=None, density=False,
 
     da_out = xr.DataArray(h_data, dims=output_dims, coords=all_coords,
                           name=output_name)
+
+    if density:
+        # correct for overcounting the bins which weren't histogrammed along
+        n_bins_bystander_dims = da_out.isel(**{bd: 0 for bd in new_dims}).size
+        da_out = da_out * n_bins_bystander_dims
+
     return da_out
 
     # we need weights to be passed through apply_func's alignment algorithm,


Test patch:
diff --git a/xhistogram/test/test_core.py b/xhistogram/test/test_core.py
index 1ba4f88..a5b0a05 100644
--- a/xhistogram/test/test_core.py
+++ b/xhistogram/test/test_core.py
@@ -27,6 +27,33 @@ def test_histogram_results_1d(block_size):
     np.testing.assert_array_equal(hist, h_na)
 
 
+@pytest.mark.parametrize('block_size', [None, 1, 2])
+def test_histogram_results_1d_density(block_size):
+    nrows, ncols = 5, 20
+    data = np.random.randn(nrows, ncols)
+    bins = np.linspace(-4, 4, 10)
+
+    h = histogram(data, bins=bins, axis=1, block_size=block_size, density=True)
+    assert h.shape == (nrows, len(bins)-1)
+
+    # make sure we get the same thing as histogram
+    hist, _ = np.histogram(data, bins=bins, density=True)
+    np.testing.assert_allclose(hist, h.sum(axis=0))
+
+    # check integral is 1
+    widths = np.diff(bins)
+    integral = np.sum(hist * widths)
+    np.testing.assert_allclose(integral, 1.0)
+
+    # now try with no axis
+    h_na = histogram(data, bins=bins, block_size=block_size, density=True)
+    np.testing.assert_array_equal(hist, h_na)
+
+    # check integral is 1
+    integral = np.sum(h_na * widths)
+    np.testing.assert_allclose(integral, 1.0)
+
+
 @pytest.mark.parametrize('block_size', [None, 1, 2])
 def test_histogram_results_1d_weighted(block_size):
     nrows, ncols = 5, 20
@@ -52,7 +79,6 @@ def test_histogram_results_1d_weighted_broadcasting(block_size):
     np.testing.assert_array_equal(2*h, h_w)
 
 
-
 def test_histogram_results_2d():
     nrows, ncols = 5, 20
     data_a = np.random.randn(nrows, ncols)
@@ -70,6 +96,61 @@ def test_histogram_results_2d():
     np.testing.assert_array_equal(hist, h)
 
 
+def test_histogram_results_2d_density():
+    nrows, ncols = 5, 20
+    data_a = np.random.randn(nrows, ncols)
+    data_b = np.random.randn(nrows, ncols)
+    nbins_a = 9
+    bins_a = np.linspace(-4, 4, nbins_a + 1)
+    nbins_b = 10
+    bins_b = np.linspace(-4, 4, nbins_b + 1)
+
+    h = histogram(data_a, data_b, bins=[bins_a, bins_b], density=True)
+    assert h.shape == (nbins_a, nbins_b)
+
+    hist, _, _ = np.histogram2d(data_a.ravel(), data_b.ravel(),
+                                bins=[bins_a, bins_b], density=True)
+    np.testing.assert_allclose(hist, h)
+
+    # check integral is 1
+    widths_a = np.diff(bins_a)
+    widths_b = np.diff(bins_b)
+    areas = np.outer(widths_a, widths_b)
+    integral = np.sum(hist * areas)
+    np.testing.assert_allclose(integral, 1.0)
+
+
+def test_histogram_results_3d_density():
+    nrows, ncols = 5, 20
+    data_a = np.random.randn(nrows, ncols)
+    data_b = np.random.randn(nrows, ncols)
+    data_c = np.random.randn(nrows, ncols)
+    nbins_a = 9
+    bins_a = np.linspace(-4, 4, nbins_a + 1)
+    nbins_b = 10
+    bins_b = np.linspace(-4, 4, nbins_b + 1)
+    nbins_c = 9
+    bins_c = np.linspace(-4, 4, nbins_c + 1)
+
+    h = histogram(data_a, data_b, data_c, bins=[bins_a, bins_b, bins_c],
+                  density=True)
+
+    assert h.shape == (nbins_a, nbins_b, nbins_c)
+
+    hist, _ = np.histogramdd((data_a.ravel(), data_b.ravel(), data_c.ravel()),
+                             bins=[bins_a, bins_b, bins_c], density=True)
+
+    np.testing.assert_allclose(hist, h)
+
+    # check integral is 1
+    widths_a = np.diff(bins_a)
+    widths_b = np.diff(bins_b)
+    widths_c = np.diff(bins_c)
+    areas = np.einsum('i,j,k', widths_a, widths_b, widths_c)
+    integral = np.sum(hist * areas)
+    np.testing.assert_allclose(integral, 1.0)
+
+
 @pytest.mark.parametrize('block_size', [None, 5, 'auto'])
 @pytest.mark.parametrize('use_dask', [False, True])
 def test_histogram_shape(use_dask, block_size):
diff --git a/xhistogram/test/test_xarray.py b/xhistogram/test/test_xarray.py
index 0018f0d..7be7449 100644
--- a/xhistogram/test/test_xarray.py
+++ b/xhistogram/test/test_xarray.py
@@ -59,6 +59,32 @@ def test_histogram_ones(ones, ndims):
         _check_result(h, d)
 
 
+@pytest.mark.parametrize('ndims', [1, 2, 3, 4])
+def test_histogram_ones_density(ones, ndims):
+    dims = ones.dims
+    if ones.ndim < ndims:
+        pytest.skip(
+            "Don't need to test when number of dimension combinations "
+            "exceeds the number of array dimensions")
+
+    # everything should be in the middle bin (index 1)
+    bins = np.array([0, 0.9, 1.1, 2])
+    bin_area = 0.2
+
+    def _check_result(h_density, d):
+        other_dims = [dim for dim in ones.dims if dim not in d]
+        if len(other_dims) > 0:
+            assert set(other_dims) <= set(h_density.dims)
+
+        # check that all integrals over pdfs at different locations are = 1
+        h_integrals = (h_density * bin_area).sum(dim='ones_bin')
+        np.testing.assert_allclose(h_integrals.values, 1.0)
+
+    for d in combinations(dims, ndims):
+        h_density = histogram(ones, bins=[bins], dim=d, density=True)
+        _check_result(h_density, d)
+
+
 # TODO: refactor this test to use better fixtures
 # (it currently has a ton of loops)
 @pytest.mark.parametrize('ndims', [1, 2, 3, 4])