Add support for flag variables. (#252)

Author: dcherian

cf_xarray/accessor.py

@@ -886,15 +886,134 @@ def __getattr__(self, attr):
         )
 
 
+def create_flag_dict(da):
+    if not da.cf.is_flag_variable:
+        raise ValueError(
+            "Comparisons are only supported for DataArrays that represent CF flag variables."
+            ".attrs must contain 'flag_values' and 'flag_meanings'"
+        )
+
+    flag_meanings = da.attrs["flag_meanings"].split(" ")
+    flag_values = da.attrs["flag_values"]
+    # TODO: assert flag_values is iterable
+    assert len(flag_values) == len(flag_meanings)
+    return dict(zip(flag_meanings, flag_values))
+
+
 class CFAccessor:
     """
     Common Dataset and DataArray accessor functionality.
     """
 
-    def __init__(self, da):
-        self._obj = da
+    def __init__(self, obj):
+        self._obj = obj
         self._all_cell_measures = None
 
+    def _assert_valid_other_comparison(self, other):
+        flag_dict = create_flag_dict(self._obj)
+        if other not in flag_dict:
+            raise ValueError(
+                f"Did not find flag value meaning [{other}] in known flag meanings: [{flag_dict.keys()!r}]"
+            )
+        return flag_dict
+
+    def __eq__(self, other):
+        """
+        Compare flag values against `other`.
+
+        `other` must be in the 'flag_meanings' attribute.
+        `other` is mapped to the corresponding value in the 'flag_values' attribute, and then
+        compared.
+        """
+        flag_dict = self._assert_valid_other_comparison(other)
+        return self._obj == flag_dict[other]
+
+    def __ne__(self, other):
+        """
+        Compare flag values against `other`.
+
+        `other` must be in the 'flag_meanings' attribute.
+        `other` is mapped to the corresponding value in the 'flag_values' attribute, and then
+        compared.
+        """
+        flag_dict = self._assert_valid_other_comparison(other)
+        return self._obj != flag_dict[other]
+
+    def __lt__(self, other):
+        """
+        Compare flag values against `other`.
+
+        `other` must be in the 'flag_meanings' attribute.
+        `other` is mapped to the corresponding value in the 'flag_values' attribute, and then
+        compared.
+        """
+        flag_dict = self._assert_valid_other_comparison(other)
+        return self._obj < flag_dict[other]
+
+    def __le__(self, other):
+        """
+        Compare flag values against `other`.
+
+        `other` must be in the 'flag_meanings' attribute.
+        `other` is mapped to the corresponding value in the 'flag_values' attribute, and then
+        compared.
+        """
+        flag_dict = self._assert_valid_other_comparison(other)
+        return self._obj <= flag_dict[other]
+
+    def __gt__(self, other):
+        """
+        Compare flag values against `other`.
+
+        `other` must be in the 'flag_meanings' attribute.
+        `other` is mapped to the corresponding value in the 'flag_values' attribute, and then
+        compared.
+        """
+        flag_dict = self._assert_valid_other_comparison(other)
+        return self._obj > flag_dict[other]
+
+    def __ge__(self, other):
+        """
+        Compare flag values against `other`.
+
+        `other` must be in the 'flag_meanings' attribute.
+        `other` is mapped to the corresponding value in the 'flag_values' attribute, and then
+        compared.
+        """
+        flag_dict = self._assert_valid_other_comparison(other)
+        return self._obj >= flag_dict[other]
+
+    def isin(self, test_elements):
+        """Test each value in the array for whether it is in test_elements.
+
+        Parameters
+        ----------
+        test_elements : array_like, 1D
+            The values against which to test each value of `element`.
+            These must be in "flag_meanings" attribute, and are mapped
+            to the corresponding value in "flag_values" before passing
+            that on to DataArray.isin.
+
+
+        Returns
+        -------
+        isin : DataArray
+            Has the same type and shape as this object, but with a bool dtype.
+        """
+        if not isinstance(self._obj, DataArray):
+            raise ValueError(
+                ".cf.isin is only supported on DataArrays that contain CF flag attributes."
+            )
+        flag_dict = create_flag_dict(self._obj)
+        mapped_test_elements = []
+        for elem in test_elements:
+            if elem not in flag_dict:
+                raise ValueError(
+                    f"Did not find flag value meaning [{elem}] in known flag meanings: [{flag_dict.keys()!r}]"
+                )
+            mapped_test_elements.append(flag_dict[elem])
+        return self._obj.isin(mapped_test_elements)
+
     def _get_all_cell_measures(self):
         """
         Get all cell measures defined in the object, adding CF pre-defined measures.
@@ -1130,7 +1249,12 @@ def make_text_section(subtitle, attr, valid_values, default_keys=None):
 
             return "\n".join(rows) + "\n"
 
-        text = "Coordinates:"
+        if isinstance(self._obj, DataArray) and self._obj.cf.is_flag_variable:
+            flag_dict = create_flag_dict(self._obj)
+            text = f"CF Flag variable with mapping:\n\t{flag_dict!r}\n\n"
+        else:
+            text = ""
+        text += "Coordinates:"
         text += make_text_section("CF Axes", "axes", coords, _AXIS_NAMES)
         text += make_text_section("CF Coordinates", "coordinates", coords, _COORD_NAMES)
         text += make_text_section(
@@ -2057,4 +2181,18 @@ def __getitem__(self, key: Union[str, List[str]]) -> DataArray:
 
         return _getitem(self, key)
 
-    pass
+    @property
+    def is_flag_variable(self):
+        """
+        Returns True if the DataArray satisfies CF conventions for flag variables.
+
+        Flag masks are not supported yet.
+        """
+        if (
+            isinstance(self._obj, DataArray)
+            and "flag_meanings" in self._obj.attrs
+            and "flag_values" in self._obj.attrs
+        ):
+            return True
+        else:
+            return False

cf_xarray/datasets.py

@@ -279,3 +279,15 @@
         }
     )
 )
+
+
+basin = xr.DataArray(
+    [1, 2, 1, 1, 2, 2, 3, 3, 3, 3],
+    dims=("time",),
+    attrs={
+        "flag_values": [1, 2, 3],
+        "flag_meanings": "atlantic_ocean pacific_ocean indian_ocean",
+        "standard_name": "region",
+    },
+    name="basin",
+)

cf_xarray/tests/test_accessor.py

@@ -17,6 +17,7 @@
 from ..datasets import (
     airds,
     anc,
+    basin,
     ds_no_attrs,
     forecast,
     mollwds,
@@ -128,6 +129,9 @@ def test_repr():
     """
     assert actual == dedent(expected)
 
