Repository: xgcm/xgcm (python)

Issue:
Weird behavior with cumsum
I've been trying to take the cummulative sum using `xgcm.cumsum()` but I'm getting a weird behavior when the data is not loaded. I've added a reproducible example below.

```
ds = xr.DataArray(np.ones(10)*.5,dims='Z',coords={'Z':np.arange(.5,10,1)}).chunk({'Z':-1}).to_dataset(name='drF')
ds.coords['Zp1'] = xr.DataArray(np.arange(11),dims='Zp1',coords={'Zp1':np.arange(11)})
grid = Grid(ds, coords={'Z': {'center': 'Z', 'outer': 'Zp1'}})

grid.cumsum(ds.drF, 'Z', boundary='extend')
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Input In [55], in <cell line: 1>()
----> 1 grid.cumsum(ds.drF, 'Z', boundary='extend')

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid.py:2245, in Grid.cumsum(self, da, axis, **kwargs)
   2196 def cumsum(self, da, axis, **kwargs):
   2197     """
   2198     Cumulatively sum a DataArray, transforming to the intermediate axis
   2199     position.
   (...)
   2243     >>> grid.max(da, ["X", "Y"], fill_value={"X": 0, "Y": 100})
   2244     """
-> 2245     return self._1d_grid_ufunc_dispatch("cumsum", da, axis, **kwargs)

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid.py:1831, in Grid._1d_grid_ufunc_dispatch(self, funcname, data, axis, to, keep_coords, metric_weighted, other_component, **kwargs)
   1828 else:
   1829     map_overlap = False
-> 1831 array = grid_ufunc(
   1832     self,
   1833     array,
   1834     axis=[(ax_name,)],
   1835     keep_coords=keep_coords,
   1836     dask=dask,
   1837     map_overlap=map_overlap,
   1838     other_component=other_component,
   1839     **remaining_kwargs,
   1840 )
   1842 if ax_metric_weighted:
   1843     metric = self.get_metric(array, ax_metric_weighted)

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid_ufunc.py:460, in GridUFunc.__call__(self, grid, axis, *args, **kwargs)
    458 map_overlap = kwargs.pop("map_overlap", self.map_overlap)
    459 pad_before_func = kwargs.pop("pad_before_func", self.pad_before_func)
--> 460 return apply_as_grid_ufunc(
    461     self.ufunc,
    462     *args,
    463     axis=axis,
    464     grid=grid,
    465     signature=self.signature,
    466     boundary_width=self.boundary_width,
    467     boundary=boundary,
    468     dask=dask,
    469     map_overlap=map_overlap,
    470     pad_before_func=pad_before_func,
    471     **kwargs,
    472 )

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid_ufunc.py:791, in apply_as_grid_ufunc(func, axis, grid, signature, boundary_width, boundary, fill_value, keep_coords, dask, map_overlap, pad_before_func, other_component, *args, **kwargs)
    780 else:  # pad after func
    781     unpadded_results = _apply(
    782         mapped_func,
    783         args,
   (...)
    789         **kwargs,
    790     )
--> 791     results = _pad_then_rechunk(
    792         unpadded_results,
    793         grid,
    794         in_core_dims,
    795         boundary_width_real_axes,
    796         boundary,
    797         fill_value,
    798         other_component,
    799     )
    801 # TODO add option to trim result if not done in ufunc
    802 # TODO loud warning if ufunc returns array of incorrect size
    803 
    804 # Restore any dimension coordinates associated with new output dims that are present in grid
    805 results_with_coords = _reattach_coords(results, grid, boundary_width, keep_coords)

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid_ufunc.py:893, in _pad_then_rechunk(args, grid, in_core_dims, boundary_width_real_axes, boundary, fill_value, other_component)
    871 def _pad_then_rechunk(
    872     args,
    873     grid,
   (...)
    878     other_component,
    879 ):
    881     padded_args = [
    882         pad(
    883             a,
   (...)
    890         for a, oc in zip(args, other_component)
    891     ]
--> 893     if any(
    894         _has_chunked_core_dims(padded_arg, core_dims)
    895         for padded_arg, core_dims in zip(padded_args, in_core_dims)
    896     ):
    897         # merge any lonely chunks on either end created by padding
    898         rechunked_padded_args = _rechunk_to_merge_in_boundary_chunks(
    899             padded_args,
    900             args,
    901             boundary_width_real_axes,
    902             grid,
    903         )
    904     else:

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid_ufunc.py:894, in <genexpr>(.0)
    871 def _pad_then_rechunk(
    872     args,
    873     grid,
   (...)
    878     other_component,
    879 ):
    881     padded_args = [
    882         pad(
    883             a,
   (...)
    890         for a, oc in zip(args, other_component)
    891     ]
    893     if any(
--> 894         _has_chunked_core_dims(padded_arg, core_dims)
    895         for padded_arg, core_dims in zip(padded_args, in_core_dims)
    896     ):
    897         # merge any lonely chunks on either end created by padding
    898         rechunked_padded_args = _rechunk_to_merge_in_boundary_chunks(
    899             padded_args,
    900             args,
    901             boundary_width_real_axes,
    902             grid,
    903         )
    904     else:

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid_ufunc.py:917, in _has_chunked_core_dims(obj, core_dims)
    915 def _has_chunked_core_dims(obj: xr.DataArray, core_dims: Sequence[str]) -> bool:
    916     # TODO what if only some of the core dimensions are chunked?
--> 917     return obj.chunks is not None and any(
    918         _is_dim_chunked(obj, dim) for dim in core_dims
    919     )

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid_ufunc.py:918, in <genexpr>(.0)
    915 def _has_chunked_core_dims(obj: xr.DataArray, core_dims: Sequence[str]) -> bool:
    916     # TODO what if only some of the core dimensions are chunked?
    917     return obj.chunks is not None and any(
--> 918         _is_dim_chunked(obj, dim) for dim in core_dims
    919     )

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid_ufunc.py:912, in _is_dim_chunked(a, dim)
    910 def _is_dim_chunked(a, dim):
    911     # TODO this func can't handle Datasets - it will error if you check multiple variables with different chunking
--> 912     return len(a.variable.chunksizes[dim]) > 1

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xarray/core/utils.py:457, in Frozen.__getitem__(self, key)
    456 def __getitem__(self, key: K) -> V:
--> 457     return self.mapping[key]

KeyError: 'Z'
```

