Repository: xgcm/xgcm (python)

Issue:
Error when interpolating `to='outer'`
This is pigging back on the issue #507 but using the version `xgcm=0.8.0` now I'm getting an error when trying to interpolate along a dimension that has `outer` coordinates. I've pasted 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.interp(ds.drF, 'Z', boundary='extend', to='outer')

--------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Input In [19], in <cell line: 5>()
      2 ds.coords['Zp1'] = xr.DataArray(np.arange(11),dims='Zp1',coords={'Zp1':np.arange(11)})
      3 grid = Grid(ds, coords={'Z': {'center': 'Z', 'outer': 'Zp1'}})
----> 5 grid.interp(ds.drF, 'Z', boundary='extend', to='outer')

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid.py:2041, in Grid.interp(self, da, axis, **kwargs)
   1989 def interp(self, da, axis, **kwargs):
   1990     """
   1991     Interpolate neighboring points to the intermediate grid point along
   1992     this axis.
   (...)
   2039     >>> grid.interp(da, ["X", "Y"], periodic={"X": True, "Y": False})
   2040     """
-> 2041     return self._1d_grid_ufunc_dispatch("interp", da, axis, **kwargs)

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid.py:1836, in Grid._1d_grid_ufunc_dispatch(self, funcname, data, axis, to, keep_coords, metric_weighted, other_component, **kwargs)
   1833 else:
   1834     map_overlap = False
-> 1836 array = grid_ufunc(
   1837     self,
   1838     array,
   1839     axis=[(ax_name,)],
   1840     keep_coords=keep_coords,
   1841     dask=dask,
   1842     map_overlap=map_overlap,
   1843     other_component=other_component,
   1844     **remaining_kwargs,
   1845 )
   1847 if ax_metric_weighted:
   1848     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:745, 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)
    742 # Maybe map function over chunked core dims using dask.array.map_overlap
    743 if map_overlap:
    744     # Disallow situations where shifting axis position would cause chunk size to change
--> 745     _check_if_length_would_change(sig)
    747     mapped_func = _map_func_over_core_dims(
    748         func,
    749         args,
   (...)
    753         out_dtypes,
    754     )
    755 else:

File ~/miniconda3/envs/xgcm/lib/python3.9/site-packages/xgcm/grid_ufunc.py:1002, in _check_if_length_would_change(signature)
    996 all_ax_positions = set(
    997     p
    998     for arg_ps in signature.in_ax_positions + signature.out_ax_positions
    999     for p in arg_ps
   1000 )
   1001 if any(pos in DISALLOWED_OVERLAP_POSITIONS for pos in all_ax_positions):
-> 1002     raise NotImplementedError(
   1003         "Cannot chunk along a core dimension for a grid ufunc which has a signature which "
   1004         f"includes one of the axis positions {DISALLOWED_OVERLAP_POSITIONS}"
   1005     )

NotImplementedError: Cannot chunk along a core dimension for a grid ufunc which has a signature which includes one of the axis positions ['inner', 'outer']
```

The error goes away again if I load the array into memory:

`grid.interp(ds.drF.load(), 'Z', boundary='extend', to='outer')`

Pull request description:
Fix gh518 bug
This fixes the bug reported in #518 and #522 (by undoing a regression introduced in #515).

Unfortunately this fix breaks `cumsum` in some cases (specifically the test `xgcm/test/test_grid_ufunc.py::TestPadAfterUFunc::test_cumsum_chunk_checking_bug` now fails on this branch). The way it breaks cumsum is pretty subtle and I think I will fix that in a complementary second PR.

EDIT: Fixing cumsum in this PR

 - [x] Closes #518
 - [x] Tests added
 - [x] Passes `pre-commit run --all-files`
 - [x] User visible changes (including notable bug fixes) are documented in `whats-new.rst`


Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 23670d5..a2b7849 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -17,6 +17,9 @@ Breaking Changes
 Internal Changes
 ~~~~~~~~~~~~~~~~
 
+- Rewrote cumsum to use a different code path from :py:func:`~xgcm.apply_as_grid_ufunc` internally,
+  which makes it less susceptible to subtle bugs like the one reported in :issue:`507`. (:pull:`558`).
+  By `Thomas Nicholas <https://github.com/tomnicholas>`_.
 
 Documentation
 ~~~~~~~~~~~~~
