Support cumsum (#451)

Author: dcherian

.pre-commit-config.yaml

@@ -10,11 +10,6 @@ repos:
         args: ["--fix", "--show-fixes"]
       - id: ruff-format
 
-  - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: "v4.0.0-alpha.8"
-    hooks:
-      - id: prettier
-
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v5.0.0
     hooks:

flox/aggregate_flox.py

@@ -276,3 +276,39 @@ def ffill(group_idx, array, *, axis, **kwargs):
 
     invert_perm = slice(None) if isinstance(perm, slice) else np.argsort(perm, kind="stable")
     return array[tuple(slc)][..., invert_perm]
+
+
+def _np_grouped_scan(group_idx, array, *, axis: int, skipna: bool, **kwargs):
+    handle_nans = not skipna and array.dtype.kind in "cfO"
+
+    group_idx, array, perm = _prepare_for_flox(group_idx, array)
+    ndim = array.ndim
+    assert axis == (ndim - 1), (axis, ndim - 1)
+
+    flag = np.concatenate((np.asarray([True], like=group_idx), group_idx[1:] != group_idx[:-1]))
+    (inv_idx,) = flag.nonzero()
+    segment_lengths = np.add.reduceat(np.ones(group_idx.shape), inv_idx, dtype=np.int64)
+
+    # TODO: set dtype to float properly for handle_nans?
+    accum = np.nancumsum(array, axis=axis)
+
+    if len(inv_idx) > 1:
+        first_group_idx = inv_idx[1]
+        # extract cumulative sum _before_ start of group
+        prev_group_cumsum = accum[..., inv_idx[1:] - 1]
+        accum[..., first_group_idx:] -= np.repeat(prev_group_cumsum, segment_lengths[1:], axis=axis)
+
+    if handle_nans:
+        mask = isnull(array)
+        accummask = np.cumsum(mask, axis=-1, dtype=np.uint64)
+        if len(inv_idx) > 1:
+            prev_group_cumsum = accummask[..., inv_idx[1:] - 1]
+            accummask[..., first_group_idx:] -= np.repeat(prev_group_cumsum, segment_lengths[1:], axis=axis)
+        accum[accummask > 0] = np.nan
+
+    invert_perm = slice(None) if isinstance(perm, slice) else np.argsort(perm, kind="stable")
+    return accum[..., invert_perm]
+
+
+cumsum = partial(_np_grouped_scan, skipna=False)
+nancumsum = partial(_np_grouped_scan, skipna=True)

flox/aggregations.py

@@ -631,7 +631,7 @@ def last(self) -> AlignedArrays:
         reduced = chunk_reduce(
             self.array,
             self.group_idx,
-            func=("nanlast",),
+            func=("last",),
             axis=-1,
             # TODO: automate?
             engine="flox",
@@ -699,6 +699,7 @@ def scan_binary_op(left_state: ScanState, right_state: ScanState, *, agg: Scan)
             fill_value=agg.identity,
         )
         result = AlignedArrays(array=final_value[..., left.group_idx.size :], group_idx=right.group_idx)
+
     else:
         raise ValueError(f"Unknown binary op application mode: {agg.mode!r}")
 
@@ -717,8 +718,7 @@ def scan_binary_op(left_state: ScanState, right_state: ScanState, *, agg: Scan)
     )
 
 
-# TODO: numpy_groupies cumsum is a broken when NaNs are present.
-# cumsum = Scan("cumsum", binary_op=np.add, reduction="sum", scan="cumsum", identity=0)
+cumsum = Scan("cumsum", binary_op=np.add, reduction="sum", scan="cumsum", identity=0)
 nancumsum = Scan("nancumsum", binary_op=np.add, reduction="nansum", scan="nancumsum", identity=0)
 # ffill uses the identity for scan, and then at the binary-op state,
 # we concatenate the blockwise-reduced values with the original block,
@@ -782,7 +782,7 @@ def scan_binary_op(left_state: ScanState, right_state: ScanState, *, agg: Scan)
     "nanquantile": nanquantile,
     "mode": mode,
     "nanmode": nanmode,
-    # "cumsum": cumsum,
+    "cumsum": cumsum,
     "nancumsum": nancumsum,
     "ffill": ffill,
     "bfill": bfill,

