Repository: xgcm/xgcm (python)

Issue:
[BUG] xgcm cant detect proper ufunc for chunked data.


**What happened**:
I just tested the new grid_ufunc branch and ran into an issue when working with chunked datasets.

**What you expected to happen**:
I expected the 'basic' grid methods like diff to work on datasets of really any chunking structure, especially when chunked along non-core dimensions

**Minimal Complete Verifiable Example**:
But when I try this:
```python
from xgcm import Grid
from xgcm.test.datasets import datasets_grid_metric
ds, coords, metrics = datasets_grid_metric('C')
ds = ds.chunk({'time':1})
grid = Grid(ds, coords=coords)
grid.diff(ds.tracer, 'X')
```
I get an error from the dispatching logic:
```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
/tmp/ipykernel_2409/1917141936.py in <module>
      5 ds = ds.chunk({'time':1})
      6 grid = Grid(ds, coords=coords)
----> 7 grid.diff(ds.tracer, 'X')

~/xgcm/xgcm/grid.py in diff(self, da, axis, **kwargs)
   1960         >>> grid.diff(da, ["X", "Y"], fill_value={"X": 0, "Y": 100})
   1961         """
-> 1962         return self._1d_grid_ufunc_dispatch("diff", da, axis, **kwargs)
   1963 
   1964     @docstrings.dedent

~/xgcm/xgcm/grid.py in _1d_grid_ufunc_dispatch(self, funcname, da, axis, to, keep_coords, metric_weighted, **kwargs)
   1767         for signature_1d, ax_name in zip(signatures, axis):
   1768 
-> 1769             grid_ufunc, remaining_kwargs = _select_grid_ufunc(
   1770                 funcname, signature_1d, module=gridops, **kwargs
   1771             )

~/xgcm/xgcm/grid.py in _select_grid_ufunc(funcname, signature, module, **kwargs)
   2336     ]
   2337     if len(name_matching_ufuncs) == 0:
-> 2338         raise NotImplementedError(
   2339             f"Could not find any pre-defined {funcname} grid ufuncs"
   2340         )

NotImplementedError: Could not find any pre-defined diff grid ufuncs
```


Pull request description:
Enable dask in gridops
This closes #429, but only works for 1D data if any core dims are chunked. I'll raise a separate issue for the 2D problem.

 - [x] Closes #429
 - [x] Tests added (just one for now, )
 - [x] Passes `pre-commit run --all-files`

Patch:
diff --git a/xgcm/grid.py b/xgcm/grid.py
index 7a9e67a..1214dc8 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -8,10 +8,16 @@ from collections import OrderedDict
 import docrep  # type: ignore
 import numpy as np
 import xarray as xr
+from dask.array import Array as Dask_Array
 
 from . import comodo, gridops
 from .duck_array_ops import _apply_boundary_condition, _pad_array, concatenate
-from .grid_ufunc import GridUFunc, _signatures_equivalent, apply_as_grid_ufunc
+from .grid_ufunc import (
+    GridUFunc,
+    _has_chunked_core_dims,
+    _signatures_equivalent,
+    apply_as_grid_ufunc,
+)
 from .metrics import iterate_axis_combinations
 
 try:
@@ -1760,6 +1766,12 @@ class Grid:
 
         signatures = self._create_1d_grid_ufunc_signatures(da, axis=axis, to=to)
 
+        # if any dims are chunked then we need dask
+        if isinstance(da.data, Dask_Array):
+            dask = "parallelized"
+        else:
+            dask = "forbidden"
+
         array = da
         # Apply 1D function over multiple axes
         # TODO This will call xarray.apply_ufunc once for each axis, but if signatures + kwargs are the same then we
@@ -1775,8 +1787,22 @@ class Grid:
                 metric = self.get_metric(array, ax_metric_weighted)
                 array = array * metric
 
