Fixed bug where numpy.histogram2d() was not propagating units. (#61)

sun-data · Jul 17, 2024 · cab7459 · cab7459
1 parent c5ace7a
commit cab7459
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/named_arrays/_scalars/scalar_named_array_functions.py b/named_arrays/_scalars/scalar_named_array_functions.py
@@ -298,6 +298,14 @@ def histogram2d(
     xedges = na.ScalarArray.empty(shape_x)
     yedges = na.ScalarArray.empty(shape_y)
 
+    unit_weights = na.unit(weights)
+    unit_x = na.unit(x)
+    unit_y = na.unit(y)
+
+    hist = hist if unit_weights is None else hist << unit_weights
+    xedges = xedges if unit_x is None else xedges << unit_x
+    yedges = yedges if unit_y is None else yedges << unit_y
+
     for i in na.ndindex(shape_orthogonal):
         min_i = min[i]
         max_i = max[i]

diff --git a/named_arrays/_scalars/tests/test_scalars.py b/named_arrays/_scalars/tests/test_scalars.py
@@ -254,6 +254,12 @@ def test_histogram2d(
                 assert np.all(result.outputs >= 0)
                 assert result.outputs.sum() > 0
 
+                unit = array.unit_normalized
+                unit_weights = na.unit_normalized(weights)
+                assert result.inputs.x.unit_normalized.is_equivalent(unit)
+                assert result.inputs.y.unit_normalized.is_equivalent(unit)
+                assert result.outputs.unit_normalized.is_equivalent(unit_weights)
+
 
 class AbstractTestAbstractScalarArray(
     AbstractTestAbstractScalar,