WIP: add xarray dataarray_accessor

[rabernat]

This is something we have discussed for a while. Today I realized it was pretty simple to implement. (Famous last words.) So I had to give it a try.

Several people (e.g. @JiaweiZhuang) have commented that they would like to be able to use xgcm with an accessor. The challenge is figuring out which Grid object to use with any given xarray.DataArray.

What I did here was copy Dask's mechanism for finding a Client object to use for computation. I used weakref to create a global set of Grid objects. When you access the .grid property of a DataArray, the accessor searches for a Grid that can align with the DataArray. The first one it finds is used for da.grid.interp and da.grid.diff methods.

Here's a simple example:

import numpy as np
import xarray as xr
import xgcm

ds = xr.Dataset(coords={'x_c': (['x_c',], np.arange(1,10)),
                        'x_g': (['x_g',], np.arange(0.5,9))})

grid = xgcm.Grid(ds, coords={'X': {'center': 'x_c', 'left': 'x_g'}})
da = np.cos(ds.x_c)
da.grid.interp('X')

There are many ways this could go wrong, particularly if there are multiple Grid objects in existence. However, in my workflows I tend to only ever have one at a time. So perhaps it's a easy, simple solution that will improve our API.

[codecov]

Codecov Report

Merging #197 into master will decrease coverage by 0.49%.
The diff coverage is 80.00%.

@@            Coverage Diff             @@
##           master     #197      +/-   ##
==========================================
- Coverage   89.73%   89.24%   -0.50%     
==========================================
  Files           4        5       +1     
  Lines         653      688      +35     
  Branches      143      146       +3     
==========================================
+ Hits          586      614      +28     
- Misses         36       41       +5     
- Partials       31       33       +2     

Impacted Files	Coverage Δ
xgcm/accessor.py	73.07% <73.07%> (ø)
xgcm/grid.py	91.46% <100.00%> (+0.18%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 03c165c...42db299. Read the comment docs.

[jbusecke]

Oh neat! I have definitely worked with many grid objects before, so I think this might be worth thinking about more. But perhaps there could be some input checks that would help sort out the “right” grid object (We discusses this in an earlier release...but the new github app has no autocomplete 😭)

[rabernat]

But perhaps there could be some input checks that would help sort out the “right” grid object

That's what the function _find_compatible_global_grid tries to do.

[jbusecke]

I see. I think this is really cool (but have to check it out in practice more).

[jbusecke]

Just linking #225 here, since it seems relevant to connect the discussions.

[github-actions]

This pull request has been marked 'stale' due to lack of recent activity. If there is no further activity, the pull request will be closed in another 30 days. Thank you for your contribution!

[rabernat]

Let's keep this alive! Not stale!

[lukelbd]

Just wanted to say this is an awesome project and having a .grid accessor would be really cool -- more elegant IMO.

Maybe ds.grid.init() could be expanded to work as an alternative to the Grid(ds, ...) call? Or, the accessor could be made callable, with ds.grid() equivalent to Grid(ds, ...).

And to handle multiple grids, maybe the accessor could be optionally indexed (e.g. ds.grid[N].interp(...)), with the "first" / "default" grid automatically selected when no index is provided (e.g. ds.grid.interp(...) -- perhaps a warning could be raised when there is more than one option). Or we could use ds.grid.add() instead of ds.grid.init(), and/or rename the accessor to .grids, if we're concerned about multiple-grid ambiguities.