Repository: xgcm/xgcm
Pull request #554 (2022-11-07 22:45:42)

[pre-commit.ci] pre-commit autoupdate
<!--pre-commit.ci start-->
updates:
- [github.com/keewis/blackdoc: v0.3.7 → v0.3.8](https://github.com/keewis/blackdoc/compare/v0.3.7...v0.3.8)
- [github.com/pre-commit/mirrors-mypy: v0.982 → v0.990](https://github.com/pre-commit/mirrors-mypy/compare/v0.982...v0.990)
<!--pre-commit.ci end-->


Patch:
diff --git a/.github/workflows/release_reminder.yaml b/.github/workflows/release_reminder.yaml
deleted file mode 100644
index 4dcd521..0000000
--- a/.github/workflows/release_reminder.yaml
+++ /dev/null
@@ -1,18 +0,0 @@
-name: Create release reminder
-on:
-  schedule:
-  # Run on the first of the month
-  - cron: "0 00 1 * *"
-
-  workflow_dispatch:
-
-jobs:
-  create-sync-issue:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v3
-    - uses: JasonEtco/create-an-issue@v2
-      env:
-        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-      with:
-        filename: .github/ISSUE_TEMPLATE/release_reminder.md
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index beaf79c..a468acd 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -16,7 +16,7 @@ repos:
     hooks:
       - id: black
   - repo: https://github.com/keewis/blackdoc
-    rev: v0.3.7
+    rev: v0.3.8
     hooks:
       - id: blackdoc
         additional_dependencies: ["black==22.3.0"]
@@ -25,7 +25,7 @@ repos:
     hooks:
       - id: flake8
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.982
+    rev: v0.990
     hooks:
       - id: mypy
         # `properies` & `asv_bench` are copied from setup.cfg.
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):