Repository: yt-project/unyt

Issue:
Incorrect units or crash in histograms when density and/or weights present
* unyt version: 3.0.3.post34+gecd71ea
* Python version: 3.12
* Operating System: Linux (Fedora 40)

### Description

When making a histogram with `density=True` I expect the histogram to have units equal to the inverse of those of the input array.

When making a histogram with `weights=...` and the weights have units, I expect the histogram to have units equal to the units of the weights (if they have any).

When both `weights=...` and `density=True` are given I expect the combination of the above.

These principles should also extend to 2d and dd histograms.

### What I Did

The basic `histogram` just doesn't apply any units to the first return value:
```
>>> np.histogram(np.arange(10) * u.cm, density=True, weights=np.arange(10) * u.s)
(array([0.        , 0.02469136, 0.04938272, 0.07407407, 0.09876543,
       0.12345679, 0.14814815, 0.17283951, 0.19753086, 0.22222222]), unyt_array([0. , 0.9, 1.8, 2.7, 3.6, 4.5, 5.4, 6.3, 7.2, 8.1, 9. ], 'cm'))
```
The `histogram2d` actually crashes when weights with units are given:
```
>>> np.histogram2d(np.arange(10) * u.cm, np.arange(10) * u.g, density=True, weights=np.arange(10) * u.s)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/txwx36/code/unyt/unyt/array.py", line 2045, in __array_function__
    return _HANDLED_FUNCTIONS[func](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/txwx36/code/unyt/unyt/_array_functions.py", line 179, in histogram2d
    counts, xbins, ybins = np.histogram2d._implementation(
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/txwx36/.venv/koman/lib64/python3.12/site-packages/numpy/lib/_twodim_base_impl.py", line 813, in histogram2d
    hist, edges = histogramdd([x, y], bins, range, density, weights)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/txwx36/code/unyt/unyt/array.py", line 2045, in __array_function__
    return _HANDLED_FUNCTIONS[func](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/txwx36/code/unyt/unyt/_array_functions.py", line 187, in histogramdd
    units = [_.units for _ in sample]
             ^^^^^^^
AttributeError: 'numpy.ndarray' object has no attribute 'units'
```

Discussion:

Patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index fa2c478..62aca69 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -158,38 +158,162 @@ def _sanitize_range(_range, units):
     return new_range.squeeze()
 
 
-@implements(np.histogram)
-def histogram(
-    a,
-    bins=10,
-    range=None,
-    *args,
-    **kwargs,
-):
-    range = _sanitize_range(range, units=[a.units])
-    counts, bins = np.histogram._implementation(
-        np.asarray(a), bins, range, *args, **kwargs
-    )
-    return counts, bins * a.units
+def _histogram(a, *, bins=10, range=None, density=None, weights=None, normed=None):
+    range = _sanitize_range(range, units=[getattr(a, "units", None)])
+    if NUMPY_VERSION >= Version("1.24"):
+        counts, bins = np.histogram._implementation(
+            np.asarray(a),
+            bins=bins,
+            range=range,
+            density=density,
+            weights=np.asarray(weights) if weights is not None else None,
+        )
+    else:
+        counts, bins = np.histogram._implementation(
+            np.asarray(a),
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=np.asarray(weights) if weights is not None else None,
+            density=density,
+        )
+    # a and/or weights could have units, only apply if present
+    # don't getattr(..., "units", NULL_UNIT) because e.g. we don't want
+    # a unyt_array if weights are not a unyt_array and not density
+    if density and hasattr(a, "units"):
+        counts /= a.units
+    if weights is not None and hasattr(weights, "units"):
+        counts *= weights.units
+    return counts, bins * getattr(a, "units", 1)
 
 
-@implements(np.histogram2d)
-def histogram2d(x, y, bins=10, range=None, *args, **kwargs):
-    range = _sanitize_range(range, units=[x.units, y.units])
-    counts, xbins, ybins = np.histogram2d._implementation(
-        np.asarray(x), np.asarray(y), bins, range, *args, **kwargs
-    )
-    return counts, xbins * x.units, ybins * y.units
+if NUMPY_VERSION >= Version("1.24"):
 
+    @implements(np.histogram)
+    def histogram(a, bins=10, range=None, density=None, weights=None):
+        return _histogram(a, bins=bins, range=range, density=density, weights=weights)
 
