Skip to content

Fix inplace arithmetic operators for numpy>=2 #271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
apdavison merged 3 commits into from
Jan 16, 2026
+67 −9
Merged

Fix inplace arithmetic operators for numpy>=2 #271

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
aec1f81
Added unit test for in-place multiplication, division and power, base…
apdavison Jan 13, 2026
529fce6
Fix for failure to handle units when performing in-place multiplicati…
apdavison Jan 15, 2026
0bb68ee
avoid unecessary copies
apdavison Jan 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 35 additions & 9 deletions quantities/quantity.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down
32 changes: 32 additions & 0 deletions quantities/tests/test_arithmetic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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)