diff --git a/quantities/quantity.py b/quantities/quantity.py index 8fc355d..35b4ef7 100644 --- a/quantities/quantity.py +++ b/quantities/quantity.py @@ -16,6 +16,9 @@ # e.g. PREFERRED = [pq.mV, pq.pA, pq.UnitQuantity('femtocoulomb', 1e-15*pq.C, 'fC')] # Intended to be overwritten in down-stream packages +_np_version = tuple(map(int, np.__version__.split(".dev")[0].split("."))) + + def validate_unit_quantity(value): try: assert isinstance(value, Quantity) @@ -318,7 +321,7 @@ def __array_prepare__(self, obj, context=None): return res def __array_wrap__(self, obj, context=None, return_scalar=False): - _np_version = tuple(map(int, np.__version__.split(".dev")[0].split("."))) + # For NumPy < 2.0 we do old behavior if _np_version < (2, 0, 0): if not isinstance(obj, Quantity): @@ -342,7 +345,6 @@ def __add__(self, other): @scale_other_units def __radd__(self, other): return np.add(other, self) - return super().__radd__(other) @with_doc(np.ndarray.__iadd__) @scale_other_units @@ -358,7 +360,6 @@ def __sub__(self, other): @scale_other_units def __rsub__(self, other): return np.subtract(other, self) - return super().__rsub__(other) @with_doc(np.ndarray.__isub__) @scale_other_units @@ -378,22 +379,40 @@ def __imod__(self, other): @with_doc(np.ndarray.__imul__) @protected_multiplication def __imul__(self, other): - return super().__imul__(other) + # the following is an inelegant fix for the removal of __array_prepare__ in NumPy 2.x + # the longer-term solution is probably to implement __array_ufunc__ + # See: + # - https://numpy.org/devdocs/release/2.0.0-notes.html#array-prepare-is-removed + # - https://numpy.org/neps/nep-0013-ufunc-overrides.html + if _np_version < (2, 0, 0): + return super().__imul__(other) + else: + cself = self.copy() + cother = other.copy() + res = super().__imul__(other) + context = (np.multiply, (cself, cother, cself), 0) + return self.__array_prepare__(res, context=context) @with_doc(np.ndarray.__rmul__) def __rmul__(self, other): return np.multiply(other, self) - return super().__rmul__(other) @with_doc(np.ndarray.__itruediv__) @protected_multiplication def __itruediv__(self, other): - return super().__itruediv__(other) + # see comment above on __imul__ + if _np_version < (2, 0, 0): + return super().__itruediv__(other) + else: + cself = self.copy() + cother = other.copy() + res = super().__itruediv__(other) + context = (np.true_divide, (cself, cother, cself), 0) + return self.__array_prepare__(res, context=context) @with_doc(np.ndarray.__rtruediv__) def __rtruediv__(self, other): return np.true_divide(other, self) - return super().__rtruediv__(other) @with_doc(np.ndarray.__pow__) @check_uniform @@ -404,7 +423,15 @@ def __pow__(self, other): @check_uniform @protected_power def __ipow__(self, other): - return super().__ipow__(other) + # see comment above on __imul__ + if _np_version < (2, 0, 0): + return super().__ipow__(other) + else: + cself = self.copy() + cother = other.copy() + res = super().__ipow__(other) + context = (np.power, (cself, cother, cself), 0) + return self.__array_prepare__(res, context=context) def __round__(self, decimals=0): return np.around(self, decimals) @@ -528,7 +555,6 @@ def sum(self, axis=None, dtype=None, out=None): @with_doc(np.nansum) def nansum(self, axis=None, dtype=None, out=None): - import numpy as np return Quantity( np.nansum(self.magnitude, axis, dtype, out), self.dimensionality diff --git a/quantities/tests/test_arithmetic.py b/quantities/tests/test_arithmetic.py index 1b6aeeb..983b0af 100644 --- a/quantities/tests/test_arithmetic.py +++ b/quantities/tests/test_arithmetic.py @@ -367,12 +367,35 @@ def test_in_place_subtraction(self): self.assertRaises(ValueError, op.isub, [1, 2, 3]*pq.m, pq.J) self.assertRaises(ValueError, op.isub, [1, 2, 3]*pq.m, 5*pq.J) + def test_in_place_multiplication(self): + velocity = 3 * pq.m/pq.s + time = 2 * pq.s + + self.assertQuantityEqual(velocity * time, 6 * pq.m) + + distance = velocity.copy() + distance *= time + self.assertQuantityEqual(distance, 6 * pq.m) + def test_division(self): molar = pq.UnitQuantity('M', 1000 * pq.mole/pq.m**3, u_symbol='M') for subtr in [1, 1.0]: q = 1*molar/(1000*pq.mole/pq.m**3) self.assertQuantityEqual((q - subtr).simplified, 0) + a = np.array([5, 10, 15]) * pq.s + b = np.array([2, 4, 6]) * pq.kg + + c = a / b + self.assertQuantityEqual(c, np.array([2.5, 2.5, 2.5]) * pq.s / pq.kg) + + def test_in_place_division(self): + a = np.array([5, 10, 15]) * pq.s + b = np.array([2, 4, 6]) * pq.kg + + a /= b + self.assertQuantityEqual(a, np.array([2.5, 2.5, 2.5]) * pq.s / pq.kg) + def test_powering(self): # test raising a quantity to a power self.assertQuantityEqual((5.5 * pq.cm)**5, (5.5**5) * (pq.cm**5)) @@ -403,3 +426,12 @@ def q_pow_r(q1, q2): def ipow(q1, q2): q1 -= q2 self.assertRaises(ValueError, ipow, 1*pq.m, [1, 2]) + + def test_inplace_powering(self): + a = 5.5 * pq.cm + a **= 5 + self.assertQuantityEqual(a, (5.5**5) * (pq.cm**5)) + + b = np.array([1, 2, 3, 4, 5]) * pq.kg + b **= 3 + self.assertQuantityEqual(b, np.array([1, 8, 27, 64, 125]) * pq.kg**3)