pycontrails 0.41.0__cp39-cp39-win_amd64.whl → 0.42.2__cp39-cp39-win_amd64.whl

pycontrails/ext/bada/__init__.py

@@ -20,6 +20,12 @@ try:
         BADAParams,
     )
 
+except ImportError as e:
+    raise ImportError(
+        'Failed to import `pycontrails-bada` extension. Install with `pip install "pycontrails-bada'
+        ' @ git+ssh://git@github.com/contrailcirrus/pycontrails-bada.git"`'
+    ) from e
+else:
     __all__ = [
         "BADA",
         "BADA3",
@@ -33,9 +39,3 @@ try:
         "bada4",
         "bada_model",
     ]
-
-except ImportError as e:
-    raise ImportError(
-        'Failed to import `pycontrails-bada` extension. Install with `pip install "pycontrails-bada'
-        ' @ git+ssh://git@github.com/contrailcirrus/pycontrails-bada.git"`'
-    ) from e

pycontrails/ext/cirium/__init__.py

@@ -5,10 +5,10 @@ from __future__ import annotations
 try:
     from pycontrails_cirium import Cirium
 
-    __all__ = ["Cirium"]
-
 except ImportError as e:
     raise ImportError(
         "Failed to import `pycontrails-cirium` extension. Install with `pip install"
         ' "pycontrails-cirium @ git+ssh://git@github.com/contrailcirrus/pycontrails-cirium.git"`'
     ) from e
+else:
+    __all__ = ["Cirium"]

pycontrails/models/cocip/cocip.py

@@ -7,6 +7,7 @@ import warnings
 from typing import Any, NoReturn, Sequence, overload
 
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 import xarray as xr
 from overrides import overrides
@@ -257,7 +258,7 @@ class Cocip(Model):
     contrail_dataset: xr.Dataset | None
 
     #: Array of np.datetime64 time steps for contrail evolution
-    timesteps: np.ndarray
+    timesteps: npt.NDArray[np.datetime64]
 
     #: Parallel copy of flight waypoints after SAC filter applied
     _sac_flight: Flight
@@ -376,8 +377,9 @@ class Cocip(Model):
         # Save humidity scaling type to output attrs
         # NOTE: Do this after _process_flight because that method automatically
         # broadcasts all numeric source params.
-        if self.params["humidity_scaling"] is not None:
-            for k, v in self.params["humidity_scaling"].description.items():
+        humidity_scaling = self.params["humidity_scaling"]
+        if humidity_scaling is not None:
+            for k, v in humidity_scaling.description.items():
                 self.source.attrs[f"humidity_scaling_{k}"] = v
 
         if isinstance(self.source, Fleet):
@@ -488,7 +490,7 @@ class Cocip(Model):
         self.source.setdefault("level", self.source.level)
         self.source.setdefault("air_pressure", self.source.air_pressure)
 