+    # Flag DataArray
+    assert "CF Flag variable" in repr(basin.cf)
+
 
 def test_axes():
     expected = dict(T=["time"], X=["lon"], Y=["lat"])
@@ -1389,3 +1393,35 @@ def test_add_canonical_attributes(override, skip, verbose, capsys):
 
         cf_da.attrs.pop("history")
         assert_identical(cf_da, cf_ds["air"])
+
+
+@pytest.mark.parametrize("op", ["ge", "gt", "eq", "ne", "le", "lt"])
+def test_flag_features(op):
+    actual = getattr(basin.cf, f"__{op}__")("atlantic_ocean")
+    expected = getattr(basin, f"__{op}__")(1)
+    assert_identical(actual, expected)
+
+
+def test_flag_isin():
+    actual = basin.cf.isin(["atlantic_ocean", "pacific_ocean"])
+    expected = basin.isin([1, 2])
+    assert_identical(actual, expected)
+
+
+def test_flag_errors():
+    with pytest.raises(ValueError):
+        basin.cf.isin(["arctic_ocean"])
+
+    with pytest.raises(ValueError):
+        basin.cf == "arctic_ocean"
+
+    ds = xr.Dataset({"basin": basin})
+    with pytest.raises(ValueError):
+        ds.cf.isin(["atlantic_ocean"])
+
+    basin.attrs.pop("flag_values")
+    with pytest.raises(ValueError):
+        basin.cf.isin(["pacific_ocean"])
+
+    with pytest.raises(ValueError):
+        basin.cf == "pacific_ocean"

ci/doc.yml

@@ -16,6 +16,7 @@ dependencies:
   - ipykernel
   - ipywidgets
   - pandas
+  - pooch
   - pydata-sphinx-theme
   - pip:
       - git+https://github.com/xarray-contrib/cf-xarray

doc/api.rst

@@ -31,6 +31,7 @@ Attributes
     DataArray.cf.cell_measures
     DataArray.cf.coordinates
     DataArray.cf.formula_terms
+    DataArray.cf.is_flag_variable
     DataArray.cf.standard_names
     DataArray.cf.plot
 
@@ -52,6 +53,24 @@ Methods
     DataArray.cf.keys
     DataArray.cf.rename_like
 
+Flag Variables
+++++++++++++++
+
+cf_xarray supports rich comparisons for `CF flag variables`_. Flag masks are not yet supported.
+
+.. autosummary::
+   :toctree: generated/
+   :template: autosummary/accessor_method.rst
+
+    DataArray.cf.__lt__
+    DataArray.cf.__le__
+    DataArray.cf.__eq__
+    DataArray.cf.__ne__
+    DataArray.cf.__ge__
+    DataArray.cf.__gt__
+    DataArray.cf.isin
+
+
 Dataset
 -------
 
@@ -92,3 +111,6 @@ Methods
     Dataset.cf.guess_coord_axis
     Dataset.cf.keys
     Dataset.cf.rename_like
+
+
+.. _`CF flag variables`: http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#flags

doc/whats-new.rst

@@ -7,6 +7,8 @@ v0.6.1 (unreleased)
 ===================
 - Support detecting pint-backed Variables with units-based criteria. By `Deepak Cherian`_.
 - Support reshaping nD bounds arrays to (n-1)D vertex arrays. By `Deepak Cherian`_.
+- Support rich comparisons  with ``DataArray.cf`` and :py:meth:`DataArray.cf.isin` for `flag variables`_.
+  By `Deepak Cherian`_ and `Julius Busecke`_
 
 v0.6.0 (June 29, 2021)
 ======================
@@ -125,3 +127,5 @@ v0.1.3
 .. _`Filipe Fernandes`: https://github.com/ocefpaf
 .. _`Julia Kent`: https://github.com/jukent
 .. _`Kristen Thyng`: https://github.com/kthyng
+.. _`Julius Busecke`: https://github.com/jbusecke
+.. _`flag variables`: http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#flags