Repository: xgcm/xgcm
Pull request #362 (2021-07-27 23:23:44)

as_grid_ufunc decorator
First step towards closing #344 - so far only does point (1) of @dcherian's [comment](https://github.com/xgcm/xgcm/issues/344#issuecomment-860760874), but the regex for the signature parser was tricky enough that I wanted to put this up as a WIP now that it passes my signature tests.

The regex parser is an adapted version of the one used by [numpy](https://github.com/numpy/numpy/blob/main/numpy/lib/function_base.py#L1846) and [dask](https://github.com/dask/dask/blob/8d49c3f1491b4aaccb6560814718f498abf63335/dask/array/gufunc.py#L18), and I'm going to copy dask's class structure for actually applying it to functions.

One question is how is the parser supposed to know what the `output_sizes` kwarg to `xr.apply_ufunc` should be? That information isn't in the signature specification @rabernat [gave](https://github.com/xgcm/xgcm/issues/344#issue-919651677). Or do we just not envisage a use case for this function where it creates new dimensions on the output?

 - [x] Tests added
 - [x] Passes `pre-commit run --all-files`
 - [ ] ~~User visible changes (including notable bug fixes) are documented in `whats-new.rst`~~
 - [x] New functions/methods are listed in `api.rst`


dcherian: > the parser supposed to know what the output_sizes kwarg to xr.apply_ufunc should be? 

You can infer the sizes of the new dimension based on a single dimension already present in `xgcm.Grid` i.e. if you know `center` is size N, the `inner` is size N-1, `outer` is N+1, etc. You might have to make up names for the new dimension though?
TomNicholas: > You can infer the sizes of the new dimension

Oh yes I guess maybe... Only if you know that X->Y though - what if the input had mutliple dimensions. Like would we ever need to deal with a function with a signature like `"(X:center),(Y:inner)->(Z:outer)"`?

> You might have to make up names for the new dimension though?

Those names will be specified in the signature, e.g. if `signature="(X:center)->(Y:outer)"` we know the output name is `Y`.
jbusecke: > Oh yes I guess maybe... Only if you know that X->Y though - what if the input had multiple dimensions? Like would we ever need to deal with a function with a signature like "(X:center),(Y:inner)->(Z:outer)"?

I couldnt think of any example that would require this at the moment. But let me sleep on it hehe.
dcherian: > Those names will be specified in the signature, e.g. if signature="(X:center)->(Y:outer)" we know the output name is Y.

I don't think so. The signature tells us that the output has one core dimension, along the Y axis at the position `outer`. But if `Grid` does not have any variable currently at position `Y: outer` we don't know the dimension name.
rabernat: > One question is how is the parser supposed to know what the `output_sizes` kwarg to `xr.apply_ufunc` should be?

I don't think it will be possible to fully populate the dimensions / coordinates without reference to the Grid dataset (the dataset used to initialize the grid, stored in `Grid._ds`). For each grid position signature (e.g. `(X:center)`) there will be at most 1 matching dimension (and possibly coordinate) in the Grid dataset.

We need to figure out the relationship between thie parser and the Grid object. But in my mind, the parser cannot be self-contained: it needs the Grid.
TomNicholas: > I don't think it will be possible to fully populate the dimensions / coordinates without reference to the Grid dataset

Okay I've obviously misunderstood. So the parser that I've written doesn't yet find the core dimensions, instead it really finds the xgcm Axes specified by the signature and the positions of those Axes. Then to get the core dims we need a second step where we consult the grid object via something like `grid.axes['Axis_name_from_parser'].coords['Axis_position_from_parser']`?
rabernat: Yes, that matches my understanding. The parser you've written is great and does just whit it needs. There is just a final step before xr.apply_ufunc...

> Then to get the core dims we need a second step where we consult the grid object via something like `grid.axes[Axis_name_from_parser].coords[Axis_position_from_parser]`?

exactly.
TomNicholas: Parsing the Axis and determining the core dims from the grid object seems fine, but moving forward there are potentially some problems with changing output sizes.

In `xarray.apply_ufunc` you can't change the size of core dimensions without "excluding" them:

```python
In [8]: def shorten(a):
   ...:     return a[1:]
   ...: 

In [9]: da = xr.DataArray(np.arange(10), dims='x')

In [10]: xr.apply_ufunc(shorten, da, input_core_dims=[['x']], output_core_dims=[['x']])
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-15-268f57ff3530> in <module>
----> 1 xr.apply_ufunc(shorten, da, input_core_dims=[['x']], output_core_dims=[['x']])

~/Documents/Work/Code/xarray/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)
   1178             join=join,
   1179             exclude_dims=exclude_dims,
-> 1180             keep_attrs=keep_attrs,
   1181         )
   1182     # feed Variables directly through apply_variable_ufunc

~/Documents/Work/Code/xarray/xarray/core/computation.py in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
    291 
    292     data_vars = [getattr(a, "variable", a) for a in args]
--> 293     result_var = func(*data_vars)
    294 
    295     if signature.num_outputs > 1:

~/Documents/Work/Code/xarray/xarray/core/computation.py in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
    779                     "dimensions specified in ``exclude_dims`` with "
    780                     "xarray.apply_ufunc are allowed to change size.".format(
--> 781                         dim, dim_sizes[dim], new_size
    782                     )
    783                 )

ValueError: size of dimension 'x' on inputs was unexpectedly changed by applied function from 10 to 9. Only dimensions specified in ``exclude_dims`` with xarray.apply_ufunc are allowed to change size.
```

Any operation that changes the Axis position from say "center" to "inner" would require changing the length of a core dimension.

You can do this

```python
In [16]: xr.apply_ufunc(shorten, da, input_core_dims=[['x']], output_core_dims=[['x']], exclude_dims={'x'})
Out[16]: 
<xarray.DataArray (x: 9)>
array([1, 2, 3, 4, 5, 6, 7, 8, 9])
Dimensions without coordinates: x
```
but I don't really understand what `exclude_dims` actually does well enough to know if this is a problem? The docs say "exclude from alignment and broadcasting entirely", but what does that mean? In what scenario will that break our use cases? One thing it will definitely do is drop any coordinates along that dimension.
rabernat: In your example, `input_core_dims` and `output_core_dims` are the same, but in xgcm they will be different. Doesn't that fix the problem?

Behold this example:
```python
import xarray as xr

def change_len_of_last_axis(a):
    return a[..., :-1]

da = xr.DataArray(np.ones((5, 10)), dims=['x', 'y'])

xr.apply_ufunc(
    change_len_of_last_axis,
    da,
    input_core_dims=['y'],
    output_core_dims=['z'],
    output_sizes={'z': da.shape[da.get_axis_num('y')] - 1}
)
```

rabernat: And just to clarify, the data model in xgcm is that _every axis position has a distinct dimension name_. Despite sharing coordinate values, `outer` and `inner` should have different dimensions. I agree this is confusing, but it helps keep things simple.
TomNicholas: Riiight. That's counter-intuitive, but I can see why it's like that. 

New dimensions are allowed to have new sizes, so that should mean this isn't a problem. 

So to summarize: the only time you need a new length on the output is when the output has an (ax:pos) pair that isn't present on the input, but so long as "every axis position has a distinct dimension name" then that new (ax:pos) pair should correspond to a core dim that isn't present on the input, so we don't need `exclude_dims`.  
TomNicholas: This now works for basic 1D cases, but if I chunk them with dask then it will complain that I'm trying to chunk along core dims. What do we want to do about that? For now we could just expose the kwargs for `dask='parallelized/'allowed'` and `allow_rechunk` as optional kwargs to the `apply_grid_ufunc` decorator?
rabernat: > This now works for basic 1D cases, but if I chunk them with dask then it will complain that I'm trying to chunk along core dims. What do we want to do about that?

This actually gets to a core issue which we need to make a decision about. Our current implementation ultimately uses xarray selections (+ concatentation for boundary conditions) in order to get the pairs of data on which to operate.

https://github.com/xgcm/xgcm/blob/95f4f33d72d2add00136e27f6b3bedecb97d4d77/xgcm/grid.py#L545-L548

This approach is agnostic about whether the "core dim" is chunked or not. It works either way, but performance is poor in both [chunked / unchunked] cases!

If we replace this with, e.g. numba-optimized ufuncs, these ufuncs will presumably require contiguous data along the core dims, thus we would not be able to use `xr.apply_ufunc(..., dask='allowed')`. In this scenario, we would lose a current feature, which is the ability to operate on chunked dimensions.

One workaround for this would be to have an extra layer of logic for chunked arrays that uses [dask.overlap](https://docs.dask.org/en/latest/generated/dask.array.overlap.overlap.html#dask.array.overlap.overlap) or its xarray equivalent. Or perhaps we always go directly to `dask.map_overlap`, which has a fast path for non-chunked dims. Not sure how this would mix with `xr.appy_ufunc`.

The question of how we handle overlap is also intimately related to the question of boundary conditions.

I'd be happy to hop on a call this week to do a little whiteboard brainstorming about this. It's tricky.
dcherian: > Or perhaps we always go directly to dask.map_overlap, which has a fast path for non-chunked dims. Not sure how this would mix with xr.appy_ufunc

`apply_ufunc(.., dask="allowed")` with `map_overlap` mapping the jitted function will work.

>  chunk boundary

Don't you mean actual boundaries? 2, 3.5, 4 are all boundary conditions at the domain edge. You can `pad` or apply any boundary condition you want and then call `map_overlap` with `boundary=None`. 

cube sphere faces will be complicated but the rest shouldn't be crazy. It'd be nice if we could add `boundary="tripolar"` and have a `Grid.pad` that does the right thing.


TomNicholas: > apply_ufunc(.., dask="allowed") with map_overlap mapping the jitted function will work.

This is what I was thinking. I'll have a go on a test case now.
rabernat: > Is there a way we can insert arbitrary values at the chunk boundary like that in dask.map_overlap?

No, but we can write our own custom function to pad the arrays and then call map_overlap.

> cube sphere faces will be complicated but the rest shouldn't be crazy. It'd be nice if we could add `boundary="tripolar"` and have a `Grid.pad` that does the right thing.

Right, I think we want to have a whole padding module that handles these things. The trick will be to avoid unnecessary dask tasks for contiguous core dims. Some numba-based solutions for periodic boundary conditions are reported here: https://github.com/ocean-eddy-cpt/gcm-filters/issues/45#issuecomment-884498140


TomNicholas: > Right, I think we want to have a whole padding module that handles these things.

Okay, I think this deserves it's own issue.


I've just pushed a test case that uses `apply_ufunc` and `map_overlap` on a chunked 1D example, but I'm getting this weird error, where the undecorated version of the function passes but the decorated one somehow doesn't. I've checked that `grid_ufunc` receives the same arguments in both cases, so I'm not really sure what's going on there...

```python
========================================================================================== FAILURES ==========================================================================================
_____________________________________________________________________________ TestGridUFunc.test_1d_dask_overlap _____________________________________________________________________________

self = <xgcm.test.test_grid_ufunc.TestGridUFunc object at 0x7fd0ceeb1340>

    def test_1d_dask_overlap(self):
        from dask.array import map_overlap
    
        def diff_center_to_left(a):
            return a - np.roll(a, shift=-1)
    
        def diff_overlap(a):
            return map_overlap(diff_center_to_left, a, depth=1, boundary=0)
    
        grid = create_1d_test_grid()
        da = np.sin(grid._ds.x_c * 2 * np.pi / 9).chunk(1)
        da.coords["x_c"] = grid._ds.x_c
    
        diffed = (da - da.roll(x_c=-1, roll_coords=False)).data
        expected = xr.DataArray(diffed, dims=["x_g"], coords={"x_g": grid._ds.x_g})
    
        # Test direct application
        result = grid_ufunc(
            diff_center_to_left,
            da,
            grid=grid,
            signature="(X:center)->(X:left)",
            dask="allowed",
        ).compute()
        assert_equal(result, expected)
    
        # Test decorator
        @as_grid_ufunc(grid, "(X:center)->(X:left)", dask="allowed")
        def diff_overlap(a):
            return map_overlap(diff_center_to_left, a, depth=1, boundary=0)
    
        result = diff_overlap(da).compute()
>       assert_equal(result, expected)
E       AssertionError: Left and right DataArray objects are not equal
E       
E       Differing values:
E       L
E           array([-3.420201e-01,  1.187823e-01,  5.240053e-01,  6.840403e-01,
E                   5.240053e-01,  1.187823e-01, -3.420201e-01, -6.427876e-01,
E                  -2.449294e-16])
E       R
E           array([-0.34202 ,  0.118782,  0.524005,  0.68404 ,  0.524005,  0.118782,
E                  -0.34202 , -0.642788, -0.642788])

xgcm/test/test_grid_ufunc.py:194: AssertionError

```
rabernat: Since numba uses a different numerical implementation than numpy, we don't necessary expect exact equality with floating point data. Maybe use `assert_allclose` instead of `assert_equal`?
TomNicholas: > Since numba uses a different numerical implementation than numpy, we don't necessary expect exact equality with floating point data. Maybe use assert_allclose instead of assert_equal?

Yes good point, but in this case the error is caused by the final value in the array being `-2.449294e-16` instead of `-0.642788`, that's what I don't understand.
jbusecke: Should `boundary` here
```python
def diff_overlap(a):
            return map_overlap(diff_center_to_left, a, depth=1, boundary=0)
```
be `'periodic'`? 
TomNicholas: > Should boundary here be 'periodic'?

That actually now passes?!

But I'm totally confused as to how - previously I had one passing with `boundary=0` being called through `grid_ufunc`, and one failing also with `boundary=0` , the only difference was that it was being called through the `as_grid_ufunc` decorator. But the `as_grid_ufunc` decorator immediately passes on to `grid_ufunc`, so they should have behaved identically...

Now now my test passes but my original mystery is not solved :sweat_smile: 
jbusecke: So the boundary applies only to the outer parts of the dask array, not the chunks? Am I understanding this right? To get the behavior of `np.roll` on the full domain, you need to have a periodic boundary? Not sure why that would work with 0?
TomNicholas: This PR now does (1), (2), and (3) from @dcherian 's [comment](https://github.com/xgcm/xgcm/issues/344#issuecomment-860760874), and there is an example test case which uses `map_overlap` within `xarray.apply_ufunc`. 

The final item is (4) 

> Do something about metrics and boundary

That seems like a can of worms (which I haven't started thinking about properly yet) - perhaps handling metrics and boundaries should be left to a second PR?
jbusecke: > That seems like a can of worms (which I haven't started thinking about properly yet) - perhaps handling metrics and boundaries should be left to a second PR?

Metrics handling at this point could be simply a multiplication onto the 'diffed'/'interped'/'ufunced' array (presuming we use the ufunc approach for the grid methods).
So if we could first implement a ufunc version of the basic operations like `.diff()` and `.interp() we still get the metrics for free?

In some higher level cases (potential vorticity?) we might want to pass the metrics into the ufunc, but I think your call is right to postpone this.
jbusecke: I think conceptually, the boundary question is harder to solve, and should be the next big ticket item?
TomNicholas: A summary of what I need to do:

```python
@as_grid_ufunc(signature="(X:center)->(X:left)")
def diff_center_to_left(a):
    return a - np.roll(a, -1)
```

`diff_centre_to_left` now has an additional `grid` argument. When `diff_centre_to_left` is called it parses the signature, looks up the axes in the grid object, then calls `xarray.apply_ufunc`. So basically the same as what `as_grid_ufunc` already does, just need to change the place where `grid` is passed (`grid` is currently passed to the decorator, but needs to be passed to the decorated function). That's the only real change to this PR needed I think.

```python
class Grid:
    def apply_grid_ufunc(self, grid_ufunc, *args, **kwargs):
        """General method that could be used to call any user-defined grid_ufunc"""
        return grid_ufunc(self, *args, **kwargs)

    def diff(self, da, axis, **kwargs):
        return self.apply_grid_ufunc(diff_center_to_left, da, axis, **kwargs)
```
rabernat: Yes that seems correct to me!
jbusecke: I really like this outline: 

> Right. It sounds like the first step will be to just get the grid_ufunc machinery working without touching the existing API. Then we look at how to refactor the implementation of Grid.diff, Grid.interp, etc to use grid_ufunc. Finally, once that is done, we see if we can accelerate the core grid_ufuncs further using numba.

Perhaps somewhere in there it might be good to discuss adding some 'performance metrics' into our tests. I think it would be useful to have some performance tests or even something like the number of dask chunks produced by a certain operation as a baseline. Having this implemented before e.g. optimizing with numba would enable us to track the improvements?
TomNicholas: > grid is currently passed to the decorator, but needs to be passed to the decorated function

I've made this change in the latest commits, which seems to work nicely. I've also made a new branch in the xgcm repo and changed the base of this PR to merge into that instead (which is why it no longer needs an approving review I believe). Unless anyone has further comments I'm going to merge this and then start changing the `Grid` class in a follow-up PR.

> Perhaps somewhere in there it might be good to discuss adding some 'performance metrics' into our tests.

I'm not sure how easy it is to write automated tests that check number of dask tasks created, but I agree we do need to at least manually check that the refactoring is doing what we hope in terms of performance. After the next PR we should hopefully have two branches to compare, and can maybe use one of Ryan's motivating examples (which was in a gist I have saved somewhere...) as a test case.
TomNicholas: While discussing with Julius I realised that it would be helpful if the grid ufuncs knew what their signatures were, so I've now promoted the return value of `@as_grid_ufunc` to be a callable `GridUFunc` class instead of a function, which possesses a `.signature` attribute. This is again inspired by how `dask.apply_gufunc` works, and opens the door to having grid ufuncs which know what `boundary` and so on that they can and can't handle.

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 41d529a..dc3e74f 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -7,16 +7,16 @@ repos:
       - id: check-yaml
 
   - repo: https://github.com/PyCQA/isort
-    rev: 5.8.0
+    rev: 5.9.3
     hooks:
       - id: isort
 
   - repo: https://github.com/python/black
-    rev: 21.5b1
+    rev: 21.7b0
     hooks:
       - id: black
   - repo: https://github.com/keewis/blackdoc
-    rev: v0.3.3
+    rev: v0.3.4
     hooks:
       - id: blackdoc
   - repo: https://gitlab.com/pycqa/flake8
diff --git a/ci/environment-core-deps.yml b/ci/environment-core-deps.yml
index 9920807..a247e24 100644
--- a/ci/environment-core-deps.yml
+++ b/ci/environment-core-deps.yml
@@ -8,6 +8,7 @@ dependencies:
   - dask
   - docrep<=0.2.7
   - numpy
+  - scipy
   - pip
   - pre-commit
   - pytest
diff --git a/ci/environment-upstream-dev.yml b/ci/environment-upstream-dev.yml
index 1a109f9..2e5cbcc 100644
--- a/ci/environment-upstream-dev.yml
+++ b/ci/environment-upstream-dev.yml
@@ -8,6 +8,7 @@ dependencies:
   - docrep<=0.2.7
   - numba
   - numpy
+  - scipy
   - pip
   - pre-commit
   - pytest
diff --git a/ci/environment.yml b/ci/environment.yml
index 3104a04..22d34ad 100644
--- a/ci/environment.yml
+++ b/ci/environment.yml
@@ -7,6 +7,7 @@ dependencies:
   - docrep<=0.2.7
   - numba
   - numpy
+  - scipy
   - pip
   - pre-commit
   - pytest
diff --git a/doc/api.rst b/doc/api.rst
index f1515fe..366d906 100644
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -13,6 +13,15 @@ Grid
 .. autoclass:: xgcm.Grid
    :members:
 
+Grid ufuncs
+====
+
+.. autosummary::
+   :toctree: generated/
+
+   apply_grid_ufunc
+   as_grid_ufunc
+
 autogenerate
 ============
 
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index a9ef774..e5fcffa 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -10,6 +10,14 @@ v0.6.0 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- Turn numpy-style ufuncs into grid-aware "grid-ufuncs" via new functions :py:meth:`~xgcm.grid_ufunc.apply_grid_ufunc`
+  and :py:meth:`~xgcm.grid_ufunc.as_grid_ufunc`. Beginning of an internal refactor to rewrite all grid computation methods
+  as grid ufuncs. (:pull:`362`, :issue:`344`)
+  By `Thomas Nicholas <https://github.com/tomnicholas>`_.
+
+- :py:meth:`~xgcm.grid.Grid.set_metrics` now enables adding metrics to a grid object (:pull:`336`, :issue:`199`).
+  By `Dianne Deauna <https://github.com/jdldeauna>`_ under the `SIParCS internship <https://www2.cisl.ucar.edu/siparcs-2021-projects#8>`_.
+
 - :py:meth:`~xgcm.grid.Grid.get_metric` refactored, and now incorporates :py:meth:`~xgcm.grid.Grid.interp_like` to allow for automatic interpolation of missing metrics from available values on surrounding positions (:pull:`345`, :pull:`354`).
   By `Dianne Deauna <https://github.com/jdldeauna>`_.[*]_
 
diff --git a/setup.py b/setup.py
index d470556..9f245a3 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@ extras_require = {
 extras_require["dev"] = extras_require["complete"] + [
     "pytest",
     "pytest-cov",
-    "flake8",
+    "scipy" "flake8",
     "black",
     "codecov",
 ]
diff --git a/xgcm/__init__.py b/xgcm/__init__.py
index 1107262..068e390 100644
--- a/xgcm/__init__.py
+++ b/xgcm/__init__.py
@@ -5,3 +5,4 @@ except ImportError:
 
 from .autogenerate import generate_grid_ds
 from .grid import Axis, Grid
+from .grid_ufunc import apply_as_grid_ufunc, as_grid_ufunc
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
new file mode 100644
index 0000000..68cf25c
--- /dev/null
+++ b/xgcm/grid_ufunc.py
@@ -0,0 +1,292 @@
+import re
+
+import xarray as xr
+
+# Modified version of `numpy.lib.function_base._parse_gufunc_signature`
+# Modifications:
+#   - Specify xgcm.Axis name and "axis positions" instead of numpy axes as (ax_name:ax_pos)
+_AXIS_NAME = r"\w+"
+_AXIS_POSITION = "(?:center|left|right|inner|outer)"
+_AXIS_NAME_POSITION_PAIR = f"{_AXIS_NAME}:{_AXIS_POSITION}"
+_AXIS_NAME_POSITION_PAIR_LIST = (
+    f"(?:{_AXIS_NAME_POSITION_PAIR}(?:,{_AXIS_NAME_POSITION_PAIR})*,?)*"
+)
+_ARGUMENT = rf"\({_AXIS_NAME_POSITION_PAIR_LIST}\)"
+_ARGUMENT_LIST = f"{_ARGUMENT}(?:,{_ARGUMENT})*"
+_SIGNATURE = f"^{_ARGUMENT_LIST}->{_ARGUMENT_LIST}$"
+
+
+def _parse_grid_ufunc_signature(signature):
+    """
+    Parse string signatures for a grid-aware universal function.
+
+    The way this parser works excludes using Axis names that match possible xgcm
+    axis positions, i.e. ['center', 'left', 'right', 'inner', 'outer'].
+
+    Arguments
+    ---------
+    signature : string
+        Grid universal function signature. Specifies the xgcm.Axis names and
+        positions for each input and output variable, e.g.,
+
+        ``"(X:center)->(X:left)"`` for ``diff_center_to_left(a)`.
+
+    Returns
+    -------
+    input_axes_names : List[Tuple[str, ...]]
+        Input Axes names parsed from the signature
+    output_axes_names : List[Tuple[str, ...]]
+        Output Axes names parsed from the signature
+    input_axes_positions : List[Tuple[str, ...]]
+        Input Axes positions parsed from the signature
+    output_axes_positions : List[Tuple[str, ...]]
+        Output Axes positions parsed from the signature
+    """
+
+    signature = signature.replace(" ", "")
+
+    if not re.match(_SIGNATURE, signature):
+        raise ValueError(f"Not a valid grid ufunc signature: {signature}")
+
+    in_txt, out_txt = signature.split("->")
+
+    in_ax_names = []
+    for arg in re.findall(_ARGUMENT, in_txt):
+        # Delete the axis positions so they aren't matched as axis names
+        only_names = re.sub(_AXIS_POSITION, "", arg)
+        in_ax_names.append(tuple(re.findall(_AXIS_NAME, only_names)))
+
+    out_ax_names = []
+    for arg in re.findall(_ARGUMENT, out_txt):
+        only_names = re.sub(_AXIS_POSITION, "", arg)
+        out_ax_names.append(tuple(re.findall(_AXIS_NAME, only_names)))
+
+    in_ax_pos = [
+        tuple(re.findall(_AXIS_POSITION, arg)) for arg in re.findall(_ARGUMENT, in_txt)
+    ]
+    out_ax_pos = [
+        tuple(re.findall(_AXIS_POSITION, arg)) for arg in re.findall(_ARGUMENT, out_txt)
+    ]
+
+    return in_ax_names, out_ax_names, in_ax_pos, out_ax_pos
+
+
+class GridUFunc:
+    """
+    Binds a numpy ufunc into a "grid-aware ufunc", meaning that when called ufunc is wrapped by `apply_as_grid_ufunc`.
+
+    Parameters
+    ----------
+    ufunc : callable
+        Function to call like `func(*args, **kwargs)` on numpy-like unlabeled
+        arrays (`.data`).
+
+        Passed directly on to `xarray.apply_ufunc`.
+    signature : string
+        Grid universal function signature. Specifies the xgcm.Axis names and
+        positions for each input and output variable, e.g.,
+
+        ``"(X:center)->(X:left)"`` for ``diff_center_to_left(a)`.
+    dask : {"forbidden", "allowed", "parallelized"}, default: "forbidden"
+        How to handle applying to objects containing lazy data in the form of
+        dask arrays. Passed directly on to `xarray.apply_ufunc`.
+
+    Returns
+    -------
+    grid_ufunc : callable
+        Class which when called consumes and produces xarray objects, whose xgcm Axis
+        names and positions must conform to the pattern specified by `signature`.
+        Calling function has an additional positional argument `grid`, of type `xgcm.Grid`,
+        so that `func`'s new signature is `func(grid, *args, **kwargs)`. This grid
+        argument is passed on to `apply_grid_ufunc`.
+
+    See Also
+    --------
+    as_grid_ufunc
+    apply_as_grid_ufunc
+    Grid.apply_as_grid_ufunc
+    """
+
+    def __init__(self, ufunc, **kwargs):
+        self.ufunc = ufunc
+        self.signature = kwargs.pop("signature", "")
+        self.dask = kwargs.pop("dask", "")
+        if kwargs:
+            raise TypeError("Unsupported keyword argument(s) provided")
+
+    def __call__(self, grid=None, *args, **kwargs):
+        return apply_as_grid_ufunc(
+            self.ufunc,
+            *args,
+            grid=grid,
+            signature=self.signature,
+            dask=self.dask,
+            **kwargs,
+        )
+
+
+def as_grid_ufunc(signature="", **kwargs):
+    """
+    Decorator which turns a numpy ufunc into a "grid-aware ufunc".
+
+    Parameters
+    ----------
+    ufunc : callable
+        Function to call like `func(*args, **kwargs)` on numpy-like unlabeled
+        arrays (`.data`).
+
+        Passed directly on to `xarray.apply_ufunc`.
+    signature : string
+        Grid universal function signature. Specifies the xgcm.Axis names and
+        positions for each input and output variable, e.g.,
+
+        ``"(X:center)->(X:left)"`` for ``diff_center_to_left(a)`.
+    dask : {"forbidden", "allowed", "parallelized"}, default: "forbidden"
+        How to handle applying to objects containing lazy data in the form of
+        dask arrays. Passed directly on to `xarray.apply_ufunc`.
+
+    Returns
+    -------
+    grid_ufunc : callable
+        Function which consumes and produces xarray objects, whose xgcm Axis
+        names and positions must conform to the pattern specified by `signature`.
+        Function has an additional positional argument `grid`, of type `xgcm.Grid`,
+        so that `func`'s new signature is `func(grid, *args, **kwargs)`. This grid
+        argument is passed on to `apply_grid_ufunc`.
+
+    See Also
+    --------
+    apply_as_grid_ufunc
+    Grid.apply_as_grid_ufunc
+    """
+    _allowedkwargs = {
+        "dask",
+    }
+    if kwargs.keys() - _allowedkwargs:
+        raise TypeError("Unsupported keyword argument(s) provided")
+
+    def _as_grid_ufunc(ufunc):
+        return GridUFunc(ufunc, signature=signature, **kwargs)
+
+    return _as_grid_ufunc
+
+
+def apply_as_grid_ufunc(
+    func, *args, grid=None, signature="", dask="forbidden", **kwargs
+):
+    """
+    Apply a function to the given arguments in a grid-aware manner.
+
+    The relationship between xgcm axes on the input and output are specified by
+    `signature`. Wraps xarray.apply_ufunc, but determines the core dimensions
+    from the grid and signature passed.
+
+    Parameters
+    ----------
+    func : callable
+        Function to call like `func(*args, **kwargs)` on numpy-like unlabeled
+        arrays (`.data`).
+
+        Passed directly on to `xarray.apply_ufunc`.
+    grid : xgcm.Grid
+        The xgcm Grid object which contains the various xgcm.Axis described by
+        `signature`.
+    signature : string
+        Grid universal function signature. Specifies the xgcm.Axis names and
+        positions for each input and output variable, e.g.,
+
+        ``"(X:center)->(X:left)"`` for ``diff_center_to_left(a)`.
+    dask : {"forbidden", "allowed", "parallelized"}, default: "forbidden"
+        How to handle applying to objects containing lazy data in the form of
+        dask arrays. Passed directly on to `xarray.apply_ufunc`.
+
+    Returns
+    -------
+    results
+        The result of the call to `xarray.apply_ufunc`, but including the coordinates
+        given by the signature, which are read from the grid. Output is either a single
+        object or a tuple of such objects.
+
+    See Also
+    --------
+    as_grid_ufunc
+    Grid.apply_as_grid_ufunc
+    """
+
+    if grid is None:
+        raise ValueError("Must provide a grid object to describe the Axes")
+
+    # Extract Axes information from signature
+    in_ax_names, out_ax_names, in_ax_pos, out_ax_pos = _parse_grid_ufunc_signature(
+        signature
+    )
+
+    # Check that input args are in correct grid positions
+    for i, (arg_ns, arg_ps, arg) in enumerate(zip(in_ax_names, in_ax_pos, args)):
+        for n, p in zip(arg_ns, arg_ps):
+            try:
+                ax_pos = grid.axes[n].coords[p]
+            except KeyError:
+                raise ValueError(f"Axis position ({n}:{p}) does not exist in grid")
+
+            if ax_pos not in arg.coords:
+                raise ValueError(
+                    f"Mismatch between signature and input argument {i}: "
+                    f"Signature specified data to lie at Axis Position ({n}:{p}), "
+                    f"but the corresponding grid coordinate {grid.axes[n].coords[p]} "
+                    f"does not appear in argument"
+                    f"{arg}"
+                )
+
+            # TODO also check that dims are the right length for their stated Axis positions on inputs?
+
+    # Determine core dimensions for apply_ufunc
+    in_core_dims = [
+        [grid.axes[n].coords[p] for n, p in zip(arg_ns, arg_ps)]
+        for arg_ns, arg_ps in zip(in_ax_names, in_ax_pos)
+    ]
+    out_core_dims = [
+        [grid.axes[n].coords[p] for n, p in zip(arg_ns, arg_ps)]
+        for arg_ns, arg_ps in zip(out_ax_names, out_ax_pos)
+    ]
+
+    all_out_core_dims = set(dim for arg in out_core_dims for dim in arg)
+
+    # Determine expected output dimension sizes from grid._ds
+    # Only required when dask='parallelized'
+    out_sizes = {out_dim: grid._ds.dims[out_dim] for out_dim in all_out_core_dims}
+
+    # Perform operation via xarray.apply_ufunc
+    results = xr.apply_ufunc(
+        func,
+        *args,
+        input_core_dims=in_core_dims,
+        output_core_dims=out_core_dims,
+        dask=dask,
+        **kwargs,
+        dask_gufunc_kwargs={"output_sizes": out_sizes},
+    )
+
+    # apply_ufunc might return multiple objects
+    if not isinstance(results, tuple):
+        results = (results,)
+
+    # Restore any coordinates associated with new output dims that are present in grid
+    # TODO should this be optional via a `keep_coords` arg?
+    results_with_coords = []
+    for res, arg_out_core_dims in zip(results, out_core_dims):
+        new_core_dim_coords = {
+            dim: grid._ds.coords[dim]
+            for dim in arg_out_core_dims
+            if dim in grid._ds.dims and dim not in res.coords
+        }
+        res = res.assign_coords(new_core_dim_coords)
+        results_with_coords.append(res)
+
+    # Return single results not wrapped in 1-element tuple, like xr.apply_ufunc
+    if len(results_with_coords) == 1:
+        (results_with_coords,) = results_with_coords
+
+    # TODO handle metrics and boundary? Or should that happen in the ufuncs themselves?
+
+    return results_with_coords


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
new file mode 100644
index 0000000..61f1aaa
--- /dev/null
+++ b/xgcm/test/test_grid_ufunc.py
@@ -0,0 +1,335 @@
+import re
+
+import numpy as np
+import pytest
+import xarray as xr
+from xarray.testing import assert_equal
+
+from xgcm.grid import Grid
+from xgcm.grid_ufunc import (
+    GridUFunc,
+    _parse_grid_ufunc_signature,
+    apply_as_grid_ufunc,
+    as_grid_ufunc,
+)
+
+
+class TestParseGridUfuncSignature:
+    @pytest.mark.parametrize(
+        "signature, exp_in_ax_names, exp_out_ax_names, exp_in_ax_pos, exp_out_ax_pos",
+        [
+            ("()->()", [()], [()], [()], [()]),
+            ("(X:center)->()", [("X",)], [()], [("center",)], [()]),
+            ("()->(X:left)", [()], [("X",)], [()], [("left",)]),
+            ("(X:center)->(X:left)", [("X",)], [("X",)], [("center",)], [("left",)]),
+            ("(X:left)->(Y:center)", [("X",)], [("Y",)], [("left",)], [("center",)]),
+            ("(X:left)->(Y:center)", [("X",)], [("Y",)], [("left",)], [("center",)]),
+            (
+                "(X:left),(X:right)->(Y:center)",
+                [("X",), ("X",)],
+                [("Y",)],
+                [("left",), ("right",)],
+                [("center",)],
+            ),
+            (
+                "(X:center)->(Y:inner),(Y:outer)",
+                [("X",)],
+                [("Y",), ("Y",)],
+                [("center",)],
+                [("inner",), ("outer",)],
+            ),
+            (
+                "(X:center,Y:center)->(Z:center)",
+                [("X", "Y")],
+                [("Z",)],
+                [("center", "center")],
+                [("center",)],
+            ),
+        ],
+    )
+    def test_parse_valid_signatures(
+        self,
+        signature,
+        exp_in_ax_names,
+        exp_out_ax_names,
+        exp_in_ax_pos,
+        exp_out_ax_pos,
+    ):
+        in_ax_names, out_ax_names, in_ax_pos, out_ax_pos = _parse_grid_ufunc_signature(
+            signature
+        )
+        assert in_ax_names == exp_in_ax_names
+        assert out_ax_names == exp_out_ax_names
+        assert in_ax_pos == exp_in_ax_pos
+        assert out_ax_pos == exp_out_ax_pos
+
+    @pytest.mark.parametrize(
+        "signature",
+        [
+            "(x:left)(y:left)->()",
+            "(x:left),(y:left)->",
+            "((x:left))->(x:left)",
+            "((x:left))->(x:left)",
+            "(x:left)->(x:left)," "(i)->(i)",
+            "(X:centre)->()",
+        ],
+    )
+    def test_invalid_signatures(self, signature):
+        with pytest.raises(ValueError):
+            _parse_grid_ufunc_signature(signature)
+
+
+def create_1d_test_grid_ds(ax_name):
+
+    grid_ds = xr.Dataset(
+        coords={
+            f"{ax_name}_c": (
+                [
+                    f"{ax_name}_c",
+                ],
+                np.arange(1, 10),
+            ),
+            f"{ax_name}_g": (
+                [
+                    f"{ax_name}_g",
+                ],
+                np.arange(0.5, 9),
+            ),
+            f"{ax_name}_r": (
+                [
+                    f"{ax_name}_r",
+                ],
+                np.arange(1.5, 10),
+            ),
+            f"{ax_name}_i": (
+                [
+                    f"{ax_name}_i",
+                ],
+                np.arange(1.5, 9),
+            ),
+        }
+    )
+
+    return grid_ds
+
+
+def create_1d_test_grid(ax_name):
+    grid_ds = create_1d_test_grid_ds(ax_name)
+    return Grid(
+        grid_ds,
+        coords={
+            f"{ax_name.upper()}": {
+                "center": f"{ax_name}_c",
+                "left": f"{ax_name}_g",
+                "right": f"{ax_name}_r",
+                "inner": f"{ax_name}_i",
+            }
+        },
+    )
+
+
+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)
+
+    return Grid(
+        ds=xr.merge([grid_ds_1, grid_ds_2]),
+        coords={
+            f"{ax_name_1.upper()}": {
+                "center": f"{ax_name_1}_c",
+                "left": f"{ax_name_1}_g",
+                "right": f"{ax_name_1}_r",
+                "inner": f"{ax_name_1}_i",
+            },
+            f"{ax_name_2.upper()}": {
+                "center": f"{ax_name_2}_c",
+                "left": f"{ax_name_2}_g",
+                "right": f"{ax_name_2}_r",
+                "inner": f"{ax_name_2}_i",
+            },
+        },
+    )
+
+
+class TestGridUFunc:
+    def test_stores_ufunc_kwarg_info(self):
+        signature = "(X:center)->(X:left)"
+
+        @as_grid_ufunc(signature)
+        def diff_center_to_left(a):
+            return a - np.roll(a, shift=-1)
+
+        assert isinstance(diff_center_to_left, GridUFunc)
+        assert diff_center_to_left.signature == signature
+
+        with pytest.raises(TypeError, match="Unsupported keyword argument"):
+
+            @as_grid_ufunc(signature, junk="useless")
+            def diff_center_to_left(a):
+                return a - np.roll(a, shift=-1)
+
+    def test_input_on_wrong_positions(self):
+        grid = create_1d_test_grid("x")
+        da = np.sin(grid._ds.x_g * 2 * np.pi / 9)
+
+        with pytest.raises(ValueError, match=re.escape("(Y:center) does not exist")):
+            apply_as_grid_ufunc(lambda x: x, da, grid=grid, signature="(Y:center)->()")
+
+        with pytest.raises(ValueError, match="coordinate x_c does not appear"):
+            apply_as_grid_ufunc(lambda x: x, da, grid=grid, signature="(X:center)->()")
+
+    def test_1d_unchanging_size_no_dask(self):
+        def diff_center_to_left(a):
+            return a - np.roll(a, shift=-1)
+
+        grid = create_1d_test_grid("x")
+        da = np.sin(grid._ds.x_c * 2 * np.pi / 9)
+        da.coords["x_c"] = grid._ds.x_c
+
+        diffed = (da - da.roll(x_c=-1, roll_coords=False)).data
+        expected = xr.DataArray(diffed, dims=["x_g"], coords={"x_g": grid._ds.x_g})
+
+        # Test direct application
+        result = apply_as_grid_ufunc(
+            diff_center_to_left, da, grid=grid, signature="(X:center)->(X:left)"
+        )
+        assert_equal(result, expected)
+
+        # Test decorator
+        @as_grid_ufunc("(X:center)->(X:left)")
+        def diff_center_to_left(a):
+            return a - np.roll(a, shift=-1)
+
+        result = diff_center_to_left(grid, da)
+        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:])
+
+        grid = create_1d_test_grid("x")
+        da = xr.DataArray(
+            np.arange(10, 19), dims=["x_c"], coords={"x_c": grid._ds.x_c}
+        ).chunk()
+
+        expected = da.interp(x_c=np.arange(1.5, 9), method="linear").rename(x_c="x_i")
+
+        # Test direct application
+        result = apply_as_grid_ufunc(
+            interp_center_to_inner,
+            da,
+            grid=grid,
+            signature="(X:center)->(X:inner)",
+            dask="parallelized",
+        ).compute()
+        assert_equal(result, expected)
+
+        # Test decorator
+        @as_grid_ufunc("(X:center)->(X:inner)", dask="parallelized")
+        def interp_center_to_inner(a):
+            return 0.5 * (a[:-1] + a[1:])
+
+        result = interp_center_to_inner(grid, da).compute()
+        assert_equal(result, expected)
+
+    def test_1d_overlap_dask_allowed(self):
+        from dask.array import map_overlap
+
+        def diff_center_to_left(a):
+            return a - np.roll(a, shift=-1)
+
+        def diff_overlap(a):
+            return map_overlap(diff_center_to_left, a, depth=1, boundary="periodic")
+
+        grid = create_1d_test_grid("x")
+        da = np.sin(grid._ds.x_c * 2 * np.pi / 9).chunk(1)
+        da.coords["x_c"] = grid._ds.x_c
+
+        diffed = (da - da.roll(x_c=-1, roll_coords=False)).data
+        expected = xr.DataArray(
+            diffed, dims=["x_g"], coords={"x_g": grid._ds.x_g}
+        ).compute()
+
+        # Test direct application
+        result = apply_as_grid_ufunc(
+            diff_center_to_left,
+            da,
+            grid=grid,
+            signature="(X:center)->(X:left)",
+            dask="allowed",
+        ).compute()
+        assert_equal(result, expected)
+
+        # Test decorator
+        @as_grid_ufunc("(X:center)->(X:left)", dask="allowed")
+        def diff_overlap(a):
+            return map_overlap(diff_center_to_left, a, depth=1, boundary="periodic")
+
+        result = diff_overlap(grid, da).compute()
+        assert_equal(result, expected)
+
+    def test_multiple_inputs(self):
+        def inner_product_left_right(a, b):
+            return np.inner(a, b)
+
+        grid = create_1d_test_grid("x")
+        a = np.sin(grid._ds.x_g * 2 * np.pi / 9)
+        a.coords["x_g"] = grid._ds.x_g
+        b = np.cos(grid._ds.x_r * 2 * np.pi / 9)
+        b.coords["x_r"] = grid._ds.x_r
+
+        expected = xr.DataArray(np.inner(a, b))
+
+        # Test direct application
+        result = apply_as_grid_ufunc(
+            inner_product_left_right,
+            a,
+            b,
+            grid=grid,
+            signature="(X:left),(X:right)->()",
+        )
+        assert_equal(result, expected)
+
+        # Test decorator
+        @as_grid_ufunc("(X:left),(X:right)->()")
+        def inner_product_left_right(a, b):
+            return np.inner(a, b)
+
+        result = inner_product_left_right(grid, a, b)
+        assert_equal(result, expected)
+
+    def test_multiple_outputs(self):
+        def diff_center_to_inner(a, axis):
+            result = a - np.roll(a, shift=1, axis=axis)
+            return np.delete(result, 0, axis)  # remove first element along axis
+
+        def grad_to_inner(a):
+            return diff_center_to_inner(a, axis=0), diff_center_to_inner(a, axis=1)
+
+        grid = create_2d_test_grid("x", "y")
+
+        a = grid._ds.x_c ** 2 + grid._ds.y_c ** 2
+
+        expected_u = 2 * grid._ds.x_i.expand_dims(dim={"y_c": len(grid._ds.y_c)})
+        expected_u.coords["y_c"] = grid._ds.y_c
+        expected_v = 2 * grid._ds.y_i.expand_dims(dim={"x_c": len(grid._ds.x_c)})
+        expected_v.coords["x_c"] = grid._ds.x_c
+
+        # Test direct application
+        u, v = apply_as_grid_ufunc(
+            grad_to_inner,
+            a,
+            grid=grid,
+            signature="(X:center,Y:center)->(X:inner,Y:center),(X:center,Y:inner)",
+        )
+        assert_equal(u.T, expected_u)
+        assert_equal(v, expected_v)
+
+        # Test decorator
+        @as_grid_ufunc("(X:center,Y:center)->(X:inner,Y:center),(X:center,Y:inner)")
+        def grad_to_inner(a):
+            return diff_center_to_inner(a, axis=0), diff_center_to_inner(a, axis=1)
+
+        u, v = grad_to_inner(grid, a)
+        assert_equal(u.T, expected_u)
+        assert_equal(v, expected_v)