fix pydata#3141 add cumsum to DatasetGroupBy

:
VladSkripniuk · Oct 19, 2019 · 28dd667 · 28dd667
1 parent 16fdac9
commit 28dd667
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py
@@ -859,3 +859,4 @@ def assign(self, **kwargs):
 
 ops.inject_reduce_methods(DatasetGroupBy)
 ops.inject_binary_ops(DatasetGroupBy)
+ops.inject_cum_methods(DatasetGroupBy)
diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py
@@ -202,4 +202,21 @@ def test_da_groupby_assign_coords():
     assert_identical(expected, actual2)
 
 
+def test_da_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)
+
+
 # TODO: move other groupby tests from test_dataset and test_dataarray over here