+            # if chunked along core dim then we need map_overlap
+            core_dim = self._get_dims_from_axis(da, ax_name)
+            if _has_chunked_core_dims(array, core_dim):
+                map_overlap = True
+                dask = "allowed"
+            else:
+                map_overlap = False
+
             array = grid_ufunc(
-                self, array, axis=[ax_name], keep_coords=keep_coords, **remaining_kwargs
+                self,
+                array,
+                axis=[(ax_name,)],
+                keep_coords=keep_coords,
+                dask=dask,
+                map_overlap=map_overlap,
+                **remaining_kwargs,
             )
 
             if ax_metric_weighted:
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
index 7c7e1d1..14bb72e 100644
--- a/xgcm/grid_ufunc.py
+++ b/xgcm/grid_ufunc.py
@@ -544,11 +544,12 @@ def apply_as_grid_ufunc(
     return results_with_coords
 
 
-def _has_chunked_core_dims(obj: xr.DataArray, core_dims: Sequence[str]) -> bool:
-    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
+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
+
 
+def _has_chunked_core_dims(obj: xr.DataArray, core_dims: Sequence[str]) -> bool:
     # TODO what if only some of the core dimensions are chunked?
     return obj.chunks is not None and any(is_dim_chunked(obj, dim) for dim in core_dims)
 


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
index 8ddad10..8d1d71c 100644
--- a/xgcm/test/test_grid_ufunc.py
+++ b/xgcm/test/test_grid_ufunc.py
@@ -81,7 +81,7 @@ class TestParseGridUfuncSignature:
             _parse_grid_ufunc_signature(signature)
 
 
-def create_1d_test_grid_ds(ax_name):
+def create_1d_test_grid_ds(ax_name, length):
 
     grid_ds = xr.Dataset(
         coords={
@@ -89,25 +89,25 @@ def create_1d_test_grid_ds(ax_name):
                 [
                     f"{ax_name}_c",
                 ],
-                np.arange(1, 10),
+                np.arange(1, length + 1),
             ),
             f"{ax_name}_g": (
                 [
                     f"{ax_name}_g",
                 ],
-                np.arange(0.5, 9),
+                np.arange(0.5, length),
             ),
             f"{ax_name}_r": (
                 [
                     f"{ax_name}_r",
                 ],
-                np.arange(1.5, 10),
+                np.arange(1.5, length + 1),
             ),
             f"{ax_name}_i": (
                 [
                     f"{ax_name}_i",
                 ],
-                np.arange(1.5, 9),
+                np.arange(1.5, length),
             ),
             f"{ax_name}_o": (
                 [
@@ -121,8 +121,8 @@ def create_1d_test_grid_ds(ax_name):
     return grid_ds
 
 
-def create_1d_test_grid(ax_name):
-    grid_ds = create_1d_test_grid_ds(ax_name)
+def create_1d_test_grid(ax_name, length=9):
+    grid_ds = create_1d_test_grid_ds(ax_name, length)
     return Grid(
         grid_ds,
         coords={
@@ -137,9 +137,9 @@ def create_1d_test_grid(ax_name):
     )
 
 
-def create_2d_test_grid(ax_name_1, ax_name_2):
-    grid_ds_1 = create_1d_test_grid_ds(ax_name_1)
-    grid_ds_2 = create_1d_test_grid_ds(ax_name_2)
+def create_2d_test_grid(ax_name_1, ax_name_2, length1=9, length2=11):
+    grid_ds_1 = create_1d_test_grid_ds(ax_name_1, length1)
+    grid_ds_2 = create_1d_test_grid_ds(ax_name_2, length2)
 
     return Grid(
         ds=xr.merge([grid_ds_1, grid_ds_2]),
@@ -542,7 +542,6 @@ class TestDaskOverlap:
             grid=grid,
             signature="(X:center)->(X:left)",
             boundary_width={"X": (1, 0)},
-            # boundary="",
             dask="allowed",
             map_overlap=True,
         ).compute()
@@ -621,7 +620,7 @@ class TestDaskOverlap:
     def test_only_some_core_dims_are_chunked(self):
         raise NotImplementedError
 
-    def test_ufunc_changes_chunksize(self):
+    def test_raise_when_ufunc_changes_chunksize(self):
         @as_grid_ufunc(
             "(X:outer)->(X:center)",
             boundary_width={"X": (1, 0)},
@@ -703,6 +702,21 @@ class TestDaskOverlap:
             )
 
 
+# TODO tests for handling dask in gri.diff etc. should eventually live in test_grid.py
+def test_grid_diff_center_to_right_1d():
+    grid = create_1d_test_grid("depth")
+    da = np.sin(grid._ds.depth_c * 2 * np.pi / 9).chunk(1)
+    da.coords["depth_c"] = grid._ds.depth_c
+
+    diffed = (da.roll(depth_c=-1, roll_coords=False) - da).data
+    expected = xr.DataArray(
+        diffed, dims=["depth_r"], coords={"depth_r": grid._ds.depth_r}
+    ).compute()
+
+    result = grid.diff(da, axis="depth", to="right").compute()
+    assert_equal(result, expected)
+
+
 class TestSignaturesEquivalent:
     def test_equivalent(self):
         sig1 = "(X:center)->(X:left)"