Add utility for modifying pop model output to be compatible with xgcm

ci/docs.yml

@@ -2,19 +2,21 @@ name: pop-tools-dev
 channels:
   - conda-forge
 dependencies:
+  - dask
   - nbsphinx
+  - netcdf4
+  - notebook
+  - numba=0.45.1
   - numpy=1.17.0
   - numpydoc=0.8.0
   - pip
+  - pooch
   - python=3.7
   - recommonmark
-  - sphinx_rtd_theme=0.4.2
-  - sphinx=1.8.2
-  - sphinx-gallery=0.2.0
-  - xarray
-  - netcdf4
-  - numba=0.45.1
   - scipy=1.3.1
+  - sphinx_rtd_theme=0.4.2
   - sphinx-copybutton=0.2.5
-  - pooch
   - watermark
+  - xarray
+  - pip:
+      - git+https://github.com/xgcm/xgcm.git

ci/environment-dev-3.6.yml

@@ -2,24 +2,36 @@ name: pop-tools-dev
 channels:
   - conda-forge
 dependencies:
+  - autopep8
+  - black
   - bottleneck
   - codecov
   - dask
-  - distributed
+  - docrep
+  - flake8
+  - isort
   - jupyterlab
+  - make
   - matplotlib
+  - nbsphinx
   - netcdf4
   - numba
-  - numpy
+  - numpydoc
   - pip
+  - pooch
   - pytest
   - pytest-cov
-  - python=3.6
-  - scipy
-  - xarray>=0.12
-  - zarr
-  - docrep
   - pytest-lazy-fixture
   - pytest-xdist
-  - pooch
+  - python=3.6
+  - pyyaml
+  - recommonmark
+  - scipy
+  - sphinx-copybutton
+  - sphinx=1.8.2
+  - sphinx_rtd_theme=0.4.2
   - watermark
+  - xarray
+  - zarr
+  - pip:
+      - git+https://github.com/xgcm/xgcm.git

ci/environment-dev-3.7.yml

@@ -7,6 +7,7 @@ dependencies:
   - bottleneck
   - codecov
   - dask
+  - docrep
   - flake8
   - isort
   - jupyterlab
@@ -15,21 +16,22 @@ dependencies:
   - nbsphinx
   - netcdf4
   - numba
-  - numpy
   - numpydoc
   - pip
+  - pooch
   - pytest
   - pytest-cov
+  - pytest-lazy-fixture
+  - pytest-xdist
   - python=3.7
+  - pyyaml
   - recommonmark
   - scipy
-  - sphinx_rtd_theme=0.4.2
-  - sphinx=1.8.2
-  - xarray>=0.12
-  - zarr
-  - docrep
-  - pytest-lazy-fixture
-  - pytest-xdist
   - sphinx-copybutton
-  - pooch
+  - sphinx=1.8.2
+  - sphinx_rtd_theme=0.4.2
   - watermark
+  - xarray
+  - zarr
+  - pip:
+      - git+https://github.com/xgcm/xgcm.git

docs/source/api.rst

@@ -35,6 +35,12 @@ Utilities
    lateral_fill
 
 
+xgcm utilities
+~~~~~~~~~~~~~~
+
+.. autosummary::
+   get_xgcm_grid
+
 .. currentmodule:: pop_tools
 
 .. autofunction:: get_grid
@@ -48,3 +54,5 @@ Utilities
 .. autofunction:: list_region_masks
 
 .. autofunction:: lateral_fill
+
+.. autofunction:: get_xgcm_grid

pop_tools/__init__.py

@@ -8,6 +8,7 @@
 from .fill import lateral_fill, lateral_fill_np_array
 from .grid import get_grid, grid_defs
 from .region_masks import list_region_masks, region_mask_3d
+from .xgcm_util import get_xgcm_grid
 
 try:
     __version__ = get_distribution(__name__).version

pop_tools/xgcm_util.py

