From a7811cd5220ca0e2d3e87cc6fc44668840935c67 Mon Sep 17 00:00:00 2001 From: Vlad Skripniuk Date: Sat, 19 Oct 2019 10:50:47 +0200 Subject: [PATCH 01/10] Add cumsum to DatasetGroupBy Fixes #3141 --- xarray/core/groupby.py | 5 +++-- xarray/tests/test_groupby.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) 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 From 5d229028defe4f0ac2a17d6498bfab5b7f118750 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 27 Apr 2022 15:21:03 +0000 Subject: [PATCH 02/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/core/groupby.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 9e79e09a082..4f439c0a610 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -971,7 +971,9 @@ def reduce_array(ar): return self.map(reduce_array, shortcut=shortcut) -class DataArrayGroupBy(DataArrayGroupByBase, DataArrayGroupByReductions, IncludeCumMethods): +class DataArrayGroupBy( + DataArrayGroupByBase, DataArrayGroupByReductions, IncludeCumMethods +): __slots__ = () From 57f6fb18d7f1890d9548c606ce675a2b631cb11c Mon Sep 17 00:00:00 2001 From: dcherian Date: Wed, 27 Apr 2022 09:27:55 -0600 Subject: [PATCH 03/10] fix --- xarray/core/groupby.py | 13 +++++++++++-- xarray/tests/test_groupby.py | 1 - 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 4f439c0a610..a3c455d5ae5 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -10,6 +10,7 @@ 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 @@ -972,7 +973,10 @@ def reduce_array(ar): class DataArrayGroupBy( - DataArrayGroupByBase, DataArrayGroupByReductions, IncludeCumMethods + DataArrayGroupByBase, + DataArrayGroupByReductions, + ImplementsArrayReduce, + IncludeCumMethods, ): __slots__ = () @@ -1111,5 +1115,10 @@ def assign(self, **kwargs): return self.map(lambda ds: ds.assign(**kwargs)) -class DatasetGroupBy(DatasetGroupByBase, DatasetGroupByReductions, IncludeCumMethods): +class DatasetGroupBy( + DatasetGroupByBase, + DatasetGroupByReductions, + ImplementsDatasetReduce, + IncludeCumMethods, +): __slots__ = () diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 526a38ae593..2ca4ba04637 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -1927,7 +1927,6 @@ def test_groupby_cumsum(): 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]}, ) From 690f1eaeec308025309ea388c4889e48f488c356 Mon Sep 17 00:00:00 2001 From: dcherian Date: Thu, 28 Apr 2022 09:24:07 -0600 Subject: [PATCH 04/10] More fix. --- xarray/tests/test_groupby.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 2ca4ba04637..955dfb92483 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -1928,11 +1928,18 @@ def test_groupby_cumsum(): { "foo": (("x",), [7, 10, 1, 2, 1, 2, 3]), }, - coords={"x": [0, 1, 2, 3, 4, 5, 6]}, + coords={ + "x": [0, 1, 2, 3, 4, 5, 6], + "group_id": ds.group_id, + }, ) - assert_identical(expected, actual) + # 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(ds.group_id).cumsum(dim="x") + 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) From 3d1beca7a3df5857a356004d0d3b87bc4cc93e43 Mon Sep 17 00:00:00 2001 From: dcherian Date: Thu, 28 Apr 2022 09:26:37 -0600 Subject: [PATCH 05/10] Add whats-new --- doc/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 4882402073c..559bdcd0ef1 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -41,6 +41,8 @@ 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`) Breaking changes ~~~~~~~~~~~~~~~~ From bcbe80160d3583a333577697be2099eb2332b9ca Mon Sep 17 00:00:00 2001 From: dcherian Date: Thu, 28 Apr 2022 09:30:41 -0600 Subject: [PATCH 06/10] [skip-ci] Add to api.rst --- doc/api.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/api.rst b/doc/api.rst index 644b86cdebb..6001ac2ab68 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -708,6 +708,7 @@ Dataset DatasetGroupBy.all DatasetGroupBy.any DatasetGroupBy.count + DatasetGroupBy.cumsum DatasetGroupBy.max DatasetGroupBy.mean DatasetGroupBy.median @@ -737,6 +738,7 @@ DataArray DataArrayGroupBy.all DataArrayGroupBy.any DataArrayGroupBy.count + DataArrayGroupBy.cumsum DataArrayGroupBy.max DataArrayGroupBy.mean DataArrayGroupBy.median From 944bd0248fddf299820ec883faace38365ab19c6 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Tue, 12 Jul 2022 15:46:43 -0600 Subject: [PATCH 07/10] Update xarray/tests/test_groupby.py Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com> --- xarray/tests/test_groupby.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index b5d389d07c3..207cb224dcd 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -1992,7 +1992,7 @@ def func(arg1, arg2, arg3=0.0): assert_identical(expected, actual) -def test_groupby_cumsum(): +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])}, From 06fc6c76386a589385d4b924deceda326a3af8b8 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Tue, 12 Jul 2022 23:50:49 +0200 Subject: [PATCH 08/10] Update xarray/core/groupby.py --- xarray/core/groupby.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index b83686c5365..3634d9e2fb7 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -1194,12 +1194,12 @@ def reduce_array(ar: DataArray) -> DataArray: # https://github.com/python/mypy/issues/9031 -class DataArrayGroupBy( +class DataArrayGroupBy( # type: ignore[misc] DataArrayGroupByBase, DataArrayGroupByReductions, ImplementsArrayReduce, IncludeCumMethods, -): # type: ignore[misc] +): __slots__ = () From aa13743e60d81032109fc87d4bc21f020fb16570 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Tue, 12 Jul 2022 23:50:55 +0200 Subject: [PATCH 09/10] Update xarray/core/groupby.py --- xarray/core/groupby.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 3634d9e2fb7..8f9ab504b51 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -1353,10 +1353,10 @@ def assign(self, **kwargs: Any) -> Dataset: # https://github.com/python/mypy/issues/9031 -class DatasetGroupBy( +class DatasetGroupBy( # type: ignore[misc] DatasetGroupByBase, DatasetGroupByReductions, ImplementsDatasetReduce, IncludeCumMethods, -): # type: ignore[misc] +): __slots__ = () From 3dd5b52414e7c41fc402b0d4a400f2857c83b9e6 Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Wed, 13 Jul 2022 00:05:57 +0200 Subject: [PATCH 10/10] Update xarray/tests/test_groupby.py --- xarray/tests/test_groupby.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 207cb224dcd..36d5b783d72 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -1997,7 +1997,7 @@ def test_groupby_cumsum() -> None: {"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") + 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]),