pycontrails 0.47.3__cp310-cp310-win_amd64.whl → 0.48.1__cp310-cp310-win_amd64.whl

pycontrails/models/cocip/cocip.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import logging
 import warnings
-from typing import Any, NoReturn, Sequence, overload
+from typing import Any, Literal, NoReturn, Sequence, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -93,9 +93,6 @@ class Cocip(Model):
         * - Ice water content
           - ``specific_cloud_ice_water_content``
           - ``ice_water_mixing_ratio``
-        * - Geopotential
-          - ``geopotential``
-          - ``geopotential_height``
 
     .. list-table:: Variable keys for single-level radiation data
         :header-rows: 1
@@ -153,13 +150,13 @@ class Cocip(Model):
 
     **Outputs**
 
-    NaN values may appear in model output. Specifically, `np.nan` values are used to indicate:
+    NaN values may appear in model output. Specifically, ``np.nan`` values are used to indicate:
 
     - Flight waypoint or contrail waypoint is not contained with the :attr:`met` domain.
     - The variable was NOT computed during the model evaluation. For example, at flight waypoints
-      not producing any persistent contrails, "radiative" variables (`rsr`, `olr`, `rf_sw`,
-      `rf_lw`, `rf_net`) are not computed. Consequently, the corresponding values in the output
-      of :meth:`eval` are NaN. One exception to this rule is found on `ef` (energy forcing)
+      not producing any persistent contrails, "radiative" variables (``rsr``, ``olr``, ``rf_sw``,
+      ``rf_lw``, ``rf_net``) are not computed. Consequently, the corresponding values in the output
+      of :meth:`eval` are NaN. One exception to this rule is found on ``ef`` (energy forcing)
       `contrail_age` predictions. For these two "cumulative" variables, waypoints not producing
       any persistent contrails are assigned 0 values.
 
@@ -212,7 +209,6 @@ class Cocip(Model):
         met_var.NorthwardWind,
         met_var.VerticalVelocity,
         (ecmwf.SpecificCloudIceWaterContent, gfs.CloudIceWaterMixingRatio),
-        (met_var.Geopotential, met_var.GeopotentialHeight),
     )
 
     #: Required single-level top of atmosphere radiation variables.
@@ -223,7 +219,7 @@ class Cocip(Model):
     )
 
     #: Minimal set of met variables needed to run the model after pre-processing.
-    #: The intention here is that ``geopotential`` and ``ciwc`` are unnecessary after
+    #: The intention here is that ``ciwc`` is unnecessary after
     #: ``tau_cirrus`` has already been calculated.
     processed_met_variables = (
         met_var.AirTemperature,
@@ -235,28 +231,33 @@ class Cocip(Model):
     )
 
     #: Additional met variables used to support outputs
-    optional_met_variables = ((ecmwf.CloudAreaFractionInLayer, gfs.TotalCloudCoverIsobaric),)
+    #: .. versionchanged:: 0.48.0
+    #: Moved Geopotential from :attr:`required_met_variables` to :attr:`optional_met_variables`
+    optional_met_variables = (
+        (met_var.Geopotential, met_var.GeopotentialHeight),
+        (ecmwf.CloudAreaFractionInLayer, gfs.TotalCloudCoverIsobaric),
+    )
 
     #: Met data is not optional
     met: MetDataset
     met_required = True
 
-    #: Radiation data formatted as a MetDataset at a single pressure level [-1]
+    #: Radiation data formatted as a :class:`MetDataset` at a single pressure level [-1]
     rad: MetDataset
 
     #: Last Flight modeled in :meth:`eval`
     source: Flight | Fleet
 
-    #: List of GeoVectorDataset contrail objects - one for each timestep
+    #: List of :class:`GeoVectorDataset` contrail objects - one for each timestep
     contrail_list: list[GeoVectorDataset]
 
     #: Contrail evolution output from model. Set to None when no contrails are formed.
     contrail: pd.DataFrame | None
 
-    #: xr.Dataset reporesentation of contrail evolution.
+    #: :class:`xr.Dataset` representation of contrail evolution.
     contrail_dataset: xr.Dataset | None
 
