Properly handle offset units for trapz (Fixes #1593)

2 files changed, 54 insertions, 2 deletions

diff --git a/pint/facets/numpy/numpy_func.py b/pint/facets/numpy/numpy_func.py
index 2a4421c..2a004a8 100644
--- a/pint/facets/numpy/numpy_func.py
+++ b/pint/facets/numpy/numpy_func.py
@@ -13,7 +13,7 @@ from inspect import signature
 from itertools import chain
 
 from ...compat import is_upcast_type, np, zero_or_nan
-from ...errors import DimensionalityError, UnitStrippedWarning
+from ...errors import DimensionalityError, OffsetUnitCalculusError, UnitStrippedWarning
 from ...util import iterable, sized
 
 HANDLED_UFUNCS = {}
@@ -729,6 +729,36 @@ for name in ["prod", "nanprod"]:
     implement_prod_func(name)
 
 
+def _base_unit_if_needed(a):
+    if a._is_multiplicative:
+        return a
+    else:
+        if a.units._REGISTRY.autoconvert_offset_to_baseunit:
+            return a.to_base_units()
+        else:
+            raise OffsetUnitCalculusError(a.units)
+
+
+@implements("trapz", "function")
+def _trapz(a, x=None, dx=1.0, **kwargs):
+    a = _base_unit_if_needed(a)
+    units = a.units
+    if x is not None:
+        if hasattr(x, "units"):
+            x = _base_unit_if_needed(x)
+            units *= x.units
+            x = x._magnitude
+        ret = np.trapz(a._magnitude, x, **kwargs)
+    else:
+        if hasattr(dx, "units"):
+            dx = _base_unit_if_needed(dx)
+            units *= dx.units
+            dx = dx._magnitude
+        ret = np.trapz(a._magnitude, dx=dx, **kwargs)
+
+    return a.units._REGISTRY.Quantity(ret, units)
+
+
 # Implement simple matching-unit or stripped-unit functions based on signature
 
 
@@ -920,7 +950,7 @@ for func_str in [
 # Handle functions with output unit defined by operation
 for func_str in ["std", "nanstd", "sum", "nansum", "cumsum", "nancumsum"]:
     implement_func("function", func_str, input_units=None, output_unit="sum")
-for func_str in ["cross", "trapz", "dot"]:
+for func_str in ["cross", "dot"]:
     implement_func("function", func_str, input_units=None, output_unit="mul")
 for func_str in ["diff", "ediff1d"]:
     implement_func("function", func_str, input_units=None, output_unit="delta")
diff --git a/pint/testsuite/test_numpy_func.py b/pint/testsuite/test_numpy_func.py
index 49caa32..4f1488c 100644
--- a/pint/testsuite/test_numpy_func.py
+++ b/pint/testsuite/test_numpy_func.py
@@ -1,3 +1,4 @@
+from contextlib import ExitStack
 from unittest.mock import patch
 
 import pytest
@@ -191,3 +192,24 @@ class TestNumPyFuncUtils(TestNumpyMethods):
             numpy_wrap("invalid", np.ones, [], {}, [])
         # TODO (#905 follow-up): test that NotImplemented is returned when upcast types
         # present
+
+    def test_trapz(self):
+        with ExitStack() as stack:
+            stack.callback(
+                setattr,
+                self.ureg,
+                "autoconvert_offset_to_baseunit",
+                self.ureg.autoconvert_offset_to_baseunit,
+            )
+            self.ureg.autoconvert_offset_to_baseunit = True
+            t = self.Q_(np.array([0.0, 4.0, 8.0]), "degC")
+            z = self.Q_(np.array([0.0, 2.0, 4.0]), "m")
+            helpers.assert_quantity_equal(
+                np.trapz(t, x=z), self.Q_(1108.6, "kelvin meter")
+            )
+
+    def test_trapz_no_autoconvert(self):
+        t = self.Q_(np.array([0.0, 4.0, 8.0]), "degC")
+        z = self.Q_(np.array([0.0, 2.0, 4.0]), "m")
+        with pytest.raises(OffsetUnitCalculusError):
+            np.trapz(t, x=z)