add support for numpy.correlate and unit test (#1990)

ikondov · web-flow · commit 9972e60ead95 · 2024-05-22T03:24:17.000+09:00
diff --git a/pint/facets/numpy/numpy_func.py b/pint/facets/numpy/numpy_func.py
@@ -779,6 +779,15 @@ def _trapz(y, x=None, dx=1.0, **kwargs):
     return y.units._REGISTRY.Quantity(ret, units)
 
 
+@implements("correlate", "function")
+def _correlate(a, v, mode="valid", **kwargs):
+    a = _base_unit_if_needed(a)
+    v = _base_unit_if_needed(v)
+    units = a.units * v.units
+    ret = np.correlate(a._magnitude, v._magnitude, mode=mode, **kwargs)
+    return a.units._REGISTRY.Quantity(ret, units)
+
+
 def implement_mul_func(func):
     # If NumPy is not available, do not attempt implement that which does not exist
     if np is None:
diff --git a/pint/testsuite/test_numpy_func.py b/pint/testsuite/test_numpy_func.py
@@ -216,6 +216,14 @@ def test_trapz_no_autoconvert(self):
         with pytest.raises(OffsetUnitCalculusError):
             np.trapz(t, x=z)
 
+    def test_correlate(self):
+        a = self.Q_(np.array([1, 2, 3]), "m")
+        v = self.Q_(np.array([0, 1, 0.5]), "s")
+        res = np.correlate(a, v, "full")
+        ref = np.array([0.5, 2.0, 3.5, 3.0, 0.0])
+        assert np.array_equal(res.magnitude, ref)
+        assert res.units == "meter * second"
+
     def test_dot(self):
         with ExitStack() as stack:
             stack.callback(
