pycontrails 0.41.0__cp311-cp311-macosx_11_0_arm64.whl → 0.42.2__cp311-cp311-macosx_11_0_arm64.whl

pycontrails/models/humidity_scaling.py

@@ -1,17 +1,20 @@
-"""Support for humidity scaling methodologies on ERA5 data."""
+"""Support for humidity scaling methodologies."""
 
 from __future__ import annotations
 
 import abc
 import dataclasses
+import functools
+import pathlib
 import warnings
 from typing import Any, NoReturn, overload
 
 import numpy as np
+import numpy.typing as npt
 import xarray as xr
 from overrides import overrides
 
-from pycontrails.core.met import MetDataset
+from pycontrails.core.met import MetDataArray, MetDataset
 from pycontrails.core.models import Model, ModelParams
 from pycontrails.core.vector import GeoVectorDataset
 from pycontrails.physics import constants, thermo, units
@@ -19,7 +22,7 @@ from pycontrails.utils.types import ArrayLike
 
 
 def _rhi_over_q(air_temperature: ArrayLike, air_pressure: ArrayLike) -> ArrayLike:
-    """Compute the quotient RHi / q."""
+    """Compute the quotient ``RHi / q``."""
     return air_pressure * (constants.R_v / constants.R_d) / thermo.e_sat_ice(air_temperature)
 
 
@@ -140,16 +143,18 @@ class HumidityScaling(Model):
                 "being applied more than once."
             )
 
+        p: np.ndarray | xr.DataArray
         if isinstance(self.source, GeoVectorDataset):
-            self.source.setdefault("air_pressure", self.source.air_pressure)
+            p = self.source.setdefault("air_pressure", self.source.air_pressure)
+        else:
+            p = self.source.data["air_pressure"]
 
         q = self.source.data["specific_humidity"]
         T = self.source.data["air_temperature"]
-        p = self.source.data["air_pressure"]
         kwargs = {k: self.get_source_param(k) for k in self.scaler_specific_keys}
 
         q, rhi = self.scale(q, T, p, **kwargs)
-        self.source.update(specific_humidity=q, rhi=rhi)  # type: ignore[arg-type]
+        self.source.update(specific_humidity=q, rhi=rhi)
 
         return self.source
 
@@ -264,7 +269,7 @@ class ExponentialBoostHumidityScaling(HumidityScaling):
         rhi /= rhi_adj
 
         # Find ISSRs
-        is_issr = rhi >= 1
+        is_issr = rhi >= 1.0
 
         # Apply boosting to ISSRs
         if isinstance(rhi, xr.DataArray):
@@ -440,7 +445,7 @@ def _calc_rhi_max(air_temperature: ArrayLike) -> ArrayLike:
         p_ice = thermo.e_sat_ice(air_temperature)
         p_liq = thermo.e_sat_liquid(air_temperature)
         return xr.where(
-            air_temperature < 235,
+            air_temperature < 235.0,
             1.67 + (1.45 - 1.67) * (air_temperature - 190.0) / (235.0 - 190.0),
             p_liq / p_ice,
         )
@@ -553,10 +558,262 @@ class HumidityScalingByLevel(HumidityScaling):
                 "attribute 'mid_troposphere_threshold'."
             )
 
-        level = air_pressure / 100
+        level = air_pressure / 100.0
         fp = [rhi_adj_stratosphere, rhi_adj_mid_troposphere]
         rhi_adj = np.interp(level, xp=xp, fp=fp)
 
         q = specific_humidity / rhi_adj
         rhi = thermo.rhi(q, air_temperature, air_pressure)
         return q, rhi
