Make np.digitize unit-aware

unyt/_array_functions.py

@@ -385,6 +385,22 @@ def histogram_bin_edges(a, *args, **kwargs):
     )
 
 
+@implements(np.digitize)
+def digitize(x, bins, right=False):
+    # convert bins to x's units before binning, reusing the same helper as the
+    # histogram functions. This both makes the comparison unit-aware (so e.g.
+    # grams and kilograms give consistent indices) and raises on dimensionally
+    # incompatible bins instead of silently comparing raw values (see #633).
+    if not hasattr(bins, "units"):
+        # let _sanitize_bins operate on an array (it divides bins by a
+        # conversion factor, which a plain list does not support)
+        bins = np.asarray(bins)
+    sanitized_bins = _sanitize_bins(x, bins)
+    return np.digitize._implementation(
+        np.asarray(x), np.asarray(sanitized_bins), right=right
+    )
+
+
 def get_units(objs):
     units = []
     for sub in objs:

unyt/tests/test_array_functions.py

@@ -148,7 +148,6 @@
     np.extract,  # works out of the box (tested)
     np.setxor1d,  # we get it for free with previously implemented functions (tested)
     np.lexsort,  # returns pure numbers
-    np.digitize,  # returns pure numbers
     np.tril_indices_from,  # returns pure numbers
     np.triu_indices_from,  # returns pure numbers
     np.imag,  # works out of the box (tested)
@@ -2097,6 +2096,30 @@ def test_searchsorted(val):
     assert res == 2
 
 
+def test_digitize_converts_bins():
+    # bins in different (but compatible) units must be converted to the data's
+    # units before binning: 1, 2, 3 cm are all far below 1 km, so they all land
+    # in the first bin (see #633)
+    res = np.digitize([1, 2, 3] * cm, [0, 1, 2, 3] * km)
+    assert_array_equal_units(res, np.array([1, 1, 1]))
+
+
+def test_digitize_right():
+    res = np.digitize([1, 2, 3] * cm, [0, 1, 2, 3] * km, right=True)
+    assert_array_equal_units(res, np.array([1, 1, 1]))
+
+
+def test_digitize_dimensionless_plain_bins():
+    # a dimensionless array with plain (unit-less) bins must still work
+    res = np.digitize([1, 2, 3] * dimensionless, [0, 1, 2, 3])
+    assert_array_equal_units(res, np.array([2, 3, 4]))
+
+
+def test_digitize_mixed_units():
+    with pytest.raises(UnitConversionError):
+        np.digitize([1, 2, 3] * g, [0, 1, 2, 3] * s)
+
+
 @pytest.mark.parametrize(
     "choicelist_gen, default",
     [