@@ -0,0 +1,169 @@
+"""
+Module for modifying pop model output to be compatible with xgcm
+
+# Ref: https://gist.github.com/rabernat/933bc785c99828352f343e48d0da6e22
+"""
+
+import numpy as np
+import xarray as xr
+
+
+def _add_pop_dims_to_dataset(ds):
+    ds_new = ds.copy()
+    ds_new['nlon_u'] = xr.Variable(
+        ('nlon_u'), np.arange(len(ds.nlon)) + 1, {'axis': 'X', 'c_grid_axis_shift': 0.5}
+    )
+    ds_new['nlat_u'] = xr.Variable(
+        ('nlat_u'), np.arange(len(ds.nlat)) + 1, {'axis': 'Y', 'c_grid_axis_shift': 0.5}
+    )
+    ds_new['nlon_t'] = xr.Variable(('nlon_t'), np.arange(len(ds.nlon)) + 0.5, {'axis': 'X'})
+    ds_new['nlat_t'] = xr.Variable(('nlat_t'), np.arange(len(ds.nlat)) + 0.5, {'axis': 'Y'})
+
+    # add metadata to z grid
+    if 'z_t' in ds_new.variables:
+        ds_new['z_t'].attrs.update({'axis': 'Z'})
+    if 'z_w' in ds_new.variables:
+        ds_new['z_w'].attrs.update({'axis': 'Z', 'c_grid_axis_shift': -0.5})
+    if 'z_w_top' in ds_new.variables:
+        ds_new['z_w_top'].attrs.update({'axis': 'Z', 'c_grid_axis_shift': -0.5})
+    if 'z_w_bot' in ds_new.variables:
+        ds_new['z_w_bot'].attrs.update({'axis': 'Z', 'c_grid_axis_shift': 0.5})
+
+    return ds_new
+
+
+def _dims_from_grid_loc(grid_loc):
+    grid_loc = str(grid_loc)
+    ndim = int(grid_loc[0])
+    x_loc_key = int(grid_loc[1])
+    y_loc_key = int(grid_loc[2])
+    z_loc_key = int(grid_loc[3])
+
+    x_loc = {1: 'nlon_t', 2: 'nlon_u', 3: 'nlon_u'}[x_loc_key]
+    y_loc = {1: 'nlat_t', 2: 'nlat_u', 3: 'nlat_t'}[y_loc_key]
+    y_loc = {1: 'nlat_t', 2: 'nlat_u'}[y_loc_key]
+    z_loc = {0: 'surface', 1: 'z_t', 2: 'z_w', 3: 'z_w_bot', 4: 'z_t_150m'}[z_loc_key]
+
+    if ndim == 3:
+        return z_loc, y_loc, x_loc
+    elif ndim == 2:
+        return y_loc, x_loc
+
+
+def _label_coord_grid_locs(ds):
+    grid_locs = {
+        'ANGLE': '2220',
+        'ANGLET': '2110',
+        'DXT': '2110',
+        'DXU': '2220',
+        'DYT': '2110',
+        'DYU': '2220',
+        'DZU': '3221',
+        'DZT': '3111',
+        'HT': '2110',
+        'HU': '2220',
+        'HTE': '2210',
+        'HTN': '2120',
+        'HUS': '2210',
+        'HUW': '2120',
+        'KMT': '2110',
+        'KMU': '2220',
+        'REGION_MASK': '2110',
+        'TAREA': '2110',
+        'TLAT': '2110',
+        'TLONG': '2110',
+        'UAREA': '2220',
+        'ULAT': '2220',
+        'ULONG': '2220',
+    }
+    ds_new = ds.copy()
+    for vname, grid_loc in grid_locs.items():
+        if vname in ds.variables:
+            ds_new[vname].attrs['grid_loc'] = grid_loc
+    return ds_new
+
+
+def relabel_pop_dims(ds):
+    """Return a new xarray dataset with distinct dimensions for variables at different
+    grid points.
+    """
+    ds_new = _label_coord_grid_locs(ds)
+    ds_new = _add_pop_dims_to_dataset(ds_new)
+    for vname in ds_new.variables:
+        if 'grid_loc' in ds_new[vname].attrs:
+            da = ds_new[vname]
+            dims_orig = da.dims
+            new_spatial_dims = _dims_from_grid_loc(da.attrs['grid_loc'])
+            if dims_orig[0] == 'time':
+                dims = ('time',) + new_spatial_dims
+            else:
+                dims = new_spatial_dims
+            ds_new[vname] = xr.Variable(dims, da.data, da.attrs, da.encoding, fastpath=True)
+    return ds_new
+
+
+def get_xgcm_grid(ds, **kwargs):
+    """Return an xgcm Grid object
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        An xarray Dataset
+    kwargs:
+       Additional keyword arguments are passed through to `xgcm.Grid` class.
+
+    Returns
+    -------
+    grid : xgcm.Grid
+        An xgcm Grid object
+
+    Examples
+    --------
+    >>> from pop_tools import get_xgcm_grid, DATASETS
+    >>> import xarray as xr
+    >>> fname = DATASETS.fetch("iron_tracer.nc")
+    >>> ds = xr.open_dataset(fname)
+    >>> ds
+    <xarray.Dataset>
+    Dimensions:    (nlat: 384, nlon: 320, time: 24, z_t: 60, z_w_bot: 60, z_w_top: 60)
+    Coordinates:
+    * time       (time) object 0249-02-01 00:00:00 ... 0251-01-01 00:00:00
+    * z_t        (z_t) float32 500.0 1500.0 2500.0 ... 487508.34 512502.8 537500.0
+        TLAT       (nlat, nlon) float64 ...
+        ULONG      (nlat, nlon) float64 ...
+        TLONG      (nlat, nlon) float64 ...
+        ULAT       (nlat, nlon) float64 ...
+    * z_w_top    (z_w_top) float32 0.0 1000.0 2000.0 ... 500004.7 525000.94
+    * z_w_bot    (z_w_bot) float32 1000.0 2000.0 3000.0 ... 525000.94 549999.06
+    Dimensions without coordinates: nlat, nlon
+    Data variables:
+        UE         (time, z_t, nlat, nlon) float32 ...
+        VN         (time, z_t, nlat, nlon) float32 ...
+        WT         (time, z_w_top, nlat, nlon) float32 ...
+        HDIFE      (time, z_t, nlat, nlon) float32 ...
+        HDIFN      (time, z_t, nlat, nlon) float32 ...
+        HDIFB      (time, z_w_bot, nlat, nlon) float32 ...
+        DIA_IMPVF  (time, z_w_bot, nlat, nlon) float32 ...
+        KPP_SRC    (time, z_t, nlat, nlon) float32 ...
+        STF        (time, nlat, nlon) float32 ...
+        SMS        (time, nlat, nlon) float32 ...
+    >>> grid = get_xgcm_grid(ds)
+    >>> grid
+    <xgcm.Grid>
+    Z Axis (periodic):
+        * center   z_t --> left
+        * right    z_w_bot --> center
+        * left     z_w_top --> center
+    Y Axis (periodic):
+        * center   nlat_t --> right
+        * right    nlat_u --> center
+    X Axis (periodic):
+        * center   nlon_t --> right
+        * right    nlon_u --> center
+
+    """
+    import xgcm
+
+    ds = relabel_pop_dims(ds)
+    grid = xgcm.Grid(ds, **kwargs)
+    return grid

requirements.txt

@@ -3,4 +3,5 @@ numba
 pyyaml
 xarray
 dask
-pooch
+pooch>=0.7.0
+xgcm

setup.cfg

@@ -17,7 +17,7 @@ collect_ignore = ['setup.py']
 
 [isort]
 known_first_party=pop_tools
-known_third_party=dask,numba,numpy,pkg_resources,pooch,pytest,setuptools,xarray,yaml
+known_third_party=dask,numba,numpy,pkg_resources,pooch,pytest,setuptools,xarray,xgcm,yaml
 multi_line_output=3
 include_trailing_comma=True
 force_grid_wrap=0

tests/test_xgcm_util.py

@@ -0,0 +1,13 @@
+import xarray as xr
+import xgcm
+
+import pop_tools
+from pop_tools import DATASETS
+
+
+def test_xgcm_grid():
+    fname = DATASETS.fetch('tend_zint_100m_Fe.nc')
+    ds = xr.open_dataset(fname)
+    grid = pop_tools.get_xgcm_grid(ds, metrics=None)
+    assert isinstance(grid, xgcm.Grid)
+    assert set(['X', 'Y', 'Z']) == set(grid.axes.keys())