Skip to content

Commit

Permalink
Support Dataset.cf.cell_measures (#117)
Browse files Browse the repository at this point in the history
* Support Dataset.cf.cell_measures

* Cleanup

* fix docstring
  • Loading branch information
dcherian authored Nov 25, 2020
1 parent b99b6b9 commit f61fb66
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 26 deletions.
38 changes: 16 additions & 22 deletions cf_xarray/accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -336,45 +336,44 @@ def _get_measure_variable(
return [da[varnames[0]]]


def _get_measure(da: Union[DataArray, Dataset], key: str) -> List[str]:
def _get_measure(obj: Union[DataArray, Dataset], key: str) -> List[str]:
"""
Translate from cell measures ("area" or "volume") to appropriate variable name.
This function interprets the ``cell_measures`` attribute on DataArrays.
Parameters
----------
da: DataArray
obj: DataArray, Dataset
DataArray belonging to the coordinate to be checked
key: str, ["area", "volume"]
key to check for.
error: bool
raise errors when key is not found or interpretable. Use False and provide default
to replicate dict.get(k, None).
default: Any
default value to return when error is False.
Returns
-------
List[str], Variable name(s) in parent xarray object that matches axis or coordinate `key`
"""
if not isinstance(da, DataArray):
raise NotImplementedError("Measures not implemented for Datasets yet.")

if "cell_measures" not in da.attrs:
raise KeyError("'cell_measures' not present in 'attrs'.")

valid_keys = _CELL_MEASURES
if key not in valid_keys:
raise KeyError(
f"cf_xarray did not understand key {key!r}. Expected one of {valid_keys!r}"
)

attr = da.attrs["cell_measures"]
measures = parse_cell_methods_attr(attr)
results: Union[str, List] = measures.get(key, [])
if isinstance(obj, DataArray):
obj = obj._to_temp_dataset()

results = set()
for var in obj.variables:
da = obj[var]
if "cell_measures" in da.attrs:
attr = da.attrs["cell_measures"]
measures = parse_cell_methods_attr(attr)
if key in measures:
results.update([measures[key]])

if isinstance(results, str):
return [results]
return results
return list(results)


#: Default mappers for common keys.
Expand Down Expand Up @@ -833,10 +832,7 @@ def describe(self):
text += f"\t{key}: {coords[key] if key in coords else []}\n"

text += "\nCell Measures:\n"
if isinstance(self._obj, Dataset):
measures = {key: "unsupported" for key in _CELL_MEASURES}
else:
measures = self.cell_measures
measures = self.cell_measures
for key in _CELL_MEASURES:
text += f"\t{key}: {measures[key] if key in measures else []}\n"

Expand Down Expand Up @@ -930,14 +926,12 @@ def cell_measures(self) -> Dict[str, List[str]]:
This is useful for checking whether a key is valid for indexing, i.e.
that the attributes necessary to allow indexing by that key exist.
However, it will only return the cell measure names.
Returns
-------
Dictionary of valid cell measure names that can be used with __getitem__ or .cf[key].
Will be ("area", "volume") or a subset thereof.
"""
assert isinstance(self._obj, DataArray), "this only works with DataArrays"

measures = {
key: apply_mapper(_get_measure, self._obj, key, error=False)
Expand Down
9 changes: 5 additions & 4 deletions cf_xarray/tests/test_accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ def test_describe(capsys):
"Axes:\n\tX: ['lon']\n\tY: ['lat']\n\tZ: []\n\tT: ['time']\n"
"\nCoordinates:\n\tlongitude: ['lon']\n\tlatitude: ['lat']"
"\n\tvertical: []\n\ttime: ['time']\n"
"\nCell Measures:\n\tarea: unsupported\n\tvolume: unsupported\n"
"\nCell Measures:\n\tarea: ['cell_area']\n\tvolume: []\n"
"\nStandard Names:\n\tair_temperature: ['air']\n\n"
)
assert actual == expected
Expand Down Expand Up @@ -58,12 +58,13 @@ def test_cell_measures():
ds["air"].attrs["cell_measures"] += " foo_measure: foo"
assert "foo_std_name" in ds.cf["air_temperature"].cf

expected = dict(area=["cell_area"])
ds["air"].attrs["cell_measures"] += " volume: foo"
expected = dict(area=["cell_area"], volume=["foo"])
actual = ds["air"].cf.cell_measures
assert actual == expected

with pytest.raises(AssertionError, match=r"this only works with DataArrays"):
popds.cf.cell_measures
actual = ds.cf.cell_measures
assert actual == expected


def test_standard_names():
Expand Down
1 change: 1 addition & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ What's New

v0.4.0 (unreleased)
===================
- Support ``Dataset.cf.cell_measures``. By `Deepak Cherian`_.
- Added ``.axes`` to return a dictionary mapping available Axis standard names to variable names of an xarray object, ``.coordinates`` for Coordinates,
``.cell_measures`` for Cell Measures, and ``.standard_names`` for all variables. `Kristen Thyng`_ and `Mattia Almansi`_.
- Changed ``get_valid_keys()`` to ``.keys()``. `Kristen Thyng`_.
Expand Down

0 comments on commit f61fb66

Please sign in to comment.