@@ -24,6 +27,9 @@ Documentation
 Bugfixes
 ~~~~~~~~
 
+- Fix bug where chunked core dims of only a single chunk triggered errors. (:pull:`558`, :issue:`518`, :issue:`522`)
+  By `Thomas Nicholas <https://github.com/tomnicholas>`_.
+
 
 v0.8.0 (2022/06/14)
 -------------------
diff --git a/xgcm/grid.py b/xgcm/grid.py
index dd7b16a..bf3cbdf 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -17,9 +17,11 @@ from .grid_ufunc import (
     _GridUFuncSignature,
     _has_chunked_core_dims,
     _maybe_unpack_vector_component,
+    _reattach_coords,
     apply_as_grid_ufunc,
 )
 from .metrics import iterate_axis_combinations
+from .padding import pad
 
 try:
     import numba  # type: ignore
@@ -1420,9 +1422,9 @@ class Grid:
     def _as_axis_kwarg_mapping(
         self,
         kwargs: Union[Any, Dict[str, Any]],
-        axes: Iterable[str] = None,
+        axes: Optional[Iterable[str]] = None,
         ax_property_name=None,
-        default_value: Any = None,
+        default_value: Optional[Any] = None,
     ) -> Dict[str, Any]:
         """Convert kwarg input into dict for each available axis
         E.g. for a grid with 2 axes for the keyword argument `periodic`
@@ -1582,7 +1584,7 @@ class Grid:
                 metric_var = self._ds[metric_varname].reset_coords(drop=True)
                 self._metrics[metric_axes].append(metric_var)
 
-    def _get_dims_from_axis(self, da, axis):
+    def _get_dims_from_axis(self, da: xr.DataArray, axis: Iterable[str]) -> List[str]:
         da = _maybe_unpack_vector_component(da)
         dim = []
         axis = _maybe_promote_str_to_list(axis)
@@ -1753,8 +1755,8 @@ class Grid:
         axis,
         to=None,
         keep_coords=False,
-        metric_weighted: Union[
-            str, Iterable[str], Dict[str, Union[str, Iterable[str]]]
+        metric_weighted: Optional[
+            Union[str, Iterable[str], Dict[str, Union[str, Iterable[str]]]]
         ] = None,
         other_component: Optional[Dict[str, xr.DataArray]] = None,
         **kwargs,
@@ -1901,11 +1903,11 @@ class Grid:
         self,
         func: Callable,
         *args: xr.DataArray,
-        axis: Sequence[Sequence[str]] = None,
+        axis: Optional[Sequence[Sequence[str]]] = None,
         signature: Union[str, _GridUFuncSignature] = "",
-        boundary_width: Mapping[str, Tuple[int, int]] = None,
-        boundary: Union[str, Mapping[str, str]] = None,
-        fill_value: Union[float, Mapping[str, float]] = None,
+        boundary_width: Optional[Mapping[str, Tuple[int, int]]] = None,
+        boundary: Optional[Union[str, Mapping[str, str]]] = None,
+        fill_value: Optional[Union[float, Mapping[str, float]]] = None,
         dask: Literal["forbidden", "parallelized", "allowed"] = "forbidden",
         map_overlap: bool = False,
         **kwargs,
@@ -2198,13 +2200,24 @@ class Grid:
         """
         return self._1d_grid_ufunc_dispatch("max", da, axis, **kwargs)
 
-    def cumsum(self, da, axis, **kwargs):
+    def cumsum(
+        self,
+        da: xr.DataArray,
+        axis: Union[str, Iterable[str]],
+        to=None,
+        boundary=None,
+        fill_value=None,
+        metric_weighted=None,
+        keep_coords: bool = False,
+    ) -> xr.DataArray:
         """
         Cumulatively sum a DataArray, transforming to the intermediate axis
         position.
 
         Parameters
         ----------
+        da: xarray.DataArray
+            Data to apply cumsum to.
         axis : str or list or tuple
             Name of the axis on which to act. Multiple axes can be passed as list or
             tuple (e.g. ``['X', 'Y']``). Functions will be executed over each axis in the
@@ -2247,7 +2260,89 @@ class Grid:
 
         >>> grid.max(da, ["X", "Y"], fill_value={"X": 0, "Y": 100})
         """
