Repository: xgcm/xgcm (python)

Issue:
[BUG] Problems with custom grid ufunc and dask options
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports: http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples: https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:
I was trying to check out how to write and apply custom grid_ufuncs and ran into an issue, that @TomNicholas and I concluded was actually a bug.

I start with a simple dataset, that is chunked, but not along any core dimensions. As such I expect `dask='parallelized'` to work as input.

```python
from xgcm.grid_ufunc import as_grid_ufunc
from xgcm.test.test_grid_ufunc import create_1d_test_grid_ds, create_1d_test_grid

ds = create_1d_test_grid_ds('X')
grid = create_1d_test_grid('X')
da = np.sin(grid._ds.X_c)
da = da.chunk({'X_c':-1})
da 


# simple forward difference
@as_grid_ufunc(
    signature="(ax0:center)->(ax0:right)",
    boundary_width={"ax0": (1, 0)},
    dask='parallelized'
)
def test_ufunc(a):
    out = a[..., :-1] - a[..., 1:]
    return out

ufunc_diff = test_ufunc(grid, da, axis='X')
ufunc_diff
```

But this returns the following error:
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_2409/2481888569.py in <module>
     19     return out
     20 
---> 21 ufunc_diff = test_ufunc(grid, da, axis='X')
     22 ufunc_diff

~/xgcm/xgcm/grid_ufunc.py in __call__(self, grid, axis, *args, **kwargs)
    158         dask = kwargs.pop("dask", self.dask)
    159         map_overlap = kwargs.pop("map_overlap", self.map_overlap)
