Skip to content

Commit

Permalink
Drop missing variables (#267)
Browse files Browse the repository at this point in the history
  • Loading branch information
malmans2 authored Dec 7, 2021
1 parent 31ceff2 commit 6f47339
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 8 deletions.
33 changes: 26 additions & 7 deletions cf_xarray/accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -1029,6 +1029,10 @@ def isin(self, test_elements):
mapped_test_elements.append(flag_dict[elem])
return self._obj.isin(mapped_test_elements)

def _drop_missing_variables(self, variables: List[str]) -> List[str]:

return [var for var in variables if var in self._obj or var in self._obj.coords]

def _get_all_cell_measures(self):
"""
Get all cell measures defined in the object, adding CF pre-defined measures.
Expand Down Expand Up @@ -1383,7 +1387,9 @@ def cell_measures(self) -> Dict[str, List[str]]:
keys = {}
for attr in all_attrs:
keys.update(parse_cell_methods_attr(attr))
measures = {key: _get_all(self._obj, key) for key in keys}
measures = {
key: self._drop_missing_variables(_get_all(self._obj, key)) for key in keys
}

return {k: sorted(set(v)) for k, v in measures.items() if v}

Expand Down Expand Up @@ -1877,9 +1883,15 @@ 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")
}
results = {}
for dim in _get_dims(self._obj, "Z"):
terms = self._obj[dim].cf.formula_terms
variables = self._drop_missing_variables(list(terms.values()))
terms = {key: val for key, val in terms.items() if val in variables}
if terms:
results[dim] = terms

return results

@property
def bounds(self) -> Dict[str, List[str]]:
Expand All @@ -1896,12 +1908,15 @@ def bounds(self) -> Dict[str, List[str]]:
keys = self.keys() | set(obj.variables)

vardict = {
key: apply_mapper(_get_bounds, obj, key, error=False) for key in keys
key: self._drop_missing_variables(
apply_mapper(_get_bounds, obj, key, error=False)
)
for key in keys
}

return {k: sorted(v) for k, v in vardict.items() if v}

def get_bounds(self, key: str) -> DataArray:
def get_bounds(self, key: str) -> Union[DataArray, Dataset]:
"""
Get bounds variable corresponding to key.
Expand All @@ -1915,7 +1930,11 @@ def get_bounds(self, key: str) -> DataArray:
DataArray
"""

return apply_mapper(_variables(_single(_get_bounds)), self._obj, key)[0]
results = self.bounds.get(key, [])
if not results:
raise KeyError(f"No results found for {key!r}.")

return self._obj[results[0] if len(results) == 1 else results]

def get_bounds_dim_name(self, key: str) -> str:
"""
Expand Down
23 changes: 22 additions & 1 deletion cf_xarray/tests/test_accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,7 @@ def test_cell_measures():
ds["air"].attrs["cell_measures"] += " volume: foo"
ds["foo"].attrs["cell_measures"] = ds["air"].attrs["cell_measures"]
expected = dict(area=["cell_area"], foo_measure=["foo"], volume=["foo"])
actual_air = ds["air"].cf.cell_measures
actual_air = ds.cf["air"].cf.cell_measures
actual_foo = ds.cf["foo_measure"].cf.cell_measures
assert actual_air == actual_foo == expected

Expand Down Expand Up @@ -1451,3 +1451,24 @@ def test_flag_errors():

with pytest.raises(ValueError):
basin.cf == "pacific_ocean"


def test_missing_variables():

# Bounds
ds = mollwds.copy(deep=True)
ds = ds.drop("lon_bounds")
assert ds.cf.bounds == {"lat": ["lat_bounds"], "latitude": ["lat_bounds"]}

with pytest.raises(KeyError, match=r"No results found for 'longitude'."):
ds.cf.get_bounds("longitude")

# Cell measures
ds = airds.copy(deep=True)
ds = ds.drop("cell_area")
assert ds.cf.cell_measures == {}

# Formula terms
ds = vert.copy(deep=True)
ds = ds.drop("ap")
assert ds.cf.formula_terms == {"lev": {"b": "b", "ps": "ps"}}

0 comments on commit 6f47339

Please sign in to comment.