+
+
+@functools.cache
+def _load_iagos_quantiles() -> npt.NDArray[np.float64]:
+    path = pathlib.Path(__file__).parent / "quantiles" / "iagos_quantiles.npy"
+    # FIXME: Recompute to avoid the divide by 100.0 here
+    return np.load(path, allow_pickle=False) / 100.0
+
+
+@functools.cache
+def _load_era5_ensemble_quantiles() -> npt.NDArray[np.float64]:
+    path = pathlib.Path(__file__).parent / "quantiles" / "era5_ensemble_quantiles.npy"
+    # FIXME: Recompute to avoid the divide by 100.0 here
+    return np.load(path, allow_pickle=False) / 100.0
+
+
+def quantile_rhi_map(era5_rhi: npt.NDArray[np.float_], member: int) -> npt.NDArray[np.float64]:
+    """Map ERA5-derived RHi to it's corresponding IAGOS quantile via histogram matching.
+
+    This matching is performed on a **single** ERA5 ensemble member.
+
+    Parameters
+    ----------
+    era5_rhi : npt.NDArray[np.float_]
+        ERA5-derived RHi values for the given ensemble member.
+    member : int
+        The ERA5 ensemble member to use. Must be in the range ``[0, 10)``.
+
+    Returns
+    -------
+    npt.NDArray[np.float64]
+        The IAGOS quantiles corresponding to the ERA5-derived RHi values.
+    """
+
+    era5_quantiles = _load_era5_ensemble_quantiles()  # shape (801, 10)
+    era5_quantiles = era5_quantiles[:, member]  # shape (801,)
+    iagos_quantiles = _load_iagos_quantiles()  # shape (801,)
+
+    return np.interp(era5_rhi, era5_quantiles, iagos_quantiles)
+
+
+def recalibrate_rhi(
+    era5_rhi_all_members: npt.NDArray[np.float_], member: int
+) -> npt.NDArray[np.float_]:
+    """Recalibrate ERA5-derived RHi values to IAGOS quantiles.
+
+    This recalibration requires values for **all** ERA5 ensemble members.
+
+    Parameters
+    ----------
+    era5_rhi_all_members : npt.NDArray[np.float_]
+        ERA5-derived RHi values for all ensemble members. This array should have shape ``(n, 10)``.
+    member : int
+        The ERA5 ensemble member to use. Must be in the range ``[0, 10)``.
+
+    Returns
+    -------
+    npt.NDArray[np.float_]
+        The recalibrated RHi values. This is an array of shape ``(n,)``.
+
+    References
+    ----------
+    :cite:`eckelCalibratedProbabilisticQuantitative1998`
+    """
+
+    n_members = 10
+    assert era5_rhi_all_members.shape[1] == n_members
+
+    # Perform histogram matching on the given ensemble member
+    recalibrated_rhi = quantile_rhi_map(era5_rhi_all_members[:, member], member)
+
+    # Perform histogram matching on all other ensemble members
+    # Add up the results into a single 'ensemble_mean_rhi' array
+    ensemble_mean_rhi: npt.NDArray[np.float_] = 0.0  # type: ignore[assignment]
+    for r in range(n_members):
+        if r == member:
+            ensemble_mean_rhi += recalibrated_rhi
+        else:
+            ensemble_mean_rhi += quantile_rhi_map(era5_rhi_all_members[:, r], r)
+
+    # Divide by the number of ensemble members to get the mean
+    ensemble_mean_rhi /= n_members
+
+    eckel_a = -0.005213832567192828
+    eckel_c = 2.7859172756970354
+
+    out = (ensemble_mean_rhi - eckel_a) + eckel_c * (recalibrated_rhi - ensemble_mean_rhi)
+    return out.astype(era5_rhi_all_members.dtype)
+
+
+@dataclasses.dataclass
+class HistogramMatchingWithEckelParams(ModelParams):
+    """Parameters for :class:`HistogramMatchingWithEckel`.
+
+    .. warning::
+        Experimental. This may change or be removed in a future release.
+    """
+
+    #: A length-10 list of ERA5 ensemble members.
+    #: Each element is a :class:`MetDataArray` holding specific humidity
+    #: values for a single ensemble member. If None, a ValueError will be
+    #: raised at model instantiation time. The order of the list must be
+    #: consistent with the order of the ERA5 ensemble members.
+    ensemble_specific_humidity: list[MetDataArray] | None = None
+
+    #: The specific member used. Must be in the range [0, 10). If None,
+    #: a ValueError will be raised at model instantiation time.
+    member: int | None = None
+
+
+class HistogramMatchingWithEckel(HumidityScaling):
+    """Scale humidity by histogram matching to IAGOS RHi quantiles.
+
+    This method also applies the Eckel scaling to the recalibrated RHi values.
+
+    Unlike other specific humidity scaling methods, this method requires met data
+    and performs interpolation at evaluation time.
+
+    .. warning::
+        Experimental. This may change or be removed in a future release.
+
+    References
+    ----------
+    :cite:`eckelCalibratedProbabilisticQuantitative1998`
+    """
+
+    name = "histogram_matching_with_eckel"
+    long_name = "IAGOS RHi histogram matching with Eckel scaling"
+    formula = "era5_quantiles -> iagos_quantiles -> recalibrated_rhi"
+    default_params = HistogramMatchingWithEckelParams
+
+    n_members = 10  # hard-coded elsewhere
+
+    def __init__(
+        self,
+        met: MetDataset | None = None,
+        params: dict[str, Any] | None = None,
+        **params_kwargs: Any,
+    ) -> None:
+        super().__init__(met, params, **params_kwargs)
+
+        # Some very crude validation
+        member = self.params["member"]
+        assert member in range(self.n_members)
+        self.member: int = member
+
+        ensemble_specific_humidity = self.params["ensemble_specific_humidity"]
+        assert len(ensemble_specific_humidity) == self.n_members
+        for member, mda in enumerate(ensemble_specific_humidity):
+            try:
+                assert mda.data["number"] == member
+            except KeyError:
+                pass
+
+        self.ensemble_specific_humidity: list[MetDataArray] = ensemble_specific_humidity
+
+    @overload
+    def eval(self, source: GeoVectorDataset, **params: Any) -> GeoVectorDataset:
+        ...
+
+    @overload
+    def eval(self, source: MetDataset, **params: Any) -> NoReturn:
+        ...
+
+    @overload
+    def eval(self, source: None = ..., **params: Any) -> NoReturn:
+        ...
+
+    def eval(
+        self, source: GeoVectorDataset | MetDataset | None = None, **params: Any
+    ) -> GeoVectorDataset | MetDataset:
+        """Scale specific humidity by histogram matching to IAGOS RHi quantiles.
+
+        This method assumes ``source`` is equipped with the following variables:
+
+        - air_temperature
+        - specific_humidity: Humidity values for the :attr:`member` ERA5 ensemble member.
+        """
+
+        self.update_params(params)
+        self.set_source(source)
+        self.source = self.require_source_type(GeoVectorDataset)
+
+        if "rhi" in self.source:
+            warnings.warn(
+                "Variable 'rhi' already found on source to be scaled. This "
+                "is unexpected and may be the result of humidity scaling "
+                "being applied more than once."
+            )
+
+        # Create a 2D array of specific humidity values for all ensemble members
+        # The specific humidity values for the current member are taken from the source
+        # This matches patterns used in other humidity scaling methods
+        # The remaining values are interpolated from the ERA5 ensemble members
+        q = self.source.data["specific_humidity"]
+        q2d = np.empty((len(self.source), self.n_members), dtype=q.dtype)
+
+        for member, mda in enumerate(self.ensemble_specific_humidity):
+            if member == self.member:
+                q2d[:, member] = q
+            else:
+                q2d[:, member] = self.source.intersect_met(mda, **self.interp_kwargs)
+
+        p = self.source.setdefault("air_pressure", self.source.air_pressure)
+        T = self.source.data["air_temperature"]
+
+        q, rhi = self.scale(q2d, T, p)
+        self.source.update(specific_humidity=q, rhi=rhi)
+
+        return self.source
+
+    @overrides
+    def scale(  # type: ignore[override]
+        self,
+        specific_humidity: npt.NDArray[np.float_],
+        air_temperature: npt.NDArray[np.float_],
+        air_pressure: npt.NDArray[np.float_],
+        **kwargs: Any,
+    ) -> tuple[npt.NDArray[np.float_], npt.NDArray[np.float_]]:
+        """Scale specific humidity values via histogram matching and Eckel scaling.
+
+        Unlike the method on the base class, the method assumes each of the input
+        arrays are :class:`np.ndarray` and not :class:`xr.DataArray` objects.
+
+        Parameters
+        ----------
+        specific_humidity : npt.NDArray[np.float_]
+            A 2D array of specific humidity values for all ERA5 ensemble members.
+            The shape of this array must be ``(n, 10)``, where ``n`` is the number
+            of observations and ``10`` is the number of ERA5 ensemble members.
+        air_temperature : npt.NDArray[np.float_]
+            A 1D array of air temperature values with shape ``(n,)``.
+        air_pressure : npt.NDArray[np.float_]
+            A 1D array of air pressure values with shape ``(n,)``.
+        kwargs: Any
+            Unused, kept for compatibility with the base class.
+
+        Returns
+        -------
+        specific_humidity : npt.NDArray[np.float_]
+            The recalibrated specific humidity values. A 1D array with shape ``(n,)``.
+        rhi : npt.NDArray[np.float_]
+            The recalibrated RHi values. A 1D array with shape ``(n,)``.
+        """
+
+        rhi_over_q = _rhi_over_q(air_temperature, air_pressure)
+        rhi = rhi_over_q[:, np.newaxis] * specific_humidity
+
+        recalibrated_rhi = recalibrate_rhi(rhi, self.member)
+        recalibrated_q = recalibrated_rhi / rhi_over_q
+
+        return recalibrated_q, recalibrated_rhi