-    #: Array of np.datetime64 time steps for contrail evolution
+    #: Array of :class:`np.datetime64` time steps for contrail evolution
     timesteps: npt.NDArray[np.datetime64]
 
     #: Parallel copy of flight waypoints after SAC filter applied
@@ -278,12 +279,8 @@ class Cocip(Model):
         # call Model init
         super().__init__(met, params=params, **params_kwargs)
 
-        shift_radiation_time = self.params["shift_radiation_time"]
         compute_tau_cirrus = self.params["compute_tau_cirrus_in_model_init"]
-        met, rad = process_met_datasets(met, rad, compute_tau_cirrus, shift_radiation_time)
-
-        self.met = met
-        self.rad = rad
+        self.met, self.rad = process_met_datasets(met, rad, compute_tau_cirrus)
 
         # initialize outputs to None
         self.contrail = None
@@ -406,14 +403,13 @@ class Cocip(Model):
         self.contrail_list = []
         self._simulate_contrail_evolution()
 
-        self._cleanup_indices()
-
         if not self.contrail_list:
             logger.debug("No contrails formed by %s", label)
             return self._fill_empty_flight_results(return_flight_list)
 
         logger.debug("Complete contrail simulation for %s", label)
 
+        self._cleanup_indices()
         self._bundle_results()
 
         if return_flight_list:
@@ -915,7 +911,7 @@ class Cocip(Model):
         self._downwash_contrail = self._create_downwash_contrail()
         buffers = {
             f"{coord}_buffer": self.params[f"met_{coord}_buffer"]
-            for coord in ["longitude", "latitude", "level"]
+            for coord in ("longitude", "latitude", "level")
         }
         logger.debug("Downselect met for start of Cocip evolution")
         met = self._downwash_contrail.downselect_met(self.met, **buffers, copy=False)
@@ -1101,12 +1097,20 @@ class Cocip(Model):
 
     @overrides
     def _cleanup_indices(self) -> None:
-        if self.params["interpolation_use_indices"]:
+        """Cleanup interpolation artifacts."""
+
+        if not self.params["interpolation_use_indices"]:
+            return
+
+        if hasattr(self, "contrail_list"):
             for contrail in self.contrail_list:
                 contrail._invalidate_indices()
-            self.source._invalidate_indices()
-            self._sac_flight._invalidate_indices()
+
+        self.source._invalidate_indices()
+        self._sac_flight._invalidate_indices()
+        if hasattr(self, "_downwash_flight"):
             self._downwash_flight._invalidate_indices()
+        if hasattr(self, "_downwash_contrail"):
             self._downwash_contrail._invalidate_indices()
 
     def _bundle_results(self) -> None:
@@ -1118,7 +1122,9 @@ class Cocip(Model):
         self.contrail = pd.concat(dfs)
 
         # add age in hours to the contrail waypoint outputs
-        self.contrail["age_hours"] = self.contrail["age"] / np.timedelta64(1, "h")
+        age_hours = np.empty_like(self.contrail["ef"])
+        np.divide(self.contrail["age"], np.timedelta64(1, "h"), out=age_hours)
+        self.contrail["age_hours"] = age_hours
 
         if self.params["verbose_outputs"]:
             # Compute dt_integration -- logic is somewhat complicated, but
@@ -1143,10 +1149,9 @@ class Cocip(Model):
 
             self.contrail = seq_index.set_index("index")
 
