Add parametric vertical coordinate generation functions

[dcherian]

Allow calculating parameterized vertical coordinates using ds.cf.decode_aux_coords() -> Dataset

CF: http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-v-coord
Iris: https://scitools.org.uk/iris/docs/latest/iris/iris/aux_factory.html

[dcherian]

I was curious so I got this working for ROMS following http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#_ocean_s_coordinate_generic_form_2

import re

dim = "s_rho"
zname = "z_rho"
ds = xr.tutorial.open_dataset('ROMS_example.nc', chunks={'ocean_time': 1})

formula_terms = ds[dim].attrs["formula_terms"]
stdname = ds[dim].attrs["standard_name"]

terms = {}
for mapping in re.sub(": ", ":", formula_terms).split(" "):
    key, value = mapping.split(":")
    # raise nice error here if value is not in ds
    terms[key] = ds[value]

if stdname == "ocean_s_coordinate_g2":
    # make sure all necessary terms are present in terms
    # (depth_c * s(k) + depth(j,i) * C(k)) / (depth_c + depth(j,i))
    S = (terms["depth_c"] * terms["s"] + terms["depth"] * terms["C"]) / (
        terms["depth_c"] + terms["depth"]
    )
    # z(n,k,j,i) = eta(n,j,i) + (eta(n,j,i) + depth(j,i)) * S(k,j,i)
    z = terms["eta"] + (terms["eta"] + terms["depth"]) * S

Zo_rho = (ds.hc * ds.s_rho + ds.Cs_r * ds.h) / (ds.hc + ds.h)
z_rho = ds.zeta + (ds.zeta + ds.h) * Zo_rho

xr.testing.assert_equal(z_rho, z)

cc @kthyng: it could be nice to implement this in cf_xarray and the xroms can just use cf_xarray to get proper values?

[kthyng]

@hetland has a history of caring about this, let's see what he thinks.

There are multiple mappings depending on the vertical set up for a simulation.

[hetland]

What you propose should work. The coordinates on the CF page you link to correspond to the two transforms listed in the ROMS Wiki. I think it makes sense for xroms to farm this out to cf-xarray, since it will be more generic.

Xarray makes working with the equations (which have different numbers of dimensions) so much easier than numpy, so as your example shows, you can just write the equation for the vertical coordinate, and xarray takes care of the broadcasting automatically. So I also think this approach will be scalable, and result in easy to maintain code.

I replaced my roms vertical coordinate code using numpy, that was almost 500 lines of code, with six lines of code using xarray. (It would be just two if I hardwired in which transformation to use).

@rsignell-usgs should also see this. He helped put the ROMS vertical coordinates into the CF-convention.

[rsignell-usgs]

Ideally we would read the CF conventions standard_name and parse the formula_terms to determine how to calculate z, for example: http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#_ocean_s_coordinate_generic_form_1

[dcherian]

yup that's what I am proposing

[dcherian]

The ocean s coordinates have been added. Leaving this issue open for the other coordinates in http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-v-coord

[kthyng]

I just compared z_rho and z_w with the code in xroms for Vtransform=2 and the values matched this code. 🎉

But, their dimensional ordering did not. I have the following in my version to reorder to convention, the type of thing we'd discussed being in cf-xarray itself (#82):

var.cf.transpose( *[dim for dim in ["T", "Z", "Y", "X"] if dim in var.reset_coords(drop=True).cf.axes] )

We have also been using these metrics on other grids, so not just z_rho and z_w, but also z_rho_u for z_rho on the u grid, etc. These are used to set up the xgcm grid object. Are you interested in including these in this function?

I additionally have added all the same but with "0" on the end indicating that they do not change in time. In the equation, zeta is set to 0. This is helpful to reduce computation when using z coordinates but not needing extreme accuracy. Are you interested in these?