diff --git a/src/earthkit/data/encoders/grib.py b/src/earthkit/data/encoders/grib.py index 63b095e21..ddf9d1c24 100644 --- a/src/earthkit/data/encoders/grib.py +++ b/src/earthkit/data/encoders/grib.py @@ -769,6 +769,10 @@ def _make_new_handle( self._update_metadata(handle, metadata, compulsory, can_infer_time) + # right now the encoder is only able to write pv for edition 2 + if "pv" in metadata and metadata.get("edition", None) != 2: + metadata["edition"] = 2 + # eccodes keys are order dependent KEY_ORDER = ("edition", "stepType") r = {k: metadata.pop(k) for k in KEY_ORDER if k in metadata} diff --git a/src/earthkit/data/field/component/level_parameters.py b/src/earthkit/data/field/component/level_parameters.py new file mode 100644 index 000000000..77f5af6f3 --- /dev/null +++ b/src/earthkit/data/field/component/level_parameters.py @@ -0,0 +1,77 @@ +# (C) Copyright 2022 ECMWF. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. +# +from abc import ABCMeta, abstractmethod + +from .level_type import LevelTypes, get_level_type + + +class LevelParameters(metaclass=ABCMeta): + _LEVEL_TYPE = None + + def level_type(self): + return self._LEVEL_TYPE + + @abstractmethod + def number_of_levels(self): + pass + + @abstractmethod + def coefficients(self): + pass + + def coefficient_names(self): + return self._LEVEL_TYPE.coefficient_names + + @abstractmethod + def coefficient_size(self): + pass + + +class HybridLevelParametersBase(LevelParameters): + _LEVEL_TYPE = LevelTypes.HYBRID.value + + +class HybridLevelParameters(HybridLevelParametersBase): + def __init__(self, A, B): + self._A = A + self._B = B + if len(A) != len(B): + raise ValueError("A and B coefficient arrays must have the same length") + + def number_of_levels(self): + return len(self._A) - 1 + + def coefficients(self): + return self._A, self._B + + def coefficient_size(self): + return 2 * (len(self._A)) + + +def create_level_parameters(level_type, coefficients) -> LevelParameters: + if coefficients is None: + raise ValueError("Coefficients must be provided for level types that require them") + + if isinstance(coefficients, LevelParameters): + if coefficients.level_type() != level_type: + raise ValueError( + f"Provided LevelParameters of type {coefficients.level_type()} do not match" + f" expected level type {level_type}" + ) + return coefficients + + level_type = get_level_type(level_type) + + if level_type == LevelTypes.HYBRID: + if len(coefficients) != 2: + raise ValueError("Hybrid level type requires two coefficient arrays (A and B)") + A, B = coefficients + return HybridLevelParameters(A, B) + + raise ValueError(f"Unsupported level type '{level_type}' or invalid coefficients") diff --git a/src/earthkit/data/field/component/level_type.py b/src/earthkit/data/field/component/level_type.py index f96282493..ba90f0731 100644 --- a/src/earthkit/data/field/component/level_type.py +++ b/src/earthkit/data/field/component/level_type.py @@ -64,6 +64,8 @@ def __init__( units: Union[str, Units], layer: bool, positive: str, + parametric: bool = False, + coefficient_names: tuple[str, ...] | None = None, ) -> None: """Initialise the LevelType object. @@ -88,6 +90,10 @@ def __init__( Whether the level type represents a layer or a single level. positive : str The positive direction of the level type. Can be either "up" or "down". + parametric : bool, optional + Whether the level type is parametric. Default is False. + coefficient_names : tuple[str, ...] | None, optional + The names of the coefficients for the level type, if parametric is True. """ self.name = name self.abbreviation = abbreviation @@ -96,6 +102,8 @@ def __init__( self.units = Units.from_any(units) self.layer = layer self.positive = positive + self.parametric = parametric + self.coefficient_names = coefficient_names self.cf = { "standard_name": self.standard_name, "long_name": self.long_name, @@ -114,6 +122,8 @@ def __eq__(self, other) -> bool: return self.name == other.name if isinstance(other, str): return self.name == other + if isinstance(other, LevelTypes): + return self.name == other.value.name return False def __repr__(self): @@ -141,7 +151,7 @@ def __hash__(self): # "layer": True, # "positive": POSITIVE_DOWN, # }, - "MODEL": { + "HYBRID": { "name": "hybrid", "abbreviation": "ml", "standard_name": "atmosphere_hybrid_sigma_pressure_coordinate", @@ -149,6 +159,8 @@ def __hash__(self): "units": "1", "layer": False, "positive": POSITIVE_DOWN, + "parametric": True, + "coefficient_names": ("A", "B"), }, "THETA": { "name": "potential_temperature", diff --git a/src/earthkit/data/field/component/vertical.py b/src/earthkit/data/field/component/vertical.py index 8ca890ec7..261cbea50 100644 --- a/src/earthkit/data/field/component/vertical.py +++ b/src/earthkit/data/field/component/vertical.py @@ -16,6 +16,7 @@ from earthkit.utils.units import Units from .component import SimpleFieldComponent, component_keys, mark_get_key +from .level_parameters import LevelParameters from .level_type import LevelType, get_level_type @@ -148,6 +149,48 @@ def level_type(self) -> str: """ pass + @mark_get_key + @abstractmethod + def parametric(self) -> bool: + """Return whether the vertical type is parametric. + + A parametric vertical type is one that is defined by a set of parameters or coefficients + (e.g. "hybrid"). This method returns True for parametric vertical types and False for others. + """ + pass + + @mark_get_key + @abstractmethod + def number_of_levels(self) -> int | None: + """Return the number of levels. + + The number of levels is only applicable for certain vertical types (e.g. "hybrid"), and + will return None for other types. + """ + pass + + @mark_get_key + @abstractmethod + def coefficients(self) -> float | tuple[float, ...] | None: + """Return the level definition coefficients. + + The level definition coefficients are only applicable for certain vertical types (e.g. "hybrid"), + and will return None for other types. The coefficients can be a single sequence or a tuple of sequences, + depending on the specific vertical type. + """ + pass + + @mark_get_key + @abstractmethod + def coefficient_names(self) -> tuple[str] | None: + """Return the names of the level definition coefficients. + + The level definition coefficient names are only applicable for certain vertical types (e.g. "hybrid"), + and will return None for other types. The names can be a single sequence or a tuple of sequences, + depending on the specific vertical type. + """ + pass + def create_vertical(d: dict) -> "Vertical": """Create a Vertical object from a dictionary. @@ -165,9 +208,14 @@ def create_vertical(d: dict) -> "Vertical": if not isinstance(d, dict): raise TypeError(f"Cannot create Vertical from {type(d)}, expected dict") - cls = Vertical - d1 = cls._normalise_create_kwargs(d, allowed_keys=("level", "layer", "level_type")) - return cls(**d1) + allowed_keys = ("level", "layer", "level_type", "coefficients") + d = Vertical._normalise_create_kwargs(d, allowed_keys=allowed_keys) + coefficients = d.pop("coefficients", None) + + if coefficients is None: + return Vertical(**d) + else: + return ParametricVertical(**d, coefficients=coefficients) class EmptyVertical(VerticalBase): @@ -225,6 +273,34 @@ def level_type(self) -> str: """ return self._type.name + def parametric(self) -> bool: + """Return whether the level type is parametric. + + The parametric information is not available for this vertical type, and this method returns None. + """ + return False + + def number_of_levels(self) -> None: + """Return the number of levels. + + The number of levels is not applicable for this vertical type, and this method returns None. + """ + return None + + def coefficients(self) -> None: + """Return the level definition coefficients. + + The level definition coefficients are not applicable for this vertical type, and this method returns None. + """ + return None + + def coefficient_names(self) -> None: + """Return the names of the level definition coefficients. + + The coefficient names are not applicable for this vertical type, and this method returns None. + """ + return None + @classmethod def from_dict(cls, d: dict) -> "VerticalBase": """Create a Vertical object from a dictionary. @@ -280,6 +356,8 @@ class Vertical(VerticalBase): metadata such as units and CF attributes. """ + _level_parameters: Optional[LevelParameters] = None + def __init__( self, level: Union[int, float] = None, @@ -313,6 +391,18 @@ def positive(self) -> Optional[str]: def level_type(self) -> str: return self._type.name + def parametric(self) -> Optional[bool]: + return self._type.parametric + + def number_of_levels(self) -> Optional[int]: + return None + + def coefficients(self) -> Optional[int]: + return None + + def coefficient_names(self) -> Optional[int]: + return None + def __print__(self) -> str: return f"{self.level} {self.units} ({self.abbreviation})" @@ -320,7 +410,7 @@ def __repr__(self) -> str: return f"{self.__class__.__name__}(level={self.level()}, units={self.units()}, level_type={self._type.name})" @classmethod - def from_dict(cls, d: dict, allow_unused=False) -> "Vertical": + def from_dict(cls, d: dict) -> "Vertical": """Create a Vertical object from a dictionary. Parameters @@ -394,12 +484,95 @@ def set(self, *args, **kwargs) -> "Vertical": specified as a string, it will be converted to a LevelType object using the :func:`get_level_type` function. """ - d = self._normalise_set_kwargs(*args, allowed_keys=("level", "layer", "level_type"), **kwargs) + d = self._normalise_set_kwargs(*args, allowed_keys=("level", "layer", "level_type", "coefficients"), **kwargs) + + current = { + "level": self._level, + "layer": self._layer, + "level_type": self._type.name, + } + + current.update(d) + return self.from_dict(current) + + +class ParametricVertical(Vertical): + """Vertical component with additional parameters. + + This class extends the Vertical class to include additional parameters that may be relevant for certain + vertical types. The additional parameters are represented as a LevelParameters object, which can + contain any number of parameters depending on the specific vertical type. + + The additional parameters can be accessed using the :meth:`number_of_levels` and :meth:`coefficients` + methods, which return the number of levels and the level definition coefficients, respectively. + These methods will return None if the additional parameters are not defined for the vertical type. + """ + + def __init__( + self, + level: Union[int, float] = None, + layer: Optional[tuple[float, float]] = None, + level_type: Optional[Union[LevelType, str]] = None, + coefficients: tuple[float] | LevelParameters | None = None, + ) -> None: + + super().__init__(level=level, layer=layer, level_type=level_type) + if not self._type.parametric: + raise ValueError(f"Level type {self._type.name} does not support parametric coefficients") + + if coefficients is None: + raise ValueError("ParametricVertical requires a valid level_parameters to be provided, got None") + from earthkit.data.field.component.level_parameters import create_level_parameters + + self._coefficients = create_level_parameters(level_type, coefficients) + + def number_of_levels(self) -> int: + return self._coefficients.number_of_levels() + + def coefficients(self) -> float | tuple[float, ...]: + return self._coefficients.coefficients() + + def coefficient_names(self) -> str | tuple[str, ...]: + return self._coefficients.coefficient_names() + + def to_dict(self): + d = super().to_dict() + d["number_of_levels"] = self.number_of_levels() + return d + + def set(self, *args, **kwargs) -> "Vertical": + """Create a new instance with updated data. + + Parameters + ---------- + args : tuple + Positional arguments containing time data. Only dictionaries are allowed. + kwargs : dict + Keyword arguments containing time data. + + Returns + ------- + Vertical + The created Vertical instance. + + + The allowed keys in the dictionaries and keyword arguments are: + + - "level" + - "layer" + - "level_type" + + The "level_type" key can be specified as a LevelType object or as a string. If + specified as a string, it will be converted to a LevelType object using + the :func:`get_level_type` function. + """ + d = self._normalise_set_kwargs(*args, allowed_keys=("level", "layer", "level_type", "coefficients"), **kwargs) current = { "level": self._level, "layer": self._layer, "level_type": self._type.name, + "coefficients": self.coefficients(), } current.update(d) diff --git a/src/earthkit/data/field/grib/vertical.py b/src/earthkit/data/field/grib/vertical.py index c2d01f6a7..4db2eac74 100644 --- a/src/earthkit/data/field/grib/vertical.py +++ b/src/earthkit/data/field/grib/vertical.py @@ -9,12 +9,39 @@ from abc import ABCMeta, abstractmethod from collections import defaultdict +from earthkit.data.field.component.level_parameters import HybridLevelParametersBase from earthkit.data.field.component.level_type import LevelTypes from .collector import GribContextCollector from .core import GribFieldComponentHandler +class GribHybridLevelParameters(HybridLevelParametersBase): + def __init__(self, handle): + self._handle = handle + + def number_of_levels(self): + coeff_num = self._handle.get("NV", default=None) + if coeff_num is not None: + return int(coeff_num / 2) - 1 + return None + + def coefficients(self): + pv = self._handle.get("pv", default=None) + if pv is not None: + import numpy as np + + coeff_num = int(len(pv) / 2) + A = np.array(pv[:coeff_num]) + B = np.array(pv[coeff_num:]) + return A, B + + return None + + def coefficient_size(self): + return 2 * (self.number_of_levels() + 1) + + class Converter: def from_grib(self, value): return value @@ -62,7 +89,7 @@ def from_grib(self, handle): pass @abstractmethod - def to_grib(self, component): + def to_grib(self, component, handle=None): pass @abstractmethod @@ -82,9 +109,9 @@ def _get(key, default=None): level = self.converter.from_grib(level) layer = None - return level, layer, self.component_type + return {"level": level, "layer": layer, "level_type": self.component_type} - def to_grib(self, component): + def to_grib(self, component, handle=None): level = self.converter.to_grib(component.level()) return { @@ -96,6 +123,37 @@ def match(self, component): return self.component_matcher.match(component) +class HybridGribLevelType(GribLevelType): + def from_grib(self, handle): + result = super().from_grib(handle) + level_parameters = GribHybridLevelParameters(handle) + result["coefficients"] = level_parameters + return result + + def to_grib(self, component, handle=None): + r = super().to_grib(component) + + # this tries to avoid writing the coefficients back to the handle if they are already + # present and correct, which can be expensive for large coefficient arrays. The check + # not robust enough and should be improved. + if handle is not None and hasattr(component, "level_parameters"): + level_parameters = component._level_parameters + if isinstance(level_parameters, GribFieldComponentHandler): + nv = handle.get("NV", default=None) + if level_parameters.coefficient_size() == nv: + return r + + coefficients = component.coefficients() + if coefficients is not None: + import numpy as np + + A, B = coefficients + r["NV"] = len(A) + len(B) + r["pv"] = np.concatenate([A, B]) + + return r + + class GribLayerType(GribVerticalType): def from_grib(self, handle): def _get(key, default=None): @@ -109,9 +167,9 @@ def _get(key, default=None): level = top layer = (top, bottom) - return level, layer, self.component_type + return {"level": level, "layer": layer, "level_type": self.component_type} - def to_grib(self, component): + def to_grib(self, component, handle=None): layer = self.converter.to_grib(component.layer()) top, bottom = layer @@ -154,7 +212,7 @@ def _get(key, default=None): GribLevelType("generalVerticalLayer", LevelTypes.GENERAL), GribLevelType("heightAboveSea", LevelTypes.HEIGHT_AS_LEVEL), GribLevelType("heightAboveGround", LevelTypes.HEIGHT_AG_LEVEL), - GribLevelType("hybrid", LevelTypes.MODEL), + HybridGribLevelType("hybrid", LevelTypes.HYBRID), GribLevelType("iceLayerOnWater", LevelTypes.ICE_LAYER_ON_WATER), GribLayerType("isobaricLayer", LevelTypes.PRESSURE), # hPa GribLevelType("isothermal", LevelTypes.TEMPERATURE), @@ -218,7 +276,7 @@ def _build_dict(handle): if t is None: from earthkit.data.field.component.level_type import get_level_type - level, layer = from_grib(handle) + # level, layer = from_grib(handle) component = get_level_type(level_type) t = register_grib_level_type( key=level_type, @@ -228,19 +286,28 @@ def _build_dict(handle): if t is None: raise ValueError(f"Unsupported level type: {level_type}") - level, layer, level_type = t.from_grib(handle) + r = t.from_grib(handle) + return r + + # level, layer, level_type = t.from_grib(handle) + + # level_parameters = None + # if level_type == LevelTypes.HYBRID: + # level_parameters = GribHybridLevelParameters(handle) - return dict( - level=level, - layer=layer, - level_type=level_type, - ) + # return dict( + # level=level, + # layer=layer, + # level_type=level_type, + # level_parameters=level_parameters, + # ) class GribVerticalContextCollector(GribContextCollector): @staticmethod def collect_keys(handler, context): component = handler.component + handle = context.get("handle") grib_level_type = _COMPONENT_TYPES.get(component.level_type()) if isinstance(grib_level_type, tuple): t = grib_level_type @@ -253,7 +320,7 @@ def collect_keys(handler, context): if grib_level_type is None: raise ValueError(f"Unknown level type: {component.level_type()}") - r = grib_level_type.to_grib(component) + r = grib_level_type.to_grib(component, handle) context.update(r) diff --git a/src/earthkit/data/field/mars/vertical.py b/src/earthkit/data/field/mars/vertical.py index a8effeceb..1e5feed2e 100644 --- a/src/earthkit/data/field/mars/vertical.py +++ b/src/earthkit/data/field/mars/vertical.py @@ -66,7 +66,7 @@ def from_mars(self, request): _TYPES = [ MarsLevelType("pl", LevelTypes.PRESSURE), - MarsLevelType("ml", LevelTypes.MODEL), + MarsLevelType("ml", LevelTypes.HYBRID), MarsLevelType("pt", LevelTypes.THETA), MarsLevelType("pv", LevelTypes.PV), MarsLevelType("sfc", LevelTypes.SURFACE), diff --git a/src/earthkit/data/field/xarray/vertical.py b/src/earthkit/data/field/xarray/vertical.py index e820e5777..eac9753d0 100644 --- a/src/earthkit/data/field/xarray/vertical.py +++ b/src/earthkit/data/field/xarray/vertical.py @@ -23,7 +23,7 @@ def __init__(self, keys, spec_type): _TYPES = [ XarrayLevelType("pl", LevelTypes.PRESSURE), - XarrayLevelType("ml", LevelTypes.MODEL), + XarrayLevelType("ml", LevelTypes.HYBRID), XarrayLevelType("sfc", LevelTypes.SURFACE), ] diff --git a/src/earthkit/data/utils/message.py b/src/earthkit/data/utils/message.py index b388dab56..fb34ccfd4 100644 --- a/src/earthkit/data/utils/message.py +++ b/src/earthkit/data/utils/message.py @@ -293,7 +293,7 @@ def set(self, name, value): try: assert self.path is None, "Only cloned handles can have values changed" - if isinstance(value, list): + if isinstance(value, (list, np.ndarray, tuple)): return eccodes.codes_set_array(self._handle, name, value) return eccodes.codes_set(self._handle, name, value) diff --git a/tests/field/test_vertical_component.py b/tests/field/test_vertical_component.py index 2c1661d62..f9b874a91 100644 --- a/tests/field/test_vertical_component.py +++ b/tests/field/test_vertical_component.py @@ -11,8 +11,13 @@ import pytest +from earthkit.data.field.component.level_parameters import HybridLevelParameters from earthkit.data.field.component.level_type import get_level_type -from earthkit.data.field.component.vertical import Vertical +from earthkit.data.field.component.vertical import ParametricVertical, Vertical + +A = (0.1, 0.2, 0.3, 0.4) +B = (0.4, 0.5, 0.6, 0.7) +hybrid_params = HybridLevelParameters(A=A, B=B) def test_vertical_component_alias_1(): @@ -29,19 +34,54 @@ def test_vertical_component_alias_1(): "level": 1000, "level_type": "pressure", }, - (1000, "pressure"), + { + "level": 1000, + "level_type": "pressure", + "number_of_levels": None, + "coefficients": None, + }, ), ( { "level": 1000, "level_type": "pressure", }, - (1000, "pressure"), + { + "level": 1000, + "level_type": "pressure", + "number_of_levels": None, + "coefficients": None, + }, + ), + ( + { + "level": 2, + "level_type": "hybrid", + "coefficients": (A, B), + }, + { + "level": 2, + "level_type": "hybrid", + "number_of_levels": 3, + "coefficients": (A, B), + }, + ), + ( + { + "level": 2, + "level_type": "hybrid", + "coefficients": hybrid_params, + }, + { + "level": 2, + "level_type": "hybrid", + "number_of_levels": 3, + "coefficients": (A, B), + }, ), ], ) def test_vertical_component_from_dict_ok(input_d, ref): - if not isinstance(input_d, list): input_d = [input_d] @@ -49,8 +89,10 @@ def test_vertical_component_from_dict_ok(input_d, ref): for d in input_d: r = Vertical.from_dict(d) - assert r.level() == ref[0] - assert r.level_type() == "pressure" + assert r.level() == ref["level"] + assert r.level_type() == ref["level_type"] + assert r.number_of_levels() == ref.get("number_of_levels") + assert r.coefficients() == ref.get("coefficients") def test_vertical_component_type(): @@ -68,6 +110,8 @@ def test_vertical_component_type(): "units": "hectopascal", "positive": "down", } + assert t.parametric is False + assert t.coefficient_names is None assert r.level() == 1000 assert r.level_type() == "pressure" @@ -80,6 +124,9 @@ def test_vertical_component_type(): "units": "hectopascal", "positive": "down", } + assert r.parametric() is False + assert r.number_of_levels() is None + assert r.coefficient_names() is None p_type = get_level_type("pressure") assert p_type == t @@ -110,8 +157,7 @@ def test_vertical_component_type(): ), ], ) -def test_vertical_component_set(input_d, ref): - +def test_vertical_component_set_1(input_d, ref): r = Vertical(level=1000, level_type="pressure") if not isinstance(input_d, list): @@ -127,3 +173,38 @@ def test_vertical_component_set(input_d, ref): # the original object is unchanged assert r.level() == 1000 assert r.level_type() == "pressure" + + +def test_vertical_component_set_hybrid_1(): + r = ParametricVertical(level=1, level_type="hybrid", coefficients=(A, B)) + + r1 = r.set({"level": 2}) + + assert r1.level() == 2 + assert r1.level_type() == "hybrid" + assert r1.number_of_levels() == 3 + A1, B1 = r1.coefficients() + assert A1 == A + assert B1 == B + + +def test_vertical_component_set_hybrid_2(): + r = ParametricVertical(level=1, level_type="hybrid", coefficients=(A, B)) + + A_new = (0.15, 0.25, 0.35, 0.45, 0.55) + B_new = (0.45, 0.55, 0.65, 0.75, 0.85) + r1 = r.set({"level": 2}, coefficients=(A_new, B_new)) + + assert r1.level() == 2 + assert r1.level_type() == "hybrid" + assert r1.number_of_levels() == 4 + A1, B1 = r1.coefficients() + assert A1 == A_new + assert B1 == B_new + + +def test_vertical_component_register(): + r = Vertical(level=1000, level_type="_my_level_type") + + assert r.level() == 1000 + assert r.level_type() == "_my_level_type" diff --git a/tests/grib/test_grib_set_vertical.py b/tests/grib/test_grib_set_vertical.py index e49aa2d18..9f13c8ead 100644 --- a/tests/grib/test_grib_set_vertical.py +++ b/tests/grib/test_grib_set_vertical.py @@ -9,7 +9,7 @@ # nor does it submit to any jurisdiction. # - +import numpy as np import pytest from grib_fixtures import load_grib_data # noqa: E402 @@ -102,3 +102,115 @@ def test_grib_set_vertical(fl_type, _kwargs, ref1, grib_ref, ref2): assert len(f_saved) == 1 for k, v in ref2.items(): assert f_saved[0].get(k) == v, f"key {k} expected {v} got {f_saved[0].get(k)}" + + +@pytest.mark.parametrize("fl_type", ["file"]) +def test_grib_set_vertical_hybrid_1(fl_type): + ds_ori, _ = load_grib_data("ml_data.grib", fl_type, folder="data") + f_ori = ds_ori[0] + + f = f_ori.set({"vertical.level": 2}) + assert f.vertical.level() == 2 + assert f.vertical.level_type() == "hybrid" + assert f.vertical.number_of_levels() == 137 + A, B = f.vertical.coefficients() + assert len(A) == 138 + assert len(B) == 138 + + +@pytest.mark.parametrize("fl_type", ["file"]) +@pytest.mark.parametrize("coeff_mode", ["object", "tuple"]) +@pytest.mark.parametrize( + "input_file, _kwargs", + [ + (("ml_data.grib", "data"), {}), + (("test4.grib", "example"), {"vertical.level_type": "hybrid", "vertical.level": 1}), + ], +) +def test_grib_set_vertical_hybrid_coefficients_1(fl_type, coeff_mode, input_file, _kwargs): + file_name, folder = input_file + ds_ori, _ = load_grib_data(file_name, fl_type, folder=folder) + f_ori = ds_ori[0] + + A = np.array([0.0, 2.000365, 4.00073, 6.001095] + [0.0] * 88) + B = np.array([0.0, 0.0, 0.0, 0.0] + [1.0] * 88) + + level_num = len(A) - 1 + + if coeff_mode == "object": + from earthkit.data.field.component.level_parameters import HybridLevelParameters + + coeff = HybridLevelParameters(A=A, B=B) + + elif coeff_mode == "tuple": + coeff = (A, B) + else: + raise ValueError(f"Invalid coeff_mode '{coeff_mode}'") + + f = f_ori.set({"vertical.coefficients": coeff}, _kwargs) + + ref1 = { + "vertical.level": 1, + "vertical.level_type": "hybrid", + "vertical.number_of_levels": level_num, + "metadata.levelist": None, + "metadata.level": None, + "metadata.levtype": None, + "metadata.typeOfLevel": None, + } + + for k, v in ref1.items(): + assert f.get(k) == v + + A1, B1 = f.vertical.coefficients() + assert np.allclose(A1, A) + assert np.allclose(B1, B) + + f = f.sync() + + ref2 = { + "vertical.level": 1, + "vertical.level_type": "hybrid", + "vertical.number_of_levels": level_num, + "metadata.levelist": 1, + "metadata.level": 1, + "metadata.levtype": "ml", + "metadata.typeOfLevel": "hybrid", + } + + for k, v in ref2.items(): + assert f.get(k) == v + + A1, B1 = f.vertical.coefficients() + assert np.allclose(A1, A) + assert np.allclose(B1, B) + + with temp_file() as tmp: + f.to_target("file", tmp) + f_saved = from_source("file", tmp).to_fieldlist() + assert len(f_saved) == 1 + for k, v in ref2.items(): + assert f_saved[0].get(k) == v, f"key {k} expected {v} got {f_saved[0].get(k)}" + + A2, B2 = f_saved[0].vertical.coefficients() + assert np.allclose(A, A2) + assert np.allclose(B, B2) + + +@pytest.mark.parametrize("fl_type", ["file"]) +def test_grib_set_vertical_hybrid_bad(fl_type): + ds_ori, _ = load_grib_data("ml_data.grib", fl_type, folder="data") + f_ori = ds_ori[0] + + from earthkit.data.field.component.level_parameters import HybridLevelParameters + + A = np.array([0.0, 2.000365, 4.00073, 6.001095] + [0.0] * 88) + B = np.array([0.0, 0.0, 0.0, 0.0] + [1.0] * 88) + + level_parameters = HybridLevelParameters(A=A, B=B) + + with pytest.raises(ValueError): + f_ori.set({"vertical.level_type": "pressure", "vertical.coefficients": (A, B)}) + + with pytest.raises(ValueError): + f_ori.set({"vertical.level_type": "pressure", "vertical.coefficients": level_parameters}) diff --git a/tests/grib/test_grib_vertical.py b/tests/grib/test_grib_vertical.py index 46749a8b6..ddb7513fb 100644 --- a/tests/grib/test_grib_vertical.py +++ b/tests/grib/test_grib_vertical.py @@ -29,6 +29,13 @@ def test_grib_vertical_1(fl_type): assert f.vertical.level() == 0 assert f.vertical.layer() is None assert f.vertical.level_type() == "surface" + assert f.vertical.number_of_levels() is None + assert f.vertical.coefficients() is None + assert f.get("vertical.level") == 0 + assert f.get("vertical.level_type") == "surface" + assert f.get("vertical.layer") is None + assert f.get("vertical.number_of_levels") is None + assert f.get("vertical.coefficients") is None @pytest.mark.parametrize("fl_type", FL_TYPES) @@ -41,6 +48,8 @@ def test_grib_vertical_2(fl_type): assert f.vertical.units() == "hPa" assert f.vertical.abbreviation() == "pl" assert f.vertical.positive() == "down" + assert f.vertical.number_of_levels() is None + assert f.vertical.coefficients() is None @pytest.mark.cache @@ -143,3 +152,49 @@ def test_grib_vertical_layer(fname, expected_values): assert np.isclose(f.vertical.layer()[0], bottom) assert np.isclose(f.vertical.layer()[1], top) assert f.vertical.level_type() == level_type + + +@pytest.mark.parametrize("fl_type", FL_TYPES) +def test_grib_vertical_hybrid(fl_type): + ds, _ = load_grib_data("ml_data.grib", fl_type, folder="data") + f = ds[0] + + assert f.vertical.level() == 1 + assert f.vertical.layer() is None + assert f.vertical.level_type() == "hybrid" + assert f.vertical.number_of_levels() == 137 + A, B = f.vertical.coefficients() + assert len(A) == 138 + assert len(B) == 138 + assert np.allclose( + A[:4], + [ + 0.0, + 2.000365, + 3.102241, + 4.666084, + ], + ) + assert np.allclose(B[:4], [0.0, 0.0, 0.0, 0.0]) + assert np.allclose(A[-4:], [22.835938, 3.757813, 0.0, 0.0]) + assert np.allclose(B[-4:], [0.9919840097, 0.9950025082, 0.9976301193, 1.0]) + + assert f.vertical.coefficient_names() == ("A", "B") + + assert f.get("vertical.level") == 1 + assert f.get("vertical.level_type") == "hybrid" + assert f.get("vertical.layer") is None + assert f.get("vertical.number_of_levels") == 137 + A_get, B_get = f.get("vertical.coefficients") + assert np.allclose( + A_get[:4], + [ + 0.0, + 2.000365, + 3.102241, + 4.666084, + ], + ) + assert np.allclose(B_get[:4], [0.0, 0.0, 0.0, 0.0]) + assert np.allclose(A_get[-4:], [22.835938, 3.757813, 0.0, 0.0]) + assert np.allclose(B_get[-4:], [0.9919840097, 0.9950025082, 0.9976301193, 1.0])