--> 160         return apply_as_grid_ufunc(
    161             self.ufunc,
    162             *args,

~/xgcm/xgcm/grid_ufunc.py in apply_as_grid_ufunc(func, axis, grid, signature, boundary_width, boundary, fill_value, keep_coords, dask, map_overlap, *args, **kwargs)
    477 
    478     # Perform operation via xarray.apply_ufunc
--> 479     results = xr.apply_ufunc(
    480         mapped_func,
    481         *rechunked_padded_args,

/srv/conda/envs/notebook/lib/python3.8/site-packages/xarray/core/computation.py in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
   1157     # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
   1158     elif any(isinstance(a, DataArray) for a in args):
-> 1159         return apply_dataarray_vfunc(
   1160             variables_vfunc,
   1161             *args,

/srv/conda/envs/notebook/lib/python3.8/site-packages/xarray/core/computation.py in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
    284 
    285     data_vars = [getattr(a, "variable", a) for a in args]
--> 286     result_var = func(*data_vars)
    287 
    288     if signature.num_outputs > 1:

/srv/conda/envs/notebook/lib/python3.8/site-packages/xarray/core/computation.py in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
    670                         for axis, dim in enumerate(core_dims, start=-len(core_dims)):
    671                             if len(data.chunks[axis]) != 1:
--> 672                                 raise ValueError(
    673                                     f"dimension {dim} on {n}th function argument to "
    674                                     "apply_ufunc with dask='parallelized' consists of "

ValueError: dimension X_c on 0th function argument to apply_ufunc with dask='parallelized' consists of multiple chunks, but is also a core dimension. To fix, either rechunk into a single dask array chunk along this dimension, i.e., ``.chunk(dict(X_c=-1))``, or pass ``allow_rechunk=True`` in ``dask_gufunc_kwargs`` but beware that this may significantly increase memory usage.
```

What seems to happen here is that the padding adds chunks along the core dimension `X_c`, which should be 'fused' with the 'main' chunk. 
This is not happening because of [this logic](https://github.com/xgcm/xgcm/blob/03d229e18264f7749619291ab4fd7a2578392698/xgcm/grid_ufunc.py#L412-L418).

A possible solution would be to apply the padding fusion always when padding is actually executed, and not just if map_overlap is specified. 

Pull request description:
Fix bug where padding changed chunking
 - [x] Closes #430
 - [x] Tests added
 - [x] Passes `pre-commit run --all-files`


Patch:
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
index c76f30e..fe0bba6 100644
--- a/xgcm/grid_ufunc.py
+++ b/xgcm/grid_ufunc.py
@@ -1,6 +1,16 @@
 import re
 import string
-from typing import TYPE_CHECKING, Any, Callable, List, Mapping, Sequence, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import numpy as np
 import xarray as xr
@@ -409,23 +419,23 @@ def apply_as_grid_ufunc(
         }
         padded_args = args
 
-    if (
-        any(
-            _has_chunked_core_dims(arg, core_dims)
-            for arg, core_dims in zip(args, in_core_dims)
-        )
-        and map_overlap
+    if any(
+        _has_chunked_core_dims(padded_arg, core_dims)
+        for padded_arg, core_dims in zip(padded_args, in_core_dims)
     ):
-        # map operation over dask chunks along core dimensions
-        from dask.array import map_overlap as dask_map_overlap  # type: ignore
-
-        # merge any lonely chunks from padding
+        # merge any lonely chunks on either end created by padding
         rechunked_padded_args = _rechunk_to_merge_in_boundary_chunks(
             padded_args,
             args,
             boundary_width_real_axes,
             grid,
         )
+    else:
+        rechunked_padded_args = padded_args
+
+    if map_overlap:
+        # map operation over dask chunks along core dimensions
+        from dask.array import map_overlap as dask_map_overlap  # type: ignore
 
         boundary_width_per_numpy_axis = {
             grid.axes[ax_name]._get_axis_dim_num(args[0]): width
@@ -466,7 +476,6 @@ def apply_as_grid_ufunc(
             )
 
     else:
-        rechunked_padded_args = padded_args
         mapped_func = func
 
     # Determine expected output dimension sizes from grid._ds
@@ -599,25 +608,36 @@ def _get_chunk_pattern_for_merging_boundary(
     da: xr.DataArray,
     original_chunks: Mapping[str, Tuple[int, ...]],
     boundary_width_real_axes: Mapping[str, Tuple[int, int]],
-):
+) -> Mapping[str, Tuple[int, ...]]:
     """Calculates the pattern of chunking needed to merge back in small chunks left on boundaries after padding"""
 
-    # Easier to work with width of boundaries in terms of str axis names rather than int axis numbers
+    # Easier to work with width of boundaries in terms of str dimension names rather than int axis numbers
     boundary_width_dims = {
         _get_dim(grid, da, ax): width for ax, width in boundary_width_real_axes.items()
     }
 
-    new_chunks = {}
+    new_chunks: Dict[str, Tuple[int, ...]] = {}
     for dim, width in boundary_width_dims.items():
         lower_boundary_width, upper_boundary_width = boundary_width_dims[dim]
-        first_chunk_width, *other_chunks_widths, last_chunk_width = original_chunks[dim]
-        new_chunks_along_dim = tuple(
-            [
-                first_chunk_width + lower_boundary_width,
-                *other_chunks_widths,
-                last_chunk_width + upper_boundary_width,
+
+        new_chunks_along_dim: Tuple[int, ...]
+        if len(original_chunks[dim]) == 1:
+            # unpadded array had only one chunk, but padding has meant new array is extended
+            original_array_length = original_chunks[dim][0]
+            new_chunks_along_dim = (
+                lower_boundary_width + original_array_length + upper_boundary_width,
+            )
+        else:
+            first_chunk_width, *other_chunks_widths, last_chunk_width = original_chunks[
+                dim
             ]
-        )
+            new_chunks_along_dim = tuple(
+                [
+                    first_chunk_width + lower_boundary_width,
+                    *other_chunks_widths,
+                    last_chunk_width + upper_boundary_width,
+                ]
+            )
         new_chunks[dim] = new_chunks_along_dim
 
     return new_chunks


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
index 079d029..8ddad10 100644
--- a/xgcm/test/test_grid_ufunc.py
+++ b/xgcm/test/test_grid_ufunc.py
@@ -237,6 +237,63 @@ class TestGridUFunc:
         result = diff_center_to_left(grid, da, axis=[("depth",)])
         assert_equal(result, expected)
 
+    def test_1d_unchanging_size_but_padded_dask_parallelized(self):
+        """
+        This test checks that the process of padding a non-chunked core dimension doesn't turn it into a chunked core
+        dimension. See GH #430.
+        """
+
+        def diff_center_to_left(a):
+            return a[..., 1:] - a[..., :-1]
+
+        grid = create_1d_test_grid("depth")
+        da = np.sin(grid._ds.depth_c * 2 * np.pi / 9).chunk()
+        da.coords["depth_c"] = grid._ds.depth_c
+
+        diffed = (da - da.roll(depth_c=1, roll_coords=False)).data
+        expected = xr.DataArray(
+            diffed, dims=["depth_g"], coords={"depth_g": grid._ds.depth_g}
+        ).compute()
+
+        # Test direct application
+        result = apply_as_grid_ufunc(
+            diff_center_to_left,
+            da,
+            axis=[("depth",)],
+            grid=grid,
+            signature="(X:center)->(X:left)",
+            boundary_width={"X": (1, 0)},
+            dask="parallelized",
+        ).compute()
+        assert_equal(result, expected)
+
+        # Test Grid method
+        result = grid.apply_as_grid_ufunc(
+            diff_center_to_left,
+            da,
+            axis=[("depth",)],
+            signature="(X:center)->(X:left)",
+            boundary_width={"X": (1, 0)},
+            dask="parallelized",
+        )
+        assert_equal(result, expected)
+
+        # Test decorator
+        @as_grid_ufunc(
+            "(X:center)->(X:left)",
+            boundary_width={"X": (1, 0)},
+            dask="parallelized",
+        )
+        def diff_center_to_left(a):
+            return a[..., 1:] - a[..., :-1]
+
+        result = diff_center_to_left(
+            grid,
+            da,
+            axis=[("depth",)],
+        ).compute()
+        assert_equal(result, expected)
+
     def test_1d_changing_size_dask_parallelized(self):
         def interp_center_to_inner(a):
             return 0.5 * (a[:-1] + a[1:])
@@ -268,7 +325,7 @@ class TestGridUFunc:
             axis=[("depth",)],
             signature="(X:center)->(X:inner)",
             dask="parallelized",
-        )
+        ).compute()
         assert_equal(result, expected)
 
         # Test decorator