pycontrails/models/issr.py

@@ -51,7 +51,7 @@ class ISSR(Model):
     >>> variables = ["air_temperature", "specific_humidity"]
     >>> pressure_levels = [200, 250, 300]
     >>> era5 = ERA5(time, variables, pressure_levels)
-    >>> met = era5.open_metdataset(xr_kwargs=dict(parallel=False))
+    >>> met = era5.open_metdataset()
 
     >>> # Instantiate and run model
     >>> scaling = ConstantHumidityScaling(rhi_adj=0.98)
@@ -136,7 +136,7 @@ class ISSR(Model):
             rhi=self.source.data.get("rhi", None),  # if rhi already known, pass it in
             rhi_threshold=self.params["rhi_threshold"],
         )
-        self.source["issr"] = issr_  # type: ignore[assignment]
+        self.source["issr"] = issr_
 
         # Tag output with additional attrs when source is MetDataset
         if isinstance(self.source, MetDataset):

pycontrails/models/pcr.py

@@ -97,7 +97,7 @@ class PCR(Model):
         sac_model.eval(self.source)
 
         pcr_ = _pcr_from_issr_and_sac(self.source.data["issr"], self.source.data["sac"])
-        self.source["pcr"] = pcr_  # type: ignore[assignment]
+        self.source["pcr"] = pcr_
         # Tag output with additional attrs when source is MetDataset
         if isinstance(self.source, MetDataset):
             attrs = {**self.source["issr"].attrs, **self.source["sac"].attrs}

