Skip to content

Commit

Permalink
Add .cf.formula_terms (#213)
Browse files Browse the repository at this point in the history
  • Loading branch information
dcherian authored Apr 24, 2021
1 parent 3252c74 commit e4fa084
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 17 deletions.
59 changes: 42 additions & 17 deletions cf_xarray/accessor.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import functools
import inspect
import itertools
import re
import warnings
from collections import ChainMap
from typing import (
Expand Down Expand Up @@ -1445,16 +1446,19 @@ def differentiate(
self, coord, *xr_args, positive_upward: bool = False, **xr_kwargs
):
"""
Differentiate an xarray object.
Parameters
----------
xr_args, xr_kwargs are passed directly to the underlying xarray function.
The following are added by cf_xarray:
positive_upward: optional, bool
Change sign of the derivative based on the ``"positive"`` attribute of ``coord``
so that positive values indicate increasing upward.
If ``positive=="down"``, then multiplied by -1.
Notes
-----
``xr_args``, ``xr_kwargs`` are passed directly to the underlying xarray function.
See Also
--------
DataArray.cf.differentiate
Expand Down Expand Up @@ -1515,6 +1519,16 @@ def __getitem__(self, key: Union[str, List[str]]) -> Union[DataArray, Dataset]:
"""
return _getitem(self, key)

@property
def formula_terms(self) -> Dict[str, Dict[str, str]]:
"""
Property that returns a dictionary
{parametric_coord_name: {standard_term_name: variable_name}}
"""
return {
dim: self._obj[dim].cf.formula_terms for dim in _get_dims(self._obj, "Z")
}

@property
def bounds(self) -> Dict[str, List[str]]:
"""
Expand Down Expand Up @@ -1719,36 +1733,29 @@ def decode_vertical_coords(self, prefix="z"):
.. warning::
Very lightly tested. Please double check the results.
"""
import re

ds = self._obj
dims = _get_dims(ds, "Z")

requirements = {
"ocean_s_coordinate_g1": {"depth_c", "depth", "s", "C", "eta"},
"ocean_s_coordinate_g2": {"depth_c", "depth", "s", "C", "eta"},
}

for dim in dims:
allterms = self.formula_terms
for dim in allterms:
suffix = dim.split("_")
zname = f"{prefix}_" + "_".join(suffix[1:])

if (
"formula_terms" not in ds[dim].attrs
or "standard_name" not in ds[dim].attrs
):
if "standard_name" not in ds[dim].attrs:
continue

formula_terms = ds[dim].attrs["formula_terms"]
stdname = ds[dim].attrs["standard_name"]

# map "standard" formula term names to actual variable names
terms = {}
for mapping in re.sub(": ", ":", formula_terms).split(" "):
key, value = mapping.split(":")
for key, value in allterms[dim].items():
if value not in ds:
raise KeyError(
f"Variable {value!r} is required to decode coordinate for {dim} but it is absent in the Dataset."
f"Variable {value!r} is required to decode coordinate for {dim!r}"
" but it is absent in the Dataset."
)
terms[key] = ds[value]

Expand Down Expand Up @@ -1776,12 +1783,30 @@ def decode_vertical_coords(self, prefix="z"):

else:
raise NotImplementedError(
f"Coordinate function for {stdname} not implemented yet. Contributions welcome!"
f"Coordinate function for {stdname!r} not implemented yet. Contributions welcome!"
)


@xr.register_dataarray_accessor("cf")
class CFDataArrayAccessor(CFAccessor):
@property
def formula_terms(self) -> Dict[str, str]:
"""
Property that returns a dictionary
{parametric_coord_name: {standard_term_name: variable_name}}
"""
da = self._obj
if "formula_terms" not in da.attrs:
var = da[_single(_get_dims)(da, "Z")[0]]
else:
var = da
terms = {}
formula_terms = var.attrs.get("formula_terms", "")
for mapping in re.sub(r"\s*:\s*", ":", formula_terms).split():
key, value = mapping.split(":")
terms[key] = value
return terms

def __getitem__(self, key: Union[str, List[str]]) -> DataArray:
"""
Index into a DataArray making use of CF attributes.
Expand Down
22 changes: 22 additions & 0 deletions cf_xarray/tests/test_accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -945,6 +945,28 @@ def test_param_vcoord_ocean_s_coord():
copy.cf.decode_vertical_coords()


def test_formula_terms():
srhoterms = {
"s": "s_rho",
"C": "Cs_r",
"eta": "zeta",
"depth": "h",
"depth_c": "hc",
}
assert romsds.cf.formula_terms == {"s_rho": srhoterms}
assert romsds["temp"].cf.formula_terms == srhoterms
assert romsds["s_rho"].cf.formula_terms == srhoterms

s_rho = romsds["s_rho"].copy(deep=True)
del s_rho.attrs["standard_name"]
del s_rho.s_rho.attrs["standard_name"] # TODO: xarray bug
assert s_rho.cf.formula_terms == srhoterms

with pytest.raises(KeyError):
# x,y,t variable
romsds["zeta"].cf.formula_terms


def test_standard_name_mapper():
da = xr.DataArray(
np.arange(6),
Expand Down
2 changes: 2 additions & 0 deletions doc/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ Attributes
DataArray.cf.axes
DataArray.cf.cell_measures
DataArray.cf.coordinates
DataArray.cf.formula_terms
DataArray.cf.standard_names
DataArray.cf.plot

Expand Down Expand Up @@ -66,6 +67,7 @@ Attributes
Dataset.cf.bounds
Dataset.cf.cell_measures
Dataset.cf.coordinates
Dataset.cf.formula_terms
Dataset.cf.standard_names

.. _dsmeth:
Expand Down
2 changes: 2 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ What's New
v0.5.2 (unreleased)
===================

- Added :py:attr:`DataArray.cf.formula_terms` and :py:attr:`Dataset.cf.formula_terms`.
By `Deepak Cherian`_.
- Added :py:attr:`Dataset.cf.bounds` to return a dictionary mapping valid keys to the variable names of their bounds. By `Mattia Almansi`_.
- :py:meth:`DataArray.cf.differentiate` and :py:meth:`Dataset.cf.differentiate` can optionally correct
sign of the derivative by interpreting the ``"positive"`` attribute. By `Deepak Cherian`_.
Expand Down

0 comments on commit e4fa084

Please sign in to comment.