The error goes away when I load the data, i.e.
```
grid.cumsum(ds.drF.load(), 'Z', boundary='extend')
```

**xgcm version and environment**:
My xgcm version is 0.7.0.



Pull request description:
Fix chunked cumsum
<!-- Feel free to remove check-list items aren't relevant to your change -->

 - [x] Closes #507
 - [x] Tests added
 - [x] Passes `pre-commit run --all-files`
 - [x] User visible changes (including notable bug fixes) are documented in `whats-new.rst`
 - [ ] ~~New functions/methods are listed in `api.rst`~~


Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 10ac3ca..3ee1b5a 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -31,6 +31,8 @@ Bugfixes
 
 - Fix formatting of the CITATION.cff file (:pull:`500`).
   By `Julius Busecke <https://github.com/jbusecke>`_.
+- Fix bug with cumsum when data chunked with dask. (:pull:`415`, :issue:`507`)
+  By `Thomas Nicholas <https://github.com/tomnicholas>`_.
 
 v0.7.0 (2022/4/20)
 -------------------
diff --git a/xgcm/grid.py b/xgcm/grid.py
index 8479e4f..dd7b16a 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -1826,9 +1826,8 @@ class Grid:
             # if chunked along core dim then we need map_overlap
             core_dim = self._get_dims_from_axis(data, ax_name)
             if _has_chunked_core_dims(data_unpacked, core_dim):