-        # ---
-        # Create contrail xr.Dataset (self.contrail_dataset)
-        # ---
-        if self.params["verbose_outputs"]:
+            # ---
+            # Create contrail xr.Dataset (self.contrail_dataset)
+            # ---
             if isinstance(self.source, Fleet):
                 self.contrail_dataset = xr.Dataset.from_dataframe(
                     self.contrail.set_index(["flight_id", "timestep", "waypoint"])
@@ -1253,9 +1258,11 @@ class Cocip(Model):
         Flight or list[Flight]
             Flight or list of Flight objects with empty variables.
         """
+        self._cleanup_indices()
 
         intersection = self.source.data.pop("_met_intersection")
-        zeros_and_nans = np.where(intersection, 0.0, np.nan)
+        zeros_and_nans = np.zeros(intersection.shape, dtype=np.float32)
+        zeros_and_nans[~intersection] = np.nan
         self.source["ef"] = zeros_and_nans.copy()
         self.source["persistent_1"] = zeros_and_nans.copy()
         self.source["cocip"] = np.sign(zeros_and_nans)
@@ -1275,8 +1282,7 @@ class Cocip(Model):
 def process_met_datasets(
     met: MetDataset,
     rad: MetDataset,
-    compute_tau_cirrus: bool | str = "auto",
-    shift_radiation_time: np.timedelta64 | None = None,
+    compute_tau_cirrus: bool | Literal["auto"] = "auto",
 ) -> tuple[MetDataset, MetDataset]:
     """Process and verify ERA5 data for :class:`Cocip` and :class:`CocipGrid`.
 
@@ -1289,18 +1295,20 @@ def process_met_datasets(
         of the ``process_met`` parameter. The same approach is also taken
         in :class:`Cocip` in version 0.27.0.
 
+    .. versionchanged:: 0.48.0
+
+        Remove the ``shift_radiation_time`` parameter. This parameter is now
+        inferred from the metadata on the `rad` instance.
+
     Parameters
     ----------
     met : MetDataset
         Met pressure-level data
     rad : MetDataset
         Rad single-level data
-    compute_tau_cirrus : bool | str
+    compute_tau_cirrus : bool | Literal["auto"]
         Whether to add ``"tau_cirrus"`` variable to pressure-level met data. If set to
         ``"auto"``, ``"tau_cirrus"`` will be computed iff the met data is dask-backed.
-    shift_radiation_time : np.timedelta64 | None
-        Shift the time dimension of radiation data to account for accumulated values.
-        If not specified, the default value from :class:`CocipGridParams` will be used.
 
     Returns
     -------
@@ -1315,35 +1323,20 @@ def process_met_datasets(
             "Specific humidity enhancement of the raw specific humidity values in "
             "the underlying met data is deprecated."
         )
-    if isinstance(compute_tau_cirrus, str) and compute_tau_cirrus != "auto":
-        raise ValueError(
-            "Parameter ``compute_tau_cirrus`` must be one of"
-            f" True, False, or 'auto'. Found {compute_tau_cirrus}."
-        )
+
+    if compute_tau_cirrus == "auto":
+        # If met data is dask-backed, compute tau_cirrus
+        compute_tau_cirrus = met.data["air_temperature"].chunks is not None
+
     if "tau_cirrus" not in met.data:
         met.ensure_vars(Cocip.met_variables)
-        if (
-            compute_tau_cirrus == "auto"
-            and met.data["specific_humidity"].chunks is not None
-            or isinstance(compute_tau_cirrus, bool)
-            and compute_tau_cirrus
-        ):
+        if compute_tau_cirrus:
             met = add_tau_cirrus(met)
     else:
         met.ensure_vars(Cocip.processed_met_variables)
 
-    # Deal with rad: check shift_radiation_time
     rad.ensure_vars(Cocip.rad_variables)
-    if "_pycontrails_modified" in rad["time"].attrs:
-        existing_shift = rad["time"].attrs["shift_radiation_time"]  # this is a string
-        if pd.Timedelta(existing_shift) != shift_radiation_time:  # compare timedeltas
-            raise ValueError(
-                "The time coordinate in MetDataset 'rad' has already been "
-                f"scaled by a CoCiP model with 'shift_radiation_time={existing_shift}'. "
-            )
-    else:
-        shift_radiation_time = shift_radiation_time or Cocip.default_params.shift_radiation_time
-        rad = _process_rad(rad, shift_radiation_time)
+    rad = _process_rad(rad)
 
     return met, rad
 
@@ -1366,7 +1359,7 @@ def add_tau_cirrus(met: MetDataset) -> MetDataset:
     return met
 
 
-def _process_rad(rad: MetDataset, shift_radiation_time: np.timedelta64) -> MetDataset:
+def _process_rad(rad: MetDataset) -> MetDataset:
     """Process radiation specific variables for model.
 
     These variables are used to calculate the reflected solar radiation (RSR),
@@ -1379,8 +1372,6 @@ def _process_rad(rad: MetDataset, shift_radiation_time: np.timedelta64) -> MetDa
     ----------
     rad : MetDataset
         Rad single-level data
-    shift_radiation_time : np.timedelta64
-        Shift the time dimension of radiation data to account for accumulated values.
 
     Returns
     -------
@@ -1392,18 +1383,67 @@ def _process_rad(rad: MetDataset, shift_radiation_time: np.timedelta64) -> MetDa
     - https://www.ecmwf.int/sites/default/files/elibrary/2015/18490-radiation-quantities-ecmwf-model-and-mars.pdf
     - https://confluence.ecmwf.int/pages/viewpage.action?pageId=155337784
     """  # noqa: E501
-    logger.debug(f"Shifting radiation time dimension by {shift_radiation_time}")
-    if not np.all(np.diff(rad.data["time"]) / 2 == -shift_radiation_time):
+    # If the time coordinate has already been shifted, early return
+    if "shift_radiation_time" in rad["time"].attrs:
+        return rad
+
+    provider = rad.provider_attr
+
+    # Only shift ECMWF data -- exit for anything else
+    # A warning is emitted upstream if the provider is not ECMWF or NCEP
+    if provider != "ECMWF":
+        return rad
+
+    dataset = rad.dataset_attr
+    product = rad.product_attr
+
+    if dataset == "HRES":
+        try:
+            radiation_accumulated = rad.attrs["radiation_accumulated"]
+        except KeyError:
+            msg = (
+                "HRES data must have a boolean 'radiation_accumulated' attribute. "
+                "This attribute is used to determine whether the radiation data "
+                "has been accumulated over the time period. This is the case for "
+                "HRES data taken from a common time of forecast with multiple "
+                "forecast steps. If this is not the case, set the "
+                "'radiation_accumulated' attribute to False."
+            )
+            raise ValueError(msg)
+        if radiation_accumulated:
+            # Keep the original attrs -- we need these later on
+            old_attrs = {k: v.attrs for k, v in rad.data.items()}
+
+            # NOTE: Taking the diff will remove the first time step
+            # This is typically what we want (forecast step 0 is all zeros)
+            # But, if the data has been downselected for a particular Flight / Fleet,
+            # we lose the first time step of the data.
+            rad.data = rad.data.diff("time", label="upper")
+
+            # Add back the original attrs
+            for k, v in rad.data.items():
+                v.attrs = old_attrs[k]
+
+        shift_radiation_time = -np.timedelta64(30, "m")
+
+    elif dataset == "ERA5" and product == "ensemble":
+        shift_radiation_time = -np.timedelta64(90, "m")
+    else:
+        shift_radiation_time = -np.timedelta64(30, "m")
+
+    # Do a final idiot check -- most likely, the time resolution of the data will
+    # agree with the shift_radiation_time. If not, emit a warning. There could be
+    # a false positive here if the data has been downsampled in time.
+    logger.debug("Shifting rad time by %s", shift_radiation_time)
+    rad_time_diff = np.diff(rad.data["time"])
+    if not np.all(rad_time_diff / 2 == -shift_radiation_time):
         warnings.warn(
             f"Shifting radiation time dimension by unexpected interval {shift_radiation_time}. "
-            "If working with ERA5 ensemble members, set "
-            "`shift_radiation_time=-np.timedelta64(90, 'm')`. Otherwise, the "
-            "expected shift is half the time difference consecutive time steps."
+            f"The rad data has metadata indicating it is {product} ECMWF data. "
+            f"This dataset should have time steps of {-2 * shift_radiation_time}."
         )
 
     rad.data = rad.data.assign_coords({"time": rad.data["time"] + shift_radiation_time})
-    msg = "Time coordinates adjusted to account for accumulation averaging"
-    rad.data["time"].attrs["_pycontrails_modified"] = msg
     rad.data["time"].attrs["shift_radiation_time"] = str(shift_radiation_time)
 
     return rad
@@ -1439,7 +1479,9 @@ def _eval_aircraft_performance(
     if aircraft_performance is None:
         raise ValueError(
             "An AircraftPerformance model parameter is required if the flight "
-            f"does not contain the following variables: {aircraft_performance_outputs}"
+            f"does not contain the following variables: {aircraft_performance_outputs}. "
+            "For example, instantiate the Cocip model with "
+            "'Cocip(..., aircraft_performance=PSFlight(...))'."
         )
 
     return aircraft_performance.eval(source=flight, copy_source=False)
@@ -1538,9 +1580,9 @@ def calc_timestep_geometry(contrail: GeoVectorDataset) -> None:
     # Finally, at the next evolution step, the previous waypoint will accrue
     # a nan value after segment_length is recalculated
 
-    segment_length[~continuous] = 0
-    sin_a[~continuous] = 0
-    cos_a[~continuous] = 0
+    segment_length[~continuous] = 0.0
+    sin_a[~continuous] = 0.0
+    cos_a[~continuous] = 0.0
 
     np.nan_to_num(segment_length, copy=False)
     np.nan_to_num(sin_a, copy=False)
@@ -1680,8 +1722,8 @@ def calc_shortwave_radiation(
     Raises
     ------
     ValueError
-        If ``rad`` does not contain ``"toa_upward_shortwave_flux"`` or ``"top_net_solar_radiation"``
-        variable.
+        If ``rad`` does not contain ``"toa_upward_shortwave_flux"`` or
+        ``"top_net_solar_radiation"`` variable.
 
     Notes
     -----
@@ -1692,30 +1734,36 @@ def calc_shortwave_radiation(
     --------
     :func:`geo.solar_direct_radiation`
     """
-
     if "sdr" in vector and "rsr" in vector:
-        return None
+        return
 
-    # calculate instantaneous theoretical solar direct radiation based on geo position and time
-    longitude = vector["longitude"]
-    latitude = vector["latitude"]
-    time = vector["time"]
-    vector["sdr"] = geo.solar_direct_radiation(longitude, latitude, time, threshold_cos_sza=0.01)
+    try:
+        sdr = vector["sdr"]
+    except KeyError:
+        # calculate instantaneous theoretical solar direct radiation based on geo position and time
+        longitude = vector["longitude"]
+        latitude = vector["latitude"]
+        time = vector["time"]
+        sdr = geo.solar_direct_radiation(longitude, latitude, time, threshold_cos_sza=0.01)
+        vector["sdr"] = sdr
 
     # GFS contains RSR (toa_upward_shortwave_flux) variable directly
-    if "toa_upward_shortwave_flux" in rad:
-        interpolate_met(rad, vector, "toa_upward_shortwave_flux", "rsr", **interp_kwargs)
+    gfs_key = "toa_upward_shortwave_flux"
+    if gfs_key in rad:
+        interpolate_met(rad, vector, gfs_key, "rsr", **interp_kwargs)
+        return
+
+    ecmwf_key = "top_net_solar_radiation"
+    if ecmwf_key not in rad:
+        msg = f"'rad' data must contain either '{gfs_key}' or '{ecmwf_key}' (ECMWF) variable."
+        raise ValueError(msg)
 
     # ECMWF contains "top_net_solar_radiation" which is SDR - RSR
-    elif "top_net_solar_radiation" in rad:
-        interpolate_met(rad, vector, "top_net_solar_radiation", **interp_kwargs)
-        vector["rsr"] = np.maximum(vector["sdr"] - vector["top_net_solar_radiation"], 0.0)
+    tnsr = interpolate_met(rad, vector, ecmwf_key, **interp_kwargs)
+    tnsr = _rad_accumulation_to_average_instantaneous(rad, ecmwf_key, tnsr)
+    vector.update({ecmwf_key: tnsr})
 
-    else:
-        raise ValueError(
-            "'rad' data must contain either 'toa_upward_shortwave_flux' or "
-            "'top_net_solar_radiation' (ECMWF) variable."
-        )
+    vector["rsr"] = np.maximum(sdr - tnsr, 0.0)
 
 
 def calc_outgoing_longwave_radiation(
@@ -1747,20 +1795,22 @@ def calc_outgoing_longwave_radiation(
         return None
 
     # GFS contains OLR (toa_upward_longwave_flux) variable directly
-    if "toa_upward_longwave_flux" in rad:
-        interpolate_met(rad, vector, "toa_upward_longwave_flux", "olr", **interp_kwargs)
+    gfs_key = "toa_upward_longwave_flux"
+    if gfs_key in rad:
+        interpolate_met(rad, vector, gfs_key, "olr", **interp_kwargs)
         return
 
     # ECMWF contains "top_net_thermal_radiation" which is -1 * OLR
-    if "top_net_thermal_radiation" in rad:
-        interpolate_met(rad, vector, "top_net_thermal_radiation", **interp_kwargs)
-        vector["olr"] = np.maximum(-vector["top_net_thermal_radiation"], 0.0)
-        return
+    ecmwf_key = "top_net_thermal_radiation"
+    if ecmwf_key not in rad:
+        msg = f"'rad' data must contain either '{gfs_key}' or '{ecmwf_key}' (ECMWF) variable."
+        raise ValueError(msg)
 
-    raise ValueError(
-        "rad data must contain either 'toa_upward_longwave_flux' "
-        "or 'top_net_thermal_radiation' (ECMWF) variable."
-    )
+    tntr = interpolate_met(rad, vector, ecmwf_key, **interp_kwargs)
+    tntr = _rad_accumulation_to_average_instantaneous(rad, ecmwf_key, tntr)
+    vector.update({ecmwf_key: tntr})
+
+    vector["olr"] = np.maximum(-tntr, 0.0)
 
 
 def calc_radiative_properties(contrail: GeoVectorDataset, params: dict[str, Any]) -> None:
@@ -1995,8 +2045,7 @@ def calc_contrail_properties(
     contrail.update(dn_dt_agg=dn_dt_agg)
     contrail.update(dn_dt_turb=dn_dt_turb)
     if radiative_heating_effects:
-        contrail.update(heat_rate=heat_rate)
-        contrail.update(d_heat_rate=d_heat_rate)
+        contrail.update(heat_rate=heat_rate, d_heat_rate=d_heat_rate)
 
 
 def calc_timestep_contrail_evolution(
@@ -2279,6 +2328,57 @@ def calc_timestep_contrail_evolution(
     return final_contrail
 
 
+def _rad_accumulation_to_average_instantaneous(
+    rad: MetDataset,
+    name: str,
+    arr: npt.NDArray[np.float_],
+) -> npt.NDArray[np.float_]:
+    """Convert from radiation accumulation to average instantaneous values.
+
+    .. versionadded:: 0.48.0
+
+    Parameters
+    ----------
+    rad : MetDataset
+        Radiation data
+    name : str
+        Variable name
+    arr : npt.NDArray[np.float_]
+        Array of values already interpolated from ``rad``
+
+    Returns
+    -------
+    npt.NDArray[np.float_]
+        Array of values converted from accumulation to average instantaneous values
+    """
+    mda = rad[name]
+    try:
+        unit = mda.attrs["units"]
+    except KeyError as e:
+        msg = (
+            f"Radiation data contains '{name}' variable "
+            "but units are not specified. Provide units in the "
+            f"rad['{name}'].attrs passed into Cocip."
+        )
+        raise KeyError(msg) from e
+
+    # The unit is already instantaneous
+    if unit == "W m**-2":
+        return arr
+
+    if unit != "J m**-2":
+        msg = f"Unexpected units '{unit}' for '{name}' variable. Expected 'J m**-2' or 'W m**-2'."
+        raise ValueError(msg)
+
+    # Convert from J m**-2 to W m**-2
+    if rad.dataset_attr == "ERA5" and rad.product_attr == "ensemble":
+        n_seconds = 3.0 * 3600.0  # 3 hour interval
+    else:
+        n_seconds = 3600.0  # 1 hour interval
+
+    return arr / n_seconds
+
+
 def _emissions_variables() -> tuple[str, ...]:
     """Return variables required for emissions calculation."""
     return (

pycontrails/models/cocip/cocip_params.py

@@ -73,10 +73,6 @@ class CocipParams(ModelParams):
     #: Constant below applies to ECMWF data.
     effective_vertical_resolution: float = 2000.0
 
-    #: Shift the time coordinates of radiation parameters for accumulated values
-    #: TODO: change this to np.timedelta64(0, "m") when we start using other datasets
-    shift_radiation_time: np.timedelta64 = -np.timedelta64(30, "m")
-
     #: Smoothing parameters for true airspeed.
     #: Only used for Flight models.
     #: Passed directly to :func:`scipy.signal.savgol_filter`.

pycontrails/models/cocip/wake_vortex.py

@@ -54,22 +54,24 @@ def max_downward_displacement(
     - :cite:`holzapfelProbabilisticTwoPhaseWake2003`
     - :cite:`schumannContrailCirrusPrediction2012`
     """
-    wingspan_arr = np.broadcast_to(wingspan, true_airspeed.shape)
-    aircraft_mass_arr = np.broadcast_to(aircraft_mass, true_airspeed.shape)
-
     rho_air = thermo.rho_d(air_temperature, air_pressure)
     n_bv = thermo.brunt_vaisala_frequency(air_pressure, air_temperature, dT_dz)
-    t_0 = effective_time_scale(wingspan, true_airspeed, aircraft_mass_arr, rho_air)
+    t_0 = effective_time_scale(wingspan, true_airspeed, aircraft_mass, rho_air)
 
     dz_max_strong = downward_displacement_strongly_stratified(
-        wingspan, true_airspeed, aircraft_mass_arr, rho_air, n_bv
+        wingspan, true_airspeed, aircraft_mass, rho_air, n_bv
     )
 
     is_weakly_stratified = n_bv * t_0 < 0.8
+    if isinstance(wingspan, np.ndarray):
+        wingspan = wingspan[is_weakly_stratified]
+    if isinstance(aircraft_mass, np.ndarray):
+        aircraft_mass = aircraft_mass[is_weakly_stratified]
+
     dz_max_weak = downward_displacement_weakly_stratified(
-        wingspan=wingspan_arr[is_weakly_stratified],
+        wingspan=wingspan,
         true_airspeed=true_airspeed[is_weakly_stratified],
-        aircraft_mass=aircraft_mass_arr[is_weakly_stratified],
+        aircraft_mass=aircraft_mass,
         rho_air=rho_air[is_weakly_stratified],
         n_bv=n_bv[is_weakly_stratified],
         dz_max_strong=dz_max_strong[is_weakly_stratified],

pycontrails/models/cocipgrid/cocip_grid.py

@@ -106,12 +106,8 @@ class CocipGrid(models.Model, cocip_time_handling.CocipTimeHandlingMixin):
         super().__init__(met, params=params, **params_kwargs)
         self.validate_time_params()
 
-        shift_radiation_time = self.params["shift_radiation_time"]
         compute_tau_cirrus = self.params["compute_tau_cirrus_in_model_init"]
-        met, rad = cocip.process_met_datasets(met, rad, compute_tau_cirrus, shift_radiation_time)
-
-        self.met = met
-        self.rad = rad
+        self.met, self.rad = cocip.process_met_datasets(met, rad, compute_tau_cirrus)
 
         # Convenience -- only used in `run_interpolators`
         self.params["_interp_kwargs"] = self.interp_kwargs
@@ -309,12 +305,12 @@ class CocipGrid(models.Model, cocip_time_handling.CocipTimeHandlingMixin):
         else:
             dt_integration_str = str(dt_integration.astype("timedelta64[s]"))
 
+        self.transfer_met_source_attrs()
         attrs: dict[str, Any] = {
             "description": self.long_name,
             "max_age": max_age_str,
             "dt_integration": dt_integration_str,
             "aircraft_type": self.get_source_param("aircraft_type"),
-            "met_source": self.met.attrs.get("met_source", "unknown"),
             "pycontrails_version": pycontrails.__version__,
             **self.source.attrs,  # type: ignore[dict-item]
         }