Use cumsum from flox

[Illviljan]

• Closes cumsum drops index coordinates #6528
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

The non-flox version reduces chunksizes significantly:

x = xr.DataArray([1, 1, 1, 1, 1], name="x").chunk()
grp_idx = xr.DataArray([-1, 0, 0, -1, 1])
with xr.set_options(use_flox=False):
    print(x.groupby(grp_idx).cumsum())
<xarray.DataArray 'x' (dim_0: 5)> Size: 40B
dask.array<getitem, shape=(5,), dtype=int64, chunksize=(2,), chunktype=numpy.ndarray>
Dimensions without coordinates: dim_0

With flox the chunksize is retained:

x = xr.DataArray([1, 1, 1, 1, 1], name="x").chunk()
grp_idx = xr.DataArray([-1, 0, 0, -1, 1])
with xr.set_options(use_flox=True):
    print(x.groupby(grp_idx).cumsum())
<xarray.DataArray 'x' (dim_0: 5)> Size: 40B
dask.array<_finalize_scan, shape=(5,), dtype=int64, chunksize=(5,), chunktype=numpy.ndarray>
Dimensions without coordinates: dim_0

Other changes:

• Changes DataArray.cumsum/Dataset.cumsum/DataTree.cumsum/DataArray.groupby.cumsum/Dataset.groupby.cumsum etc.
• Coordinates are now retained

Notes
groupby_scan was added in: https://github.com/xarray-contrib/flox/releases/tag/v0.9.9
cumsum was added in: https://github.com/xarray-contrib/flox/releases/tag/v0.10.5

[Illviljan]

Writing down these flox issues before I forget:

ds = xr.Dataset(
    {
        "foo": (
            ("test", "time"),
            [[7, 2, 0, 1, 2, np.nan], [1, 1, 1, 1, 1, 1], [2, 2, 2, 2, 2, 2]],
        )
    },
    coords={
        "time": [0, 1 / 6, 2 / 6, 3 / 6, 4 / 6, 5 / 6],
        "test": ["a", "b", "b"],
        "group_idx": ("time", [0, 0, 1, 1, 2, 2]),
        "group_idx2": ("time", [0, 1, 1, 1, 1, 1]),
    },
)

# group_idx along 1 dim and cumsum dim along another fails with flox:
ds.groupby("group_idx").cumsum("test") 

# cumsum along multple dims fails with flox:
ds.groupby("group_idx").cumsum(...)

[Illviljan]

I wanted to give pytest accept a final try within pixi. But I'm still unable to get it to correctly indent in win10

    pixi shell
    python xarray/util/generate_aggregations.py
    pytest --doctest-modules xarray/core/_aggregations.py --accept
    pytest --doctest-modules xarray/core/_aggregations.py

[dcherian]

Sorry I've lost track here. Are there upstream fixes required in flox? My other concern is that resample().cumsum() is probably quite bad because I haven't implemented method="cohorts" for scans upstream

[Illviljan]

No I think this PR is finished. Mostly waiting on the green CI and maybe a release to give it chance to sit on main for a while.

I haven't dug into the performance of resample().cumsum(), is it a stopper for you? Are you expecting it to be worse than what we had previously?
At least with groupby().cumsum() it was unusable for me, I just do a early .compute() to avoid the dask arrays.