Add cumsum to DatasetGroupBy (#6525)

* Add cumsum to DatasetGroupBy Fixes #3141 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * More fix. * Add whats-new * [skip-ci] Add to api.rst * Update xarray/tests/test_groupby.py Co-authored-by: Illviljan <[email protected]> * Update xarray/core/groupby.py * Update xarray/core/groupby.py * Update xarray/tests/test_groupby.py Co-authored-by: Vlad Skripniuk <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Maximilian Roos <[email protected]> Co-authored-by: Illviljan <[email protected]> Co-authored-by: Anderson Banihirwe <[email protected]>
pydata · Jul 20, 2022 · dabd977 · dabd977
1 parent 9b54b44
commit dabd977
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 2 deletions.
diff --git 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
@@ -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 <https://github.com/TomNicholas>`_.
+- Add :py:meth:`core.groupby.DatasetGroupBy.cumsum` and :py:meth:`core.groupby.DataArrayGroupBy.cumsum`.
+  By `Vladislav Skripniuk <https://github.com/VladSkripniuk>`_ and `Deepak Cherian <https://github.com/dcherian>`_. (: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 <https://github.com/TomNicholas>`_.

diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py
@@ -25,10 +25,12 @@
 from ._reductions import DataArrayGroupByReductions, DatasetGroupByReductions
 from .alignment import align
 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
@@ -1187,7 +1189,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__ = ()
 
 
@@ -1341,5 +1348,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
@@ -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