-                # cumsum is a special case because it can't be correctly applied chunk-wise with map_overlap (it would need blockwise instead)
-                # TODO double check that this dispatches to the dask version of cumsum?
-                # TODO could we use the same approach with the diff etc?
+                # cumsum is a special case because it can't be correctly applied chunk-wise with map_overlap
+                # (it would need blockwise instead)
                 map_overlap = True if funcname != "cumsum" else False
                 dask = "allowed"
             else:
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
index c908b8d..576c83f 100644
--- a/xgcm/grid_ufunc.py
+++ b/xgcm/grid_ufunc.py
@@ -791,7 +791,7 @@ def apply_as_grid_ufunc(
         results = _pad_then_rechunk(
             unpadded_results,
             grid,
-            in_core_dims,
+            out_core_dims,
             boundary_width_real_axes,
             boundary,
             fill_value,
@@ -799,9 +799,9 @@ def apply_as_grid_ufunc(
         )
 
     # TODO add option to trim result if not done in ufunc
-    # TODO loud warning if ufunc returns array of incorrect size
 
     # Restore any dimension coordinates associated with new output dims that are present in grid
+    # Also throws loud warning if ufunc returns array of incorrect size
     results_with_coords = _reattach_coords(results, grid, boundary_width, keep_coords)
 
     # Return single results not wrapped in 1-element tuple, like xr.apply_ufunc does
@@ -909,7 +909,7 @@ def _pad_then_rechunk(
 
 def _is_dim_chunked(a, dim):
     # TODO this func can't handle Datasets - it will error if you check multiple variables with different chunking
-    return len(a.variable.chunksizes[dim]) > 1
+    return len(a.variable.chunksizes[dim]) > 0
 
 
 def _has_chunked_core_dims(obj: xr.DataArray, core_dims: Sequence[str]) -> bool:


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
index 0199725..5ea4268 100644
--- a/xgcm/test/test_grid_ufunc.py
+++ b/xgcm/test/test_grid_ufunc.py
@@ -799,6 +799,53 @@ class TestPadAfterUFunc:
         )
         assert_equal(result, expected)
 
+    def test_cumsum_dask(self):
+        def cumsum_center_to_left(a):
+            return np.cumsum(a, axis=-1)[..., :-1]
+
+        grid = create_1d_test_grid("depth")
+        da = grid._ds.depth_c**2
+        da.coords["depth_c"] = grid._ds.depth_c
+        da = da.chunk({"depth_c": 3})
+
+        cum = da.cumsum(dim="depth_c").roll(depth_c=1, roll_coords=False)
+        cum[0] = 0
+        expected = xr.DataArray(
+            cum.data, dims=["depth_g"], coords={"depth_g": grid._ds.depth_g}
+        )
+
+        result = apply_as_grid_ufunc(
+            cumsum_center_to_left,
+            da,
+            axis=[("depth",)],
+            grid=grid,
+            signature="(X:center)->(X:left)",
+            boundary_width={"X": (1, 0)},
+            boundary="fill",
+            fill_value=0,
+            pad_before_func=False,
+            dask="allowed",
+            map_overlap=False,
+        )
+        assert_equal(result, expected)
+
+    def test_cumsum_chunk_checking_bug(self):
+        # see issue #507
+        ds = (
+            xr.DataArray(
+                np.ones(10) * 0.5, dims="Z", coords={"Z": np.arange(0.5, 10, 1)}
+            )
+            .chunk({"Z": -1})
+            .to_dataset(name="drF")
+        )
+        ds.coords["Zp1"] = xr.DataArray(
+            np.arange(11), dims="Zp1", coords={"Zp1": np.arange(11)}
+        )
+        grid = Grid(ds, coords={"Z": {"center": "Z", "outer": "Zp1"}})
+
+        grid.cumsum(ds.drF, "Z", boundary="periodic")
+        grid.cumsum(ds.drF, "Z", boundary="extend")
+
 
 class TestDaskNoOverlap:
     def test_chunked_non_core_dims(self):