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?

[dcherian]

I think these may all be out of scope since there is no s_u (for e.g.)

The time-invariant version could be constructed by creating a new ds with zeta = 0 and then running ds.cf.decode_vertical_coords()

[kthyng]

I think these may all be out of scope since there is no s_u (for e.g.)

Ok I will do this in xroms then.

The time-invariant version could be constructed by creating a new ds with zeta = 0 and then running ds.cf.decode_vertical_coords()

Ok I will try this.

But what about the dimensional ordering being off? Is this important for cf-xarray or should I keep this in xroms too?

[stefraynaud]

@dcherian there is an invariable part in the sigma formula. A dedicated class or a cache system like in the example below would be nice to take advantage of it.

https://github.com/VACUMM/xoa/blob/233540257af6932c781585e33b75b173c1949ce7/xoa/sigma.py#L210

As for the transposition problem, I personally chose to push presumed horizontal dimensions to the right by making the them appear first in the computations. Then, the computed vertical variable is transposed to match the first physical variable that have the same dimensions:

https://github.com/VACUMM/xoa/blob/233540257af6932c781585e33b75b173c1949ce7/xoa/sigma.py#L615

[kthyng]

@dcherian Is this in cf-xarray or just the proof of concept code shown in this issue? I'm interested in moving this forward — what would be the next steps? I personally care most about getting typical ocean-related vertical coordinates to work: ROMS and HYCOM, for example.

[dcherian]

I added the ROMS functions: https://cf-xarray.readthedocs.io/en/latest/parametricz.html

PRs for the rest are very welcome!

[kthyng]

Ok, I think the presently-available ROMS decoding might not work when both s_rho and s_w are present in the Dataset — does that sound right? The error is KeyError: 'ocean_s_coordinate' but actually ds.cf.standard_names gives 'ocean_s_coordinate': ['s_rho', 's_w'].

[kthyng]

Never mind, I think I interpreted the error incorrectly. I need to add the info for standard_name = "ocean_s_coordinate"... I'll see what I can do!

[kthyng]

@dcherian more information! So in two ROMS models I have checked, the variable s_rho has standard_name ocean_s_coordinate but actually it should be either the *_g1 or the _*g2 variety that you already have coded in depending on the value of the ROMS variable Vtransform. Would you be ok with a check in cf-xarray to try to align these variables together or would that not work because it is too specific to ROMS (though I would make it so it doesn't break other code)?

[dcherian]

I prefer xroms handle such things and fix the standard name.

[kthyng]

Ok got it.

[kthyng]

For future reference, ocean_sigma_coordinate is now another vertical decoding option in cf-xarray!

[dcherian]

@ocefpaf has code for some of the remaining ocean stuff here: https://github.com/ioos/odvc/blob/main/odvc/formulas.py

[kthyng]

I'm just using this capability again! I am excited to use it with FVCOM in particular but still can't read in those files with xarray. I will write a work-around soon.