diff --git a/cf_xarray/accessor.py b/cf_xarray/accessor.py index 44efd292..f7028998 100644 --- a/cf_xarray/accessor.py +++ b/cf_xarray/accessor.py @@ -2190,12 +2190,15 @@ def bounds_to_vertices( ) return obj - def decode_vertical_coords(self, prefix="z"): + def decode_vertical_coords(self, *, outnames=None, prefix=None): """ Decode parameterized vertical coordinates in place. Parameters ---------- + outnames : dict, optional + Keys of outnames are the input sigma/s coordinate variable name and + the values are the name to use for the associated vertical coordinate. prefix : str, optional Prefix for newly created z variables. E.g. ``s_rho`` becomes ``z_rho`` @@ -2209,7 +2212,8 @@ def decode_vertical_coords(self, prefix="z"): Will only decode when the ``formula_terms`` and ``standard_name`` attributes are set on the parameter (e.g ``s_rho`` ) - Currently only supports ``ocean_s_coordinate_g1`` and ``ocean_s_coordinate_g2``. + Currently only supports ``ocean_s_coordinate_g1``, ``ocean_s_coordinate_g2``, + and ``ocean_sigma_coordinate``. .. warning:: Very lightly tested. Please double check the results. @@ -2223,12 +2227,28 @@ def decode_vertical_coords(self, prefix="z"): requirements = { "ocean_s_coordinate_g1": {"depth_c", "depth", "s", "C", "eta"}, "ocean_s_coordinate_g2": {"depth_c", "depth", "s", "C", "eta"}, + "ocean_sigma_coordinate": {"sigma", "eta", "depth"}, } allterms = self.formula_terms for dim in allterms: - suffix = dim.split("_") - zname = f"{prefix}_" + "_".join(suffix[1:]) + if prefix is None: + assert ( + outnames is not None + ), "if prefix is None, outnames must be provided" + # set outnames here + try: + zname = outnames[dim] + except KeyError: + raise KeyError("Your `outnames` need to include a key of `dim`.") + + else: + warnings.warn( + "`prefix` is being deprecated; use `outnames` instead.", + DeprecationWarning, + ) + suffix = dim.split("_") + zname = f"{prefix}_" + "_".join(suffix[1:]) if "standard_name" not in ds[dim].attrs: continue @@ -2254,8 +2274,9 @@ def decode_vertical_coords(self, prefix="z"): terms["depth_c"] * terms["s"] + (terms["depth"] - terms["depth_c"]) * terms["C"] ) + # z(n,k,j,i) = S(k,j,i) + eta(n,j,i) * (1 + S(k,j,i) / depth(j,i)) - ds.coords[zname] = S + terms["eta"] * (1 + S / terms["depth"]) + ztemp = S + terms["eta"] * (1 + S / terms["depth"]) elif stdname == "ocean_s_coordinate_g2": # make sure all necessary terms are present in terms @@ -2263,14 +2284,21 @@ def decode_vertical_coords(self, prefix="z"): 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) - ds.coords[zname] = terms["eta"] + (terms["eta"] + terms["depth"]) * S + ztemp = terms["eta"] + (terms["eta"] + terms["depth"]) * S + + elif stdname == "ocean_sigma_coordinate": + # z(n,k,j,i) = eta(n,j,i) + sigma(k)*(depth(j,i)+eta(n,j,i)) + ztemp = terms["eta"] + terms["sigma"] * (terms["depth"] + terms["eta"]) else: raise NotImplementedError( f"Coordinate function for {stdname!r} not implemented yet. Contributions welcome!" ) + ds.coords[zname] = ztemp + @xr.register_dataarray_accessor("cf") class CFDataArrayAccessor(CFAccessor): diff --git a/cf_xarray/datasets.py b/cf_xarray/datasets.py index f77c9f29..7e24055e 100644 --- a/cf_xarray/datasets.py +++ b/cf_xarray/datasets.py @@ -16,6 +16,29 @@ ds_no_attrs[variable].attrs = {} +# POM dataset +pomds = xr.Dataset() +pomds["sigma"] = ( + # fmt: off + "sigma", + [-0.983333, -0.95 , -0.916667, -0.883333, -0.85 , -0.816667, + -0.783333, -0.75 , -0.716667, -0.683333, -0.65 , -0.616667, + -0.583333, -0.55 , -0.516667, -0.483333, -0.45 , -0.416667, + -0.383333, -0.35 , -0.316667, -0.283333, -0.25 , -0.216667, + -0.183333, -0.15 , -0.116667, -0.083333, -0.05 , -0.016667], + # fmt: on + { + "units": "sigma_level", + "long_name": "Sigma Stretched Vertical Coordinate at Nodes", + "positive": "down", + "standard_name": "ocean_sigma_coordinate", + "formula_terms": "sigma: sigma eta: zeta depth: depth", + } +) +pomds["depth"] = 175.0 +pomds["zeta"] = ("ocean_time", [-0.155356, -0.127435]) + + popds = xr.Dataset() popds.coords["TLONG"] = ( ("nlat", "nlon"), diff --git a/cf_xarray/tests/test_accessor.py b/cf_xarray/tests/test_accessor.py index a17cf5ba..d41f0343 100644 --- a/cf_xarray/tests/test_accessor.py +++ b/cf_xarray/tests/test_accessor.py @@ -24,6 +24,7 @@ forecast, mollwds, multiple, + pomds, popds, romsds, vert, @@ -1029,32 +1030,47 @@ def test_param_vcoord_ocean_s_coord(): romsds.hc + romsds.h ) expected = romsds.zeta + (romsds.zeta + romsds.h) * Zo_rho - romsds.cf.decode_vertical_coords() + romsds.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) assert_allclose( romsds.z_rho.reset_coords(drop=True), expected.reset_coords(drop=True) ) romsds.s_rho.attrs["standard_name"] = "ocean_s_coordinate_g1" Zo_rho = romsds.hc * (romsds.s_rho - romsds.Cs_r) + romsds.Cs_r * romsds.h + expected = Zo_rho + romsds.zeta * (1 + Zo_rho / romsds.h) - romsds.cf.decode_vertical_coords() + romsds.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) assert_allclose( romsds.z_rho.reset_coords(drop=True), expected.reset_coords(drop=True) ) - romsds.cf.decode_vertical_coords(prefix="ZZZ") + romsds.cf.decode_vertical_coords(outnames={"s_rho": "ZZZ_rho"}) assert "ZZZ_rho" in romsds.coords copy = romsds.copy(deep=True) del copy["zeta"] with pytest.raises(KeyError): - copy.cf.decode_vertical_coords() + copy.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) copy = romsds.copy(deep=True) copy.s_rho.attrs["formula_terms"] = "s: s_rho C: Cs_r depth: h depth_c: hc" with pytest.raises(KeyError): + copy.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) + + +def test_param_vcoord_ocean_sigma_coordinate(): + expected = pomds.zeta + pomds.sigma * (pomds.depth + pomds.zeta) + pomds.cf.decode_vertical_coords(outnames={"sigma": "z"}) + assert_allclose(pomds.z.reset_coords(drop=True), expected.reset_coords(drop=True)) + + copy = pomds.copy(deep=True) + del copy["zeta"] + with pytest.raises(AssertionError): copy.cf.decode_vertical_coords() + with pytest.raises(KeyError): + copy.cf.decode_vertical_coords(outnames={}) + def test_formula_terms(): srhoterms = { diff --git a/doc/parametricz.md b/doc/parametricz.md index 77833a1e..aab8a780 100644 --- a/doc/parametricz.md +++ b/doc/parametricz.md @@ -21,7 +21,7 @@ xr.set_options(display_expand_data=False) # Parametric Vertical Coordinates -`cf_xarray` supports decoding [parametric vertical coordinates](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-vertical-coordinate) encoded in the `formula_terms` attribute using {py:meth}`Dataset.cf.decode_vertical_coords`. Right now, only the two ocean s-coordinates are supported, but support for the [rest](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-v-coord) should be easy to add (Pull Requests are very welcome!). +`cf_xarray` supports decoding [parametric vertical coordinates](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-vertical-coordinate) encoded in the `formula_terms` attribute using {py:meth}`Dataset.cf.decode_vertical_coords`. Right now, only the two ocean s-coordinates and `ocean_sigma_coordinate` are supported, but support for the [rest](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-v-coord) should be easy to add (Pull Requests are very welcome!). ## Decoding parametric coordinates ```{code-cell} @@ -33,7 +33,7 @@ romsds Now we decode the vertical coordinates **in-place**. Note the new `z_rho` variable. `cf_xarray` sees that `s_rho` has a `formula_terms` attribute, looks up the right formula using `s_rho.attrs["standard_name"]` and computes a new vertical coordinate variable. ```{code-cell} -romsds.cf.decode_vertical_coords() # adds new z_rho variable +romsds.cf.decode_vertical_coords(outnames={'s_rho': 'z_rho'}) # adds new z_rho variable romsds.z_rho ``` diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 608cc9a5..f911cb56 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -2,6 +2,12 @@ What's New ---------- + +v 0.7.1 (unreleased) +==================== +- added another type of vertical coordinate to decode: ``ocean_sigma_coordinate``. By `Kristen Thyng`_. + + v0.7.0 (January 24, 2022) ========================= - Many improvements to autoguessing for plotting. By `Deepak Cherian`_