-        core_columns = ["longitude", "latitude", "altitude", "time"]
+        core_columns = ("longitude", "latitude", "altitude", "time")
         for col in core_columns:
             if np.isnan(self.source[col]).any():
                 raise ValueError(
@@ -532,23 +534,20 @@ class Cocip(Model):
         # STEP 4: Begin met interpolation
         # Unfortunately we use both "u_wind" and "eastward_wind" to refer to the
         # same variable, so the logic gets a bit more complex.
-        scale_humidity = (self.params["humidity_scaling"] is not None) and (
-            "specific_humidity" not in self.source
-        )
-        variables = {
-            "air_temperature": "air_temperature",
-            "specific_humidity": "specific_humidity",
-            "eastward_wind": "u_wind",
-            "northward_wind": "v_wind",
-        }
-        for met_variable, fl_variable in variables.items():
-            interpolate_met(met, self.source, met_variable, fl_variable, **self.interp_kwargs)
+        humidity_scaling = self.params["humidity_scaling"]
+        scale_humidity = humidity_scaling is not None and "specific_humidity" not in self.source
+        verbose_outputs = self.params["verbose_outputs"]
+
+        interpolate_met(met, self.source, "air_temperature", **self.interp_kwargs)
+        interpolate_met(met, self.source, "specific_humidity", **self.interp_kwargs)
+        interpolate_met(met, self.source, "eastward_wind", "u_wind", **self.interp_kwargs)
+        interpolate_met(met, self.source, "northward_wind", "v_wind", **self.interp_kwargs)
 
         if scale_humidity:
-            self.params["humidity_scaling"].eval(self.source, copy_source=False)
+            humidity_scaling.eval(self.source, copy_source=False)
 
         # if humidity_scaling isn't defined, add rhi to source for verbose_outputs
-        elif self.params["verbose_outputs"]:
+        elif verbose_outputs:
             self.source["rhi"] = thermo.rhi(
                 self.source["specific_humidity"],
                 self.source["air_temperature"],
@@ -556,7 +555,7 @@ class Cocip(Model):
             )
 
         # Cache extra met properties for post-analysis
-        if self.params["verbose_outputs"]:
+        if verbose_outputs:
             interpolate_met(met, self.source, "tau_cirrus", **self.interp_kwargs)
 
             # handle ECMWF/GFS ciwc variables
@@ -568,6 +567,9 @@ class Cocip(Model):
             self.source["rho_air"] = thermo.rho_d(
                 self.source["air_temperature"], self.source.air_pressure
             )
+            self.source["sdr"] = geo.solar_direct_radiation(
+                self.source["longitude"], self.source["latitude"], self.source["time"]
+            )
 
         # STEP 5: Calculate segment-specific properties if they are not already attached
         if "true_airspeed" not in self.source:
@@ -579,7 +581,8 @@ class Cocip(Model):
             )
         if "segment_length" not in self.source:
             self.source["segment_length"] = self.source.segment_length()
-        max_ = np.nanmax(self.source["segment_length"])
+
+        max_ = self.source.max_distance_gap
         lim_ = self.params["max_seg_length_m"]
         if max_ > 0.9 * lim_:
             warnings.warn(
@@ -867,8 +870,9 @@ class Cocip(Model):
             **self.interp_kwargs,
         )
 
