Support hybrid level coefficients

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}

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")

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,14 +151,16 @@ def __hash__(self):
     #     "layer": True,
     #     "positive": POSITIVE_DOWN,
     # },
-    "MODEL": {
+    "HYBRID": {
         "name": "hybrid",
         "abbreviation": "ml",
         "standard_name": "atmosphere_hybrid_sigma_pressure_coordinate",
         "long_name": "hybrid level",
         "units": "1",
         "layer": False,
         "positive": POSITIVE_DOWN,
+        "parametric": True,
+        "coefficient_names": ("A", "B"),
     },
     "THETA": {
         "name": "potential_temperature",

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,14 +391,26 @@ 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})"
 
     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)