diff --git a/src/blosc2/__init__.py b/src/blosc2/__init__.py index 12a2a908..c820e199 100644 --- a/src/blosc2/__init__.py +++ b/src/blosc2/__init__.py @@ -433,6 +433,8 @@ def _raise(exc): full_like, save, stack, + cumulative_sum, + cumulative_prod, ) from .embed_store import EmbedStore, estore_from_cframe from .dict_store import DictStore diff --git a/src/blosc2/ndarray.py b/src/blosc2/ndarray.py index 3a33c54c..00b2def6 100644 --- a/src/blosc2/ndarray.py +++ b/src/blosc2/ndarray.py @@ -6370,3 +6370,69 @@ def mygen(i): for a in myarrs: out += (broadcast_to(a, shape),) return out + + +def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False): + """Array API cumulative_sum. + + Parameters + ---------- + x : array-like or NDArray + Input array. + axis : int or None + Axis along which to compute the cumulative sum. For ndim > 1, this + must be provided. Default: None (flatten). + dtype : dtype or None + Accumulator/output dtype. If None, follow type-promotion rules. + include_initial : bool, default False + If True, include the additive identity (zero) as the first value along + the reduced axis (output length N+1). + + Returns + ------- + NDArray + Array containing the cumulative sums. + """ + # Accept blosc2.NDArray or array-like + if isinstance(x, NDArray): + arr = x[:] # read into NumPy array + else: + arr = np.asarray(x) + + # Use NumPy's Array API compatible cumulative_sum function + res = np.cumulative_sum(arr, axis=axis, dtype=dtype, include_initial=include_initial) + + return blosc2.asarray(res) + + +def cumulative_prod(x, /, *, axis=None, dtype=None, include_initial=False): + """Array API cumulative_prod. + + Parameters + ---------- + x : array-like or NDArray + Input array. + axis : int or None + Axis along which to compute the cumulative product. For ndim > 1, this + must be provided. Default: None (flatten). + dtype : dtype or None + Accumulator/output dtype. If None, follow type-promotion rules. + include_initial : bool, default False + If True, include the multiplicative identity (one) as the first value along + the reduced axis (output length N+1). + + Returns + ------- + NDArray + Array containing the cumulative products. + """ + # Accept blosc2.NDArray or array-like + if isinstance(x, NDArray): + arr = x[:] # read into NumPy array + else: + arr = np.asarray(x) + + # Use NumPy's Array API compatible cumulative_prod function + res = np.cumulative_prod(arr, axis=axis, dtype=dtype, include_initial=include_initial) + + return blosc2.asarray(res)