Add .cf.formula_terms (#213)

Author: dcherian

cf_xarray/accessor.py

@@ -1,6 +1,7 @@
 import functools
 import inspect
 import itertools
+import re
 import warnings
 from collections import ChainMap
 from typing import (
@@ -1445,16 +1446,19 @@ def differentiate(
         self, coord, *xr_args, positive_upward: bool = False, **xr_kwargs
     ):
         """
+        Differentiate an xarray object.
+
         Parameters
         ----------
-        xr_args, xr_kwargs are passed directly to the underlying xarray function.
-        The following are added by cf_xarray:
-
         positive_upward: optional, bool
             Change sign of the derivative based on the ``"positive"`` attribute of ``coord``
             so that positive values indicate increasing upward.
             If ``positive=="down"``, then multiplied by -1.
 
+        Notes
+        -----
+        ``xr_args``, ``xr_kwargs`` are passed directly to the underlying xarray function.
+
         See Also
         --------
         DataArray.cf.differentiate
@@ -1515,6 +1519,16 @@ def __getitem__(self, key: Union[str, List[str]]) -> Union[DataArray, Dataset]:
         """
         return _getitem(self, key)
 
+    @property
+    def formula_terms(self) -> Dict[str, Dict[str, str]]:
+        """
+        Property that returns a dictionary
+            {parametric_coord_name: {standard_term_name: variable_name}}
+        """
+        return {
+            dim: self._obj[dim].cf.formula_terms for dim in _get_dims(self._obj, "Z")
+        }
+
     @property
     def bounds(self) -> Dict[str, List[str]]:
         """
@@ -1719,36 +1733,29 @@ def decode_vertical_coords(self, prefix="z"):
         .. warning::
            Very lightly tested. Please double check the results.
         """
-        import re
-
         ds = self._obj
-        dims = _get_dims(ds, "Z")
 
         requirements = {
             "ocean_s_coordinate_g1": {"depth_c", "depth", "s", "C", "eta"},
             "ocean_s_coordinate_g2": {"depth_c", "depth", "s", "C", "eta"},
         }
 
-        for dim in dims:
+        allterms = self.formula_terms
+        for dim in allterms:
             suffix = dim.split("_")
             zname = f"{prefix}_" + "_".join(suffix[1:])
 
-            if (
-                "formula_terms" not in ds[dim].attrs
-                or "standard_name" not in ds[dim].attrs
-            ):
+            if "standard_name" not in ds[dim].attrs:
                 continue
-
-            formula_terms = ds[dim].attrs["formula_terms"]
             stdname = ds[dim].attrs["standard_name"]
 
             # map "standard" formula term names to actual variable names
             terms = {}
-            for mapping in re.sub(": ", ":", formula_terms).split(" "):
-                key, value = mapping.split(":")
+            for key, value in allterms[dim].items():
                 if value not in ds:
                     raise KeyError(
-                        f"Variable {value!r} is required to decode coordinate for {dim} but it is absent in the Dataset."
+                        f"Variable {value!r} is required to decode coordinate for {dim!r}"
+                        " but it is absent in the Dataset."
                     )
                 terms[key] = ds[value]
 
@@ -1776,12 +1783,30 @@ def decode_vertical_coords(self, prefix="z"):
 
             else:
                 raise NotImplementedError(
-                    f"Coordinate function for {stdname} not implemented yet. Contributions welcome!"
+                    f"Coordinate function for {stdname!r} not implemented yet. Contributions welcome!"
                 )
 
 
 @xr.register_dataarray_accessor("cf")
 class CFDataArrayAccessor(CFAccessor):
+    @property
+    def formula_terms(self) -> Dict[str, str]:
+        """
+        Property that returns a dictionary
+            {parametric_coord_name: {standard_term_name: variable_name}}
+        """
+        da = self._obj
+        if "formula_terms" not in da.attrs:
+            var = da[_single(_get_dims)(da, "Z")[0]]
+        else:
+            var = da
+        terms = {}
+        formula_terms = var.attrs.get("formula_terms", "")
+        for mapping in re.sub(r"\s*:\s*", ":", formula_terms).split():
+            key, value = mapping.split(":")
+            terms[key] = value
+        return terms
+
     def __getitem__(self, key: Union[str, List[str]]) -> DataArray:
         """
         Index into a DataArray making use of CF attributes.

cf_xarray/tests/test_accessor.py

@@ -945,6 +945,28 @@ def test_param_vcoord_ocean_s_coord():
         copy.cf.decode_vertical_coords()
 
 
+def test_formula_terms():
+    srhoterms = {
+        "s": "s_rho",
+        "C": "Cs_r",
+        "eta": "zeta",
+        "depth": "h",
+        "depth_c": "hc",
+    }
+    assert romsds.cf.formula_terms == {"s_rho": srhoterms}
+    assert romsds["temp"].cf.formula_terms == srhoterms
+    assert romsds["s_rho"].cf.formula_terms == srhoterms
+
+    s_rho = romsds["s_rho"].copy(deep=True)
+    del s_rho.attrs["standard_name"]
+    del s_rho.s_rho.attrs["standard_name"]  # TODO: xarray bug
+    assert s_rho.cf.formula_terms == srhoterms
+
+    with pytest.raises(KeyError):
+        # x,y,t variable
+        romsds["zeta"].cf.formula_terms
+
+
 def test_standard_name_mapper():
     da = xr.DataArray(
         np.arange(6),

doc/api.rst

@@ -30,6 +30,7 @@ Attributes
     DataArray.cf.axes
     DataArray.cf.cell_measures
     DataArray.cf.coordinates
+    DataArray.cf.formula_terms
     DataArray.cf.standard_names
     DataArray.cf.plot
 
@@ -66,6 +67,7 @@ Attributes
     Dataset.cf.bounds
     Dataset.cf.cell_measures
     Dataset.cf.coordinates
+    Dataset.cf.formula_terms
     Dataset.cf.standard_names
 
 .. _dsmeth:

doc/whats-new.rst

@@ -6,6 +6,8 @@ What's New
 v0.5.2 (unreleased)
 ===================
 
+- Added :py:attr:`DataArray.cf.formula_terms` and :py:attr:`Dataset.cf.formula_terms`.
+  By `Deepak Cherian`_.
 - Added :py:attr:`Dataset.cf.bounds` to return a dictionary mapping valid keys to the variable names of their bounds. By `Mattia Almansi`_.
 - :py:meth:`DataArray.cf.differentiate` and :py:meth:`Dataset.cf.differentiate` can optionally correct
   sign of the derivative by interpreting the ``"positive"`` attribute. By `Deepak Cherian`_.