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 2 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 @@ -2633,6 +2633,54 @@ def grid_mapping_names(self) -> dict[str, list[str]]:
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():
# default to using existing short name or standard name
name = getattr(values, "name", standard_name)
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(
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

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

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