Adds support for parametric vertical coordinate (#528)

* Adds initial parametric functions

* Updates parametric module and adds tests

* Adds function to help transpose outputs

* Removes hardcoded dims

* Fixes optional argument bug and handling case insensitive terms

* Moves from function to class based implementation of transforms

* Fixes failing test

* Removes extras spaces in text

* Fixes formatting

* Resolves mypy errors

* Removes redundant code

* Fixes global variable case

* Updates _derive_ocean_stdname

* Moves to dataclass

* Fixes passing variables to class constructor

* Removes redundant code

* Fixes handling unknown standard names

Co-authored-by: Deepak Cherian <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Removes squeeze and fixes constants

* Adds entry to CITATIONS.cff

* cleanup

* tweak docs

---------

Co-authored-by: Deepak Cherian <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Deepak Cherian <[email protected]>

Author: 3 people

CITATION.cff

@@ -83,6 +83,10 @@ authors:
   - family-names: Haëck
     given-names: Clément
     affiliation: Laboratoire d'Océanographie et du Climat (LOCEAN), Paris
+  - family-names: Boutte
+    given-names: Jason
+    orcid: 'https://orcid.org/0009-0009-3996-3772'
+    affiliation: Lawrence Livermore National Laboratory
 identifiers:
   - type: doi
     value: 10.5281/zenodo.4749735

cf_xarray/accessor.py

@@ -30,7 +30,7 @@
 from xarray.core.rolling import Coarsen, Rolling
 from xarray.core.weighted import Weighted
 
-from . import sgrid
+from . import parametric, sgrid
 from .criteria import (
     _DSG_ROLES,
     _GEOMETRY_TYPES,
@@ -2754,13 +2754,8 @@ def decode_vertical_coords(self, *, outnames=None, prefix=None):
         """
         ds = self._obj
 
-        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:
             if prefix is None:
                 assert (
@@ -2782,6 +2777,7 @@ def decode_vertical_coords(self, *, outnames=None, prefix=None):
                 suffix = dim.split("_")
                 zname = f"{prefix}_" + "_".join(suffix[1:])
 
+            # never touched, if standard name is missing it's not included in allterms
             if "standard_name" not in ds[dim].attrs:
                 continue
             stdname = ds[dim].attrs["standard_name"]
@@ -2790,46 +2786,23 @@ def decode_vertical_coords(self, *, outnames=None, prefix=None):
             terms = {}
             for key, value in allterms[dim].items():
                 if value not in ds:
+                    # is this ever hit, if variable is missing it's missing in decoded allterms
                     raise KeyError(
                         f"Variable {value!r} is required to decode coordinate for {dim!r}"
                         " but it is absent in the Dataset."
                     )
-                terms[key] = ds[value]
-
-            absent_terms = requirements[stdname] - set(terms)
-            if absent_terms:
-                raise KeyError(f"Required terms {absent_terms} absent in dataset.")
-
-            if stdname == "ocean_s_coordinate_g1":
-                # S(k,j,i) = depth_c * s(k) + (depth(j,i) - depth_c) * C(k)
-                S = (
-                    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))
-                ztemp = S + terms["eta"] * (1 + S / terms["depth"])
+                # keys should be case insensitive
+                terms[key.lower()] = ds[value]
 
-            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)
-                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:
+            try:
+                transform = parametric.TRANSFORM_FROM_STDNAME[stdname]
+            except KeyError:
+                # Should occur since stdname is check before
                 raise NotImplementedError(
-                    f"Coordinate function for {stdname!r} not implemented yet. Contributions welcome!"
-                )
+                    f"Coordinate function for {stdname!r} not implmented yet. Contributions welcome!"
+                ) from None
 
-            ds.coords[zname] = ztemp
+            ds.coords[zname] = transform.from_terms(terms)
 
 
 @xr.register_dataarray_accessor("cf")