Increased support for cf_role (#327)

1. Add cf_roles attribute
2. Add to .cf.keys()
3. Add to repr

Closes #305

Author: dcherian

cf_xarray/accessor.py

@@ -1239,7 +1239,7 @@ def __repr__(self):
         coords = self._obj.coords
         dims = self._obj.dims
 
-        def make_text_section(subtitle, attr, valid_values, default_keys=None):
+        def make_text_section(subtitle, attr, valid_values=None, default_keys=None):
 
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore")
@@ -1260,11 +1260,12 @@ def make_text_section(subtitle, attr, valid_values, default_keys=None):
             vardict = {key: vardict[key] for key in ordered_keys if key in vardict}
 
             # Keep only valid values (e.g., coords or data_vars)
-            vardict = {
-                key: set(value).intersection(valid_values)
-                for key, value in vardict.items()
-                if set(value).intersection(valid_values)
-            }
+            if valid_values is not None:
+                vardict = {
+                    key: set(value).intersection(valid_values)
+                    for key, value in vardict.items()
+                    if set(value).intersection(valid_values)
+                }
 
             # Star for keys with dims only, tab otherwise
             rows = [
@@ -1293,6 +1294,11 @@ def make_text_section(subtitle, attr, valid_values, default_keys=None):
             text = f"CF Flag variable with mapping:\n\t{flag_dict!r}\n\n"
         else:
             text = ""
+
+        if self.cf_roles:
+            text += make_text_section("CF Roles", "cf_roles")
+            text += "\n"
+
         text += "Coordinates:"
         text += make_text_section("CF Axes", "axes", coords, _AXIS_NAMES)
         text += make_text_section("CF Coordinates", "coordinates", coords, _COORD_NAMES)
@@ -1337,6 +1343,7 @@ def keys(self) -> set[str]:
         varnames = list(self.axes) + list(self.coordinates)
         varnames.extend(list(self.cell_measures))
         varnames.extend(list(self.standard_names))
+        varnames.extend(list(self.cf_roles))
 
         return set(varnames)
 
@@ -1461,6 +1468,33 @@ def standard_names(self) -> dict[str, list[str]]:
 
         return {k: sorted(v) for k, v in vardict.items()}
 
+    @property
+    def cf_roles(self) -> dict[str, list[str]]:
+        """
+        Returns a dictionary mapping cf_role names to variable names.
+
+        Returns
+        -------
+        dict
+            Dictionary mapping cf_role names to variable names.
+
+        References
+        ----------
+        Please refer to the CF conventions document : http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#coordinates-metadata
+        """
+        if isinstance(self._obj, Dataset):
+            variables = self._obj.variables
+        elif isinstance(self._obj, DataArray):
+            variables = self._obj.coords
+
+        vardict: dict[str, list[str]] = {}
+        for k, v in variables.items():
+            if "cf_role" in v.attrs:
+                role = v.attrs["cf_role"]
+                vardict[role] = vardict.setdefault(role, []) + [k]
+
+        return {k: sorted(v) for k, v in vardict.items()}
+
     def get_associated_variable_names(
         self, name: Hashable, skip_bounds: bool = False, error: bool = True
     ) -> dict[str, list[str]]:

cf_xarray/datasets.py

@@ -483,3 +483,12 @@
         },
     }
 )
+
+
+dsg = xr.Dataset(
+    {"foo": (("trajectory", "profile"), [[1, 2, 3], [1, 2, 3]])},
+    coords={
+        "profile": ("profile", [0, 1, 2], {"cf_role": "profile_id"}),
+        "trajectory": ("trajectory", [0, 1], {"cf_role": "trajectory_id"}),
+    },
+)

cf_xarray/tests/test_accessor.py

@@ -21,6 +21,7 @@
     anc,
     basin,
     ds_no_attrs,
+    dsg,
     forecast,
     mollwds,
     multiple,
@@ -168,6 +169,32 @@ def test_repr():
     """
     assert actual == dedent(expected)
 
+    # CF roles
+    actual = dsg.cf.__repr__()
+    expected = """
+    - CF Roles: * profile_id: ['profile']
+                * trajectory_id: ['trajectory']
+
+    Coordinates:
+    - CF Axes:   X, Y, Z, T: n/a
+
+    - CF Coordinates:   longitude, latitude, vertical, time: n/a
+
+    - Cell Measures:   area, volume: n/a
+
+    - Standard Names:   n/a
+
+    - Bounds:   n/a
+
+    Data Variables:
+    - Cell Measures:   area, volume: n/a
+
+    - Standard Names:   n/a
+
+    - Bounds:   n/a
+    """
+    assert actual == dedent(expected)
+
 
 def test_axes():
     expected = dict(T=["time"], X=["lon"], Y=["lat"])
@@ -1579,3 +1606,16 @@ def test_pickle():
     ds = da.to_dataset()
     pickle.loads(pickle.dumps(da.cf))
     pickle.loads(pickle.dumps(ds.cf))
+
+
+def test_cf_role():
+    for name in ["profile_id", "trajectory_id"]:
+        assert name in dsg.cf.keys()
+
+    assert dsg.cf.cf_roles == {
+        "profile_id": ["profile"],
+        "trajectory_id": ["trajectory"],
+    }
+
+    dsg.foo.cf.plot(x="profile_id")
+    dsg.foo.cf.plot(x="trajectory_id")

doc/api.rst

@@ -33,6 +33,7 @@ Attributes
 
     DataArray.cf.axes
     DataArray.cf.cell_measures
+    DataArray.cf.cf_roles
     DataArray.cf.coordinates
     DataArray.cf.formula_terms
     DataArray.cf.is_flag_variable
@@ -90,6 +91,7 @@ Attributes
     Dataset.cf.axes
     Dataset.cf.bounds
     Dataset.cf.cell_measures
+    Dataset.cf.cf_roles
     Dataset.cf.coordinates
     Dataset.cf.formula_terms
     Dataset.cf.standard_names

doc/dsg.md

@@ -30,5 +30,17 @@ ds = xr.Dataset(
     {"temp": ("x", np.arange(10))},
     coords={"cast": ("x", np.arange(10), {"cf_role": "profile_id"})}
 )
+ds.cf
+```
+
+Access `"cast"` using it's `cf_role`
+
+```{code-cell}
 ds.cf["profile_id"]
 ```
+
+Find all `cf_role` variables using {py:attr}`Dataset.cf.cf_roles` and {py:attr}`DataArray.cf.cf_roles`
+
+```{code-cell}
+ds.cf.cf_roles
+```

doc/whats-new.rst

@@ -5,6 +5,7 @@ What's New
 
 v0.7.3 (unreleased)
 ===================
+- Increased support for ``cf_role`` variables. Added :py:attr:`Dataset.cf.cf_roles` By `Deepak Cherian`_.
 
 v0.7.2 (April 5, 2022)
 ======================