Repository: xgcm/xgcm
Pull request #368 (2021-09-10 20:05:08)

Grid ufunc refactor - proof of principle
Next step towards #344:
- Defines various grid ufuncs in the file `gridops.py`
- Moves the `GCMDataset` class that was in that file to a new file `gcmdataset.py` (do we even still need that code? @jbusecke said that was some of your old code @rabernat ?)
- Starts work on a new `Grid._grid_ufunc_dispatch` method that aims to eventually replace the existing `Grid._grid_func` method and everything that calls below (including `Grid._neighbour_binary_func` etc.).
- Sketches out a mechanism to search through all the grid ufuncs defined in `gridops.py` ~~(but the searching doesn't actually work yet, still a WIP)~~ this now works, next step is to add support for a `boundary` kwarg so that some of the tests actually pass.


Builds on top of #362, so actually there are far fewer changes than this looks like (I don't know why this still lists `test_grid_ufunc.py` as a new change...)

 - [x] Continues work towards #344
 - [x] Tests added
 - [x] Passes `pre-commit run --all-files`
 - [ ] User visible changes (including notable bug fixes) are documented in `whats-new.rst`
 - [ ] New functions/methods are listed in `api.rst`


rabernat: > * do we even still need that code?

No. It's very old and unused.
TomNicholas: > No. It's very old and unused.

Okay I'll just delete the whole class then.
TomNicholas: So I'm adding the `boundary` kwargs, and I would have liked to have a `gridops.py` module that looked like this:

```python
@as_grid_ufunc(signature="(X:center)->(X:left)", boundary="fill")
def diff_center_to_left(a):
    return a - np.roll(a, -1)


@as_grid_ufunc(signature="(X:center)->(X:left)", boundary="extend")
def diff_center_to_left(a):
    return a - np.roll(a, -1)
```
where you have identically-named functions distinguished only by their keyword arguments to the decorator. 

However I don't think this is possible, as despite the members of this module being different instances of the `GridUFunc` class, python refers to both these functions via their name (`diff_center_to_left`), and so the second one ends up overwriting the first, and I have to give each function different names even if they are distinguished by their decorator arguments, i.e.

```python
@as_grid_ufunc(signature="(X:center)->(X:left)", boundary="fill")
def diff_center_to_left_fill(a):
    return a - np.roll(a, -1)


@as_grid_ufunc(signature="(X:center)->(X:left)", boundary="extend")
def diff_center_to_left_extend(a):
    return a - np.roll(a, -1)
```
rabernat: > so the second one ends up overwriting the first, and I have to give each function different names even if they are distinguished by their decorator arguments

This might be suggesting our abstraction is wrong.

First let's talk carefully about what we mean by _boundary conditions_. I distinguish between two different types of boundary:
- The **global array boundary**: obviously our array has a finite size (Ny, Nx). We need to decide what to do at the boundaries of that array. _This is what xgcm's current "boundary" kwarg is about._ Those boundaries may or may not have any physical meaning. For example, on a lat lon grid, the X-axis array boundary has zero physical meaning--it just happens to be where we sliced the sphere. This is ultimately a topological issue. Our implementation of [face connections](https://xgcm.readthedocs.io/en/latest/grid_topology.html) has some very fancy treatment for the global array boundaries, e.g. for cubed-sphere grids.
- The **physical boundary**: ocean models (and probably other FV models) have actual boundaries, where there is no fluid, i.e. land. We often want to do something distinct on these boundaries, like apply a no-flux boundary condition. These may or may not correspond with the global array boundary. This is discussed in #324. Also closely related to #293. Currently _not supported in xgcm_.

It's very useful to think about what real GCMs do. GCMs parallelize the grids over many MPI processes, each operating on a "tile." Each title only knows about its local data, so the concept of _global array boundaries_ is not relevant. ([Example](https://mitgcm.readthedocs.io/en/latest/software_arch/software_arch.html#specifying-a-domain-decomposition))
Instead, each tile has an "overlap" or "halo" region that is duplicated with its neighbors, plus a land mask.

![grid](https://mitgcm.readthedocs.io/en/latest/_images/size_h.png)

**I would like to suggest that `as_grid_ufunc` should be unaware of _global array boundaries_ but possibly aware of _physical boundaries_.** Global array boundaries can be handled separately. For example, when we call `Grid.diff`, we might imagine taking the following steps.

1. Build halo region around the array
2. Apply `diff_center_to_left` array_ufunc
3. Remove halo region

TomNicholas: So there are a two issues here - we're not quite talking about the same thing:

1) **How should grid ufuncs with the same signature be distinguished?** This is what my comment was about, rather than specifically about the `boundary` kwarg, because we will have this naming issue whether or not it's `metric` or `fill_value` or `boundary` that distinguishes the grid ufuncs. (It is possible though that the correct abstraction doesn't have any kwargs other than `signature` in the `@as_grid_ufunc` decorator, and everything else is either always attached to `grid` or handled before/after the ufunc...)

2) **Which keyword arguments do grid ufuncs need to be aware of?** This is what your comment discusses @rabernat . My understanding is that you're basically saying that we might need something like `@as_grid_ufunc(signature, physical_boundary={'X': 'dirichlet'})`, but we shouldn't need `@as_grid_ufunc(signature, boundary='fill')`, because that should instead be handled by code that wraps the application of the grid ufunc? In that model is the current `boundary` kwarg any different to the `boundary` kwarg that `dask.array.map_blocks` accepts? 



I've actually realised that I have a more immediate design problem to solve first though: **how do I check that two signatures are functionally equivalent?** i.e. `"(X:center)->(X:right)"` and `"(Y:center)->(Y:right)"` are equivalent, in the same way that for normal ufuncs `(i)->(i)` and `(j)->(j)` are equivalent.
TomNicholas: > For example, when we call Grid.diff, we might imagine taking the following steps.
> 
> - Build halo region around the array
> - Apply diff_center_to_left array_ufunc
> - Remove halo region

Isn't this literally what `dask.array.map_overlap` does? If so then can we just use that? 

If we do then the `boundary` question mostly becomes whether the call to `map_overlap` should be part of the grid ufunc or not. There are a few options I can see:

1) **Call `dask.array.map_overlap` inside each grid ufunc.** This is currently supported, and there is actually a test for this case in #362. The main disadvantage is potentially duplicated code within many (if not all) of the pre-defined ufuncs in `gridops.py`. It would also mean that `boundary` needs to be an argument to each grid ufunc, which I think conflicts with your recommendation above @rabernat .

2) **Automatically wrap each grid ufunc with `dask.array.map_overlap`.** If we know all pre-defined grid ufuncs should have the same behaviour and options with respect to internal boundaries, then we can take the supplied numpy ufunc, wrap it in `dask.array.map_overlap`, and then feed the resulting function into `xarray.apply_ufunc`. We could do this within the `apply_as_grid_ufunc` call, so it would happen within the `as_grid_ufunc` decorator, and not be mentioned at all in the grid ufunc definition. The disadvantage is that this approach relies on the assumption that the underlying numpy implementation of the grid ufunc is independent of the choice of `boundary={'fill', 'extend', 'extrapolate'}`. It might also complicate the implementation of different physical boundaries.

3) **If xarray implements `map_overlap` then that might end up being a better way?** See https://github.com/pydata/xarray/issues/3147. I don't think this is the correct solution for us though. 

One problem with this idea of using `dask.array.map_overlap` though is that while xgcm's `'fill'` is essentially dask's constant value option, and xgcm's `extend` is essentially dask's `'nearest'` options, it doesn't look like `dask.array.map_overlap` has any equivalent of xgcm's `'extrapolate'` option.

rabernat: > Isn't this literally what `dask.array.map_overlap` does? If so then can we just use that?

