Add sgrid axes parsing (#421)

* Add sgrid axes parsing

* Update cf_xarray/sgrid.py

* Fix repr

* Add more sgrid examples

* Fix ROMS detection

* Propagate the grid attribute too.

* Fix.

* Fix repr

* Cleanup formatting.py

* Fix repr

* Add docs.

* Update docs.

* Add whats-new

* last edits

Author: dcherian

cf_xarray/__init__.py

@@ -1,3 +1,4 @@
+from . import sgrid  # noqa
 from .accessor import CFAccessor  # noqa
 from .coding import (  # noqa
     decode_compress_to_multi_index,

cf_xarray/accessor.py

@@ -30,7 +30,9 @@
 from xarray.core.rolling import Coarsen, Rolling
 from xarray.core.weighted import Weighted
 
+from . import sgrid
 from .criteria import (
+    _DSG_ROLES,
     cf_role_criteria,
     coordinate_criteria,
     grid_mapping_var_criteria,
@@ -40,8 +42,9 @@
     _format_coordinates,
     # _format_conventions,
     _format_data_vars,
+    _format_dsg_roles,
     _format_flags,
-    _format_roles,
+    _format_sgrid,
     _maybe_panel,
 )
 from .helpers import _guess_bounds_1d, _guess_bounds_2d, bounds_to_vertices
@@ -313,6 +316,11 @@ def _get_axis_coord(obj: DataArray | Dataset, key: str) -> list[str]:
                     units = getattr(var.data, "units", None)
                     if units in expected:
                         results.update((coord,))
+
+    if key in _AXIS_NAMES and "grid_topology" in obj.cf.cf_roles:
+        sgrid_axes = sgrid.parse_axes(obj)
+        results.update((search_in | set(obj.dims)) & sgrid_axes[key])
+
     return list(results)
 
 
@@ -474,7 +482,7 @@ def _get_coords(obj: DataArray | Dataset, key: Hashable) -> list[Hashable]:
     One or more of ('X', 'Y', 'Z', 'T', 'longitude', 'latitude', 'vertical', 'time',
     'area', 'volume'), or arbitrary measures, or standard names present in .coords
     """
-    return [k for k in _get_all(obj, key) if k in obj.coords]
+    return [k for k in _get_all(obj, key) if k in obj.coords or k in obj.dims]
 
 
 def _variables(func: F) -> F:
@@ -1435,12 +1443,22 @@ def _generate_repr(self, rich=False):
                 _format_flags(self, rich), title="Flag Variable", rich=rich
             )
 
-        if self.cf_roles:
-            yield _maybe_panel(
-                _format_roles(self, dims, rich),
-                title="Discrete Sampling Geometry",
-                rich=rich,
-            )
+        roles = self.cf_roles
+        if roles:
+            if any(role in roles for role in _DSG_ROLES):
+                yield _maybe_panel(
+                    _format_dsg_roles(self, dims, rich),
+                    title="Discrete Sampling Geometry",
+                    rich=rich,
+                )
+
+            if "grid_topology" in self.cf_roles:
+                axes = sgrid.parse_axes(self._obj)
+                yield _maybe_panel(
+                    _format_sgrid(self, axes, rich),
+                    title="SGRID",
+                    rich=rich,
+                )
 
         yield _maybe_panel(
             _format_coordinates(self, dims, coords, rich),
@@ -1642,6 +1660,7 @@ def get_associated_variable_names(
             3. "cell_measures"
             4. "coordinates"
             5. "grid_mapping"
+            6. "grid"
         to a list of variable names referred to in the appropriate attribute
 
         Parameters
@@ -1654,15 +1673,18 @@ def get_associated_variable_names(
         Returns
         -------
         names : dict
-            Dictionary with keys "ancillary_variables", "cell_measures", "coordinates", "bounds".
+            Dictionary with keys "ancillary_variables", "cell_measures", "coordinates", "bounds",
+            "grid_mapping", "grid".
         """
         keys = [
             "ancillary_variables",
             "cell_measures",
             "coordinates",
             "bounds",
             "grid_mapping",
+            "grid",
         ]
+
         coords: dict[str, list[Hashable]] = {k: [] for k in keys}
         attrs_or_encoding = ChainMap(self._obj[name].attrs, self._obj[name].encoding)
 
@@ -1704,6 +1726,9 @@ def get_associated_variable_names(
                 if dbounds:
                     coords["bounds"].append(dbounds)
 
+        if "grid" in attrs_or_encoding:
+            coords["grid"] = [attrs_or_encoding["grid"]]
+
         if "grid_mapping" in attrs_or_encoding:
             coords["grid_mapping"] = [attrs_or_encoding["grid_mapping"]]
 

cf_xarray/criteria.py

@@ -11,13 +11,13 @@
 
 from typing import Any, Mapping, MutableMapping, Tuple
 
+_DSG_ROLES = ["timeseries_id", "profile_id", "trajectory_id"]
+
 cf_role_criteria: Mapping[str, Mapping[str, str]] = {
     k: {"cf_role": k}
     for k in (
         # CF Discrete sampling geometry
-        "timeseries_id",
-        "profile_id",
-        "trajectory_id",
+        *_DSG_ROLES,
         # SGRID
         "grid_topology",
         # UGRID

cf_xarray/datasets.py

@@ -655,3 +655,48 @@ def _create_inexact_bounds():
         "trajectory": ("trajectory", [0, 1], {"cf_role": "trajectory_id"}),
     },
 )
+
+
+sgrid_roms = xr.Dataset()
+sgrid_roms["grid"] = xr.DataArray(
+    0,
+    attrs=dict(
+        cf_role="grid_topology",
+        topology_dimension=2,
+        node_dimensions="xi_psi eta_psi",
+        face_dimensions="xi_rho: xi_psi (padding: both) eta_rho: eta_psi (padding: both)",
+        edge1_dimensions="xi_u: xi_psi eta_u: eta_psi (padding: both)",
+        edge2_dimensions="xi_v: xi_psi (padding: both) eta_v: eta_psi",
+        node_coordinates="lon_psi lat_psi",
+        face_coordinates="lon_rho lat_rho",
+        edge1_coordinates="lon_u lat_u",
+        edge2_coordinates="lon_v lat_v",
+        vertical_dimensions="s_rho: s_w (padding: none)",
+    ),
+)
+sgrid_roms["u"] = (("xi_u", "eta_u"), np.ones((2, 2)), {"grid": "grid"})
+
+sgrid_delft = xr.Dataset()
+sgrid_delft["grid"] = xr.DataArray(
+    0,
+    attrs=dict(
+        cf_role="grid_topology",
+        topology_dimension=2,
+        node_dimensions="inode jnode",
+        face_dimensions="icell: inode (padding: none) jcell: jnode (padding: none)",
+        node_coordinates="node_lon node_lat",
+    ),
+)
+
+
+sgrid_delft3 = xr.Dataset()
+sgrid_delft3["grid"] = xr.DataArray(
+    0,
+    attrs=dict(
+        cf_role="grid_topology",
+        topology_dimension=3,
+        node_dimensions="inode jnode knode",
+        volume_dimensions="iface: inode (padding: none) jface: jnode (padding: none) kface: knode (padding: none)",
+        node_coordinates="node_lon node_lat node_elevation",
+    ),
+)

cf_xarray/formatting.py

@@ -1,4 +1,5 @@
 import warnings
+from functools import partial
 from typing import Dict, Hashable, Iterable, List
 
 STAR = " * "
@@ -35,6 +36,7 @@ def make_text_section(
     subtitle: str,
     attr: str,
     dims=None,
+    valid_keys=None,
     valid_values=None,
     default_keys=None,
     rich: bool = False,
@@ -46,10 +48,16 @@ def make_text_section(
         dims = []
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        try:
-            vardict: Dict[str, Iterable[Hashable]] = getattr(accessor, attr, {})
-        except ValueError:
-            vardict = {}
+        if isinstance(attr, str):
+            try:
+                vardict: Dict[str, Iterable[Hashable]] = getattr(accessor, attr, {})
+            except ValueError:
+                vardict = {}
+        else:
+            assert isinstance(attr, dict)
+            vardict = attr
+    if valid_keys:
+        vardict = {k: v for k, v in vardict.items() if k in valid_keys}
 
     # Sort keys if there aren't extra keys,
     # preserve default keys order otherwise.
@@ -142,53 +150,72 @@ def _format_flags(accessor, rich):
     return _print_rows("Flag Meanings", rows, rich)
 
 
-def _format_roles(accessor, dims, rich):
-    yield make_text_section(accessor, "CF Roles", "cf_roles", dims=dims, rich=rich)
+def _format_dsg_roles(accessor, dims, rich):
+    from .criteria import _DSG_ROLES
+
+    yield make_text_section(
+        accessor,
+        "CF Roles",
+        "cf_roles",
+        dims=dims,
+        valid_keys=_DSG_ROLES,
+        rich=rich,
+    )
 
 
 def _format_coordinates(accessor, dims, coords, rich):
     from .accessor import _AXIS_NAMES, _CELL_MEASURES, _COORD_NAMES
 
-    yield make_text_section(
-        accessor, "CF Axes", "axes", dims, coords, _AXIS_NAMES, rich=rich
+    section = partial(
+        make_text_section, accessor=accessor, dims=dims, valid_values=coords, rich=rich
     )
-    yield make_text_section(
-        accessor, "CF Coordinates", "coordinates", dims, coords, _COORD_NAMES, rich=rich
-    )
-    yield make_text_section(
-        accessor,
-        "Cell Measures",
-        "cell_measures",
-        dims,
-        coords,
-        _CELL_MEASURES,
-        rich=rich,
-    )
-    yield make_text_section(
-        accessor, "Standard Names", "standard_names", dims, coords, rich=rich
+
+    yield section(subtitle="CF Axes", attr="axes", default_keys=_AXIS_NAMES)
+    yield section(
+        subtitle="CF Coordinates", attr="coordinates", default_keys=_COORD_NAMES
     )
-    yield make_text_section(accessor, "Bounds", "bounds", dims, coords, rich=rich)
-    yield make_text_section(
-        accessor, "Grid Mappings", "grid_mapping_names", dims, coords, rich=rich
+    yield section(
+        subtitle="Cell Measures", attr="cell_measures", default_keys=_CELL_MEASURES
     )
+    yield section(subtitle="Standard Names", attr="standard_names")
+    yield section(subtitle="Bounds", attr="bounds")
+    yield section(subtitle="Grid Mappings", attr="grid_mapping_names")
 
 
 def _format_data_vars(accessor, data_vars, rich):
     from .accessor import _CELL_MEASURES
 
-    yield make_text_section(
-        accessor,
-        "Cell Measures",
-        "cell_measures",
-        None,
-        data_vars,
-        _CELL_MEASURES,
+    section = partial(
+        make_text_section,
+        accessor=accessor,
+        dims=None,
+        valid_values=data_vars,
         rich=rich,
     )
+
+    yield section(
+        subtitle="Cell Measures", attr="cell_measures", default_keys=_CELL_MEASURES
+    )
+    yield section(subtitle="Standard Names", attr="standard_names")
+    yield section(subtitle="Bounds", attr="bounds")
+    yield section(subtitle="Grid Mappings", attr="grid_mapping_names")
+
+
+def _format_sgrid(accessor, axes, rich):
     yield make_text_section(
-        accessor, "Standard Names", "standard_names", None, data_vars, rich=rich
+        accessor,
+        "CF role",
+        "cf_roles",
+        valid_keys=["grid_topology"],
+        rich=rich,
     )
-    yield make_text_section(accessor, "Bounds", "bounds", None, data_vars, rich=rich)
+
     yield make_text_section(
-        accessor, "Grid Mappings", "grid_mapping_names", None, data_vars, rich=rich
+        accessor,
+        "Axes",
+        axes,
+        accessor._obj.dims,
+        valid_values=accessor._obj.dims,
+        default_keys=axes.keys(),
+        rich=rich,
     )

cf_xarray/sgrid.py

@@ -0,0 +1,41 @@
+SGRID_DIM_ATTRS = [
+    "face_dimensions",
+    "volume_dimensions",
+    # the following are optional and should be redundant with the above
+    # at least for dimension names
+    # "face1_dimensions",
+    # "face2_dimensions",
+    # "face3_dimensions",
+    "edge1_dimensions",
+    "edge2_dimensions",
+    # "edge3_dimensions",
+]
+
+
+def parse_axes(ds):
+    import re
+
+    (gridvar,) = ds.cf.cf_roles["grid_topology"]
+    grid = ds[gridvar]
+    pattern = re.compile("\\s?(.*?):\\s*(.*?)\\s+(?:\\(padding:(.+?)\\))?")
+    ndim = grid.attrs["topology_dimension"]
+    axes_names = ["X", "Y", "Z"][:ndim]
+    axes = dict(
+        zip(
+            axes_names,
+            ({k} for k in grid.attrs["node_dimensions"].split(" ")),
+        )
+    )
+    for attr in SGRID_DIM_ATTRS:
+        if attr in grid.attrs:
+            matches = re.findall(pattern, grid.attrs[attr] + "\n")
+            assert len(matches) == ndim, matches
+            for ax, match in zip(axes_names, matches):
+                axes[ax].update(set(match[:2]))
+
+    if ndim == 2 and "vertical_dimensions" in grid.attrs:
+        matches = re.findall(pattern, grid.attrs["vertical_dimensions"] + "\n")
+        assert len(matches) == 1
+        axes["Z"] = set(matches[0][:2])
+
+    return axes