Automatic grid generation

cc @jbusecke, @naomi-henderson

All of the basic functionality in xgcm now could easily be applied to generic gridded datasets (e.g. satellite, model data without auxiliary coordinates, etc.), if we had the ability to automatically generate other axis variables. (This may also be necessary for @naomi-henderson to move forward with pangeo-data/pangeo-discussion#1.)

What would this look like? Imagine we have an xarray dataset that looks like this

import xarray as xr
import numpy as np
ds = xr.DataArray(np.random.rand(180, 360),
                         dims=['lat', 'lon'],
                         coords={'lon': np.arange(0, 360)+0.5,
                                 'lat': np.arange(-90, 90)+0.5}
                        ).to_dataset(name='somedata')
ds

<xarray.Dataset>
Dimensions:   (lat: 180, lon: 360)
Coordinates:
  * lon       (lon) float64 0.5 1.5 2.5 3.5 4.5 5.5 6.5 7.5 8.5 9.5 10.5 ...
  * lat       (lat) float64 -89.5 -88.5 -87.5 -86.5 -85.5 -84.5 -83.5 -82.5 ...
Data variables:
    somedata  (lat, lon) float64 0.2838 0.3383 0.1441 0.8124 0.6643 0.6623 ...

This is a dataset with the values located at the midpoints between lat and lon.

It would be amazing to be able to say this

grid = xgcm.autogenerate_axis(ds, axes={'X': 'lon', 'Y': 'lat'})

There are lots of potential options we would want to include, such as the coordinate positions to generate (i.e. left, right, inner, outer). I’m also assuming we would want to treat the original coordinates as the cell centers, but maybe that’s not always the case.

This could be implemented via a lower level function called autogenerate_axis that just works on one axis at a time (the way the Grid object works now). I would recommend creating a new module called autogenerate for these functions.

A reminder of the lexicon of grid positions

 Center
 |------o-------|------o-------|------o-------|------o-------|
       [0]            [1]            [2]            [3]

 Left
 |------o-------|------o-------|------o-------|------o-------|
[0]            [1]            [2]            [3]

 Right
 |------o-------|------o-------|------o-------|------o-------|
               [0]            [1]            [2]            [3]

 Inner
 |------o-------|------o-------|------o-------|------o-------|
               [0]            [1]            [2]

 Outer
 |------o-------|------o-------|------o-------|------o-------|
[0]            [1]            [2]            [3]            [4]

[jbusecke]

I took a shot at a very basic implementation (#74). Based on the example above from @rabernat it would work something like this:

import xgcm as xg
import xarray as xr
import numpy as np

# create a 'typical observational' dataset
ds = xr.DataArray(np.random.rand(180, 360*2,10*2),
                         dims=['lat', 'lon','z'],
                         coords={'lon': np.arange(-180, 180, 0.5)+0.25,
                                 'lat': np.arange(-90, 90)+0.5,
                                 'z': np.arange(0,10,0.5)+0.25}
                        ).to_dataset(name='somedata')
ds

>>> <xarray.Dataset>
>>> Dimensions:   (lat: 180, lon: 720, z: 20)
>>> Coordinates:
>>>   * lon       (lon) float64 -179.8 -179.2 -178.8 -178.2 -177.8 -177.2 -176.8 ...
>>>   * lat       (lat) float64 -89.5 -88.5 -87.5 -86.5 -85.5 -84.5 -83.5 -82.5 ...
>>>   * z         (z) float64 0.25 0.75 1.25 1.75 2.25 2.75 3.25 3.75 4.25 4.75 ...
>>> Data variables:
>>>     somedata  (lat, lon, z) float64 0.2058 0.06754 0.89 0.9464 0.02511 ...

# autogenerate_ds recovers the 'original' coordinates
new_ds = xg.autogenerate_ds(ds,axes={'X':'lon','Y':'lat','Z':'z'},
                            position='center')

print(new_ds['lon'].data[0:3])
print(new_ds['lon'].data[-4:-1])
print(new_ds['lon'+'_inferred'].data[0:3])
print(new_ds['lon'+'_inferred'].data[-4:-1])

print(new_ds['lat'].data[0:3])
print(new_ds['lat'].data[-4:-1])
print(new_ds['lat'+'_inferred'].data[0:3])
print(new_ds['lat'+'_inferred'].data[-4:-1])

print(new_ds['z'].data[0:3])
print(new_ds['z'].data[-4:-1])
print(new_ds['z'+'_inferred'].data[0:3])
print(new_ds['z'+'_inferred'].data[-4:-1])

>>> [-179.75 -179.25 -178.75]
>>> [ 178.25  178.75  179.25]
>>> [-180.  -179.5 -179. ]
>>> [ 178.   178.5  179. ]
>>> [-89.5 -88.5 -87.5]
>>> [ 86.5  87.5  88.5]
>>> [-90. -89. -88.]
>>> [ 86.  87.  88.]
>>> [ 0.25  0.75  1.25]
>>> [ 8.25  8.75  9.25]
>>> [ 0.   0.5  1. ]
>>> [ 8.   8.5  9. ]

Some elements are still a bit clunky, particularly how add_slice is implemented in xgcm.grid, but it seems to work.

Still a lot more to implement (calculate area based on lat/lon coordinates etc), but I wanted to see how everybody feels about this basic stuff.

[rabernat]

It would be great to hear from @naomi-henderson what numerical approach they used in their past work with generic gridded datasets (like CMIP5). This paper on moisture budgets for example is the basis for one of our Pangeo use cases.