flox/xarray.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Hashable, Iterable, Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import pandas as pd
@@ -249,7 +249,7 @@ def xarray_reduce(
                 grouper_dims.append(d)
 
     if isinstance(obj, xr.Dataset):
-        ds = obj
+        ds = cast(xr.Dataset, obj)
     else:
         ds = obj._to_temp_dataset()
 
@@ -295,7 +295,7 @@ def xarray_reduce(
         not set(grouper_dims).issubset(set(variable.dims)) for variable in ds.data_vars.values()
     )
     if needs_broadcast:
-        ds_broad = xr.broadcast(ds, *by_da, exclude=exclude_dims)[0]
+        ds_broad = cast(xr.Dataset, xr.broadcast(ds, *by_da, exclude=exclude_dims)[0])
     else:
         ds_broad = ds
 

tests/test_core.py

@@ -1974,18 +1974,56 @@ def test_nanlen_string(dtype, engine) -> None:
     assert_equal(expected, actual)
 
 
-def test_cumsum() -> None:
-    array = np.array([1, 1, 1], dtype=np.uint64)
+@pytest.mark.parametrize(
+    "array",
+    [
+        np.array([1, 1, 1, 2, 3, 4, 5], dtype=np.uint64),
+        np.array([1, 1, 1, 2, np.nan, 4, 5], dtype=np.float64),
+    ],
+)
+@pytest.mark.parametrize("func", ["cumsum", "nancumsum"])
+def test_cumsum_simple(array, func) -> None:
     by = np.array([0] * array.shape[-1])
-    expected = np.nancumsum(array, axis=-1)
+    expected = getattr(np, func)(array, axis=-1)
 
-    actual = groupby_scan(array, by, func="nancumsum", axis=-1)
-    assert_equal(expected, actual)
+    actual = groupby_scan(array, by, func=func, axis=-1)
+    assert_equal(actual, expected)
+
+    if has_dask:
+        da = dask.array.from_array(array, chunks=2)
+        actual = groupby_scan(da, by, func=func, axis=-1)
+        assert_equal(actual, expected)
+
+
+def test_cumsum() -> None:
+    array = np.array(
+        [
+            [1, 2, np.nan, 4, 5],
+            [3, np.nan, 4, 6, 6],
+        ]
+    )
+    by = [0, 1, 1, 0, 1]
+
+    expected = np.array(
+        [
+            [1, 2, np.nan, 5, np.nan],
+            [3, np.nan, np.nan, 9, np.nan],
+        ]
+    )
+    actual = groupby_scan(array, by, func="cumsum", axis=-1)
+    assert_equal(actual, expected)
+    if has_dask:
+        da = dask.array.from_array(array, chunks=2)
+        actual = groupby_scan(da, by, func="cumsum", axis=-1)
+        assert_equal(actual, expected)
 
+    expected = np.array([[1, 2, 2, 5, 7], [3, 0, 4, 9, 10]], dtype=np.float64)
+    actual = groupby_scan(array, by, func="nancumsum", axis=-1)
+    assert_equal(actual, expected)
     if has_dask:
         da = dask.array.from_array(array, chunks=2)
         actual = groupby_scan(da, by, func="nancumsum", axis=-1)
-        assert_equal(expected, actual)
+        assert_equal(actual, expected)
 
 
 @pytest.mark.parametrize(

tests/test_properties.py

@@ -56,10 +56,11 @@ def bfill(array, axis, dtype=None):
 
 
 NUMPY_SCAN_FUNCS: dict[str, Callable] = {
+    "cumsum": np.cumsum,
     "nancumsum": np.nancumsum,
     "ffill": ffill,
     "bfill": bfill,
-}  # "cumsum": np.cumsum,
+}
 
 
 def not_overflowing_array(array: np.ndarray[Any, Any]) -> bool:
@@ -210,7 +211,7 @@ def test_groupby_reduce_numpy_vs_other(data, array, func: str) -> None:
     array=chunked_arrays(arrays=numeric_like_arrays),
     func=st.sampled_from(tuple(NUMPY_SCAN_FUNCS)),
 )
-def test_scans(data, array: dask.array.Array, func: str) -> None:
+def test_scans_against_numpy(data, array: dask.array.Array, func: str) -> None:
     if "cum" in func:
         assume(not_overflowing_array(np.asarray(array)))