-        return self._1d_grid_ufunc_dispatch("cumsum", da, axis, **kwargs)
+
+        if isinstance(axis, str):
+            axis = [axis]
+        to = self._as_axis_kwarg_mapping(to)
+
+        if isinstance(metric_weighted, str):
+            metric_weighted = (metric_weighted,)
+        metric_weighted = self._as_axis_kwarg_mapping(metric_weighted)
+
+        data = da
+        axes = [self.axes[ax_name] for ax_name in axis]
+        for ax in axes:
+            pos, dim = ax._get_position_name(da)
+
+            ax_metric_weighted = metric_weighted[ax.name]
+            if ax_metric_weighted:
+                metric = self.get_metric(data, ax_metric_weighted)
+                data = data * metric
+
+            # first use xarray's cumsum method
+            data = data.cumsum(dim=dim)
+
+            ax_to = to[ax.name]
+            if ax_to is None:
+                ax_to = ax._default_shifts[pos]
+
+            # now pad / trim the data as necessary
+            # here we enumerate all the valid possible shifts
+            if (pos == "center" and ax_to == "right") or (
+                pos == "left" and ax_to == "center"
+            ):
+                # do nothing, this is the default for how cumsum works
+                ax_boundary_width = {ax.name: (0, 0)}
+            elif (pos == "center" and ax_to == "left") or (
+                pos == "right" and ax_to == "center"
+            ):
+                data = data.isel(**{dim: slice(0, -1)})
+                ax_boundary_width = {ax.name: (1, 0)}
+            elif (pos == "center" and ax_to == "inner") or (
+                pos == "outer" and ax_to == "center"
+            ):
+                data = data.isel(**{dim: slice(0, -1)})
+                ax_boundary_width = {ax.name: (0, 0)}
+            elif (pos == "center" and ax_to == "outer") or (
+                pos == "inner" and ax_to == "center"
+            ):
+                ax_boundary_width = {ax.name: (1, 0)}
+            else:
+                raise ValueError(
+                    f"From `{pos}` to `{ax_to}` is not a valid position "
+                    f"shift for cumsum operation along axis {ax}."
+                )
+
+            padded = pad(
+                data=data,
+                grid=self,
+                boundary_width=ax_boundary_width,
+                boundary=boundary,
+                fill_value=fill_value,
+            )
+
+            # get dim with position to
+            new_dim_name = ax.coords[ax_to]
+            renamed = padded.rename(**{dim: new_dim_name})
+
+            # drop all coords to avoid conflicts when attaching new ones
+            coordless = renamed.drop_vars(renamed.coords)
+
+            reattached = _reattach_coords(
+                [coordless],
+                grid=self,
+                boundary_width=ax_boundary_width,
+                keep_coords=keep_coords,
+            )[0]
+
+            ax_metric_weighted = metric_weighted[ax.name]
+            if ax_metric_weighted:
+                metric = self.get_metric(reattached, ax_metric_weighted)
+                reattached = reattached / metric
+
+            data = reattached
+
+        return data
 
     def _apply_vector_function(self, function, vector, **kwargs):
         if not (len(vector) == 2 and isinstance(vector, dict)):
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
index 576c83f..781de4c 100644
--- a/xgcm/grid_ufunc.py
+++ b/xgcm/grid_ufunc.py
@@ -448,7 +448,7 @@ class GridUFunc:
 
     def __call__(
         self,
-        grid: "Grid" = None,
+        grid: Optional["Grid"] = None,
         *args: xr.DataArray,
         axis: Sequence[str],
         **kwargs,
@@ -473,7 +473,9 @@ class GridUFunc:
 
 
 def as_grid_ufunc(
-    signature: str = "", boundary_width: Mapping[str, Tuple[int, int]] = None, **kwargs
+    signature: str = "",
+    boundary_width: Optional[Mapping[str, Tuple[int, int]]] = None,
+    **kwargs,
 ) -> Callable:
     """
     Decorator which turns a numpy ufunc into a "grid-aware ufunc".
@@ -555,18 +557,18 @@ def as_grid_ufunc(
 def apply_as_grid_ufunc(
     func: Callable,
     *args: Union[xr.DataArray, Dict[str, xr.DataArray]],
-    axis: Sequence[Sequence[str]] = None,
-    grid: "Grid" = None,
+    axis: Optional[Sequence[Sequence[str]]] = None,
+    grid: Optional["Grid"] = None,
     signature: Union[str, _GridUFuncSignature] = "",
-    boundary_width: Mapping[str, Tuple[int, int]] = None,
-    boundary: Union[str, Mapping[str, str]] = None,
-    fill_value: Union[float, Mapping[str, float]] = None,
+    boundary_width: Optional[Mapping[str, Tuple[int, int]]] = None,
+    boundary: Optional[Union[str, Mapping[str, str]]] = None,
+    fill_value: Optional[Union[float, Mapping[str, float]]] = None,
     keep_coords: bool = True,
     dask: Literal["forbidden", "parallelized", "allowed"] = "forbidden",
     map_overlap: bool = False,
     pad_before_func: bool = True,
-    other_component: Union[
-        Dict[str, xr.DataArray], Sequence[Dict[str, xr.DataArray]]
+    other_component: Optional[
+        Union[Dict[str, xr.DataArray], Sequence[Dict[str, xr.DataArray]]]
     ] = None,
     **kwargs,
 ) -> List[Any]:
@@ -909,7 +911,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]) > 0
+    return len(a.variable.chunksizes[dim]) > 1
 
 
 def _has_chunked_core_dims(obj: xr.DataArray, core_dims: Sequence[str]) -> bool:
@@ -1111,7 +1113,7 @@ def _reattach_coords(
     results_with_coords = []
     for res in results:
 
-        # padding strips all coordinates (inlcuding dimension coordinates).
+        # padding strips all coordinates (including dimension coordinates).
         # Here we centrally restore them from the grid._ds.
         all_matching_coords = {
             coord: da_coord
diff --git a/xgcm/padding.py b/xgcm/padding.py
index 6444e84..3fb80af 100644
--- a/xgcm/padding.py
+++ b/xgcm/padding.py
@@ -72,7 +72,7 @@ def _pad_face_connections(
     padding_width: Dict[str, Tuple[int, int]],
     padding: Dict[str, str],
     fill_value: Dict[str, float],
-    other_component: Dict[str, xr.DataArray] = None,
+    other_component: Optional[Dict[str, xr.DataArray]] = None,
 ):
     facedim = grid._facedim
     connections = grid._connections
@@ -342,8 +342,8 @@ def pad(
     grid: Grid,
     boundary_width: Optional[Dict[str, Tuple[int, int]]],
     boundary: Union[str, Mapping[str, str], None] = None,
-    fill_value: Union[float, Mapping[str, float]] = None,
-    other_component: Dict[str, xr.DataArray] = None,
+    fill_value: Optional[Union[float, Mapping[str, float]]] = None,
+    other_component: Optional[Dict[str, xr.DataArray]] = None,
 ):
     """
     Pads the boundary of given arrays along given Axes, according to information in Axes.boundary.


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
index 5ea4268..6efab8d 100644
--- a/xgcm/test/test_grid_ufunc.py
+++ b/xgcm/test/test_grid_ufunc.py
@@ -117,7 +117,7 @@ class TestParseSignatureFromTypeHints:
         def ufunc(
             a: Annotated[np.ndarray, "X:center"]
         ) -> Annotated[np.ndarray, "X:center"]:
-            ...
+            return np.array([])
 
         assert str(ufunc.signature) == "(X:center)->(X:center)"
 
@@ -125,7 +125,7 @@ class TestParseSignatureFromTypeHints:
         def ufunc(
             a: Annotated[np.ndarray, "X:center,Y:center"]
         ) -> Annotated[np.ndarray, "X:center"]:
-            ...
+            return np.array([])
 
         assert str(ufunc.signature) == "(X:center,Y:center)->(X:center)"
 
@@ -134,7 +134,7 @@ class TestParseSignatureFromTypeHints:
             a: Annotated[np.ndarray, "X:left"],
             b: Annotated[np.ndarray, "Y:right"],
         ) -> Annotated[np.ndarray, "X:center"]:
-            ...
+            return np.array([])
 
         assert str(ufunc.signature) == "(X:left),(Y:right)->(X:center)"
 
@@ -142,7 +142,7 @@ class TestParseSignatureFromTypeHints:
         def ufunc(
             a: Annotated[np.ndarray, "X:center"]
         ) -> Annotated[np.ndarray, "X:left,Y:right"]:
-            ...
+            return np.array([])
 
         assert str(ufunc.signature) == "(X:center)->(X:left,Y:right)"
 
@@ -150,7 +150,7 @@ class TestParseSignatureFromTypeHints:
         def ufunc(
             a: Annotated[np.ndarray, "X:center"]
         ) -> Tuple[Annotated[np.ndarray, "X:left"], Annotated[np.ndarray, "Y:right"]]:
-            ...
+            return np.array([]), np.array([])
 
         assert str(ufunc.signature) == "(X:center)->(X:left),(Y:right)"
 
@@ -164,7 +164,7 @@ class TestParseSignatureFromTypeHints:
             def ufunc(
                 a: Annotated[np.ndarray, "nonsense"]  # type: ignore
             ) -> Annotated[np.ndarray, "X:center"]:
-                ...
+                return np.array([])
 
         with pytest.raises(ValueError, match="Not a valid grid ufunc signature"):
 
