new vertical coord decode and reordering (#315)

* new vertical coord decode and reordering The vertical decoding has been added for `ocean_sigma_coordinate` and the free surface variable has been expanded dimensionally for all the calculations so that the results preserve the standard ordering of time, vertical, lat/Y, lon/X. Three of the tests do not work anymore, though. Expanding dimensionally apparently bumps the attributes off of s_rho, for one. For another, the values are very similar but do not match anymore — not sure why. Interested to see what @dcherian thinks! * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update doc/whats-new.rst Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com> * Update cf_xarray/accessor.py Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com> * preferentially use outnames over prefix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * added in transpose statement that does not work yet * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * removed reordering stuff that did not work * provide a default name now * I think the tests are back how they were * trying to follow dcherian suggestions * small precommit change * Update cf_xarray/tests/test_accessor.py Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com> * small fix from review * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update cf_xarray/accessor.py * Update cf_xarray/tests/test_accessor.py * fixed tests I think Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>
xarray-contrib · Mar 29, 2022 · 228ee4a · 228ee4a
1 parent 0908033
commit 228ee4a
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 12 deletions.
diff --git 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,23 +2274,31 @@ 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
                 # (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)
-                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
@@ -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
@@ -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
@@ -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
@@ -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`_