ENH: (NEP 18) implement and test np.in1d and np.interp

unyt/_array_functions.py

@@ -756,6 +756,12 @@ def isin(element, test_elements, *args, **kwargs):
     )
 
 
+@implements(np.in1d)
+def in1d(ar1, ar2, *args, **kwargs):
+    _validate_units_consistency((ar1, ar2))
+    return np.isin._implementation(np.asarray(ar1), np.asarray(ar2), *args, **kwargs)
+
+
 @implements(np.place)
 def place(arr, mask, vals, *args, **kwargs) -> None:
     _validate_units_consistency_v2(arr.units, vals)
@@ -938,3 +944,18 @@ def tensordot(a, b, *args, **kwargs):
 def unwrap(p, *args, **kwargs):
     ret_units = p.units
     return np.unwrap._implementation(p.view(np.ndarray), *args, **kwargs) * ret_units
+
+
+@implements(np.interp)
+def interp(x, xp, fp, *args, **kwargs):
+    _validate_units_consistency((x, xp))
+
+    # return array type should match fp's
+    # so, the fallback multiplier is 1 instead of NULL_UNITS
+    # This avoid leaking a dimensionless unyt_array if reference data
+    # is a pure np.ndarray
+    ret_units = getattr(fp, "units", 1)
+    return (
+        np.interp(np.asarray(x), np.asarray(xp), np.asarray(fp), *args, **kwargs)
+        * ret_units
+    )

unyt/tests/test_array_functions.py

@@ -183,8 +183,6 @@
 # this set represents all functions that need inspection, tests, or both
 # it is always possible that some of its elements belong in NOOP_FUNCTIONS
 TODO_FUNCTIONS = {
-    np.in1d,
-    np.interp,
     np.ix_,
     np.linalg.svd,
 }
@@ -1412,6 +1410,18 @@ def test_isin():
     assert np.isin(1 * cm, a)
 
 
+def test_in1d_mixed_units():
+    a = [1, 2, 3] * cm
+    with pytest.raises(UnitInconsistencyError):
+        np.in1d([1, 2], a)
+
+
+def test_in1d():
+    a = [1, 2, 3] * cm
+    b = [1, 2] * cm
+    assert np.all(np.in1d(b, a))
+
+
 def test_place_mixed_units():
     arr = np.arange(6).reshape(2, 3) * cm
     with pytest.raises(UnitInconsistencyError):
@@ -1662,3 +1672,30 @@ def test_unwrap():
     res = np.unwrap(phase)
     assert type(res) is unyt_array
     assert res.units == rad
+
+
+def test_interp():
+    _x = np.array([1.1, 2.2, 3.3])
+    _xp = np.array([1, 2, 3])
+    _fp = np.array([4, 8, 12])
+
+    # any of the three input array-like might be unitful
+    # let's test all relevant combinations
+    # return type should match fp's
+
+    with pytest.raises(UnitInconsistencyError):
+        np.interp(_x * cm, _xp, _fp)
+
+    with pytest.raises(UnitInconsistencyError):
+        res = np.interp(_x, _xp * cm, _fp)
+
+    res = np.interp(_x * cm, _xp * cm, _fp)
+    assert type(res) is np.ndarray
+
+    res = np.interp(_x, _xp, _fp * K)
+    assert type(res) is unyt_array
+    assert res.units == K
+
+    res = np.interp(_x * cm, _xp * cm, _fp * K)
+    assert type(res) is unyt_array
+    assert res.units == K