diff --git a/doc/api.rst b/doc/api.rst index 1036b476c83..840fa32bf43 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -736,6 +736,7 @@ Dataset DatasetGroupBy.all DatasetGroupBy.any DatasetGroupBy.count + DatasetGroupBy.cumsum DatasetGroupBy.max DatasetGroupBy.mean DatasetGroupBy.median @@ -765,6 +766,7 @@ DataArray DataArrayGroupBy.all DataArrayGroupBy.any DataArrayGroupBy.count + DataArrayGroupBy.cumsum DataArrayGroupBy.max DataArrayGroupBy.mean DataArrayGroupBy.median diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 9f6f3622f71..6dfea45e673 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -123,6 +123,10 @@ New Features - Allow passing chunks in ``**kwargs`` form to :py:meth:`Dataset.chunk`, :py:meth:`DataArray.chunk`, and :py:meth:`Variable.chunk`. (:pull:`6471`) By `Tom Nicholas `_. +- Add :py:meth:`core.groupby.DatasetGroupBy.cumsum` and :py:meth:`core.groupby.DataArrayGroupBy.cumsum`. + By `Vladislav Skripniuk `_ and `Deepak Cherian `_. (:pull:`3147`, :pull:`6525`, :issue:`3141`) +- Expose `inline_array` kwarg from `dask.array.from_array` in :py:func:`open_dataset`, :py:meth:`Dataset.chunk`, + :py:meth:`DataArray.chunk`, and :py:meth:`Variable.chunk`. (:pull:`6471`) - Expose the ``inline_array`` kwarg from :py:func:`dask.array.from_array` in :py:func:`open_dataset`, :py:meth:`Dataset.chunk`, :py:meth:`DataArray.chunk`, and :py:meth:`Variable.chunk`. (:pull:`6471`) By `Tom Nicholas `_. diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 5fa78ae76de..8f9ab504b51 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -24,10 +24,12 @@ from . import dtypes, duck_array_ops, nputils, ops from ._reductions import DataArrayGroupByReductions, DatasetGroupByReductions from .arithmetic import DataArrayGroupbyArithmetic, DatasetGroupbyArithmetic +from .common import ImplementsArrayReduce, ImplementsDatasetReduce from .concat import concat from .formatting import format_array_flat from .indexes import create_default_index_implicit, filter_indexes_from_coords from .npcompat import QUANTILE_METHODS, ArrayLike +from .ops import IncludeCumMethods from .options import _get_keep_attrs from .pycompat import integer_types from .types import T_Xarray @@ -1192,7 +1194,12 @@ def reduce_array(ar: DataArray) -> DataArray: # https://github.com/python/mypy/issues/9031 -class DataArrayGroupBy(DataArrayGroupByBase, DataArrayGroupByReductions): # type: ignore[misc] +class DataArrayGroupBy( # type: ignore[misc] + DataArrayGroupByBase, + DataArrayGroupByReductions, + ImplementsArrayReduce, + IncludeCumMethods, +): __slots__ = () @@ -1346,5 +1353,10 @@ def assign(self, **kwargs: Any) -> Dataset: # https://github.com/python/mypy/issues/9031 -class DatasetGroupBy(DatasetGroupByBase, DatasetGroupByReductions): # type: ignore[misc] +class DatasetGroupBy( # type: ignore[misc] + DatasetGroupByBase, + DatasetGroupByReductions, + ImplementsDatasetReduce, + IncludeCumMethods, +): __slots__ = () diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index fc3e1434684..97bed870dd4 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -2002,3 +2002,31 @@ def func(arg1, arg2, arg3=0.0): expected = xr.Dataset({"foo": ("time", [3.0, 3.0, 3.0]), "time": times}) actual = ds.resample(time="D").map(func, args=(1.0,), arg3=1.0) assert_identical(expected, actual) + + +def test_groupby_cumsum() -> None: + ds = xr.Dataset( + {"foo": (("x",), [7, 3, 1, 1, 1, 1, 1])}, + coords={"x": [0, 1, 2, 3, 4, 5, 6], "group_id": ("x", [0, 0, 1, 1, 2, 2, 2])}, + ) + actual = ds.groupby("group_id").cumsum(dim="x") # type: ignore[attr-defined] # TODO: move cumsum to generate_reductions.py + expected = xr.Dataset( + { + "foo": (("x",), [7, 10, 1, 2, 1, 2, 3]), + }, + coords={ + "x": [0, 1, 2, 3, 4, 5, 6], + "group_id": ds.group_id, + }, + ) + # TODO: Remove drop_vars when GH6528 is fixed + # when Dataset.cumsum propagates indexes, and the group variable? + assert_identical(expected.drop_vars(["x", "group_id"]), actual) + + actual = ds.foo.groupby("group_id").cumsum(dim="x") + expected.coords["group_id"] = ds.group_id + expected.coords["x"] = np.arange(7) + assert_identical(expected.foo, actual) + + +# TODO: move other groupby tests from test_dataset and test_dataarray over here