-@implements(np.histogramdd)
-def histogramdd(sample, bins=10, range=None, *args, **kwargs):
-    units = [_.units for _ in sample]
-    range = _sanitize_range(range, units=units)
-    counts, bins = np.histogramdd._implementation(
-        [np.asarray(_) for _ in sample], bins, range, *args, **kwargs
+else:
+
+    @implements(np.histogram)
+    def histogram(a, bins=10, range=None, normed=None, weights=None, density=None):
+        return _histogram(
+            a, bins=bins, range=range, normed=normed, weights=weights, density=density
+        )
+
+
+def _histogram2d(x, y, *, bins=10, range=None, density=None, weights=None, normed=None):
+    range = _sanitize_range(
+        range, units=[getattr(x, "units", None), getattr(y, "units", None)]
     )
-    return counts, tuple(_bin * u for _bin, u in zip(bins, units))
+    if NUMPY_VERSION >= Version("1.24"):
+        counts, xbins, ybins = np.histogram2d._implementation(
+            np.asarray(x),
+            np.asarray(y),
+            bins=bins,
+            range=range,
+            density=density,
+            weights=np.asarray(weights) if weights is not None else None,
+        )
+    else:
+        counts, xbins, ybins = np.histogram2d._implementation(
+            np.asarray(x),
+            np.asarray(y),
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=np.asarray(weights) if weights is not None else None,
+            density=density,
+        )
+    # x, y and/or weights could have units, only apply if present
+    # don't getattr(..., "units", NULL_UNIT) because e.g. we don't want
+    # a unyt_array if weights are not a unyt_array and not density
+    if density:
+        if hasattr(x, "units"):
+            counts /= x.units
+        if hasattr(y, "units"):
+            counts /= y.units
+    if weights is not None and hasattr(weights, "units"):
+        counts *= weights.units
+    return counts, xbins * getattr(x, "units", 1), ybins * getattr(y, "units", 1)
+
+
+if NUMPY_VERSION >= Version("1.24"):
+
+    @implements(np.histogram2d)
+    def histogram2d(x, y, bins=10, range=None, density=None, weights=None):
+        return _histogram2d(
+            x, y, bins=bins, range=range, density=density, weights=weights
+        )
+
+else:
+
+    @implements(np.histogram2d)
+    def histogram2d(x, y, bins=10, range=None, normed=None, weights=None, density=None):
+        return _histogram2d(
+            x,
+            y,
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=weights,
+            density=density,
+        )
+
+
+def _histogramdd(
+    sample, *, bins=10, range=None, density=None, weights=None, normed=None
+):
+    range = _sanitize_range(range, units=[getattr(_, "units", None) for _ in sample])
+    if NUMPY_VERSION >= Version("1.24"):
+        counts, bins = np.histogramdd._implementation(
+            [np.asarray(_) for _ in sample],
+            bins=bins,
+            range=range,
+            density=density,
+            weights=np.asarray(weights) if weights is not None else None,
+        )
+    else:
+        counts, bins = np.histogramdd._implementation(
+            [np.asarray(_) for _ in sample],
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=np.asarray(weights) if weights is not None else None,
+            density=density,
+        )
+    # sample(s) and/or weights could have units, only apply if present
+    # don't getattr(..., "units", NULL_UNIT) because e.g. we don't want
+    # a unyt_array if weights are not a unyt_array and not density
+    if density:
+        for s in sample:
+            counts /= getattr(s, "units", 1)
+    counts *= getattr(weights, "units", 1)
+    return counts, tuple(_bin * getattr(s, "units", 1) for _bin, s in zip(bins, sample))
+
+
+if NUMPY_VERSION >= Version("1.24"):
+
+    @implements(np.histogramdd)
+    def histogramdd(sample, bins=10, range=None, density=None, weights=None):
+        return _histogramdd(
+            sample, bins=bins, range=range, density=density, weights=weights
+        )
+
+else:
+
+    @implements(np.histogramdd)
+    def histogramdd(
+        sample, bins=10, range=None, normed=None, weights=None, density=None
+    ):
+        return _histogramdd(
+            sample,
+            bins=bins,
+            range=range,
+            normed=normed,
+            weights=weights,
+            density=density,
+        )
 
 
 @implements(np.histogram_bin_edges)
@@ -535,7 +659,8 @@ def copyto(dst, src, *args, **kwargs):
 
 @implements(np.prod)
 def prod(a, *args, **kwargs):
-    return np.prod._implementation(np.asarray(a), *args, **kwargs) * a.units**a.size
+    res = np.prod._implementation(np.asarray(a), *args, **kwargs)
+    return res * a.units ** (a.size // res.size)
 
 
 @implements(np.var)


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index bfd900a..a039939 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -646,51 +646,197 @@ def test_linalg_tensordot():
     assert_array_equal_units(res, ref)
 
 
-def test_histogram():
-    rng = np.random.default_rng()
-    arr = rng.normal(size=1000) * cm
-    counts, bins = np.histogram(arr, bins=10, range=(arr.min(), arr.max()))
-    assert type(counts) is np.ndarray
-    assert bins.units == arr.units
-
-
-def test_histogram_implicit_units():
-    # see https://github.com/yt-project/unyt/issues/465
-    rng = np.random.default_rng()
-    arr = rng.normal(size=1000) * cm
-    counts, bins = np.histogram(arr, bins=10, range=(arr.min().value, arr.max().value))
-    assert type(counts) is np.ndarray
-    assert bins.units == arr.units
-
-
-def test_histogram2d():
-    rng = np.random.default_rng()
-    x = rng.normal(size=100) * cm
-    y = rng.normal(loc=10, size=100) * s
-    counts, xbins, ybins = np.histogram2d(x, y)
-    assert counts.ndim == 2
-    assert xbins.units == x.units
-    assert ybins.units == y.units
-
-
-def test_histogramdd():
-    rng = np.random.default_rng()
-    x = rng.normal(size=100) * cm
-    y = rng.normal(size=100) * s
-    z = rng.normal(size=100) * g
-    counts, (xbins, ybins, zbins) = np.histogramdd((x, y, z))
-    assert counts.ndim == 3
-    assert xbins.units == x.units
-    assert ybins.units == y.units
-    assert zbins.units == z.units
-
+class TestHistograms:
+    def test_histogram(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        counts, bins = np.histogram(arr, bins=10, range=(arr.min(), arr.max()))
+        assert type(counts) is np.ndarray
+        assert bins.units == arr.units
+
+    def test_histogram_implicit_units(self):
+        # see https://github.com/yt-project/unyt/issues/465
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        counts, bins = np.histogram(
+            arr, bins=10, range=(arr.min().value, arr.max().value)
+        )
+        assert type(counts) is np.ndarray
+        assert bins.units == arr.units
+
+    def test_histogram_with_density(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        density, bins = np.histogram(
+            arr, bins=10, range=(arr.min(), arr.max()), density=True
+        )
+        assert type(density) is unyt_array
+        assert density.units == arr.units**-1
+        assert bins.units == arr.units
+
+    def test_histogram_with_weights(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        w = rng.uniform(size=1000) * g
+        wcounts, wbins = np.histogram(
+            arr, bins=10, range=(arr.min(), arr.max()), weights=w
+        )
+        assert type(wcounts) is unyt_array
+        assert wcounts.units == w.units
+        assert wbins.units == arr.units
+
+    def test_histogram_with_weights_and_density(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        w = rng.uniform(size=1000) * g
+        wdensity, wdbins = np.histogram(
+            arr, bins=10, range=(arr.min(), arr.max()), density=True, weights=w
+        )
+        assert type(wdensity) is unyt_array
+        assert wdensity.units == w.units / arr.units
+        assert wdbins.units == arr.units
+
+    def test_histogram_with_weights_and_dimless_arr(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        w = rng.uniform(size=1000) * g
+        wcounts2, wbins2 = np.histogram(arr.to_value(arr.units), weights=w)
+        assert type(wcounts2) is unyt_array
+        assert wcounts2.units == w.units
+        assert not hasattr(wbins2, "units")
+
+    def test_histogram2d(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        counts, xbins, ybins = np.histogram2d(x, y)
+        assert counts.ndim == 2
+        assert xbins.units == x.units
+        assert ybins.units == y.units
+
+    def test_histogram2d_with_density(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        density, xbins, ybins = np.histogram2d(x, y, density=True)
+        assert density.ndim == 2
+        assert type(density) is unyt_array
+        assert density.units == (x.units * y.units) ** -1
+        assert xbins.units == x.units
+        assert ybins.units == y.units
+
+    def test_histogram2d_with_weights(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        w = rng.uniform(size=100) * g
+        wcounts, xwbins, ywbins = np.histogram2d(x, y, weights=w)
+        assert wcounts.ndim == 2
+        assert type(wcounts) is unyt_array
+        assert wcounts.units == w.units
+        assert xwbins.units == x.units
+        assert ywbins.units == y.units
+
+    def test_histogram2d_with_weights_and_density(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        w = rng.uniform(size=100) * g
+        wdensity, xwdbins, ywdbins = np.histogram2d(x, y, weights=w, density=True)
+        assert wdensity.ndim == 2
+        assert type(wdensity) is unyt_array
+        assert wdensity.units == w.units / (x.units * y.units)
+        assert xwdbins.units == x.units
+        assert ywdbins.units == y.units
+
+    def test_histogram2d_with_weights_and_dimless_arr(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        w = rng.uniform(size=100) * g
+        wcounts2, xwbins2, ywbins2 = np.histogram2d(
+            x.to_value(x.units), y.to_value(y.units), weights=w
+        )
+        assert type(wcounts2) is unyt_array
+        assert wcounts2.units == w.units
+        assert not hasattr(xwbins2, "units")
+        assert not hasattr(ywbins2, "units")
+
+    def test_histogramdd(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(size=100) * s
+        z = rng.normal(size=100) * g
+        counts, (xbins, ybins, zbins) = np.histogramdd((x, y, z))
+        assert counts.ndim == 3
+        assert xbins.units == x.units
+        assert ybins.units == y.units
+        assert zbins.units == z.units
+
+    def test_histogramdd_with_density(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        z = rng.normal(size=100) * g
+        density, (xbins, ybins, zbins) = np.histogramdd((x, y, z), density=True)
+        assert density.ndim == 3
+        assert type(density) is unyt_array
+        assert density.units == (x.units * y.units * z.units) ** -1
+        assert xbins.units == x.units
+        assert ybins.units == y.units
+        assert zbins.units == z.units
+
+    def test_histogramdd_with_weights(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        z = rng.normal(size=100) * g
+        w = rng.uniform(size=100) * K
+        wcounts, (xwbins, ywbins, zwbins) = np.histogramdd((x, y, z), weights=w)
+        assert wcounts.ndim == 3
+        assert type(wcounts) is unyt_array
+        assert wcounts.units == w.units
+        assert xwbins.units == x.units
+        assert ywbins.units == y.units
+        assert zwbins.units == z.units
+
+    def test_histogramdd_with_weights_and_density(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        z = rng.normal(size=100) * g
+        w = rng.uniform(size=100) * K
+        wdensity, (xwdbins, ywdbins, zwdbins) = np.histogramdd(
+            (x, y, z), weights=w, density=True
+        )
+        assert wdensity.ndim == 3
+        assert type(wdensity) is unyt_array
+        assert wdensity.units == w.units / (x.units * y.units * z.units)
+        assert xwdbins.units == x.units
+        assert ywdbins.units == y.units
+        assert zwdbins.units == z.units
+
+    def test_histogramdd_with_weights_and_dimless_arr(self):
+        rng = np.random.default_rng()
+        x = rng.normal(size=100) * cm
+        y = rng.normal(loc=10, size=100) * s
+        z = rng.normal(size=100) * g
+        w = rng.uniform(size=100) * K
+        wcounts2, (xwbins2, ywbins2, zwbins2) = np.histogramdd(
+            (x.to_value(x.units), y.to_value(y.units), z.to_value(z.units)), weights=w
+        )
+        assert type(wcounts2) is unyt_array
+        assert wcounts2.units == w.units
+        assert not hasattr(xwbins2, "units")
+        assert not hasattr(ywbins2, "units")
+        assert not hasattr(zwbins2, "units")
 
-def test_histogram_bin_edges():
-    rng = np.random.default_rng()
-    arr = rng.normal(size=1000) * cm
-    bins = np.histogram_bin_edges(arr)
-    assert type(bins) is unyt_array
-    assert bins.units == arr.units
+    def test_histogram_bin_edges(self):
+        rng = np.random.default_rng()
+        arr = rng.normal(size=1000) * cm
+        bins = np.histogram_bin_edges(arr)
+        assert type(bins) is unyt_array
+        assert bins.units == arr.units
 
 
 def test_concatenate():
@@ -1291,6 +1437,18 @@ def test_scalar_reducer(func, expected_units):
     assert y.units == expected_units
 
 
+@pytest.mark.parametrize("func", [np.prod, np.nanprod])
+def test_prod_with_axis(func):
+    x = [
+        [1, 2, 3],
+        [4, 5, 6],
+        [7, 8, 9],
+    ] * cm
+    y = func(x, axis=0)
+    assert type(y) is unyt_array
+    assert y.units == cm**3
+
+
 @pytest.mark.parametrize(
     "func",
     [