diff --git a/astropy/units/quantity_helper/function_helpers.py b/astropy/units/quantity_helper/function_helpers.py index 380253ec7343..11dc6d9812e1 100644 --- a/astropy/units/quantity_helper/function_helpers.py +++ b/astropy/units/quantity_helper/function_helpers.py @@ -133,7 +133,7 @@ # trapz was renamed to trapezoid SUBCLASS_SAFE_FUNCTIONS |= {np.trapezoid} if not NUMPY_LT_2_1: - SUBCLASS_SAFE_FUNCTIONS |= {np.unstack} + SUBCLASS_SAFE_FUNCTIONS |= {np.unstack, np.cumulative_prod, np.cumulative_sum} # Implemented as methods on Quantity: # np.ediff1d is from setops, but we support it anyway; the others diff --git a/astropy/units/tests/test_quantity_non_ufuncs.py b/astropy/units/tests/test_quantity_non_ufuncs.py index c0adcac72068..e380dd8a091f 100644 --- a/astropy/units/tests/test_quantity_non_ufuncs.py +++ b/astropy/units/tests/test_quantity_non_ufuncs.py @@ -671,6 +671,10 @@ def test_sum(self): def test_cumsum(self): self.check(np.cumsum) + @pytest.mark.skipif(NUMPY_LT_2_1, reason="np.cumulative_sum is new in NumPy 2.1") + def test_cumulative_sum(self): + self.check(np.cumulative_sum, axis=1) + def test_any(self): with pytest.raises(TypeError): np.any(self.q) @@ -713,6 +717,11 @@ def test_cumproduct(self): with pytest.raises(u.UnitsError): np.cumproduct(self.q) # noqa: NPY003, NPY201 + @pytest.mark.skipif(NUMPY_LT_2_1, reason="np.cumulative_prod is new in NumPy 2.1") + def test_cumulative_prod(self): + with pytest.raises(u.UnitsError): + np.cumulative_prod(self.q, axis=1) + class TestUfuncLike(InvariantUnitTestSetup): def test_ptp(self): diff --git a/astropy/utils/masked/tests/test_function_helpers.py b/astropy/utils/masked/tests/test_function_helpers.py index 4761967782f4..8468802f1e0a 100644 --- a/astropy/utils/masked/tests/test_function_helpers.py +++ b/astropy/utils/masked/tests/test_function_helpers.py @@ -891,6 +891,31 @@ def test_array_equiv(self): assert np.array_equiv(self.mb, np.stack([self.mb, self.mb])) +class TestArrayAPI: + @classmethod + def setup_class(self): + self.a = np.tile(np.arange(5.0), 2).reshape(2, 5) + self.mask_a = np.array([[False] * 5, [True] * 4 + [False]]) + self.ma = Masked(self.a, mask=self.mask_a) + + def check(self, func, *args, **kwargs): + out = func(self.ma, *args, **kwargs) + expected = func(self.a, *args, **kwargs) + assert type(out) is MaskedNDArray + assert out.dtype.kind == "f" + assert_array_equal(out.unmasked, expected) + assert_array_equal(out.mask, self.mask_a) + assert not np.may_share_memory(out.mask, self.mask_a) + + @pytest.mark.skipif(NUMPY_LT_2_1, reason="np.cumulative_prod is new in NumPy 2.1") + def test_cumulative_prod(self): + self.check(np.cumulative_prod, axis=0) + + @pytest.mark.skipif(NUMPY_LT_2_1, reason="np.cumulative_sum is new in NumPy 2.1") + def test_cumulative_sum(self): + self.check(np.cumulative_sum, axis=0) + + class TestOuterLikeFunctions(MaskedArraySetup): def test_outer(self): result = np.outer(self.ma, self.mb)