Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Assign new variables by standard names #516

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 48 additions & 0 deletions cf_xarray/accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -2719,6 +2719,54 @@
results[v].append(k)
return results

def assign(self, **standard_variables: Any) -> Dataset:
"""Assign new variables by standard_name.

Parameters
----------
standard_variables : mapping
A mapping from netCDF Climate and Forecast compliant standard names
(http://cfconventions.org/standard-names.html), to the new values.
The new values are passed to ``:meth:Dataset.assign``.

Returns
-------
dataset : Dataset
A new dataset with the new variables added, replacing any existing
variable with the same standard name. New variables are labelled
according to their input (short) names, or standard names if short
names are missing, with added trailing underscores in the event of
duplicates.

See Also
--------
DataSet.assign
"""

# initalize dictionary mapping short names to the new variables
variables = {}

# for each standard name and value pair
for standard_name, values in standard_variables.items():

Check warning on line 2750 in cf_xarray/accessor.py

View check run for this annotation

Codecov / codecov/patch

cf_xarray/accessor.py#L2750

Added line #L2750 was not covered by tests
# default to using existing short name or standard name
name = getattr(values, "name", standard_name)

Check warning on line 2752 in cf_xarray/accessor.py

View check run for this annotation

Codecov / codecov/patch

cf_xarray/accessor.py#L2752

Added line #L2752 was not covered by tests
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Something like

apply_mapper(_get_all, self._obj, key, error=False, default=[key])

should be used here. So

mapped_name, = apply_mapper(_get_all, self._obj, standard_name, error=False, default=[standard_name])

will do the translation to actual name in the dataset, if possible. Use this to create a new dictionary that you can then pass to self._obj.assign

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the feedback, I will be looking at this next week.


# add trailing underscores until we find a free slot
while name in self._obj.data_vars or name in variables:
emit_user_level_warning(

Check warning on line 2756 in cf_xarray/accessor.py

View check run for this annotation

Codecov / codecov/patch

cf_xarray/accessor.py#L2755-L2756

Added lines #L2755 - L2756 were not covered by tests
f"found existing variable {name}, using {name}_ instead",
UserWarning,
)
name += "_"

# map name to values (without coords, see #513, xarray#6447)
values = xr.as_variable(values)
values.attrs.update(standard_name=standard_name)
variables[name] = values

Check warning on line 2765 in cf_xarray/accessor.py

View check run for this annotation

Codecov / codecov/patch

cf_xarray/accessor.py#L2763-L2765

Added lines #L2763 - L2765 were not covered by tests

# assign new variables and return a new dataset
return self._obj.assign(variables)

Check warning on line 2768 in cf_xarray/accessor.py

View check run for this annotation

Codecov / codecov/patch

cf_xarray/accessor.py#L2768

Added line #L2768 was not covered by tests

def decode_vertical_coords(self, *, outnames=None, prefix=None):
"""
Decode parameterized vertical coordinates in place.
Expand Down