Support custom vocabularies/criteria (#234)

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

Author: kthyng

cf_xarray/accessor.py

@@ -65,6 +65,13 @@
 ATTRS["time"] = ATTRS["T"]
 ATTRS["vertical"] = ATTRS["Z"]
 
+OPTIONS: MutableMapping[str, Any] = {"custom_criteria": []}
+
+
+def set_options(custom_criteria):
+    OPTIONS["custom_criteria"] = always_iterable(custom_criteria, allowed=(tuple, list))
+
+
 # Type for Mapper functions
 Mapper = Callable[[Union[DataArray, Dataset], str], List[str]]
 
@@ -170,6 +177,55 @@ def _get_groupby_time_accessor(var: Union[DataArray, Dataset], key: str) -> List
         return []
 
 
+def _get_custom_criteria(
+    obj: Union[DataArray, Dataset], key: str, criteria=None
+) -> List[str]:
+    """
+    Translate from axis, coord, or custom name to variable name optionally
+    using ``custom_criteria``
+
+    Parameters
+    ----------
+    obj : DataArray, Dataset
+    key : str
+        key to check for.
+    criteria : dict, optional
+        Criteria to use to map from variable to attributes describing the
+        variable. An example is coordinate_criteria which maps coordinates to
+        their attributes and attribute values. If user has defined
+        custom_criteria, this will be used by default.
+
+    Returns
+    -------
+    List[str], Variable name(s) in parent xarray object that matches axis, coordinate, or custom `key`
+
+    """
+
+    if isinstance(obj, DataArray):
+        obj = obj._to_temp_dataset()
+
+    if criteria is None:
+        if not OPTIONS["custom_criteria"]:
+            return []
+        criteria = OPTIONS["custom_criteria"]
+
+    if criteria is not None:
+        criteria = always_iterable(criteria, allowed=(tuple, list, set))
+
+    criteria = ChainMap(*criteria)
+
+    results: Set = set()
+    if key in criteria:
+        for criterion, patterns in criteria[key].items():
+            for var in obj.variables:
+                if re.match(patterns, obj[var].attrs.get(criterion, "")):
+                    results.update((var,))
+                # also check name specifically since not in attributes
+                elif criterion == "name" and re.match(patterns, var):
+                    results.update((var,))
+    return list(results)
+
+
 def _get_axis_coord(var: Union[DataArray, Dataset], key: str) -> List[str]:
     """
     Translate from axis or coord name to variable name
@@ -314,7 +370,12 @@ def _get_all(obj: Union[DataArray, Dataset], key: str) -> List[str]:
     One or more of ('X', 'Y', 'Z', 'T', 'longitude', 'latitude', 'vertical', 'time',
     'area', 'volume'), or arbitrary measures, or standard names
     """
-    all_mappers = (_get_axis_coord, _get_measure, _get_with_standard_name)
+    all_mappers = (
+        _get_custom_criteria,
+        _get_axis_coord,
+        _get_measure,
+        _get_with_standard_name,
+    )
     results = apply_mapper(all_mappers, obj, key, error=False, default=None)
     return results
 
@@ -586,6 +647,8 @@ def check_results(names, key):
         measures = []
         warnings.warn("Ignoring bad cell_measures attribute.", UserWarning)
 
+    custom_criteria = ChainMap(*OPTIONS["custom_criteria"])
+
     varnames: List[Hashable] = []
     coords: List[Hashable] = []
     successful = dict.fromkeys(key, False)
@@ -602,6 +665,11 @@ def check_results(names, key):
             successful[k] = bool(measure)
             if measure:
                 varnames.extend(measure)
+        elif k in custom_criteria:
+            names = _get_all(obj, k)
+            check_results(names, k)
+            successful[k] = bool(names)
+            varnames.extend(names)
         else:
             stdnames = set(_get_with_standard_name(obj, k))
             objcoords = set(obj.coords)

cf_xarray/tests/test_accessor.py

@@ -1237,6 +1237,82 @@ def test_cmip6_attrs():
     assert da.cf.axes["Y"] == ["nlat"]
 
 
+def test_custom_criteria():
+    my_custom_criteria = {
+        "ssh": {
+            "standard_name": "sea_surface_elev*|sea_surface_height",
+            "name": "sea_surface_elevation$",  # variable name
+        },
+        "salt": {
+            "standard_name": "salinity",
+            "name": "sal*",
+        },
+        "wind_speed": {
+            "standard_name": "wind_speed$",
+        },
+    }
+    my_custom_criteria2 = {"temp": {"name": "temperature"}}
+    cf_xarray.accessor.set_options(my_custom_criteria)
+    my_custom_criteria_list = [my_custom_criteria, my_custom_criteria2]
+    my_custom_criteria_tuple = (my_custom_criteria, my_custom_criteria2)
+
+    # Match by name regex match
+    ds = xr.Dataset()
+    ds["salinity"] = ("dim", np.arange(10))
+    assert_identical(ds.cf["salt"], ds["salinity"])
+
+    # Match by standard_name regex match
+    ds = xr.Dataset()
+    ds["elev"] = ("dim", np.arange(10), {"standard_name": "sea_surface_elevBLAH"})
+    assert_identical(ds.cf["ssh"], ds["elev"])
+
+    # Match by standard_name exact match
+    ds = xr.Dataset()
+    ds["salinity"] = ("dim", np.arange(10), {"standard_name": "salinity"})
+    assert_identical(ds.cf["salt"], ds["salinity"])
+
+    # If not exact name, won't match
+    ds = xr.Dataset()
+    ds["sea_surface_elevation123"] = ("dim", np.arange(10))
+    # Since this will not match, this should error
+    with pytest.raises(KeyError):
+        ds.cf["ssh"]
+
+    # will select only one variable here since exact match
+    ds = xr.Dataset()
+    ds["winds"] = ("dim", np.arange(10), {"standard_name": "wind_speed"})
+    ds["gusts"] = ("dim", np.arange(10), {"standard_name": "wind_speed_of_gust"})
+    assert_identical(ds.cf["wind_speed"], ds["winds"])
+
+    # Match by exact name
+    ds = xr.Dataset()
+    ds["sea_surface_elevation"] = ("dim", np.arange(10))
+    ds["sea_surface_height"] = (
+        "dim",
+        np.arange(10),
+        {"standard_name": "sea_surface_elevBLAH"},
+    )
+    # Since there are two variables, this should error
+    with pytest.raises(KeyError):
+        ds.cf["ssh"]
+    # But the following should work instead given the two ssh variables
+    assert_identical(
+        ds.cf[["ssh"]], ds[["sea_surface_elevation", "sea_surface_height"]]
+    )
+
+    # test criteria list of dicts
+    cf_xarray.accessor.set_options(my_custom_criteria_list)
+    ds = xr.Dataset()
+    ds["temperature"] = ("dim", np.arange(10))
+    assert_identical(ds.cf["temp"], ds["temperature"])
+
+    # test criteria tuple of dicts
+    cf_xarray.accessor.set_options(my_custom_criteria_tuple)
+    ds = xr.Dataset()
+    ds["temperature"] = ("dim", np.arange(10))
+    assert_identical(ds.cf["temp"], ds["temperature"])
+
+
 def test_cf_standard_name_table_version():
 
     url = (

cf_xarray/utils.py

@@ -58,8 +58,8 @@ def invert_mappings(*mappings):
     return merged
 
 
-def always_iterable(obj: Any) -> Iterable:
-    return [obj] if not isinstance(obj, (tuple, list, set, dict)) else obj
+def always_iterable(obj: Any, allowed=(tuple, list, set, dict)) -> Iterable:
+    return [obj] if not isinstance(obj, allowed) else obj
 
 
 def parse_cf_standard_name_table(source=None):

doc/whats-new.rst

@@ -10,6 +10,7 @@ v0.5.3 (unreleased)
 - Begin adding support for units with a unit registry for pint arrays. :pr:`197`.
   By `Jon Thielen`_ and `Justus Magin`_.
 - :py:meth:`Dataset.cf.rename_like` also updates the ``bounds`` and ``cell_measures`` attributes. By `Mattia Almansi`_.
+- Support of custom vocabularies/criteria: user can input criteria for identifying variables by their name and attributes to be able to refer to them by custom names like `ds.cf["ssh"]`. :pr:`234`. By `Kristen Thyng`_ and `Deepak Cherian`_.
 
 v0.5.2 (May 11, 2021)
 =====================