@@ -172,7 +172,7 @@ class TestParseSignatureFromTypeHints:
             def ufunc(
                 a: Annotated[np.ndarray, "X:Mars"]
             ) -> Annotated[np.ndarray, "X:center"]:
-                ...
+                return np.array([])
 
     @pytest.mark.xfail(reason="signature regex will assume nonsense==no inputs")
     def test_invalid_return_arg_annotation(self):
@@ -184,7 +184,7 @@ class TestParseSignatureFromTypeHints:
             def ufunc(
                 a: Annotated[np.ndarray, "X:center"]
             ) -> Annotated[np.ndarray, "X:Venus"]:
-                ...
+                return np.array([])
 
     def test_both_sig_kwarg_and_hints_given(self):
         with pytest.raises(
@@ -195,7 +195,7 @@ class TestParseSignatureFromTypeHints:
             def ufunc(
                 a: Annotated[np.ndarray, "X:center"]
             ) -> Annotated[np.ndarray, "X:left"]:
-                ...
+                return np.array([])
 
     def test_type_hint_as_numpy_ndarray(self):
 
@@ -829,13 +829,14 @@ class TestPadAfterUFunc:
         )
         assert_equal(result, expected)
 
-    def test_cumsum_chunk_checking_bug(self):
-        # see issue #507
+    @pytest.mark.parametrize("chunksize", [-1, 3])
+    def test_cumsum_chunk_checking_bug(self, chunksize):
+        # for chunksize=-1 see issue #507
         ds = (
             xr.DataArray(
                 np.ones(10) * 0.5, dims="Z", coords={"Z": np.arange(0.5, 10, 1)}
             )
-            .chunk({"Z": -1})
+            .chunk({"Z": chunksize})
             .to_dataset(name="drF")
         )
         ds.coords["Zp1"] = xr.DataArray(
@@ -1110,6 +1111,32 @@ class TestMapOverlapGridops:
         result = grid.diff(da, axis="depth", to="right").compute()
         assert_equal(result, expected)
 
+    def test_single_chunk_core_dims_center_to_outer(self):
+        # test for issue #518
+        z_coord = xr.DataArray(np.arange(0.5, 10, 1), dims="Z")
+        zp1_coord = xr.DataArray(np.arange(11), dims="Zp1")
+        ds = (
+            xr.DataArray(np.linspace(1, 10, num=10), dims="Z", coords={"Z": z_coord})
+            .chunk({"Z": -1})
+            .to_dataset(name="drF")
+        )
+        ds.coords["Zp1"] = zp1_coord
+        grid = Grid(ds, coords={"Z": {"center": "Z", "outer": "Zp1"}})
+
+        expected_values = np.concatenate(
+            (np.array([1.0]), np.linspace(1.5, 9.5, num=9), np.array([10.0]))
+        )
+        expected = (
+            xr.DataArray(expected_values, dims="Zp1", coords={"Zp1": zp1_coord})
+            .chunk({"Zp1": -1})
+            .to_dataset(name="drF")
+        )
+
+        result = grid.interp(ds.drF, "Z", boundary="extend", to="outer").to_dataset(
+            name="drF"
+        )
+        assert_equal(result, expected)
+
 
 class TestSignaturesEquivalent:
     def test_equivalent(self):