diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index e97499f06b4..9e79e09a082 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -13,6 +13,7 @@ from .concat import concat from .formatting import format_array_flat from .indexes import create_default_index_implicit, filter_indexes_from_coords +from .ops import IncludeCumMethods from .options import _get_keep_attrs from .pycompat import integer_types from .utils import ( @@ -970,7 +971,7 @@ def reduce_array(ar): return self.map(reduce_array, shortcut=shortcut) -class DataArrayGroupBy(DataArrayGroupByBase, DataArrayGroupByReductions): +class DataArrayGroupBy(DataArrayGroupByBase, DataArrayGroupByReductions, IncludeCumMethods): __slots__ = () @@ -1108,5 +1109,5 @@ def assign(self, **kwargs): return self.map(lambda ds: ds.assign(**kwargs)) -class DatasetGroupBy(DatasetGroupByBase, DatasetGroupByReductions): +class DatasetGroupBy(DatasetGroupByBase, DatasetGroupByReductions, IncludeCumMethods): __slots__ = () diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index b4b93d1dba3..526a38ae593 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -1918,4 +1918,23 @@ def func(arg1, arg2, arg3=0.0): assert_identical(expected, actual) +def test_groupby_cumsum(): + 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") + expected = xr.Dataset( + { + "foo": (("x",), [7, 10, 1, 2, 1, 2, 3]), + "group_id": (("x",), [0, 0, 1, 1, 2, 2, 2]), + }, + coords={"x": [0, 1, 2, 3, 4, 5, 6]}, + ) + assert_identical(expected, actual) + + actual = ds.foo.groupby(ds.group_id).cumsum(dim="x") + assert_identical(expected.foo, actual) + + # TODO: move other groupby tests from test_dataset and test_dataarray over here