diff --git a/pint/facets/numpy/numpy_func.py b/pint/facets/numpy/numpy_func.py index 6bccd409f..b79700f9f 100644 --- 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 index 979b6ee25..190ecddfc 100644 --- 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., 3.5, 3., 0.]) + assert np.array_equal(res.magnitude, ref) + assert res.units == "meter * second" + def test_dot(self): with ExitStack() as stack: stack.callback(