-        if self.params["humidity_scaling"] is not None:
-            self.params["humidity_scaling"].eval(contrail_1, copy_source=False)
+        humidity_scaling = self.params["humidity_scaling"]
+        if humidity_scaling is not None:
+            humidity_scaling.eval(contrail_1, copy_source=False)
         else:
             contrail_1["air_pressure"] = contrail_1.air_pressure
             contrail_1["rhi"] = thermo.rhi(
@@ -1274,15 +1278,6 @@ class Cocip(Model):
             Flight or list of Flight objects with empty variables.
         """
 
-        empty = np.full(shape=len(self.source), fill_value=np.nan)
-
-        self.source["sdr"] = empty.copy()
-        self.source["rsr"] = empty.copy()
-        self.source["olr"] = empty.copy()
-        self.source["rf_sw"] = empty.copy()
-        self.source["rf_lw"] = empty.copy()
-        self.source["rf_net"] = empty.copy()
-
         intersection = self.source.data.pop("_met_intersection")
         zeros_and_nans = np.where(intersection, 0.0, np.nan)
         self.source["ef"] = zeros_and_nans.copy()
@@ -1830,9 +1825,9 @@ def calc_radiative_properties(contrail: GeoVectorDataset, params: dict[str, Any]
 
 def calc_contrail_properties(
     contrail: GeoVectorDataset,
-    effective_vertical_resolution: float | np.ndarray,
-    wind_shear_enhancement_exponent: float | np.ndarray,
-    sedimentation_impact_factor: float | np.ndarray,
+    effective_vertical_resolution: float | npt.NDArray[np.float_],
+    wind_shear_enhancement_exponent: float | npt.NDArray[np.float_],
+    sedimentation_impact_factor: float | npt.NDArray[np.float_],
     radiative_heating_effects: bool,
 ) -> None:
     """Calculate geometric and ice-related properties of contrail.
@@ -1859,10 +1854,12 @@ def calc_contrail_properties(
     ----------
     contrail : GeoVectorDataset
         Grid points with many precomputed keys.
-    effective_vertical_resolution, wind_shear_enhancement_exponent : float | np.ndarray
-        Passed into :func:`wind_shear.wind_shear_enhancement_factor` function.
-    sedimentation_impact_factor: float | np.ndarray
-        Passed into `contrail_properties.vertical_diffusivity` function.
+    effective_vertical_resolution : float | npt.NDArray[np.float_]
+        Passed into :func:`wind_shear.wind_shear_enhancement_factor`.
+    wind_shear_enhancement_exponent : float | npt.NDArray[np.float_]
+        Passed into :func:`wind_shear.wind_shear_enhancement_factor`.
+    sedimentation_impact_factor: float | npt.NDArray[np.float_]
+        Passed into `contrail_properties.vertical_diffusivity`.
     radiative_heating_effects: bool
         Include radiative heating effects on contrail cirrus properties.
     """
@@ -2141,8 +2138,9 @@ def calc_timestep_contrail_evolution(
 
     interpolate_met(met, contrail_2, "specific_humidity", **interp_kwargs)
 
-    if params["humidity_scaling"] is not None:
-        params["humidity_scaling"].eval(contrail_2, copy_source=False)
+    humidity_scaling = params["humidity_scaling"]
+    if humidity_scaling is not None:
+        humidity_scaling.eval(contrail_2, copy_source=False)
     else:
         contrail_2["air_pressure"] = contrail_2.air_pressure
         contrail_2["rhi"] = thermo.rhi(

pycontrails/models/cocip/cocip_params.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 import dataclasses
 
 import numpy as np
+import numpy.typing as npt
 
 from pycontrails.core.models import ModelParams
 from pycontrails.models.aircraft_performance import AircraftPerformance
@@ -15,6 +16,39 @@ from pycontrails.models.emissions.emissions import EmissionsParams
 from pycontrails.models.humidity_scaling import HumidityScaling
 
 
+def _radius_threshold_um() -> npt.NDArray[np.float32]:
+    return np.array([5.0, 9.5, 23.0, 190.0, 310.0], dtype=np.float32)
+
+
+def _habit_distributions() -> npt.NDArray[np.float32]:
+    return np.array(
+        [
+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
+            [0.0, 0.3, 0.0, 0.0, 0.0, 0.0, 0.7, 0.0],
+            [0.0, 0.3, 0.0, 0.0, 0.3, 0.0, 0.4, 0.0],
+            [0.0, 0.5, 0.0, 0.0, 0.15, 0.35, 0.0, 0.0],
+            [0.0, 0.45, 0.45, 0.1, 0.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.03, 0.97, 0.0, 0.0, 0.0],
+        ],
+        dtype=np.float32,
+    )
+
+
+def _habits() -> npt.NDArray[np.str_]:
+    return np.array(
+        [
+            "Sphere",
+            "Solid column",
+            "Hollow column",
+            "Rough aggregate",
+            "Rosette-6",
+            "Plate",
+            "Droxtal",
+            "Myhre",
+        ]
+    )
+
+
 @dataclasses.dataclass
 class CocipParams(ModelParams):
     """Model parameters required by the CoCiP models."""
@@ -147,45 +181,18 @@ class CocipParams(ModelParams):
     #: Radius threshold for regime bins, [:math:`\mu m`]
     #: This is the row index label for ``habit_distributions``.
     #: See Table 2 in :cite:`schumannEffectiveRadiusIce2011`.
-    radius_threshold_um: np.ndarray = dataclasses.field(
-        default_factory=lambda: np.array([5.0, 9.5, 23.0, 190.0, 310.0], dtype=np.float32)
-    )
+    radius_threshold_um: np.ndarray = dataclasses.field(default_factory=_radius_threshold_um)
 
     #: Particle habit (shape) types.
     #: This is the column index label for ``habit_distributions``.
     #: See Table 2 in :cite:`schumannEffectiveRadiusIce2011`.
-    habits: np.ndarray = dataclasses.field(
-        default_factory=lambda: np.array(
-            [
-                "Sphere",
-                "Solid column",
-                "Hollow column",
-                "Rough aggregate",
-                "Rosette-6",
-                "Plate",
-                "Droxtal",
-                "Myhre",
-            ]
-        )
-    )
+    habits: np.ndarray = dataclasses.field(default_factory=_habits)
 
     #: Mix of ice particle habits in each radius regime.
     #: Rows indexes are ``radius_threshold_um`` elements.
     #: Columns indexes are ``habits`` particle habit type.
     #: See Table 2 from :cite:`schumannEffectiveRadiusIce2011`.
-    habit_distributions: np.ndarray = dataclasses.field(
-        default_factory=lambda: np.array(
-            [
-                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
-                [0.0, 0.3, 0.0, 0.0, 0.0, 0.0, 0.7, 0.0],
-                [0.0, 0.3, 0.0, 0.0, 0.3, 0.0, 0.4, 0.0],
-                [0.0, 0.5, 0.0, 0.0, 0.15, 0.35, 0.0, 0.0],
-                [0.0, 0.45, 0.45, 0.1, 0.0, 0.0, 0.0, 0.0],
-                [0.0, 0.0, 0.0, 0.03, 0.97, 0.0, 0.0, 0.0],
-            ],
-            dtype=np.float32,
-        )
-    )
+    habit_distributions: np.ndarray = dataclasses.field(default_factory=_habit_distributions)
 
     #: Scale shortwave radiative forcing.
     #: Primarily used to support uncertainty estimation.

pycontrails/models/cocip/cocip_uncertainty.py

@@ -7,11 +7,12 @@ from dataclasses import asdict, dataclass
 from typing import Any, ClassVar
 
 import numpy as np
+import numpy.typing as npt
 from scipy import stats
 from scipy.stats.distributions import rv_frozen
 
+from pycontrails.models.cocip import cocip_params
 from pycontrails.models.cocip.cocip_params import CocipParams
-from pycontrails.models.humidity_scaling import ExponentialBoostHumidityScaling
 
 logger = logging.getLogger(__name__)
 
@@ -32,7 +33,7 @@ class habit_dirichlet(rv_frozen):
     def __init__(self, C: float = 96.0):
         self.C = C
 
-    def rvs(self, *args: Any, **kwds: Any) -> np.ndarray:
+    def rvs(self, *args: Any, **kwds: Any) -> npt.NDArray[np.float32]:
         """Generate sample set of habit distributions.
 
         Sampled using dirichlet distribution.
@@ -46,20 +47,25 @@ class habit_dirichlet(rv_frozen):
 
         Returns
         -------
-        np.ndarray
+        npt.NDArray[np.float32]
             Sampled habit weight distribution with the same shape
             as :attr:`CocipParams().habit_distributions`
         """
-        if args or ("size" in kwds and kwds["size"] is not None and kwds["size"] > 1):
+        if args or (kwds.get("size") is not None and kwds["size"] > 1):
             raise ValueError("habit_dirichlet distribution only supports creating one rv at a time")
 
-        default_distributions = CocipParams().habit_distributions
-        alpha_i = [0.5 + self.C * G_i for G_i in default_distributions]
-        distr = [stats.dirichlet(a) for a in alpha_i]
+        default_distributions = cocip_params._habit_distributions()
+        alpha_i = 0.5 + self.C * default_distributions
+
+        # In the first distribution, we assume all ice particles are droxtals
+        # There is no way to quantify the uncertainty in this assumption
+        # Consequently, we leave the first distribution in default_distributions
+        # alone, and only perturb the rest
+        distr_list = [stats.dirichlet(a) for a in alpha_i[1:]]
 
         habit_weights = default_distributions.copy()
-        for i in range(1, habit_weights.shape[0]):
-            habit_weights[i] = distr[i].rvs(**kwds)
+        for i, distr in enumerate(distr_list, start=1):
+            habit_weights[i] = distr.rvs(**kwds)
 
         return habit_weights
 
@@ -81,17 +87,19 @@ class CocipUncertaintyParams(CocipParams):
     >>> import scipy.stats
     >>> from pycontrails.models.cocip import CocipUncertaintyParams
 
-    # Override the `rhi_adj` field from `CocipParams` with a value in [0.97, 0.99]
-    >>> distr = scipy.stats.uniform(loc=0.97, scale=0.02)
-    >>> params = CocipUncertaintyParams(seed=123, rhi_adj_uncertainty=distr)
-    >>> params.rhi_adj
-    0.98364703
+    >>> # Override the 'initial_wake_vortex_depth' field from
+    >>> # CocipParams with a uniform value in [0.4, 0.6]
+    >>> distr = scipy.stats.uniform(loc=0.4, scale=0.2)
+    >>> params = CocipUncertaintyParams(seed=123, initial_wake_vortex_depth_uncertainty=distr)
+    >>> params.initial_wake_vortex_depth
+    0.41076420
 
-    # Once seeded, calling the class again gives a new value
-    >>> params = CocipUncertaintyParams(rhi_adj_uncertainty=distr)
-    >>> params.rhi_adj
-    0.98310570
+    >>> # Once seeded, calling the class again gives a new value
+    >>> params = CocipUncertaintyParams(initial_wake_vortex_depth=distr)
+    >>> params.initial_wake_vortex_depth
+    0.43526372
 
+    >>> # To retain the default value, set the uncertainty to None
     >>> params = CocipUncertaintyParams(rf_lw_enhancement_factor_uncertainty=None)
     >>> params.rf_lw_enhancement_factor
     1.0
@@ -106,23 +114,6 @@ class CocipUncertaintyParams(CocipParams):
     #: Reseed the random generator defined in ``__post_init__``
     seed: int | None = None
 
-    # This might need fixing!
-    # Whenever we overhaul / redo an uncertainty analysis, we may need to change
-    # how it interacts with a humidity scaler
-    humidity_scaling: ExponentialBoostHumidityScaling = ExponentialBoostHumidityScaling()
-
-    #: Parameters for specific humidity and RHi enhancement
-    #: This assumes Cocip.Params.humidity_scaling is an
-    #: :class:`ExponentialBoostHumidityScaling`` instance.
-    rhi_adj_uncertainty: rv_frozen | None = stats.norm(
-        loc=ExponentialBoostHumidityScaling.default_params.rhi_adj, scale=0.1
-    )
-    rhi_boost_exponent_uncertainty: rv_frozen = stats.triang(
-        loc=1.0,
-        c=ExponentialBoostHumidityScaling.default_params.rhi_boost_exponent - 1.0,
-        scale=1.0,
-    )
-
     #: Schumann takes ``wind_shear_enhancement_exponent`` = 0.5 and discusses the case of 0 and 2/3
     #: as possibilities.
     #: With a value of 0, wind shear is not enhanced.
@@ -172,7 +163,7 @@ class CocipUncertaintyParams(CocipParams):
     #: where :math:`\text{G}_{i}` is the approximate habit weight distributions
     #: defined in :attr:`CocipParams().habit_distributions`.
     #: Higher values of :math:`\text{C}` correspond to higher confidence in initial estimates.
-    habit_distributions_uncertainty: rv_frozen | None = habit_dirichlet(C=96)
+    habit_distributions_uncertainty: rv_frozen | None = habit_dirichlet(C=96.0)
 
     def __post_init__(self) -> None:
         """Override values of model parameters according to ranges."""
@@ -221,7 +212,7 @@ class CocipUncertaintyParams(CocipParams):
 
         return out
 
-    def rvs(self, size: None | int = None) -> dict[str, float | np.ndarray]:
+    def rvs(self, size: None | int = None) -> dict[str, float | npt.NDArray[np.float_]]:
         """Call each distribution's `rvs` method to generate random parameters.
 
         Seed calls to `rvs` with class variable `rng`.
@@ -233,7 +224,7 @@ class CocipUncertaintyParams(CocipParams):
 
         Returns
         -------
-        dict[str, float | np.ndarray]
+        dict[str, float | npt.NDArray[np.float_]]
             Dictionary of random parameters. Dictionary keys consists of names of parameters in
             `CocipParams` to be overridden by random value.
 
@@ -244,26 +235,24 @@ class CocipUncertaintyParams(CocipParams):
         >>> params = CocipUncertaintyParams(seed=456)
         >>> pprint(params.rvs())
         {'habit_distributions': array([[0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,
-               0.0000000e+00, 0.0000000e+00, 1.0000000e+00, 0.0000000e+00],
-              [3.6062175e-04, 2.9157758e-01, 6.5901526e-04, 1.3415117e-03,
-               5.0444747e-03, 2.6956137e-02, 6.7330402e-01, 7.5663364e-04],
-              [2.3834489e-03, 3.6896354e-01, 1.3977488e-03, 4.7844747e-04,
-               3.0474603e-01, 4.8765671e-03, 3.1715181e-01, 2.4031256e-06],
-              [4.0944587e-03, 5.1772928e-01, 4.8383059e-05, 1.3166944e-03,
-               1.0002965e-01, 3.7468293e-01, 3.9641498e-04, 1.7021718e-03],
-              [1.5192166e-05, 3.8423434e-01, 5.0536448e-01, 9.4547205e-02,
-               1.1378267e-03, 3.3812344e-03, 4.6305601e-03, 6.6891452e-03],
-              [1.1620003e-03, 2.3161094e-03, 1.6705607e-04, 2.0091899e-02,
-               9.7001791e-01, 1.5312615e-03, 3.9347797e-03, 7.7895907e-04]],
-             dtype=float32),
-        'initial_wake_vortex_depth': 0.4752143043559352,
-        'nvpm_ei_n_enhancement_factor': 1.0094726146454185,
-        'rf_lw_enhancement_factor': 0.9718354129386728,
-        'rf_sw_enhancement_factor': 1.016450567639041,
-        'rhi_adj': 1.0293074966694602,
-        'rhi_boost_exponent': 1.3308280817638443,
-        'sedimentation_impact_factor': 0.6433086271594898,
-        'wind_shear_enhancement_exponent': 0.31735944665388505}
+                0.0000000e+00, 0.0000000e+00, 1.0000000e+00, 0.0000000e+00],
+               [1.5554131e-02, 2.1363135e-01, 7.7715185e-03, 1.7690966e-02,
+                3.1576434e-03, 3.2992734e-06, 7.4009895e-01, 2.0921326e-03],
+               [3.8193921e-03, 2.1235342e-01, 3.3554080e-04, 5.2846869e-04,
+                3.1945917e-01, 4.8709914e-04, 4.6250960e-01, 5.0730183e-04],
+               [5.7327619e-04, 4.7781631e-01, 4.2596990e-03, 6.7235163e-04,
+                1.4447135e-01, 3.6184600e-01, 1.0150939e-02, 2.1006212e-04],
+               [1.5397545e-02, 4.0522218e-01, 4.2781001e-01, 1.4331797e-01,
+                7.1088417e-04, 9.4511814e-04, 3.3900745e-03, 3.2062260e-03],
+               [7.9063961e-04, 3.0336906e-03, 7.7571563e-04, 2.0577813e-02,
+                9.4205803e-01, 4.3379897e-03, 3.6786550e-03, 2.4747452e-02]],
+              dtype=float32),
+         'initial_wake_vortex_depth': 0.39805019708566847,
+         'nvpm_ei_n_enhancement_factor': 0.9371878437312526,
+         'rf_lw_enhancement_factor': 1.1017491252832377,
+         'rf_sw_enhancement_factor': 0.99721639115012,
+         'sedimentation_impact_factor': 0.5071779847244678,
+         'wind_shear_enhancement_exponent': 0.34100931239701004}
         """
         return {
             param: distr.rvs(size=size, random_state=self.rng)