diff --git a/CITATION.cff b/CITATION.cff index 1ca4a1c2..f4a53fa4 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -83,6 +83,10 @@ authors: - family-names: Haëck given-names: Clément affiliation: Laboratoire d'Océanographie et du Climat (LOCEAN), Paris + - family-names: Boutte + given-names: Jason + orcid: 'https://orcid.org/0009-0009-3996-3772' + affiliation: Lawrence Livermore National Laboratory identifiers: - type: doi value: 10.5281/zenodo.4749735 diff --git a/cf_xarray/accessor.py b/cf_xarray/accessor.py index ed45da18..7775582d 100644 --- a/cf_xarray/accessor.py +++ b/cf_xarray/accessor.py @@ -30,7 +30,7 @@ from xarray.core.rolling import Coarsen, Rolling from xarray.core.weighted import Weighted -from . import sgrid +from . import parametric, sgrid from .criteria import ( _DSG_ROLES, _GEOMETRY_TYPES, @@ -2754,13 +2754,8 @@ def decode_vertical_coords(self, *, outnames=None, prefix=None): """ ds = self._obj - requirements = { - "ocean_s_coordinate_g1": {"depth_c", "depth", "s", "C", "eta"}, - "ocean_s_coordinate_g2": {"depth_c", "depth", "s", "C", "eta"}, - "ocean_sigma_coordinate": {"sigma", "eta", "depth"}, - } - allterms = self.formula_terms + for dim in allterms: if prefix is None: assert ( @@ -2782,6 +2777,7 @@ def decode_vertical_coords(self, *, outnames=None, prefix=None): suffix = dim.split("_") zname = f"{prefix}_" + "_".join(suffix[1:]) + # never touched, if standard name is missing it's not included in allterms if "standard_name" not in ds[dim].attrs: continue stdname = ds[dim].attrs["standard_name"] @@ -2790,46 +2786,23 @@ def decode_vertical_coords(self, *, outnames=None, prefix=None): terms = {} for key, value in allterms[dim].items(): if value not in ds: + # is this ever hit, if variable is missing it's missing in decoded allterms raise KeyError( f"Variable {value!r} is required to decode coordinate for {dim!r}" " but it is absent in the Dataset." ) - terms[key] = ds[value] - - absent_terms = requirements[stdname] - set(terms) - if absent_terms: - raise KeyError(f"Required terms {absent_terms} absent in dataset.") - - if stdname == "ocean_s_coordinate_g1": - # S(k,j,i) = depth_c * s(k) + (depth(j,i) - depth_c) * C(k) - S = ( - terms["depth_c"] * terms["s"] - + (terms["depth"] - terms["depth_c"]) * terms["C"] - ) - - # z(n,k,j,i) = S(k,j,i) + eta(n,j,i) * (1 + S(k,j,i) / depth(j,i)) - ztemp = S + terms["eta"] * (1 + S / terms["depth"]) + # keys should be case insensitive + terms[key.lower()] = ds[value] - elif stdname == "ocean_s_coordinate_g2": - # make sure all necessary terms are present in terms - # (depth_c * s(k) + depth(j,i) * C(k)) / (depth_c + depth(j,i)) - S = (terms["depth_c"] * terms["s"] + terms["depth"] * terms["C"]) / ( - terms["depth_c"] + terms["depth"] - ) - - # z(n,k,j,i) = eta(n,j,i) + (eta(n,j,i) + depth(j,i)) * S(k,j,i) - ztemp = terms["eta"] + (terms["eta"] + terms["depth"]) * S - - elif stdname == "ocean_sigma_coordinate": - # z(n,k,j,i) = eta(n,j,i) + sigma(k)*(depth(j,i)+eta(n,j,i)) - ztemp = terms["eta"] + terms["sigma"] * (terms["depth"] + terms["eta"]) - - else: + try: + transform = parametric.TRANSFORM_FROM_STDNAME[stdname] + except KeyError: + # Should occur since stdname is check before raise NotImplementedError( - f"Coordinate function for {stdname!r} not implemented yet. Contributions welcome!" - ) + f"Coordinate function for {stdname!r} not implmented yet. Contributions welcome!" + ) from None - ds.coords[zname] = ztemp + ds.coords[zname] = transform.from_terms(terms) @xr.register_dataarray_accessor("cf") diff --git a/cf_xarray/parametric.py b/cf_xarray/parametric.py new file mode 100644 index 00000000..06b3d7fd --- /dev/null +++ b/cf_xarray/parametric.py @@ -0,0 +1,821 @@ +from abc import ABC, abstractmethod +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np +import xarray as xr +from xarray import DataArray + +OCEAN_STDNAME_MAP = { + "altitude": { + "zlev": "altitude", + "eta": "sea_surface_height_above_geoid", + "depth": "sea_floor_depth_below_geoid", + }, + "height_above_geopotential_datum": { + "zlev": "height_above_geopotential_datum", + "eta": "sea_surface_height_above_geopotential_datum", + "depth": "sea_floor_depth_below_geopotential_datum", + }, + "height_above_reference_ellipsoid": { + "zlev": "height_above_reference_ellipsoid", + "eta": "sea_surface_height_above_reference_ellipsoid", + "depth": "sea_floor_depth_below_reference_ellipsoid", + }, + "height_above_mean_sea_level": { + "zlev": "height_above_mean_sea_level", + "eta": "sea_surface_height_above_mean_sea_level", + "depth": "sea_floor_depth_below_mean_sea_level", + }, +} + + +def _derive_ocean_stdname(*, zlev=None, eta=None, depth=None): + """Derive standard name for computer ocean coordinates. + + Uses the concatentation of formula terms `zlev`, `eta`, and `depth` + standard names to compare against formula term and standard names + from a table. This can occur with any combination e.g. `zlev`, or + `zlev` + `depth`. If a match is found the standard name for the + computed value is returned. + + Parameters + ---------- + zlev : dict + Attributes for `zlev` variable. + eta : dict + Attributes for `eta` variable. + depth : dict + Attributes for `depth` variable. + + Returns + ------- + str + Standard name for the computer value. + + Raises + ------ + ValueError + If `kwargs` is empty, missing values for `kwargs` keys, or could not derive the standard name. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#table-computed-standard-names + """ + found_stdname = None + search_term = "" + search_vars = {"zlev": zlev, "eta": eta, "depth": depth} + + for x, y in sorted(search_vars.items(), key=lambda x: x[0]): + if y is None: + continue + + try: + search_term = f"{search_term}{y['standard_name']}" + except TypeError: + raise ValueError( + f"The values for {', '.join(sorted(search_vars.keys()))} cannot be `None`." + ) from None + except KeyError: + raise ValueError( + f"The standard name for the {x!r} variable is not available." + ) from None + + for x, y in OCEAN_STDNAME_MAP.items(): + check_term = "".join( + [ + y[i] + for i, j in sorted(search_vars.items(), key=lambda x: x[0]) + if j is not None + ] + ) + + if search_term == check_term: + found_stdname = x + + break + + if found_stdname is None: + stdnames = ", ".join( + [ + y["standard_name"] + for _, y in sorted(search_vars.items(), key=lambda x: x[0]) + if y is not None + ] + ) + + raise ValueError( + f"Could not derive standard name from combination of {stdnames}." + ) + + return found_stdname + + +class ParametricVerticalCoordinate(ABC): + @classmethod + @abstractmethod + def from_terms(cls, terms: dict): + pass + + @abstractmethod + def decode(self): + pass + + @property + @abstractmethod + def computed_standard_name(self): + pass + + +@dataclass +class AtmosphereLnPressure(ParametricVerticalCoordinate): + """Atmosphere natural log pressure coordinate. + + Standard name: atmosphere_ln_pressure_coordinate + + Parameters + ---------- + p0 : xr.DataArray + Reference pressure. + lev : xr.DataArray + Vertical dimensionless coordinate. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#atmosphere-natural-log-pressure-coordinate + """ + + p0: DataArray + lev: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + p = self.p0 * np.exp(-self.lev) + + return p.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self): + """Computes coordinate standard name.""" + return "air_pressure" + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "p0", "lev")) + + +@dataclass +class AtmosphereSigma(ParametricVerticalCoordinate): + """Atmosphere sigma coordinate. + + Standard name: atmosphere_sigma_coordinate + + Parameters + ---------- + sigma : xr.DataArray + Vertical dimensionless coordinate. + ps : xr.DataArray + Horizontal surface pressure. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_atmosphere_sigma_coordinate + """ + + sigma: DataArray + ps: DataArray + ptop: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + p = self.ptop + self.sigma * (self.ps - self.ptop) + + return p.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + return "air_pressure" + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "sigma", "ps", "ptop")) + + +@dataclass +class AtmosphereHybridSigmaPressure(ParametricVerticalCoordinate): + """Atmosphere hybrid sigma pressure coordinate. + + Standard name: atmosphere_hybrid_sigma_pressure_coordinate + + Parameters + ---------- + b : xr.DataArray + Component of hybrid coordinate. + ps : xr.DataArray + Horizontal surface pressure. + p0 : xr.DataArray + Reference pressure. + a : xr.DataArray + Component of hybrid coordinate. + ap : xr.DataArray + Component of hybrid coordinate. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_atmosphere_hybrid_sigma_pressure_coordinate + """ + + b: DataArray + ps: DataArray + p0: DataArray + a: DataArray + ap: DataArray + + def __post_init__(self): + if self.a is None and self.ap is None: + raise KeyError( + "Optional terms 'a', 'ap' are absent in the dataset, atleast one must be present." + ) + + if self.a is not None and self.ap is not None: + raise ValueError( + "Both optional terms 'a' and 'ap' are present in the dataset, please drop one of them." + ) + + if self.a is not None and self.p0 is None: + raise KeyError( + "Optional term 'a' is present but 'p0' is absent in the dataset." + ) + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + if self.a is None: + p = self.ap + self.b * self.ps # type: ignore[unreachable] + else: + p = self.a * self.p0 + self.b * self.ps + + return p.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + return "air_pressure" + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "b", "ps", optional=("p0", "a", "ap"))) + + +@dataclass +class AtmosphereHybridHeight(ParametricVerticalCoordinate): + """Atmosphere hybrid height coordinate. + + Standard name: atmosphere_hybrid_height_coordinate + + Parameters + ---------- + a : xr.DataArray + Height. + b : xr.DataArray + Dimensionless. + orog : xr.DataArray + Height of the surface above the datum. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#atmosphere-hybrid-height-coordinate + """ + + a: DataArray + b: DataArray + orog: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + z = self.a + self.b * self.orog + + return z.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + orog_stdname = self.orog.attrs["standard_name"] + + if orog_stdname == "surface_altitude": + out_stdname = "altitude" + elif orog_stdname == "surface_height_above_geopotential_datum": + out_stdname = "height_above_geopotential_datum" + else: + raise ValueError( + f"Unknown standard name for hybrid height coordinate: {orog_stdname!r}" + ) + + return out_stdname + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "a", "b", "orog")) + + +@dataclass +class AtmosphereSleve(ParametricVerticalCoordinate): + """Atmosphere smooth level vertical (SLEVE) coordinate. + + Standard name: atmosphere_sleve_coordinate + + Parameters + ---------- + a : xr.DataArray + Dimensionless coordinate whcih defines hybrid level. + b1 : xr.DataArray + Dimensionless coordinate whcih defines hybrid level. + b2 : xr.DataArray + Dimensionless coordinate whcih defines hybrid level. + ztop : xr.DataArray + Height above the top of the model above datum. + zsurf1 : xr.DataArray + Large-scale component of the topography. + zsurf2 : xr.DataArray + Small-scale component of the topography. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_atmosphere_smooth_level_vertical_sleve_coordinate + """ + + a: DataArray + b1: DataArray + b2: DataArray + ztop: DataArray + zsurf1: DataArray + zsurf2: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + z = self.a * self.ztop + self.b1 * self.zsurf1 + self.b2 * self.zsurf2 + + return z.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + ztop_stdname = self.ztop.attrs["standard_name"] + + if ztop_stdname == "altitude_at_top_of_atmosphere_model": + out_stdname = "altitude" + elif ( + ztop_stdname == "height_above_geopotential_datum_at_top_of_atmosphere_model" + ): + out_stdname = "height_above_geopotential_datum" + else: + raise ValueError(f"Unknown standard name: {out_stdname!r}") + + return out_stdname + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "a", "b1", "b2", "ztop", "zsurf1", "zsurf2")) + + +@dataclass +class OceanSigma(ParametricVerticalCoordinate): + """Ocean sigma coordinate. + + Standard name: ocean_sigma_coordinate + + Parameters + ---------- + sigma : xr.DataArray + Vertical dimensionless coordinate. + eta : xr.DataArray + Height of the sea surface (positive upwards) relative to the datum. + depth : xr.DataArray + Distance (positive value) from the datum to the sea floor. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_ocean_sigma_coordinate + """ + + sigma: DataArray + eta: DataArray + depth: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + z = self.eta + self.sigma * (self.depth + self.eta) + + return z.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + out_stdname = _derive_ocean_stdname(eta=self.eta.attrs, depth=self.depth.attrs) + + return out_stdname + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "sigma", "eta", "depth")) + + +@dataclass +class OceanS(ParametricVerticalCoordinate): + """Ocean s-coordinate. + + Standard name: ocean_s_coordinate + + Parameters + ---------- + s : xr.DataArray + Dimensionless coordinate. + eta : xr.DataArray + Height of the sea surface (positive upwards) relative to the datum. + depth : xr.DataArray + Distance (positive value) from the datum to the sea floor. + a : xr.DataArray + Constant controlling stretch. + b : xr.DataArray + Constant controlling stretch. + depth_c : xr.DataArray + Constant controlling stretch. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_ocean_s_coordinate + """ + + s: DataArray + eta: DataArray + depth: DataArray + a: DataArray + b: DataArray + depth_c: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + C = (1 - self.b) * np.sinh(self.a * self.s) / np.sinh(self.a) + self.b * ( + np.tanh(self.a * (self.s + 0.5)) / 2 * np.tanh(0.5 * self.a) - 0.5 + ) + + z = ( + self.eta * (1 + self.s) + + self.depth_c * self.s + + (self.depth - self.depth_c) * C + ) + + return z.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + return _derive_ocean_stdname(eta=self.eta.attrs, depth=self.depth.attrs) + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "s", "eta", "depth", "a", "b", "depth_c")) + + +@dataclass +class OceanSG1(ParametricVerticalCoordinate): + """Ocean s-coordinate, generic form 1. + + Standard name: ocean_s_coordinate_g1 + + Parameters + ---------- + s : xr.DataArray + Dimensionless coordinate. + C : xr.DataArray + Dimensionless vertical coordinate stretching function. + eta : xr.DataArray + Height of the ocean surface (positive upwards) relative to the ocean datum. + depth : xr.DataArray + Distance from ocean datum to sea floor (positive value). + depth_c : xr.DataArray + Constant (positive value) is a critical depth controlling the stretching. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_ocean_s_coordinate_generic_form_1 + """ + + s: DataArray + c: DataArray + eta: DataArray + depth: DataArray + depth_c: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + S = self.depth_c * self.s + (self.depth - self.depth_c) * self.c + + z = S + self.eta * (1 + self.s / self.depth) + + return z.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + return _derive_ocean_stdname(eta=self.eta.attrs, depth=self.depth.attrs) + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "s", "c", "eta", "depth", "depth_c")) + + +@dataclass +class OceanSG2(ParametricVerticalCoordinate): + """Ocean s-coordinate, generic form 2. + + Standard name: ocean_s_coordinate_g2 + + Parameters + ---------- + s : xr.DataArray + Dimensionless coordinate. + C : xr.DataArray + Dimensionless vertical coordinate stretching function. + eta : xr.DataArray + Height of the ocean surface (positive upwards) relative to the ocean datum. + depth : xr.DataArray + Distance from ocean datum to sea floor (positive value). + depth_c : xr.DataArray + Constant (positive value) is a critical depth controlling the stretching. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_ocean_s_coordinate_generic_form_2 + """ + + s: DataArray + c: DataArray + eta: DataArray + depth: DataArray + depth_c: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + S = (self.depth_c * self.s + self.depth * self.c) / (self.depth_c + self.depth) + + z = self.eta + (self.eta + self.depth) * S + + return z.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + return _derive_ocean_stdname(eta=self.eta.attrs, depth=self.depth.attrs) + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "s", "c", "eta", "depth", "depth_c")) + + +@dataclass +class OceanSigmaZ(ParametricVerticalCoordinate): + """Ocean sigma over z coordinate. + + Standard name: ocean_sigma_z_coordinate + + Parameters + ---------- + sigma : xr.DataArray + Coordinate defined only for `nsigma` layers nearest the ocean surface. + eta : xr.DataArray + Height of the ocean surface (positive upwards) relative to ocean datum. + depth : xr.DataArray + Distance from ocean datum to sea floor (positive value). + depth_c : xr.DataArray + Constant. + nsigma : xr.DataArray + Layers nearest the ocean surface. + zlev : xr.DataArray + Coordinate defined only for `nlayer - nsigma` where `nlayer` is the size of the vertical coordinate. + + Notes + ----- + The description of this type of parametric vertical coordinate is defective in version 1.8 and earlier versions of the standard, in that it does not state what values the vertical coordinate variable should contain. Therefore, in accordance with the rules, all versions of the standard before 1.9 are deprecated for datasets that use the "ocean sigma over z" coordinate. + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_ocean_sigma_over_z_coordinate + """ + + sigma: DataArray + eta: DataArray + depth: DataArray + depth_c: DataArray + nsigma: DataArray + zlev: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + z_sigma = self.eta + self.sigma * ( + np.minimum(self.depth_c, self.depth) + self.eta + ) + + z = xr.where(np.isnan(self.sigma), self.zlev, z_sigma) + + return z.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + return _derive_ocean_stdname( + eta=self.eta.attrs, depth=self.depth.attrs, zlev=self.zlev.attrs + ) + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls( + **get_terms(terms, "sigma", "eta", "depth", "depth_c", "nsigma", "zlev") + ) + + +@dataclass +class OceanDoubleSigma(ParametricVerticalCoordinate): + """Ocean double sigma coordinate. + + Standard name: ocean_double_sigma_coordinate + + Parameters + ---------- + sigma : xr.DataArray + Dimensionless coordinate. + depth : xr.DataArray + Distance (positive value) from datum to the sea floor. + z1 : xr.DataArray + Constant with units of length. + z2 : xr.DataArray + Constant with units of length. + a : xr.DataArray + Constant with units of length. + href : xr.DataArray + Constant with units of length. + k_c : xr.DataArray + + References + ---------- + Please refer to the CF conventions document : + 1. https://cfconventions.org/cf-conventions/cf-conventions.html#_ocean_double_sigma_coordinate + """ + + sigma: DataArray + depth: DataArray + z1: DataArray + z2: DataArray + a: DataArray + href: DataArray + k_c: DataArray + + def decode(self) -> xr.DataArray: + """Decode coordinate. + + Returns + ------- + xr.DataArray + Decoded parametric vertical coordinate. + """ + f = 0.5 * (self.z1 + self.z2) + 0.5 * (self.z1 - self.z2) * np.tanh( + 2 * self.a / (self.z1 - self.z2) * (self.depth - self.href) + ) + + z = xr.where( + self.sigma.k <= self.k_c, + self.sigma * f, + f + (self.sigma - 1) * (self.depth - f), + ) + + return z.assign_attrs(standard_name=self.computed_standard_name) + + @property + def computed_standard_name(self) -> str: + """Computes coordinate standard name.""" + return _derive_ocean_stdname(depth=self.depth.attrs) + + @classmethod + def from_terms(cls, terms: dict): + """Create coordinate from terms.""" + return cls(**get_terms(terms, "sigma", "depth", "z1", "z2", "a", "href", "k_c")) + + +TRANSFORM_FROM_STDNAME = { + "atmosphere_ln_pressure_coordinate": AtmosphereLnPressure, + "atmosphere_sigma_coordinate": AtmosphereSigma, + "atmosphere_hybrid_sigma_pressure_coordinate": AtmosphereHybridSigmaPressure, + "atmosphere_hybrid_height_coordinate": AtmosphereHybridHeight, + "atmosphere_sleve_coordinate": AtmosphereSleve, + "ocean_sigma_coordinate": OceanSigma, + "ocean_s_coordinate": OceanS, + "ocean_s_coordinate_g1": OceanSG1, + "ocean_s_coordinate_g2": OceanSG2, + "ocean_sigma_z_coordinate": OceanSigmaZ, + "ocean_double_sigma_coordinate": OceanDoubleSigma, +} + + +def get_terms( + terms: dict[str, DataArray], *required, optional: Sequence[str] | None = None +) -> dict[str, DataArray]: + if optional is None: + optional = [] + + selected_terms = {} + + for term in required + tuple(optional): + da = None + + try: + da = terms[term] + except KeyError: + if term not in optional: + raise KeyError( + f"Required term {term} is absent in the dataset." + ) from None + + selected_terms[term] = da + + return selected_terms # type: ignore[return-value] diff --git a/cf_xarray/tests/test_accessor.py b/cf_xarray/tests/test_accessor.py index 444771fa..36f95016 100644 --- a/cf_xarray/tests/test_accessor.py +++ b/cf_xarray/tests/test_accessor.py @@ -31,7 +31,6 @@ forecast, mollwds, multiple, - pomds, popds, romsds, rotds, @@ -1306,52 +1305,23 @@ def test_Z_vs_vertical_ROMS() -> None: ) -def test_param_vcoord_ocean_s_coord() -> None: - romsds.s_rho.attrs["standard_name"] = "ocean_s_coordinate_g2" - Zo_rho = (romsds.hc * romsds.s_rho + romsds.Cs_r * romsds.h) / ( - romsds.hc + romsds.h - ) - expected = romsds.zeta + (romsds.zeta + romsds.h) * Zo_rho - romsds.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) - assert_allclose( - romsds.z_rho.reset_coords(drop=True), expected.reset_coords(drop=True) - ) - - romsds.s_rho.attrs["standard_name"] = "ocean_s_coordinate_g1" - Zo_rho = romsds.hc * (romsds.s_rho - romsds.Cs_r) + romsds.Cs_r * romsds.h - - expected = Zo_rho + romsds.zeta * (1 + Zo_rho / romsds.h) - romsds.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) - assert_allclose( - romsds.z_rho.reset_coords(drop=True), expected.reset_coords(drop=True) - ) - - romsds.cf.decode_vertical_coords(outnames={"s_rho": "ZZZ_rho"}) - assert "ZZZ_rho" in romsds.coords - - copy = romsds.copy(deep=False) - del copy["zeta"] - with pytest.raises(KeyError): - copy.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) - - copy = romsds.copy(deep=False) - copy.s_rho.attrs["formula_terms"] = "s: s_rho C: Cs_r depth: h depth_c: hc" - with pytest.raises(KeyError): - copy.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) +def test_decode_vertical_coords() -> None: + with pytest.raises( + AssertionError, match="if prefix is None, outnames must be provided" + ): + romsds.cf.decode_vertical_coords() + # needs standard names on `eta` and `depth` to derive computed standard name + romsds.h.attrs["standard_name"] = "sea_floor_depth_below_ geopotential_datum" + romsds.zeta.attrs["standard_name"] = "sea_surface_height_above_ geopotential_datum" -def test_param_vcoord_ocean_sigma_coordinate() -> None: - expected = pomds.zeta + pomds.sigma * (pomds.depth + pomds.zeta) - pomds.cf.decode_vertical_coords(outnames={"sigma": "z"}) - assert_allclose(pomds.z.reset_coords(drop=True), expected.reset_coords(drop=True)) + with pytest.warns(DeprecationWarning): + romsds.cf.decode_vertical_coords(prefix="z_rho") - copy = pomds.copy(deep=False) - del copy["zeta"] - with pytest.raises(AssertionError): - copy.cf.decode_vertical_coords() + romsds_less_h = romsds.drop_vars(["h"]) - with pytest.raises(KeyError): - copy.cf.decode_vertical_coords(outnames={}) + with pytest.raises(KeyError, match="Required term depth is absent in the dataset."): + romsds_less_h.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"}) def test_formula_terms() -> None: diff --git a/cf_xarray/tests/test_parametric.py b/cf_xarray/tests/test_parametric.py new file mode 100644 index 00000000..9708b3be --- /dev/null +++ b/cf_xarray/tests/test_parametric.py @@ -0,0 +1,615 @@ +import numpy as np +import pytest +import xarray as xr +from xarray.testing import assert_allclose + +from cf_xarray import parametric + +ps = xr.DataArray(np.ones((2, 2, 2)), dims=("time", "lat", "lon"), name="ps") + +p0 = xr.DataArray( + 10.0, + dims=(), + coords={}, + name="p0", +) + +a = xr.DataArray([0, 1, 2], dims=("lev",), name="a") + +b = xr.DataArray([6, 7, 8], dims=("lev",), name="b") + +sigma = xr.DataArray( + [0, 1, 2], dims=("lev",), name="sigma", coords={"k": (("lev",), [0, 1, 2])} +) + +eta = xr.DataArray( + np.ones((2, 2, 2)), + dims=("time", "lat", "lon"), + name="eta", + attrs={"standard_name": "sea_surface_height_above_geoid"}, +) + +depth = xr.DataArray( + np.ones((2, 2)), + dims=("lat", "lon"), + name="depth", + attrs={"standard_name": "sea_floor_depth_below_geoid"}, +) + +depth_c = xr.DataArray(30.0, dims=(), coords={}, name="depth_c") + +s = xr.DataArray([0, 1, 2], dims=("lev"), name="s") + + +def test_atmosphere_ln_pressure_coordinate(): + lev = xr.DataArray( + [0, 1, 2], + dims=("lev",), + name="lev", + ) + + transform = parametric.AtmosphereLnPressure.from_terms( + { + "p0": p0, + "lev": lev, + } + ) + + output = transform.decode() + + expected = xr.DataArray([10.0, 3.678794, 1.353353], dims=("lev",), name="p") + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "air_pressure" + + +def test_atmosphere_sigma_coordinate(): + ptop = xr.DataArray(0.98692327, dims=(), coords={}, name="ptop") + + transform = parametric.AtmosphereSigma.from_terms( + { + "sigma": sigma, + "ps": ps, + "ptop": ptop, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [[0.986923, 0.986923], [0.986923, 0.986923]], + [[0.986923, 0.986923], [0.986923, 0.986923]], + ], + [ + [[1.0, 1.0], [1.0, 1.0]], + [[1.0, 1.0], [1.0, 1.0]], + ], + [ + [[1.013077, 1.013077], [1.013077, 1.013077]], + [[1.013077, 1.013077], [1.013077, 1.013077]], + ], + ], + dims=("lev", "time", "lat", "lon"), + coords={"k": (("lev",), [0, 1, 2])}, + name="p", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "air_pressure" + + +def test_atmosphere_hybrid_sigma_pressure_coordinate(): + ap = xr.DataArray([3, 4, 5], dims=("lev",), name="ap") + + transform = parametric.AtmosphereHybridSigmaPressure.from_terms( + { + "b": b, + "ps": ps, + "a": a, + "p0": p0, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [[6.0, 6.0], [6.0, 6.0]], + [[6.0, 6.0], [6.0, 6.0]], + ], + [ + [[17.0, 17.0], [17.0, 17.0]], + [[17.0, 17.0], [17.0, 17.0]], + ], + [ + [[28.0, 28.0], [28.0, 28.0]], + [[28.0, 28.0], [28.0, 28.0]], + ], + ], + dims=("lev", "time", "lat", "lon"), + name="p", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "air_pressure" + + transform = parametric.AtmosphereHybridSigmaPressure.from_terms( + { + "b": b, + "ps": ps, + "ap": ap, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [[9.0, 9.0], [9.0, 9.0]], + [[9.0, 9.0], [9.0, 9.0]], + ], + [ + [[11.0, 11.0], [11.0, 11.0]], + [[11.0, 11.0], [11.0, 11.0]], + ], + [ + [[13.0, 13.0], [13.0, 13.0]], + [[13.0, 13.0], [13.0, 13.0]], + ], + ], + dims=("lev", "time", "lat", "lon"), + name="p", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "air_pressure" + + +def test_atmosphere_hybrid_height_coordinate(): + orog = xr.DataArray( + np.zeros((2, 2, 2)), + dims=("time", "lat", "lon"), + attrs={"standard_name": "surface_altitude"}, + ) + + transform = parametric.AtmosphereHybridHeight.from_terms( + { + "a": a, + "b": b, + "orog": orog, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [[0.0, 0.0], [0.0, 0.0]], + [[0.0, 0.0], [0.0, 0.0]], + ], + [ + [[1.0, 1.0], [1.0, 1.0]], + [[1.0, 1.0], [1.0, 1.0]], + ], + [ + [[2.0, 2.0], [2.0, 2.0]], + [[2.0, 2.0], [2.0, 2.0]], + ], + ], + dims=("lev", "time", "lat", "lon"), + name="p", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "altitude" + + +def test_atmosphere_sleve_coordinate(): + b1 = xr.DataArray([0, 0, 1], dims=("lev",), name="b1") + + b2 = xr.DataArray([1, 1, 0], dims=("lev",), name="b2") + + ztop = xr.DataArray( + 30.0, + dims=(), + coords={}, + name="ztop", + attrs={"standard_name": "altitude_at_top_of_atmosphere_model"}, + ) + + zsurf1 = xr.DataArray(np.ones((2, 2, 2)), dims=("time", "lat", "lon")) + + zsurf2 = xr.DataArray(np.ones((2, 2, 2)), dims=("time", "lat", "lon")) + + transform = parametric.AtmosphereSleve.from_terms( + { + "a": a, + "b1": b1, + "b2": b2, + "ztop": ztop, + "zsurf1": zsurf1, + "zsurf2": zsurf2, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [[1.0, 1.0], [1.0, 1.0]], + [[1.0, 1.0], [1.0, 1.0]], + ], + [ + [[31.0, 31.0], [31.0, 31.0]], + [[31.0, 31.0], [31.0, 31.0]], + ], + [ + [[61.0, 61.0], [61.0, 61.0]], + [[61.0, 61.0], [61.0, 61.0]], + ], + ], + dims=("lev", "time", "lat", "lon"), + name="z", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "altitude" + + +def test_ocean_sigma_coordinate(): + transform = parametric.OceanSigma.from_terms( + { + "sigma": sigma, + "eta": eta, + "depth": depth, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + ], + [ + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + ], + ], + dims=("time", "lat", "lon", "lev"), + name="z", + coords={"k": (("lev",), [0, 1, 2])}, + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "altitude" + + +def test_ocean_s_coordinate(): + _a = xr.DataArray(1, dims=(), coords={}, name="a") + + _b = xr.DataArray(1, dims=(), coords={}, name="b") + + transform = parametric.OceanS.from_terms( + { + "s": s, + "eta": eta, + "depth": depth, + "a": _a, + "b": _b, + "depth_c": depth_c, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [ + [12.403492, 40.434874, 70.888995], + [12.403492, 40.434874, 70.888995], + ], + [ + [12.403492, 40.434874, 70.888995], + [12.403492, 40.434874, 70.888995], + ], + ], + [ + [ + [12.403492, 40.434874, 70.888995], + [12.403492, 40.434874, 70.888995], + ], + [ + [12.403492, 40.434874, 70.888995], + [12.403492, 40.434874, 70.888995], + ], + ], + ], + dims=("time", "lat", "lon", "lev"), + name="z", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "altitude" + + +def test_ocean_s_coordinate_g1(): + C = xr.DataArray([0, 1, 2], dims=("lev",), name="C") + + transform = parametric.OceanSG2.from_terms( + { + "s": s, + "c": C, + "eta": eta, + "depth": depth, + "depth_c": depth_c, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + ], + [ + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + ], + ], + dims=("time", "lat", "lon", "lev"), + name="z", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "altitude" + + +def test_ocean_s_coordinate_g2(): + C = xr.DataArray([0, 1, 2], dims=("lev",), name="C") + + transform = parametric.OceanSG2.from_terms( + { + "s": s, + "c": C, + "eta": eta, + "depth": depth, + "depth_c": depth_c, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + ], + [ + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + [ + [1.0, 3.0, 5.0], + [1.0, 3.0, 5.0], + ], + ], + ], + dims=("time", "lat", "lon", "lev"), + name="z", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "altitude" + + +def test_ocean_sigma_z_coordinate(): + zlev = xr.DataArray( + [0, 1, np.nan], dims=("lev",), name="zlev", attrs={"standard_name": "altitude"} + ) + + _sigma = xr.DataArray([np.nan, np.nan, 3], dims=("lev",), name="sigma") + + transform = parametric.OceanSigmaZ.from_terms( + { + "sigma": _sigma, + "eta": eta, + "depth": depth, + "depth_c": depth_c, + "nsigma": 10, + "zlev": zlev, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [ + [[0.0, 0.0], [0.0, 0.0]], + [[0.0, 0.0], [0.0, 0.0]], + ], + [ + [[1.0, 1.0], [1.0, 1.0]], + [[1.0, 1.0], [1.0, 1.0]], + ], + [ + [[7.0, 7.0], [7.0, 7.0]], + [[7.0, 7.0], [7.0, 7.0]], + ], + ], + dims=("lev", "time", "lat", "lon"), + name="z", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "altitude" + + +def test_ocean_double_sigma_coordinate(): + k_c = xr.DataArray( + 1, + dims=(), + coords={}, + name="k_c", + ) + + href = xr.DataArray( + 20.0, + dims=(), + coords={}, + name="href", + ) + + z1 = xr.DataArray( + 10.0, + dims=(), + coords={}, + name="z1", + ) + + z2 = xr.DataArray( + 30.0, + dims=(), + coords={}, + name="z2", + ) + + a = xr.DataArray( + 2.0, + dims=(), + coords={}, + name="a", + ) + + transform = parametric.OceanDoubleSigma.from_terms( + { + "sigma": sigma, + "depth": depth, + "z1": z1, + "z2": z2, + "a": a, + "href": href, + "k_c": k_c, + } + ) + + output = transform.decode() + + expected = xr.DataArray( + [ + [[0.0, 0.0], [0.0, 0.0]], + [[10.010004, 10.010004], [10.010004, 10.010004]], + [[1.0, 1.0], [1.0, 1.0]], + ], + dims=("lev", "lat", "lon"), + coords={"k": (("lev",), [0, 1, 2])}, + name="z", + ) + + assert_allclose(output, expected) + + assert output.attrs["standard_name"] == "altitude" + + +@pytest.mark.parametrize( + "input,expected", + [ + ({"zlev": {"standard_name": "altitude"}}, "altitude"), + ( + { + "zlev": {"standard_name": "altitude"}, + "eta": {"standard_name": "sea_surface_height_above_geoid"}, + }, + "altitude", + ), + ( + { + "zlev": {"standard_name": "altitude"}, + "eta": {"standard_name": "sea_surface_height_above_geoid"}, + "depth": {"standard_name": "sea_floor_depth_below_geoid"}, + }, + "altitude", + ), + ( + { + "eta": {"standard_name": "sea_surface_height_above_geoid"}, + "depth": {"standard_name": "sea_floor_depth_below_geoid"}, + }, + "altitude", + ), + ], +) +def test_derive_ocean_stdname(input, expected): + output = parametric._derive_ocean_stdname(**input) + + assert output == expected + + +def test_derive_ocean_stdname_no_standard_name(): + with pytest.raises( + ValueError, match="The standard name for the 'zlev' variable is not available." + ): + parametric._derive_ocean_stdname(zlev={}) + + +def test_derive_ocean_stdname_no_match(): + with pytest.raises( + ValueError, + match="Could not derive standard name from combination of not in any list.", + ): + parametric._derive_ocean_stdname(zlev={"standard_name": "not in any list"}) diff --git a/doc/parametricz.md b/doc/parametricz.md index 77f67924..897ae838 100644 --- a/doc/parametricz.md +++ b/doc/parametricz.md @@ -28,7 +28,7 @@ xr.set_options(display_expand_data=False) 3. {py:attr}`Dataset.cf.formula_terms` ``` -`cf_xarray` supports decoding [parametric vertical coordinates](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-vertical-coordinate) encoded in the `formula_terms` attribute using {py:meth}`Dataset.cf.decode_vertical_coords`. Right now, only the two ocean s-coordinates and `ocean_sigma_coordinate` are supported, but support for the [rest](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-v-coord) should be easy to add (Pull Requests are very welcome!). +`cf_xarray` supports decoding [parametric vertical coordinates](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-vertical-coordinate) encoded in the `formula_terms` attribute using {py:meth}`Dataset.cf.decode_vertical_coords`. ## Decoding parametric coordinates