pycontrails/models/sac.py

@@ -139,11 +139,11 @@ class SAC(Model):
         sac_ = sac(rh, rh_crit_sac)
 
         # Attaching some intermediate artifacts onto the source
-        self.source["G"] = G  # type: ignore[assignment]
-        self.source["T_sat_liquid"] = T_sat_liquid_  # type: ignore[assignment]
-        self.source["rh"] = rh  # type: ignore[assignment]
-        self.source["rh_critical_sac"] = rh_crit_sac  # type: ignore[assignment]
-        self.source["sac"] = sac_  # type: ignore[assignment]
+        self.source["G"] = G
+        self.source["T_sat_liquid"] = T_sat_liquid_
+        self.source["rh"] = rh
+        self.source["rh_critical_sac"] = rh_crit_sac
+        self.source["sac"] = sac_
 
         # Tag output with additional attrs when source is MetDataset
         if isinstance(self.source, MetDataset):
@@ -438,6 +438,8 @@ def T_critical_sac(
     # We only apply Newton's method at points with rh bounded below 1 (scipy will
     # raise an error if Newton's method is not converging well).
     filt = relative_humidity < 0.999
+    if not np.any(filt):
+        return T_LM
 
     U_filt = relative_humidity[filt]
     T_LM_filt = T_LM[filt]

pycontrails/physics/geo.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import numpy as np
+import numpy.typing as npt
 import xarray as xr
 
 from pycontrails.physics import constants, units
@@ -32,11 +33,12 @@ def haversine(lons0: ArrayLike, lats0: ArrayLike, lons1: ArrayLike, lats1: Array
 
     Notes
     -----
-    This formula does not take into account the non-spheroidal (ellipsoidal) shape of the Earth
+    This formula does not take into account the non-spheroidal (ellipsoidal) shape of the Earth.
+    Originally referenced from https://andrew.hedges.name/experiments/haversine/.
 
     References
     ----------
-    Source: https://andrew.hedges.name/experiments/haversine/
+    - :cite:`CalculateDistanceBearing`
 
     See Also
     --------
@@ -833,7 +835,7 @@ def advect_level(
     return (level * 100.0 + (dt_s * velocity)) / 100.0
 
 
-def _dt_to_float_seconds(dt: np.ndarray | np.timedelta64, dtype: np.dtype) -> np.ndarray:
+def _dt_to_float_seconds(dt: np.ndarray | np.timedelta64, dtype: npt.DTypeLike) -> np.ndarray:
     """Convert a time delta to seconds as a float with specified ``dtype`` precision.
 
     Parameters