Not quite.
- [map_overlap](https://docs.dask.org/en/latest/array-overlap.html#dask.array.map_overlap) as is is not sufficiently flexible to handle the really global array boundary conditions. It supports the following boundary conditions:
  >  ‘reflect’, ‘periodic’, ‘nearest’, ‘none’, or any constant value like 0 or np.nan. 
  
  It can't deal with cubed sphere, tripolar grid, etc. So we may have to vendor this function and augment it with special options.
- Xgcm should work without Dask! It has always supported an eager evaluation mode.

However, we can overlook these in the short term if we want to just move forward with an implementation.
rabernat: > **how do I check that two signatures are functionally equivalent?** i.e. `"(X:center)->(X:right)"` and `"(Y:center)->(Y:right)"` are equivalent, in the same way that for normal ufuncs `(i)->(i)` and `(j)->(j)` are equivalent.

It seems like a regex could handle this, no?
TomNicholas: > It can't deal with cubed sphere, tripolar grid, etc. 

Hmm okay.

> we may have to vendor this function and augment it with special options.

I haven't heard the word "vendor" used in this context before - what exactly do you mean? Make our own version of `map_overlap`?

I think if we had our own function to perform a similar job (e.g. `map_over_faces` or something) then we would still have to choose between my (1) and (2) above.

> However, we can overlook these in the short term if we want to just move forward with an implementation.

This might be a good idea - use `dask.array.map_overlap` to get the basic test cases with just `'fill'` and `'extend'` working, and then re-evaluate.

> It seems like a regex could handle this, no?

Probably, but so far my solution has been to call `_parse_grid_ufunc_signature` (so re-using my previous regex in a way) and then replace dummy variables. There is probably a neater regex way though.

https://github.com/xgcm/xgcm/pull/368/files#diff-3af42ae01706abd10c8c0a6185e2598b72bacfe3ef33eb162ae8da1a9d61a257R307-R335
dcherian: I think you want to build `Grid.pad` first so then making the halo is a call to `Grid.pad` (accounts for tripole, cubed-sphere, periodic etc.); then pass to `map_overlap(..., boundary=None)` so map_overlap only takes care of communicating between blocks; then remove the halo region.
jbusecke: I like the suggestion by @dcherian above. This would take care of the 'data_boundaries' (which are not physical) and then 'internal_boundaries' e.g. actual land in the data should be handled internally (at the numpy level) in the ufunc. The latter would be a new feature (to be implemented after this refactor?) and we could discuss if we want to implement that with an additional keyword or with another subset of ufuncs.
TomNicholas: >I think you want to build Grid.pad first so then making the halo is a call to Grid.pad (accounts for tripole, cubed-sphere, periodic etc.); then pass to map_overlap(..., boundary=None) so map_overlap only takes care of communicating between blocks; then remove the halo region.

Thanks @dcherian , that was a really helpful suggestion.

So in the latest commit I added a prototype of a `grid.pad` method, which is called within `apply_as_grid_ufunc`, and pads the given dataarrays by a number of elements specified by a new `boundary_width` argument to the `@as_grid_ufunc` decorator. (Thanks to @jbusecke for a very helpful pair programming session on that on Friday.)

We also identified the simplest existing failing test case that needs to pass in order to prove that this approach works: `xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_left]` - I was trying to debug via some other tests before.

The padding seems work for basic cases (though I still need to write tests that will isolate that method), but the test above still doesn't pass, instead failing with an error I told it to raise when the ufunc output hasn't been trimmed:

```python
ValueError: conflicting sizes for dimension 'XG': length 101 on <this-array> and length 100 on {'XG': 'XG'} - does your grid ufunc correctly trim off the same number of elements which were added by padding using boundary_width={'X': (1, 0)}?
```

I think the problem is that @jbusecke the `boundary_width` kwarg we added to `diff_center_to_left` in `gridops.py` isn't correct: when `diff` is called with `boundary='fill'` by the test, this ends up padding, then applying the ufunc, but the ufunc does not automatically trim the new elements back off the data, so the result is the incorrect length.

```python
@as_grid_ufunc(signature="(X:center)->(X:left)", boundary_width={"X": (1, 0)})
def diff_center_to_left(a):
    return a - np.roll(a, -1)
```

If this particular ufunc isn't supposed to have any elements padded when applying boundary conditions, then I need another example ufunc which does. In general it would be helpful if I were actually sure what the ufuncs in `gridops.py` for the various `diff_*` variations were actually supposed to be!
rabernat: > but the ufunc does not automatically trim the new elements back off the data, so the result is the incorrect length.

So the implementation needs to change
```python
@as_grid_ufunc(signature="(X:center)->(X:left)", boundary_width={"X": (1, 0)})
def diff_center_to_left(a):
    return a[..., 1:] - a[..., :-1]
```
TomNicholas: Thanks @rabernat - I expect that will work.

> So the implementation needs to change

There is possibly another option, where I add a `trim_after` argument to`@as_grid_ufunc`, that defaults to false. That would allow people to define ufuncs which don't trim the padded elements, and trim for them instead.
jbusecke: > There is possibly another option, where I add a trim_after argument to@as_grid_ufunc, that defaults to false. That would allow people to define ufuncs which don't trim the padded elements, and trim for them instead.

I am actually not against that. It would enable both types of workflow and might make these 'basic' ufuncs cleaner to read?
dcherian: > a[..., 1:] - a[..., :-1]

I think this is actually better from a memory perspective two, since we create two slices instead of one extra copy with `np.roll`

:+1: for a `trim_after` argument, though the analogous argument in `map_overlap` is just `trim`
rabernat: My general point is that, if padding is handled separately, our grid ufuncs should **never** roll arrays or do anything else that requires global array boundary communication. They should always assume they are operating on a "local tile" of data.

I agree that `trim` may be needed, depending on the particular implementation. Perhaps this can be optional, for the grid ufunc writer to choose?
TomNicholas: > My general point is that, if padding is handled separately, our grid ufuncs should never roll arrays or do anything else that requires global array boundary communication. 

I think that now all our basic ufuncs are just going to look like slices, e.g. `a[..., 1:] - a[..., :-1]`. The global array communication, if needed, is being handled (in a dask-aware way) by `DataArray.pad` now.

> Perhaps this can be optional, for the grid ufunc writer to choose?

Yes I meant as an additional option for someone writing a custom ufunc.
TomNicholas: I realised there was a major omission in #362 - `apply_as_grid_ufunc` had no `axis` argument! #371 fixes that, so this PR now will need to build on top of #371, which needs to be merged before this one.
TomNicholas: For some reason now the trimming works in 1D but has an off-by-one error in 2D :confused: 

```
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_outer] PASSED                                                                                                                     [ 10%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_inner] PASSED                                                                                                                     [ 20%]
xgcm/test/test_grid.py::test_grid_ops[periodic_1d_left] PASSED                                                                                                                         [ 30%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_left] PASSED                                                                                                                      [ 40%]
xgcm/test/test_grid.py::test_grid_ops[periodic_1d_right] PASSED                                                                                                                        [ 50%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_right] PASSED                                                                                                                     [ 60%]
xgcm/test/test_grid.py::test_grid_ops[periodic_2d_left] FAILED                                                                                                                         [ 70%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_2d_left] FAILED                                                                                                                      [ 80%]
xgcm/test/test_grid.py::test_grid_ops[xperiodic_2d_left] FAILED                                                                                                                        [ 90%]
xgcm/test/test_grid.py::test_grid_ops[yperiodic_2d_left] FAILED                                                                                                                        [100%]


FAILED xgcm/test/test_grid.py::test_grid_ops[periodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} - ...
FAILED xgcm/test/test_grid.py::test_grid_ops[nonperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'}...
FAILED xgcm/test/test_grid.py::test_grid_ops[xperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} -...
FAILED xgcm/test/test_grid.py::test_grid_ops[yperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} -...
```


TomNicholas: For some reason now the trimming works in 1D but has an off-by-one error in 2D :confused: 

```
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_outer] PASSED                                                                                                                     [ 10%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_inner] PASSED                                                                                                                     [ 20%]
xgcm/test/test_grid.py::test_grid_ops[periodic_1d_left] PASSED                                                                                                                         [ 30%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_left] PASSED                                                                                                                      [ 40%]
xgcm/test/test_grid.py::test_grid_ops[periodic_1d_right] PASSED                                                                                                                        [ 50%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_right] PASSED                                                                                                                     [ 60%]
xgcm/test/test_grid.py::test_grid_ops[periodic_2d_left] FAILED                                                                                                                         [ 70%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_2d_left] FAILED                                                                                                                      [ 80%]
xgcm/test/test_grid.py::test_grid_ops[xperiodic_2d_left] FAILED                                                                                                                        [ 90%]
xgcm/test/test_grid.py::test_grid_ops[yperiodic_2d_left] FAILED                                                                                                                        [100%]


FAILED xgcm/test/test_grid.py::test_grid_ops[periodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} - ...
FAILED xgcm/test/test_grid.py::test_grid_ops[nonperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'}...
FAILED xgcm/test/test_grid.py::test_grid_ops[xperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} -...
FAILED xgcm/test/test_grid.py::test_grid_ops[yperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} -...
```


TomNicholas: For some reason now the trimming works in 1D but has an off-by-one error in 2D :confused: 

```
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_outer] PASSED                                                                                                                     [ 10%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_inner] PASSED                                                                                                                     [ 20%]
xgcm/test/test_grid.py::test_grid_ops[periodic_1d_left] PASSED                                                                                                                         [ 30%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_left] PASSED                                                                                                                      [ 40%]
xgcm/test/test_grid.py::test_grid_ops[periodic_1d_right] PASSED                                                                                                                        [ 50%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_right] PASSED                                                                                                                     [ 60%]
xgcm/test/test_grid.py::test_grid_ops[periodic_2d_left] FAILED                                                                                                                         [ 70%]
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_2d_left] FAILED                                                                                                                      [ 80%]
xgcm/test/test_grid.py::test_grid_ops[xperiodic_2d_left] FAILED                                                                                                                        [ 90%]
xgcm/test/test_grid.py::test_grid_ops[yperiodic_2d_left] FAILED                                                                                                                        [100%]


FAILED xgcm/test/test_grid.py::test_grid_ops[periodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} - ...
FAILED xgcm/test/test_grid.py::test_grid_ops[nonperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'}...
FAILED xgcm/test/test_grid.py::test_grid_ops[xperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} -...
FAILED xgcm/test/test_grid.py::test_grid_ops[yperiodic_2d_left] - ValueError: conflicting sizes for dimension 'YG': length 199 on <this-array> and length 200 on {'XC': 'XC', 'YG': 'YG'} -...
```


TomNicholas: This now passes for 2d data too (despite not knowing what is wrong with https://github.com/xgcm/xgcm/pull/371#issuecomment-926220770)...

```
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_outer] PASSED
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_inner] PASSED
xgcm/test/test_grid.py::test_grid_ops[periodic_1d_left] PASSED
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_left] PASSED
xgcm/test/test_grid.py::test_grid_ops[periodic_1d_right] PASSED
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_1d_right] PASSED
xgcm/test/test_grid.py::test_grid_ops[periodic_2d_left] PASSED
xgcm/test/test_grid.py::test_grid_ops[nonperiodic_2d_left] PASSED
xgcm/test/test_grid.py::test_grid_ops[xperiodic_2d_left] PASSED
xgcm/test/test_grid.py::test_grid_ops[yperiodic_2d_left] PASSED
```

but only up to a transposition. Specifically, I had to add 

```python
if axis_name == "Y":
    assert da_diff.T.equals(da_diff_ax)
else:
    assert da_diff.equals(da_diff_ax)
```
to the test to get all the `[*periodic_2d_left]` tests to pass. This is because `xarray.apply_ufunc` always transposes so that the axes being acted over are the last axes in the data.

Now I'm not sure what to do here because I could:
1) Write some code that attempts to transpose the result back - unfortunately I think this is not a trivial problem, given that a grid ufunc can create and remove arbitrary axes. It might not always be clear which dim has been shifted to create which new dim...
2) Change the test - I've not changed any tests at all yet, and don't want to. But it is interesting that this would pass if xarray's assert functions didn't care about dimension order, as you suggested in https://github.com/pydata/xarray/issues/5733 @jbusecke .
3) Leave it for now - this particular test only compares `grid.diff` to `axis.diff`, and the latter is still using the old code path. I expect that when everything has been refactored internally then this particular tests will pass (though any other tests elsewhere that apply over multidimensional data will presumably fail with the same problem).


