Repository navigation
Fix NaN-aware reductions and NumPy 2 output units - #276
Conversation
apdavison
left a comment
There was a problem hiding this comment.
Thanks for this, it's a nice fix, and it simplifies the in-place operators as well. I've checked it against master on NumPy 2.0 and 2.5, and apart from the bug fixes I found no differences in behaviour. One request before merging:
Raise an Exception when a ufunc would change the units of a view. With this PR, writing into a view of a Quantity gives the view new units, but the base array keeps its old ones:
a = [1., 2., 3.] * pq.m
v = a[:2]
np.multiply(v, v, out=v)
v.units # m**2
a # [1., 4., 3.] m, data changed but units didn'tOn master the units are also wrong, so this isn't a regression. But the in-place operators already prevent it: v *= v raises ValueError: can not modify units of a view of a Quantity. Could __array_prepare__ apply the same check when the result is a view of a Quantity and its dimensionality would change? The check in protected_multiplication could be reused. A test alongside test_ufunc_output_units would be good too.
np.nanstd([1., 2.] * pq.ms)currently raises because its internal sum passeskeepdims. Simply accepting that argument is insufficient on NumPy 2: the variance calculation uses in-place ufunc outputs, whose dimensionality is not updated, so the final standard deviation can have the wrong units.This forwards
keepdimsthroughQuantity.sumand applies the existing dimensionality handlers to NumPy 2 ufunc outputs even when the result is already a Quantity. The duplicated in-place multiply/divide/power repair code can then be removed; power still copies its exponent to handle an exponent that aliases the magnitude buffer.clippasses the output magnitude to its magnitude-only calculation, as the other reduction wrappers do.Fixes #223. This uses the existing
__array_wrap__path rather than adding a new__array_ufunc__implementation.Validation on Python 3.10 / NumPy 1.24.4 and Python 3.12 / NumPy 2.0.2 and 2.4.4: full suite 182 passed, 3 expected failures, 24 subtests passed on each. README doctests pass; mypy passes on NumPy 1.24 and 2.4. Regression coverage includes real/complex inputs, axes, kept dimensions, caller-owned outputs, unchanged input, degrees of freedom, chained output-unit changes, and an aliased scalar exponent. The new tests fail on the base revision. The existing NumPy 2.4 pickle deprecation warning remains.
AI-assisted implementation and validation. No claim of support for additional ufuncs or all NumPy reduction keyword options.
Review revision: direct ufunc outputs now raise
ValueErrorwhen they would change the units of a view of a Quantity, using the same ownership rule as the in-place operators. Unit-preserving view outputs and newly allocated results remain supported. Added checks for multiply, divide, square, sqrt, and distinct output views; all three local NumPy environments now pass 183 tests, 3 expected failures, 28 subtests. README doctests and mypy checks also pass.NumPy 2 calls this existing wrapping hook after numerical computation: the exception prevents silently assigning inconsistent units, but it does not roll back values already written to
out. This revision does not add transactional output handling or a new__array_ufunc__implementation.