TomNicholas: Actually I think I've got it - I don't need to deal with the general problem of keeping dimension order, only the (much easier) case of one-to-one functions like `diff`, `interp` etc. That way I can transpose back for those functions defined on `grid`, but the fact that `apply_as_grid_ufunc` will stll rearrange dims shouldn't really matter.
TomNicholas: [I tested this with dask in this notebook](https://gist.github.com/TomNicholas/47be490cbeae03686a24de87f505f364), but found that confusingly applying these two functions behave differently with respect to dask:

```python
def diff_center_to_left(a):
    return a[..., 1:] - a[..., :-1]
```

```python
def manual_diff(array):
    return array - np.roll(array, 1, axis=-1)
```
The former complains about chunking along a core dim
```
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.
```
whilst the latter doesn't. I don't really understand why that is. Currently this PR uses the former definition in the `gridops.py` file, which is why you see the same behaviour when calliing `grid.diff` in the notebook.

Using bare `apply_ufunc` behaves the same way as `apply_as_grid_ufunc` so I don't think it's my code causing any difference in behaviour.

Additionally I tried the former function but with `dask='allowed'` instead of `'parallelized'` - that runs but gives a surprisingly complex task graph :confused: 

So it appears that a grid ufunc defined with `np.roll` is most dask-streamlined, but I thought we were expecting the opposite?

Note: This is all before I attempt to call `map_overlap` anywhere - currently all that happens is a `dask.pad` (if `boundary_width` is nonzero) and then an `xarray.apply_ufunc`.
rabernat: > this would be a good time for a review / input before I try to move any further forward I think

I just put this on my calendar in my next free spot, in the morning of Oct. 4.
dcherian: Is the graph just as complex if you called the second one with `dask="allowed"`?  It should  be...

The graph for `dask="parallelized"` is neat because it's just applying it blockwise; but that  can only work if there's a single chunk along that axis.

You're getting a chunk error for the second one because pad inserts an extra chunk IIRC with the padded  values.

Anecdotally, I know dask slicing can get weirdly complex, but don't really understand why. Isn't that why we're doing this refactor in the first place?  

I think should try `dask="allowed"` with `map_overlap` and the first function to get a real sense of the before/after improvement (if any).
TomNicholas: > Is the graph just as complex if you called the second one with dask="allowed"? It should be...

You actually get a 3rd different graph, of moderate complexity :confused: I've updated the gist with it.

> Anecdotally, I know dask slicing can get weirdly complex, but don't really understand why. Isn't that why we're doing this refactor in the first place?



> I think should try dask="allowed" with map_overlap and the first function to get a real sense of the before/after improvement (if any).

I can do that - I can define the grid ufunc to call `map_overlap` within it, like this:

```python
from dask.array import map_overlap

def diff_center_to_left(a):
    return a[..., 1:] - a[..., :-1]

def diff_overlap(a):
    return map_overlap(
        diff_center_to_left,
        a,
        depth={'X': (1, 0)},
        boundary=None,
        trim=True,
        meta=np.array(())
    )

bar_manual_diff = grid_test.apply_as_grid_ufunc(
    diff_overlap, 
    ds_test.bar, 
    axis='X', 
    signature="(X:center)->(X:left)", 
    boundary_width={"X": (0, 0)},
    boundary="fill",
    dask="allowed",
)
bar_manual_diff
```
which gives a nice streamlined graph. (I am a bit confused by why I needed to put `trim=True` and `boundary_width={'X': (0,0)}` to make this compute an array of the correct length - I thought I would need to put `trim=True` and `boundary_width={'X': (1,0)}`...)

Still, I think the before / after graph comparison (so using this `map_overlap` code above) looks like this

(before)
![image](https://user-images.githubusercontent.com/35968931/135655416-f517d874-f58d-4722-9e29-c3819a0ea5bb.png)

(after)
![image](https://user-images.githubusercontent.com/35968931/135655384-8a82cb23-2c09-4b3c-a6ec-8512ecd163ac.png)

rabernat: Merging this for now. We can keep iterating on the dask optimizations in future PRs.

Patch:
diff --git a/xgcm/grid.py b/xgcm/grid.py
index 4b36943..5bf68b7 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -1,4 +1,5 @@
 import functools
+import inspect
 import itertools
 import operator
 import warnings
@@ -8,8 +9,9 @@ import docrep
 import numpy as np
 import xarray as xr
 
-from . import comodo
+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 .metrics import iterate_axis_combinations
 
 try:
@@ -1041,6 +1043,13 @@ class Axis:
         return da.get_axis_num(coord_name)
 
 
+_XGCM_BOUNDARY_KWARG_TO_XARRAY_PAD_KWARG = {
+    "periodic": "wrap",
+    "fill": "constant",
+    "extend": "edge",
+}
+
+
 class Grid:
     """
     An object with multiple :class:`xgcm.Axis` objects representing different
@@ -1504,6 +1513,7 @@ class Grid:
             The direction in which to shift the array (can be ['center','left','right','inner','outer']).
             If not specified, default will be used.
             Optionally a dict with seperate values for each axis can be passed (see example)
+            Optionally a dict with seperate values for each axis can be passed (see example)
         boundary : None or str or dict, optional
             A flag indicating how to handle boundaries:
 
@@ -1558,6 +1568,241 @@ class Grid:
 
         return out
 
+    def pad(self, *arrays, boundary_width=None, boundary=None, fill_value=None):
+        """
+        Pads the boundary of given arrays along given Axes, according to information in Axes.boundary.
+
+        Parameters
+        ----------
+        arrays : Sequence[xarray.DataArray]
+            Arrays to pad according to boundary and boundary_width.
+        boundary_width : Dict[str: Tuple[int, int]
+            The widths of the boundaries at the edge of each array.
+            Supplied in a mapping of the form {axis_name: (lower_width, upper_width)}.
+        boundary : {None, 'fill', 'extend', 'extrapolate', dict}, optional
+            A flag indicating how to handle boundaries:
+            * None: Do not apply any boundary conditions. Raise an error if
+              boundary conditions are required for the operation.
+            * 'fill':  Set values outside the array boundary to fill_value
+              (i.e. a Dirichlet boundary condition.)
+            * 'extend': Set values outside the array to the nearest array
+              value. (i.e. a limited form of Neumann boundary condition.)
+            * 'extrapolate': Set values by extrapolating linearly from the two
+              points nearest to the edge
+            Optionally a dict mapping axis name to separate values for each axis
+            can be passed.
+        fill_value : {float, dict}, optional
+            The value to use in boundary conditions with `boundary='fill'`.
+            Optionally a dict mapping axis name to separate values for each axis
+            can be passed. Default is 0.
+        """
+        # TODO accept a general padding function like numpy.pad does as an argument to boundary
+
+        if not boundary_width:
+            raise ValueError("Must provide the widths of the boundaries")
+
+        if boundary and isinstance(boundary, str):
+            boundary = {ax_name: boundary for ax_name in self.axes.keys()}
+        if fill_value is None:
+            fill_value = 0.0
+        if isinstance(fill_value, float):
+            fill_value = {ax_name: fill_value for ax_name in self.axes.keys()}
+
+        padded = []
+        for da in arrays:
+            new_da = da
+            for ax, widths in boundary_width.items():
+                axis = self.axes[ax]
+                _, dim = axis._get_position_name(da)
+
+                # Use default boundary for axis unless overridden
+                if boundary:
+                    ax_boundary = boundary[ax]
+                else:
+                    ax_boundary = axis.boundary
+
+                if ax_boundary == "extrapolate":
+                    # TODO implement extrapolation
+                    raise NotImplementedError
+                elif ax_boundary is None:
+                    # TODO this is necessary, but also seems inconsistent with the docstring, which says that None = "no boundary condition"
+                    ax_boundary = "periodic"
+
+                # TODO avoid repeatedly calling xarray pad
+                try:
+                    mode = _XGCM_BOUNDARY_KWARG_TO_XARRAY_PAD_KWARG[ax_boundary]
+                except KeyError:
+                    raise ValueError(
+                        f"{ax_boundary} is not a supported type of boundary"
+                    )
+
+                if mode == "constant":
+                    new_da = new_da.pad(
+                        {dim: widths}, mode, constant_values=fill_value[ax]
+                    )
+                else:
+                    new_da = new_da.pad({dim: widths}, mode)
+
+            padded.append(new_da)
+
+        return padded
+
+    def _grid_ufunc_dispatch(self, funcname, da, axis, to=None, **kwargs):
+        """
+        Calls appropriate grid ufuncs on data, along the specified axes.
+
+        Parameters
+        ----------
+        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
+            given order.
+        to : str or dict, optional
+            The direction in which to shift the array (can be ['center','left','right','inner','outer']).
+            Can be passed as a single str to use for all axis, or as a dict with separate values for each axis.
+            If not specified, the `default_shifts` stored in each Axis object will be used for that axis.
+        """
+
+        if isinstance(axis, str):
+            axis = [axis]
+        if to is None:
+            to = {ax: None for ax in axis}
+        elif isinstance(to, str):
+            to = {ax: to for ax in axis}
+
+        signature = self._create_grid_ufunc_signature(da, axis=axis, to=to)
+
+        grid_ufunc, remaining_kwargs = _select_grid_ufunc(
+            funcname, signature, module=gridops, **kwargs
+        )
+
+        result = grid_ufunc(self, da, axis=[axis], **remaining_kwargs)
+
+        return self._transpose_to_keep_same_dim_order(da, result, axis)
+
+    def _create_grid_ufunc_signature(self, da, axis, to):
+        """
+        Create a signature to pass to apply_grid_ufunc from data, list of input axes, and list of target axis positions.
+        """
+
+        from_ax_positions = []
+        to_ax_positions = []
+        for ax_name in axis:
+            ax = self.axes[ax_name]
+
+            from_pos, dim = ax._get_position_name(da)
+            from_ax_positions.append(from_pos)
+
+            to_pos = to[ax_name]
+            if to_pos is None:
+                to_pos = ax._default_shifts[from_pos]
+            to_ax_positions.append(to_pos)
+
+        input_arg_signature = ",".join(
+            [f"{name}:{from_pos}" for name, from_pos in zip(axis, from_ax_positions)]
+        )
+        output_arg_signature = ",".join(
+            [f"{name}:{to_pos}" for name, to_pos in zip(axis, to_ax_positions)]
+        )
+
+        return f"({input_arg_signature})->({output_arg_signature})"
+
+    def _transpose_to_keep_same_dim_order(self, da, result, axis):
+        """Reorder DataArray dimensions to match the original input."""
+
+        initial_dims = da.dims
+
+        shifted_dims = {}
+        for ax_name in axis:
+            ax = self.axes[ax_name]
+
+            _, old_dim = ax._get_position_name(da)
+            _, new_dim = ax._get_position_name(result)
+            shifted_dims[old_dim] = new_dim
+
+        output_dims_but_in_original_order = [
+            shifted_dims[dim] if dim in shifted_dims else dim for dim in initial_dims
+        ]
+
+        return result.transpose(*output_dims_but_in_original_order)
+
+    def apply_as_grid_ufunc(
+        self,
+        func,
+        *args,
+        signature="",
+        boundary_width=None,
+        boundary=None,
+        fill_value=None,
+        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`.
+        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)`.
+        boundary_width : Dict[str: Tuple[int, int]
+            The widths of the boundaries at the edge of each array.
+            Supplied in a mapping of the form {axis_name: (lower_width, upper_width)}.
+        boundary : {None, 'fill', 'extend', 'extrapolate', dict}, optional
+            A flag indicating how to handle boundaries:
+            * None: Do not apply any boundary conditions. Raise an error if
+              boundary conditions are required for the operation.
+            * 'fill':  Set values outside the array boundary to fill_value
+              (i.e. a Dirichlet boundary condition.)
+            * 'extend': Set values outside the array to the nearest array
+              value. (i.e. a limited form of Neumann boundary condition.)
+            * 'extrapolate': Set values by extrapolating linearly from the two
+              points nearest to the edge
+            Optionally a dict mapping axis name to separate values for each axis
+            can be passed.
+        fill_value : {float, dict}, optional
+            The value to use in boundary conditions with `boundary='fill'`.
+            Optionally a dict mapping axis name to separate values for each axis
+            can be passed. Default is 0.
+        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
+        --------
+        apply_as_grid_ufunc
+        as_grid_ufunc
+        """
+        return apply_as_grid_ufunc(
+            func,
+            *args,
+            grid=self,
+            signature=signature,
+            boundary_width=boundary_width,
+            boundary=boundary,
+            fill_value=fill_value,
+            dask=dask,
+            **kwargs,
+        )
+
     @docstrings.dedent
     def interp(self, da, axis, **kwargs):
         """
@@ -1614,7 +1859,7 @@ class Grid:
 
         >>> grid.diff(da, ["X", "Y"], fill_value={"X": 0, "Y": 100})
         """
-        return self._grid_func("diff", da, axis, **kwargs)
+        return self._grid_ufunc_dispatch("diff", da, axis, **kwargs)
 
     @docstrings.dedent
     def min(self, da, axis, **kwargs):
@@ -1978,6 +2223,56 @@ class Grid:
         return self._apply_vector_function(Axis.diff, vector, **kwargs)
 
 
+def _select_grid_ufunc(funcname, signature, module, **kwargs):
+    # TODO to select via other kwargs (e.g. boundary) the signature of this function needs to be generalised
+
+    def is_grid_ufunc(obj):
+        return isinstance(obj, GridUFunc)
+
+    # This avoids defining a list of functions in gridops.py
+    all_predefined_ufuncs = inspect.getmembers(module, is_grid_ufunc)
+
+    name_matching_ufuncs = [
+        f for name, f in all_predefined_ufuncs if name.startswith(funcname)
+    ]
+    if len(name_matching_ufuncs) == 0:
+        raise NotImplementedError(
+            f"Could not find any pre-defined {funcname} grid ufuncs"
+        )
+
+    signature_matching_ufuncs = [
+        f
+        for f in name_matching_ufuncs
+        if _signatures_equivalent(f.signature, signature)
+    ]
+    if len(signature_matching_ufuncs) == 0:
+        raise NotImplementedError(
+            f"Could not find any pre-defined {funcname} grid ufuncs with signature {signature}"
+        )
+
+    matching_ufuncs = signature_matching_ufuncs
+
+    # TODO select via any other kwargs (such as boundary? metrics?) once implemented
+    all_kwargs = kwargs.copy()
+    # boundary = kwargs.pop("boundary", None)
+    # if boundary:
+    #    matching_ufuncs = [uf for uf in matching_ufuncs if uf.boundary == boundary]
+
+    if len(matching_ufuncs) > 1:
+        # TODO include kwargs used to match in this error message
+        raise ValueError(
+            f"Function {funcname} with signature='{signature}' and kwargs={all_kwargs} is an ambiguous selection"
+        )
+    elif len(matching_ufuncs) == 0:
+        raise NotImplementedError(
+            f"Could not find any pre-defined {funcname} grid ufuncs with signature='{signature}' and kwargs"
+            f"={all_kwargs}"
+        )
+    else:
+        # Exactly 1 matching function
+        return matching_ufuncs[0], kwargs
+
+
 def raw_interp_function(data_left, data_right):
     # linear, centered interpolation
     # TODO: generalize to higher order interpolation
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
index 68cf25c..117e16e 100644
--- a/xgcm/grid_ufunc.py
+++ b/xgcm/grid_ufunc.py
@@ -1,4 +1,5 @@
 import re
+import string
 
 import xarray as xr
 
@@ -79,14 +80,15 @@ class GridUFunc:
     ----------
     ufunc : callable
         Function to call like `func(*args, **kwargs)` on numpy-like unlabeled
-        arrays (`.data`).
-
-        Passed directly on to `xarray.apply_ufunc`.
+        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)`.
+    boundary_width : Dict[str: Tuple[int, int], optional
+        The widths of the boundaries at the edge of each array.
+        Supplied in a mapping of the form {axis_name: (lower_width, upper_width)}.
     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`.
@@ -97,8 +99,9 @@ class GridUFunc:
         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`.
+        and another additional positional argument `axis`, of type Sequence[Tuple[str]],
+        so that `func`'s new signature is `func(grid, *args, axis, **kwargs)`.
+        The grid and axis arguments are passed on to `apply_grid_ufunc`.
 
     See Also
     --------
@@ -110,22 +113,30 @@ class GridUFunc:
     def __init__(self, ufunc, **kwargs):
         self.ufunc = ufunc
         self.signature = kwargs.pop("signature", "")
-        self.dask = kwargs.pop("dask", "")
+        self.boundary_width = kwargs.pop("boundary_width", None)
+        self.dask = kwargs.pop("dask", "forbidden")
         if kwargs:
             raise TypeError("Unsupported keyword argument(s) provided")
 
-    def __call__(self, grid=None, *args, **kwargs):
+    def __repr__(self):
+        return f"GridUFunc(ufunc={self.ufunc}, signature='{self.signature}', boundary_width='{self.boundary_width}', dask='{self.dask})'"
+
+    def __call__(self, grid, *args, axis, boundary=None, **kwargs):
+        dask = kwargs.pop("dask", self.dask)
         return apply_as_grid_ufunc(
             self.ufunc,
             *args,
+            axis=axis,
             grid=grid,
             signature=self.signature,
-            dask=self.dask,
+            boundary_width=self.boundary_width,
+            boundary=boundary,
+            dask=dask,
             **kwargs,
         )
 
 
-def as_grid_ufunc(signature="", **kwargs):
+def as_grid_ufunc(signature="", boundary_width=None, **kwargs):
     """
     Decorator which turns a numpy ufunc into a "grid-aware ufunc".
 
@@ -133,14 +144,15 @@ def as_grid_ufunc(signature="", **kwargs):
     ----------
     ufunc : callable
         Function to call like `func(*args, **kwargs)` on numpy-like unlabeled
-        arrays (`.data`).
-
-        Passed directly on to `xarray.apply_ufunc`.
+        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)`.
+    boundary_width : Dict[str: Tuple[int, int], optional
+        The widths of the boundaries at the edge of each array.
+        Supplied in a mapping of the form {axis_name: (lower_width, upper_width)}.
     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`.
@@ -151,8 +163,9 @@ def as_grid_ufunc(signature="", **kwargs):
         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`.
+        and another additional positional argument `axis`, of type Sequence[Tuple[str]],
+        so that `func`'s new signature is `func(grid, *args, axis, boundary=None, **kwargs)`.
+        The grid and axis arguments are passed on to `apply_grid_ufunc`.
 
     See Also
     --------
@@ -166,13 +179,24 @@ def as_grid_ufunc(signature="", **kwargs):
         raise TypeError("Unsupported keyword argument(s) provided")
 
     def _as_grid_ufunc(ufunc):
-        return GridUFunc(ufunc, signature=signature, **kwargs)
+        return GridUFunc(
+            ufunc, signature=signature, boundary_width=boundary_width, **kwargs
+        )
 
     return _as_grid_ufunc
 
 
 def apply_as_grid_ufunc(
-    func, *args, grid=None, signature="", dask="forbidden", **kwargs
+    func,
+    *args,
+    axis,
+    grid=None,
+    signature="",
+    boundary_width=None,
+    boundary=None,
+    fill_value=None,
+    dask="forbidden",
+    **kwargs,
 ):
     """
     Apply a function to the given arguments in a grid-aware manner.
@@ -188,14 +212,45 @@ def apply_as_grid_ufunc(
         arrays (`.data`).
 
         Passed directly on to `xarray.apply_ufunc`.
+    args : xarray.DataArray
+        One or more xarray objects to apply the function to.
+    axis : Sequence[Tuple[str]]
+        Names of xgcm.Axes on which to act, for each array in args. Multiple axes can be passed as a sequence (e.g. ``['X', 'Y']``).
+        Function will be executed over all Axes simultaneously, and each Axis must be present in the Grid.
     grid : xgcm.Grid
-        The xgcm Grid object which contains the various xgcm.Axis described by
-        `signature`.
+        The xgcm Grid object which contains the various xgcm.Axis named in the axis kwarg, with positions matching the
+         first half of the `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)`.
+        Grid universal function signature. Specifies the relationship between xgcm.Axis positions before and after the
+        operation for each input and output variable, e.g.,
+
+        ``signature="(X:center)->(X:left)"`` for ``func=diff_center_to_left(a)`.
+
+        The axis names in the signature are dummy variables, so do not have to present in the Grid. Instead, these dummy
+        variables will be identified with the actual named Axes in the `axis` kwarg in order of appearance. For
+        instance, ``"(Z:center)->(Z:left)"`` is equivalent to ``"(X:center)->(X:left)"`` - both choices of `signature`
+        require only that there is exactly one xgcm.Axis name in `axis` which exists in Grid and starts on position
+        `center`.
+    boundary_width : Dict[str: Tuple[int, int]
+        The widths of the boundaries at the edge of each array.
+        Supplied in a mapping of the form {dummy_axis_name: (lower_width, upper_width)}.
+        The axis names here are again dummy variables, each of which must be present in the signature.
+    boundary : {None, 'fill', 'extend', 'extrapolate', dict}, optional
+        A flag indicating how to handle boundaries:
+        * None: Do not apply any boundary conditions. Raise an error if
+          boundary conditions are required for the operation.
+        * 'fill':  Set values outside the array boundary to fill_value
+          (i.e. a Dirichlet boundary condition.)
+        * 'extend': Set values outside the array to the nearest array
+          value. (i.e. a limited form of Neumann boundary condition.)
+        * 'extrapolate': Set values by extrapolating linearly from the two
+          points nearest to the edge
+        Optionally a dict mapping axis name to separate values for each axis
+        can be passed.
+    fill_value : {float, dict}, optional
+        The value to use in boundary conditions with `boundary='fill'`.
+        Optionally a dict mapping axis name to separate values for each axis
+        can be passed. Default is 0.
     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`.
@@ -211,18 +266,37 @@ def apply_as_grid_ufunc(
     --------
     as_grid_ufunc
     Grid.apply_as_grid_ufunc
+    xarray.apply_ufunc
     """
 
     if grid is None:
         raise ValueError("Must provide a grid object to describe the Axes")
 
+    if len(args) != len(axis):
+        raise ValueError(
+            "Number of entries in `axis` does not match the number of data arguments supplied"
+        )
+
     # Extract Axes information from signature
-    in_ax_names, out_ax_names, in_ax_pos, out_ax_pos = _parse_grid_ufunc_signature(
-        signature
+    (
+        in_dummy_ax_names,
+        out_dummy_ax_names,
+        in_ax_pos,
+        out_ax_pos,
+    ) = _parse_grid_ufunc_signature(signature)
+
+    dummy_to_real_axes_mapping = _identify_dummy_axes_with_real_axes(
+        in_dummy_ax_names, axis
     )
 
+    # Determine names of output axes from names in signature
+    # TODO what if we need to add a new core dim to the output that does match an input axis? Where do we get the name from?
+    out_ax_names = [
+        [dummy_to_real_axes_mapping[ax] for ax in arg] for arg in out_dummy_ax_names
+    ]
+
     # 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 i, (arg_ns, arg_ps, arg) in enumerate(zip(axis, in_ax_pos, args)):
         for n, p in zip(arg_ns, arg_ps):
             try:
                 ax_pos = grid.axes[n].coords[p]
@@ -243,18 +317,44 @@ def apply_as_grid_ufunc(
     # 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)
+        for arg_ns, arg_ps in zip(axis, 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)
+    # Pad arrays according to internal boundary condition information
+    if boundary and not boundary_width:
+        raise ValueError(
+            "To apply a boundary condition you must provide the widths of the boundaries"
+        )
+    if boundary_width:
+        # convert dummy axes names in boundary_width to match real names of given axes
+        boundary_width_real_axes = {
+            dummy_to_real_axes_mapping[ax]: width
+            for ax, width in boundary_width.items()
+        }
+
+        args = grid.pad(
+            *args,
+            boundary_width=boundary_width_real_axes,
+            boundary=boundary,
+            fill_value=fill_value,
+        )
 
     # 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}
+    # TODO does padding change this?
+    out_sizes = {
+        out_dim: grid._ds.dims[out_dim] for arg in out_core_dims for out_dim in arg
+    }
+    # TODO allow users to specify new output dtypes
+    out_dtypes = [a.dtype for a in args]
+
+    # TODO Map operation over dask chunks?
+    # def mapped_func(*a, **kw):
+    #    return map_overlap(func, *a, **kw, depths=boundary_depths, boundary=None, trim=True)
 
     # Perform operation via xarray.apply_ufunc
     results = xr.apply_ufunc(
@@ -265,8 +365,12 @@ def apply_as_grid_ufunc(
         dask=dask,
         **kwargs,
         dask_gufunc_kwargs={"output_sizes": out_sizes},
+        output_dtypes=out_dtypes,
     )
 
+    # TODO add option to trim result if not done in ufunc
+    # TODO loud warning if ufunc returns array of incorrect size
+
     # apply_ufunc might return multiple objects
     if not isinstance(results, tuple):
         results = (results,)
@@ -280,13 +384,86 @@ def apply_as_grid_ufunc(
             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)
+
+        try:
+            res = res.assign_coords(new_core_dim_coords)
+        except ValueError as err:
+            if boundary_width and str(err).startswith("conflicting sizes"):
+                # TODO make this error more informative?
+                raise ValueError(
+                    f"{str(err)} - does your grid ufunc correctly trim off the same number of elements "
+                    f"which were added by padding using boundary_width={boundary_width}?"
+                )
+            else:
+                raise
         results_with_coords.append(res)
 
-    # Return single results not wrapped in 1-element tuple, like xr.apply_ufunc
+    # Return single results not wrapped in 1-element tuple, like xr.apply_ufunc does
     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
+
+
+def _identify_dummy_axes_with_real_axes(sig_in_dummy_ax_names, axis):
+    """Create a mapping between the dummy axis names in the signature and the real axis names of the data passed."""
+
+    if len(axis) != len(sig_in_dummy_ax_names):
+        raise ValueError(
+            "Number of entries in `axis` does not match the number of variables in the input signature"
+        )
+    for i, (arg_axes, dummy_arg_axes) in enumerate(zip(axis, sig_in_dummy_ax_names)):
+        if len(arg_axes) != len(dummy_arg_axes):
+            raise ValueError(
+                f"Number of Axes in `axis` entry number {i} does not match the number of Axes in that entry in the input signature"
+            )
+
+    # We can't just use set because we need these two lists to retain their ordering relative to one another
+    unique_dummy_axes = list(
+        dict.fromkeys(ax for arg in sig_in_dummy_ax_names for ax in arg)
+    )
+    unique_real_axes = list(dict.fromkeys(ax for arg in axis for ax in arg))
+
+    if len(unique_dummy_axes) != len(unique_real_axes):
+        raise ValueError(
+            f"Found {len(unique_dummy_axes)} unique input axes in signature but {len(unique_real_axes)} "
+            f"real unique input axes were supplied to the grid ufunc when called"
+        )
+
+    return dict(zip(unique_dummy_axes, unique_real_axes))
+
+
+_REPLACEMENT_DUMMY_INDEX_NAMES = [f"__{char}" for char in string.ascii_letters]
+
+
+def _signatures_equivalent(sig1, sig2):
+    """
+    Axes names in signatures are dummy variables, so an exact string match is not required.
+
+    Our comparison strategy is to instead work through both signatures left to right, replacing all occurrences
+    of each dummy index with names drawn from a common list. If after this process the replaced names are not
+    identical, the signatures must not be equivalent. Axes positions do have to match exactly.
+    """
+    sig1_in, sig1_out, _, _ = _parse_grid_ufunc_signature(sig1)
+    sig2_in, sig2_out, _, _ = _parse_grid_ufunc_signature(sig2)
+
+    all_unique_sig1_indices = set([i for arg in sig1_in for i in arg])
+    all_unique_sig2_indices = set([i for arg in sig2_in for i in arg])
+
+    if len(all_unique_sig1_indices) != len(all_unique_sig2_indices):
+        return False
+
+    sig1_replaced = sig1
+    sig2_replaced = sig2
+    for dummy1, dummy2, common_replacement in zip(
+        all_unique_sig1_indices, all_unique_sig2_indices, _REPLACEMENT_DUMMY_INDEX_NAMES
+    ):
+        sig1_replaced = sig1_replaced.replace(dummy1, common_replacement)
+        sig2_replaced = sig2_replaced.replace(dummy2, common_replacement)
+
+    if sig1_replaced == sig2_replaced:
+        return True
+    else:
+        return False
diff --git a/xgcm/gridops.py b/xgcm/gridops.py
index 79562fb..95ec301 100644
--- a/xgcm/gridops.py
+++ b/xgcm/gridops.py
@@ -1,331 +1,66 @@
-# python 3 compatiblity
+from .grid_ufunc import as_grid_ufunc
 
-import numpy as np
-import xarray as xr
+"""
+This module is intended to only contain "grid ufuncs".
 
-# make some functions for taking divergence
-from dask import array as da
+(It is however fine to add other non-"grid ufunc" functions to this list, as xgcm.grid._select_grid_ufunc will ignore
+anything that does not return an instance of the GridUFunc class.)
 
+If adding a new function to this list, make sure the function name starts with the name of the xgcm.Grid method you
+want it to be called from. e.g. `diff_centre_to_left_second_order` will be added to the list of functions to search
+through when `xgcm.Grid.diff()` is called. See xgcm.grid_ufunc._select_grid_ufunc for more details.
+"""
 
-def _append_to_name(array, append):
-    try:
-        return array.name + "_" + append
-    except TypeError:
-        return append
 
+def diff_forward(a):
+    return a[..., 1:] - a[..., :-1]
 
-class GCMDataset(object):
-    """Representation of GCM (General Circulation Model) output data, numerical
-    grid information, and operations related to finite-volume analysis.
-    """
 
-    # without these variables in the dataset, can't function
-    needed_vars = [
-        "Z",
-        "Zp1",
-        "Zl",
-        "Zu",
-        "X",
-        "Xp1",
-        "Y",
-        "Yp1",
-        "XC",
-        "YC",
-        "XG",
-        "YG",
-        "drF",
-        "drC",
-        "dxC",
-        "dxG",
-        "dyC",
-        "dyG",
-    ]
+@as_grid_ufunc(signature="(X:center)->(X:left)", boundary_width={"X": (1, 0)})
+def diff_center_to_left(a):
+    return diff_forward(a)
 
-    def __init__(self, ds):
-        """Initialize GCM object.
 
-        Parameters
-        ----------
-        ds : xr.Dataset
-        """
+@as_grid_ufunc(signature="(X:left)->(X:center)", boundary_width={"X": (0, 1)})
+def diff_left_to_center(a):
+    return diff_forward(a)
 
-        # check that needed variables are present
-        for v in self.needed_vars:
-            if v not in ds:
-                raise KeyError("Needed variable %s not found in dataset" % v)
 
-        self.ds = ds
+@as_grid_ufunc(signature="(X:center)->(X:right)", boundary_width={"X": (0, 1)})
+def diff_center_to_right(a):
+    return diff_forward(a)
 
-    # silly functions
 
-    def _get_coords_from_dims(self, dims, replace=None):
-        """Utility function for quickly fetching coordinates from parent
-        dataset.
-        """
-        dims = list(dims)
-        if replace:
-            for k in replace:
-                dims[dims.index(k)] = replace[k]
-        return {dim: self.ds[dim] for dim in dims}, dims
+@as_grid_ufunc(signature="(X:right)->(X:center)", boundary_width={"X": (1, 0)})
+def diff_right_to_center(a):
+    return diff_forward(a)
 
-    def _get_hfac_for_array(self, array):
-        """Figure out the correct hfac given array dimensions."""
-        hfac = None
-        if "X" in array.dims and "Y" in array.dims and "HFacC" in self.ds:
-            hfac = self.ds.HFacC
-        if "Xp1" in array.dims and "Y" in array.dims and "HFacW" in self.ds:
-            hfac = self.ds.HFacW
-        if "X" in array.dims and "Yp1" in array.dims and "HFacW" in self.ds:
-            hfac = self.ds.HFacS
-        return hfac
 
-    # Vertical Differences, Derivatives, and Interpolation
+@as_grid_ufunc(signature="(X:center)->(X:outer)", boundary_width={"X": (1, 1)})
+def diff_center_to_outer(a):
+    return diff_forward(a)
 
-    def pad_zl_to_zp1(self, array, fill_value=0.0, zlname="Zl", zp1name="Zp1"):
-        """Pad an array located at zl points such that it is located at
-        zp1 points. An additional fill value is required for the bottom point.
 
-        Parameters
-        ----------
-        array : xr.DataArray
-            The array to difference. Must have the coordinate zp1.
-        fill_value : number, optional
-            The value to be used at the bottom point.
-        zlname : str, optional
-            The variable name for the zl point
-        zp1name : str, optional
-            The variable name for the zp1 point
+# TODO this actually makes the array end up smaller, but boundary_width={"X": (-1, -1)} is not the correct kwarg value.
+# TODO rename `boundary_width` argument to `pad_width` to better reflect this possibility?
+@as_grid_ufunc(signature="(X:outer)->(X:center)", boundary_width={"X": (0, 0)})
+def diff_outer_to_center(a):
+    return diff_forward(a)
 
-        Returns
-        -------
-        padded : xr.DataArray
-            Padded array with vertical coordinate zp1.
-        """
-        coords, dims = self._get_coords_from_dims(array.dims)
-        zdim = dims.index(zlname)
-        # shape of the new array to concat at the bottom
-        shape = list(array.shape)
-        shape[zdim] = 1
-        # replace Zl with the bottom level
-        coords[zlname] = np.atleast_1d(self.ds[zp1name][-1].data)
-        # an array of zeros at the bottom
-        # need different behavior for numpy vs dask
-        if array.chunks:
-            chunks = list(array.data.chunks)
-            chunks[zdim] = (1,)
-            zarr = fill_value * da.ones(shape, dtype=array.dtype, chunks=chunks)
-            zeros = xr.DataArray(zarr, coords, dims).chunk()
-        else:
-            zarr = np.zeros(shape, array.dtype)
-            zeros = xr.DataArray(zarr, coords, dims)
-        newarray = xr.concat([array, zeros], dim=zlname).rename({zlname: zp1name})
-        if newarray.chunks:
-            # this assumes that there was only one chunk in the vertical to begin with
-            # how can we do that better
-            return newarray.chunk({zp1name: len(newarray[zp1name])})
-        else:
-            return newarray
 
-    def diff_zp1_to_z(self, array, zname="Z", zp1name="Zp1"):
-        """Take the vertical difference of an array located at zp1 points, resulting
-        in a new array at z points.
+@as_grid_ufunc(signature="(X:center)->(X:inner)", boundary_width={"X": (0, 0)})
+def diff_center_to_inner(a):
+    return diff_forward(a)
 
-        Parameters
-        ----------
-        array : xr.DataArray
-            The array to difference. Must have the coordinate zp1.
-        zname : str, optional
-            The variable name for the z point
-        zp1name : str, optional
-            The variable name for the zp1 point
 
-        Returns
-        -------
-        diff : xr.DataArray
-            A new array with vertical coordinate z.
-        """
-        a_up = array.isel(**{zp1name: slice(None, -1)})
-        a_dn = array.isel(**{zp1name: slice(1, None)})
-        a_diff = a_up.data - a_dn.data
-        # dimensions and coords of new array
-        coords, dims = self._get_coords_from_dims(array.dims, replace={zp1name: zname})
-        return xr.DataArray(
-            a_diff, coords, dims, name=_append_to_name(array, "diff_zp1_to_z")
-        )
+@as_grid_ufunc(signature="(X:inner)->(X:center)", boundary_width={"X": (1, 1)})
+def diff_inner_to_center(a):
+    return diff_forward(a)
 
-    def diff_zl_to_z(self, array, fill_value=0.0):
-        """Take the vertical difference of an array located at zl points, resulting
-        in a new array at z points. A fill value is required to provide the bottom
-        boundary condition for ``array``.
 
-        Parameters
-        ----------
-        array : xr.DataArray
-            The array to difference. Must have the coordinate zl.
-        fill_value : number, optional
-            The value to be used at the bottom point. The default (0) is the
-            appropriate choice for vertical fluxes.
+@as_grid_ufunc(signature="(X:left)->(X:inner)")
+def diff_left_to_inner(a):
+    raise NotImplementedError
 
-        Returns
-        -------
-        diff : xr.DataArray
-            A new array with vertical coordinate z.
-        """
-        array_zp1 = self.pad_zl_to_zp1(array, fill_value)
-        array_diff = self.diff_zp1_to_z(array_zp1)
-        return array_diff.rename(_append_to_name(array, "_diff_zl_to_z"))
 
-    def diff_z_to_zp1(self, array):
-        """Take the vertical difference of an array located at z points, resulting
-        in a new array at zp1 points, but missing the upper and lower point.
-
-        Parameters
-        ----------
-        array : xr.DataArray
-            The array to difference. Must have the coordinate z.
-
-        Returns
-        -------
-        diff : xr.DataArray
-            A new array with vertical coordinate zp1.
-        """
-        a_up = array.isel(Z=slice(None, -1))
-        a_dn = array.isel(Z=slice(1, None))
-        a_diff = a_up.data - a_dn.data
-        # dimensions and coords of new array
-        coords, dims = self._get_coords_from_dims(array.dims, replace={"Z": "Zp1"})
-        # trim vertical
-        coords["Zp1"] = coords["Zp1"][1:-1]
-        return xr.DataArray(
-            a_diff, coords, dims, name=_append_to_name(array, "diff_z_to_zp1")
-        )
-
-    def derivative_zp1_to_z(self, array):
-        """Take the vertical derivative of an array located at zp1 points, resulting
-        in a new array at z points.
-
-        Parameters
-        ----------
-        array : xr.DataArray
-            The array to differentiate. Must have the coordinate zp1.
-
-        Returns
-        -------
-        deriv : xr.DataArray
-            A new array with vertical coordinate z.
-        """
-        a_diff = self.diff_zp1_to_z(array)
-        dz = self.ds.drF
-        return a_diff / dz
-
-    def derivative_zl_to_z(self, array, fill_value=0.0):
-        """Take the vertical derivative of an array located at zl points, resulting
-        in a new array at z points. A fill value is required to provide the bottom
-        boundary condition for ``array``.
-
-        Parameters
-        ----------
-        array : xr.DataArray
-            The array to differentiate. Must have the coordinate zl.
-        fill_value : number, optional
-            The assumed value at the bottom point. The default (0) is the
-            appropriate choice for vertical fluxes.
-
-        Returns
-        -------
-        deriv : xr.DataArray
-            A new array with vertical coordinate z.
-        """
-
-        a_diff = self.diff_zl_to_z(array, fill_value)
-        dz = self.ds.drF
-        return a_diff / dz
-
-    def derivative_z_to_zp1(self, array):
-        """Take the vertical derivative of an array located at z points, resulting
-        in a new array at zp1 points, but missing the upper and lower point.
-
-        Parameters
-        ----------
-        array : xr.DataArray
-            The array to differentiate. Must have the coordinate z.
-
-        Returns
-        -------
-        diff : xr.DataArray
-            A new array with vertical coordinate zp1.
-        """
-        a_diff = self.diff_z_to_zp1(array)
-        dz = self.ds.drC[1:-1]
-        return a_diff / dz
-
-    # Vertical Integrals
-    # if the array to integrate is 1D or 2D, don't multiply by hFac
-    # but what if it is 3D, how do we decide what to do?
-    # what if points are missing? xray should take care of that
-    # how do we pick which hFac to use? look at dims
-    def integrate_z(self, array, average=False):
-        """Integrate ``array`` in vertical dimension, accounting for vertical
-        grid geometry.
-
-        Parameters
-        ----------
-        array : xr.DataArray
-            The array to integrate. Must have the dimension Z.
-        average : bool, optional
-            If ``True``, return an average instead of an integral.
-
-        Returns
-        -------
-        integral : xr.DataArray
-            The vertical integral of ``array``.
-        """
-        if "Z" not in array.dims:
-            raise ValueError("Can only integrate arrays on Z grid")
-        dz = self.ds.drF
-        # look at horizontal dimensions and try to find an hfac
-        hfac = self._get_hfac_for_array(array)
-        if hfac is not None:
-            # brodcast hfac against dz
-            dz *= hfac
-        a_int = (array * dz).sum(dim="Z")
-        if average:
-            return a_int / dz.sum(dim="Z")
-        else:
-            return a_int
-
-    # Horizontal Differences, Derivatives, and Interpolation
-
-    # doesn't actually need parent ds
-    # this could go in xray
-    def roll(self, array, n, dim):
-        """Clone of numpy.roll for xr.DataArrays."""
-        left = array.isel(**{dim: slice(None, -n)})
-        right = array.isel(**{dim: slice(-n, None)})
-        return xr.concat([right, left], dim=dim)
-
-    def diff_xp1_to_x(self, array):
-        """Difference DataArray ``array`` in the x direction.
-        Assumes that ``array`` is located at the xp1 point."""
-        left = array
-        right = self.roll(array, -1, "Xp1")
-        if array.chunks:
-            right = right.chunk(array.chunks)
-        diff = right.data - left.data
-        coords, dims = self._get_coords_from_dims(array.dims, replace={"Xp1": "X"})
-        return xr.DataArray(diff, coords, dims).rename(
-            _append_to_name(array, "diff_xp1_to_x")
-        )
-
-    def diff_yp1_to_y(self, array):
-        """Difference DataArray ``array`` in the y direction.
-        Assumes that ``array`` is located at the yp1 point."""
-        left = array
-        right = self.roll(array, -1, "Yp1")
-        if array.chunks:
-            right = right.chunk(array.chunks)
-        diff = right.data - left.data
-        coords, dims = self._get_coords_from_dims(array.dims, replace={"Yp1": "Y"})
-        return xr.DataArray(diff, coords, dims).rename(
-            _append_to_name(array, "_diff_yp1_to_y")
-        )
+# TODO fill out all the other ufuncs for interp etc...


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
index 61f1aaa..99ae956 100644
--- a/xgcm/test/test_grid_ufunc.py
+++ b/xgcm/test/test_grid_ufunc.py
@@ -5,10 +5,11 @@ import pytest
 import xarray as xr
 from xarray.testing import assert_equal
 
-from xgcm.grid import Grid
+from xgcm.grid import Grid, _select_grid_ufunc
 from xgcm.grid_ufunc import (
     GridUFunc,
     _parse_grid_ufunc_signature,
+    _signatures_equivalent,
     apply_as_grid_ufunc,
     as_grid_ufunc,
 )
@@ -118,7 +119,7 @@ def create_1d_test_grid(ax_name):
     return Grid(
         grid_ds,
         coords={
-            f"{ax_name.upper()}": {
+            f"{ax_name}": {
                 "center": f"{ax_name}_c",
                 "left": f"{ax_name}_g",
                 "right": f"{ax_name}_r",
@@ -135,13 +136,13 @@ def create_2d_test_grid(ax_name_1, ax_name_2):
     return Grid(
         ds=xr.merge([grid_ds_1, grid_ds_2]),
         coords={
-            f"{ax_name_1.upper()}": {
+            f"{ax_name_1}": {
                 "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()}": {
+            f"{ax_name_2}": {
                 "center": f"{ax_name_2}_c",
                 "left": f"{ax_name_2}_g",
                 "right": f"{ax_name_2}_r",
@@ -169,29 +170,49 @@ class TestGridUFunc:
                 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)->()")
+        grid = create_1d_test_grid("depth")
+        da = np.sin(grid._ds.depth_g * 2 * np.pi / 9)
+
+        with pytest.raises(ValueError, match=re.escape("(depth:outer) does not exist")):
+            apply_as_grid_ufunc(
+                lambda x: x, da, axis=[("depth",)], grid=grid, signature="(X:outer)->()"
+            )
+
+        with pytest.raises(ValueError, match="coordinate depth_c does not appear"):
+            apply_as_grid_ufunc(
+                lambda x: x,
+                da,
+                axis=[("depth",)],
+                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
+        grid = create_1d_test_grid("depth")
+        da = np.sin(grid._ds.depth_c * 2 * np.pi / 9)
+        da.coords["depth_c"] = grid._ds.depth_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})
+        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}
+        )
 
         # Test direct application
         result = apply_as_grid_ufunc(
-            diff_center_to_left, da, grid=grid, signature="(X:center)->(X:left)"
+            diff_center_to_left,
+            da,
+            axis=[("depth",)],
+            grid=grid,
+            signature="(X:center)->(X:left)",
+        )
+        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)"
         )
         assert_equal(result, expected)
 
@@ -200,36 +221,49 @@ class TestGridUFunc:
         def diff_center_to_left(a):
             return a - np.roll(a, shift=-1)
 
-        result = diff_center_to_left(grid, da)
+        result = diff_center_to_left(grid, da, axis=[("depth",)])
         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")
+        grid = create_1d_test_grid("depth")
         da = xr.DataArray(
-            np.arange(10, 19), dims=["x_c"], coords={"x_c": grid._ds.x_c}
+            np.arange(10, 19), dims=["depth_c"], coords={"depth_c": grid._ds.depth_c}
         ).chunk()
 
-        expected = da.interp(x_c=np.arange(1.5, 9), method="linear").rename(x_c="x_i")
+        expected = da.interp(depth_c=np.arange(1.5, 9), method="linear").rename(
+            depth_c="depth_i"
+        )
 
         # Test direct application
         result = apply_as_grid_ufunc(
             interp_center_to_inner,
             da,
+            axis=[("depth",)],
             grid=grid,
             signature="(X:center)->(X:inner)",
             dask="parallelized",
         ).compute()
         assert_equal(result, expected)
 
+        # Test Grid method
+        result = grid.apply_as_grid_ufunc(
+            interp_center_to_inner,
+            da,
+            axis=[("depth",)],
+            signature="(X:center)->(X:inner)",
+            dask="parallelized",
+        )
+        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()
+        result = interp_center_to_inner(grid, da, axis=[("depth",)]).compute()
         assert_equal(result, expected)
 
     def test_1d_overlap_dask_allowed(self):
@@ -241,42 +275,93 @@ class TestGridUFunc:
         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
+        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 - da.roll(x_c=-1, roll_coords=False)).data
+        diffed = (da - da.roll(depth_c=-1, roll_coords=False)).data
         expected = xr.DataArray(
-            diffed, dims=["x_g"], coords={"x_g": grid._ds.x_g}
+            diffed, dims=["depth_g"], coords={"depth_g": grid._ds.depth_g}
         ).compute()
 
         # Test direct application
         result = apply_as_grid_ufunc(
-            diff_center_to_left,
+            diff_overlap,
             da,
+            axis=[("depth",)],
             grid=grid,
             signature="(X:center)->(X:left)",
             dask="allowed",
         ).compute()
         assert_equal(result, expected)
 
+        # Test Grid method
+        result = grid.apply_as_grid_ufunc(
+            diff_overlap,
+            da,
+            axis=[("depth",)],
+            signature="(X:center)->(X:left)",
+            dask="allowed",
+        )
+        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()
+        result = diff_overlap(
+            grid,
+            da,
+            axis=[("depth",)],
+        ).compute()
+        assert_equal(result, expected)
+
+    @pytest.mark.xfail(reason="Need to fix PR #371")
+    def test_apply_along_one_axis(self):
+        grid = create_2d_test_grid("lon", "lat")
+
+        def diff_center_to_left(a):
+            return a - np.roll(a, shift=-1)
+
+        da = grid._ds.lat_c ** 2 + grid._ds.lon_c ** 2
+
+        diffed = (da - da.roll(lon_c=-1, roll_coords=False)).data
+        expected = xr.DataArray(
+            diffed,
+            dims=["lat_c", "lon_g"],
+            coords={"lat_c": grid._ds.lat_c, "lon_g": grid._ds.lon_g},
+        )
+
+        # Test direct application
+        result = apply_as_grid_ufunc(
+            diff_center_to_left,
+            da,
+            axis=[("lon",)],
+            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, axis=[("lon",)])
         assert_equal(result, expected)
 
+    # TODO test a function with padding
+
     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
+        grid = create_1d_test_grid("depth")
+        a = np.sin(grid._ds.depth_g * 2 * np.pi / 9)
+        a.coords["depth_g"] = grid._ds.depth_g
+        b = np.cos(grid._ds.depth_r * 2 * np.pi / 9)
+        b.coords["depth_r"] = grid._ds.depth_r
 
         expected = xr.DataArray(np.inner(a, b))
 
@@ -285,17 +370,28 @@ class TestGridUFunc:
             inner_product_left_right,
             a,
             b,
+            axis=[("depth",), ("depth",)],
             grid=grid,
             signature="(X:left),(X:right)->()",
         )
         assert_equal(result, expected)
 
+        # Test Grid method
+        result = grid.apply_as_grid_ufunc(
+            inner_product_left_right,
+            a,
+            b,
+            axis=[("depth",), ("depth",)],
+            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)
+        result = inner_product_left_right(grid, a, b, axis=[("depth",), ("depth",)])
         assert_equal(result, expected)
 
     def test_multiple_outputs(self):
@@ -306,30 +402,137 @@ class TestGridUFunc:
         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")
+        grid = create_2d_test_grid("lon", "lat")
 
-        a = grid._ds.x_c ** 2 + grid._ds.y_c ** 2
+        a = grid._ds.lon_c ** 2 + grid._ds.lat_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
+        expected_u = 2 * grid._ds.lon_i.expand_dims(dim={"lat_c": len(grid._ds.lat_c)})
+        expected_u.coords["lat_c"] = grid._ds.lat_c
+        expected_v = 2 * grid._ds.lat_i.expand_dims(dim={"lon_c": len(grid._ds.lon_c)})
+        expected_v.coords["lon_c"] = grid._ds.lon_c
 
         # Test direct application
         u, v = apply_as_grid_ufunc(
             grad_to_inner,
             a,
+            axis=[("lon", "lat")],
             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 Grid method
+        u, v = grid.apply_as_grid_ufunc(
+            grad_to_inner,
+            a,
+            axis=[("lon", "lat")],
+            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)
+        u, v = grad_to_inner(grid, a, axis=[("lon", "lat")])
         assert_equal(u.T, expected_u)
         assert_equal(v, expected_v)
+
+
+class TestSignaturesEquivalent:
+    def test_equivalent(self):
+        sig1 = "(X:center)->(X:left)"
+        sig2 = "(X:center)->(X:left)"
+        assert _signatures_equivalent(sig1, sig2)
+
+        sig3 = "(Y:center)->(Y:left)"
+        assert _signatures_equivalent(sig1, sig3)
+
+    def test_not_equivalent(self):
+        sig1 = "(X:center)->(X:left)"
+        sig2 = "(X:center)->(X:center)"
+        assert not _signatures_equivalent(sig1, sig2)
+
+        sig3 = "(X:center)->(Y:left)"
+        assert not _signatures_equivalent(sig1, sig3)
+
+        sig4 = "(X:center,X:center)->(X:left)"
+        assert not _signatures_equivalent(sig1, sig4)
+
+    def test_no_indices(self):
+        sig = "()->()"
+        assert _signatures_equivalent(sig, sig)
+
+
+class GridOpsMockUp:
+    """
+    Container that stores some mocked-up grid ufuncs to look through.
+    Intended to be used as if it were the gridops.py module file.
+    """
+
+    @staticmethod
+    @as_grid_ufunc(signature="(X:center)->(X:left)")
+    def diff_center_to_left(a):
+        return a - np.roll(a, -1)
+
+    @staticmethod
+    @as_grid_ufunc(signature="(X:center)->(X:right)")
+    def diff_center_to_right_fill(a):
+        return np.roll(a, 1) - a
+
+    @staticmethod
+    @as_grid_ufunc(signature="(X:center)->(X:right)")
+    def diff_center_to_right_extend(a):
+        return np.roll(a, 1) - a
+
+    @staticmethod
+    @as_grid_ufunc(signature="()->()")
+    def pass_through_kwargs(**kwargs):
+        return kwargs
+
+
+class TestGridUFuncDispatch:
+    def test_select_ufunc(self):
+        gridufunc, _ = _select_grid_ufunc(
+            "diff", "(X:center)->(X:left)", module=GridOpsMockUp
+        )
+        assert gridufunc is GridOpsMockUp.diff_center_to_left
+
+    def test_select_ufunc_equivalent_signature(self):
+        gridufunc, _ = _select_grid_ufunc(
+            "diff", "(Y:center)->(Y:left)", module=GridOpsMockUp
+        )
+        assert gridufunc is GridOpsMockUp.diff_center_to_left
+
+        with pytest.raises(NotImplementedError):
+            _select_grid_ufunc("diff", "(X:center)->(Y:left)", module=GridOpsMockUp)
+
+    def test_select_ufunc_wrong_signature(self):
+        with pytest.raises(NotImplementedError):
+            _select_grid_ufunc("diff", "(X:center)->(X:center)", module=GridOpsMockUp)
+
+    @pytest.mark.xfail(reason="currently no need for this")
+    def test_select_ufunc_by_kwarg(self):
+        gridufunc, _ = _select_grid_ufunc(
+            "diff", "(X:center)->(X:right)", module=GridOpsMockUp, boundary="fill"
+        )
+        assert gridufunc is GridOpsMockUp.diff_center_to_right_fill
+
+        with pytest.raises(NotImplementedError):
+            _select_grid_ufunc(
+                "diff",
+                "(X:center)->(X:right)",
+                module=GridOpsMockUp,
+                boundary="nonsense",
+            )
+
+    @pytest.mark.xfail
+    def test_pass_through_other_kwargs(self):
+        # TODO put this in test_grid.py instead?
+        gridufunc, _ = _select_grid_ufunc(
+            "pass", "()->()", module=GridOpsMockUp, boundary="fill"
+        )
+        assert gridufunc(a=1) == {"a": 1}