pycontrails 0.54.4__cp311-cp311-win_amd64.whl → 0.54.6__cp311-cp311-win_amd64.whl

pycontrails/models/cocip/cocip.py

@@ -68,12 +68,26 @@ class Cocip(Model):
     -----
     **Inputs**
 
-    The required meteorology variables depend on the data source (e.g. ECMWF, GFS).
+    The required meteorology variables depend on the data source. :class:`Cocip`
+    supports data-source-specific variables from ECMWF models (HRES, ERA5) and the NCEP GFS, plus
+    a generic set of model-agnostic variables.
 
     See :attr:`met_variables` and :attr:`rad_variables` for the list of required variables
     to the ``met`` and ``rad`` parameters, respectively.
     When an item in one of these arrays is a :class:`tuple`, variable keys depend on data source.
 
+    A warning will be raised if meteorology data is from a source not currently supported by
+    a pycontrails datalib. In this case it is the responsibility of the user to ensure that
+    meteorology data is formatted correctly. The warning can be suppressed with a context manager:
+
+    .. code-block:: python
+        :emphasize-lines: 2,3
+
+        import warnings
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=UserWarning, message="Unknown provider")
+            cocip = Cocip(met, rad, ...)
+
     The current list of required variables (labelled by ``"standard_name"``):
 
     .. list-table:: Variable keys for pressure level data
@@ -82,24 +96,31 @@ class Cocip(Model):
         * - Parameter
           - ECMWF
           - GFS
+          - Generic
         * - Air Temperature
           - ``air_temperature``
           - ``air_temperature``
+          - ``air_temperature``
         * - Specific Humidity
           - ``specific_humidity``
           - ``specific_humidity``
+          - ``specific_humidity``
         * - Eastward wind
           - ``eastward_wind``
           - ``eastward_wind``
+          - ``eastward_wind``
         * - Northward wind
           - ``northward_wind``
           - ``northward_wind``
+          - ``northward_wind``
         * - Vertical velocity
           - ``lagrangian_tendency_of_air_pressure``
           - ``lagrangian_tendency_of_air_pressure``
+          - ``lagrangian_tendency_of_air_pressure``
         * - Ice water content
           - ``specific_cloud_ice_water_content``
           - ``ice_water_mixing_ratio``
+          - ``mass_fraction_of_cloud_ice_in_air``
 
     .. list-table:: Variable keys for single-level radiation data
         :header-rows: 1
@@ -107,12 +128,15 @@ class Cocip(Model):
         * - Parameter
           - ECMWF
           - GFS
+          - Generic
         * - Top solar radiation
           - ``top_net_solar_radiation``
           - ``toa_upward_shortwave_flux``
+          - ``toa_net_downward_shortwave_flux``
         * - Top thermal radiation
           - ``top_net_thermal_radiation``
           - ``toa_upward_longwave_flux``
+          - ``toa_outgoing_longwave_flux``
 
     **Modifications**
 
@@ -214,14 +238,26 @@ class Cocip(Model):
         met_var.EastwardWind,
         met_var.NorthwardWind,
         met_var.VerticalVelocity,
-        (ecmwf.SpecificCloudIceWaterContent, gfs.CloudIceWaterMixingRatio),
+        (
+            met_var.MassFractionOfCloudIceInAir,
+            ecmwf.SpecificCloudIceWaterContent,
+            gfs.CloudIceWaterMixingRatio,
+        ),
     )
 
     #: Required single-level top of atmosphere radiation variables.
     #: Variable keys depend on data source (e.g. ECMWF, GFS).
     rad_variables = (
-        (ecmwf.TopNetSolarRadiation, gfs.TOAUpwardShortwaveRadiation),
-        (ecmwf.TopNetThermalRadiation, gfs.TOAUpwardLongwaveRadiation),
+        (
+            met_var.TOANetDownwardShortwaveFlux,
+            ecmwf.TopNetSolarRadiation,
+            gfs.TOAUpwardShortwaveRadiation,
+        ),
+        (
+            met_var.TOAOutgoingLongwaveFlux,
+            ecmwf.TopNetThermalRadiation,
+            gfs.TOAUpwardLongwaveRadiation,
+        ),
     )
 
     #: Minimal set of met variables needed to run the model after pre-processing.
@@ -242,7 +278,11 @@ class Cocip(Model):
     #:   Moved Geopotential from :attr:`met_variables` to :attr:`optional_met_variables`
     optional_met_variables = (
         (met_var.Geopotential, met_var.GeopotentialHeight),
-        (ecmwf.CloudAreaFractionInLayer, gfs.TotalCloudCoverIsobaric),
+        (
+            met_var.CloudAreaFractionInAtmosphereLayer,
+            ecmwf.CloudAreaFractionInLayer,
+            gfs.TotalCloudCoverIsobaric,
+        ),
     )
 
     #: Met data is not optional
@@ -391,7 +431,7 @@ class Cocip(Model):
         # which is the positive direction for level
         logger.debug("Downselect met for Cocip initialization")
         level_buffer = 0, self.params["met_level_buffer"][1]
-        met = self.source.downselect_met(self.met, level_buffer=level_buffer, copy=False)
+        met = self.source.downselect_met(self.met, level_buffer=level_buffer)
         met = add_tau_cirrus(met)
 
         # Prepare flight for model
@@ -575,10 +615,12 @@ class Cocip(Model):
         if verbose_outputs:
             interpolate_met(met, self.source, "tau_cirrus", **interp_kwargs)
 
-            # handle ECMWF/GFS ciwc variables
+            # handle ECMWF/GFS/generic ciwc variables
             if (key := "specific_cloud_ice_water_content") in met:  # noqa: SIM114
                 interpolate_met(met, self.source, key, **interp_kwargs)
-            elif (key := "ice_water_mixing_ratio") in met:
+            elif (key := "ice_water_mixing_ratio") in met:  # noqa: SIM114
+                interpolate_met(met, self.source, key, **interp_kwargs)
+            elif (key := "mass_fraction_of_cloud_ice_in_air") in met:
                 interpolate_met(met, self.source, key, **interp_kwargs)
 
             self.source["rho_air"] = thermo.rho_d(
@@ -976,9 +1018,9 @@ class Cocip(Model):
             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)
+        met = self._downwash_contrail.downselect_met(self.met, **buffers)
         met = add_tau_cirrus(met)
-        rad = self._downwash_contrail.downselect_met(self.rad, **buffers, copy=False)
+        rad = self._downwash_contrail.downselect_met(self.rad, **buffers)
 
         calc_continuous(self._downwash_contrail)
         calc_timestep_geometry(self._downwash_contrail)
@@ -1135,11 +1177,11 @@ class Cocip(Model):
             & (self._downwash_flight["time"] <= lookahead),
             copy=False,
         )
-        vector = GeoVectorDataset(
+        vector = GeoVectorDataset._from_fastpath(
             {
                 key: np.concatenate((latest_contrail[key], future_contrails[key]))
                 for key in ("longitude", "latitude", "level", "time")
-            }
+            },
         )
 
         # compute time buffer to ensure downselection extends to time_end
@@ -1152,7 +1194,7 @@ class Cocip(Model):
             max(np.timedelta64(0, "ns"), time_end - vector["time"].max()),
         )
 
-        return vector.downselect_met(met, **buffers, copy=False)
+        return vector.downselect_met(met, **buffers)
 
     def _create_downwash_contrail(self) -> GeoVectorDataset:
         """Get Contrail representation of downwash flight."""
@@ -1180,7 +1222,7 @@ class Cocip(Model):
             "persistent": self._downwash_flight["persistent_1"],
         }
 
-        contrail = GeoVectorDataset(downwash_contrail_data, copy=True)
+        contrail = GeoVectorDataset._from_fastpath(downwash_contrail_data).copy()
         contrail["formation_time"] = contrail["time"].copy()
         contrail["age"] = contrail["formation_time"] - contrail["time"]
 
@@ -1587,8 +1629,7 @@ def _process_rad(rad: MetDataset) -> MetDataset:
                 rad.data["time"].attrs["shift_radiation_time"] = "variable"
             return rad
 
-        else:
-            shift_radiation_time = -np.timedelta64(30, "m")
+        shift_radiation_time = -np.timedelta64(30, "m")
 
     elif dataset == "ERA5" and product == "ensemble":
         shift_radiation_time = -np.timedelta64(90, "m")
@@ -1893,8 +1934,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_net_downward_shortwave_flux"``,
+        ``"toa_upward_shortwave_flux"`` or ``"top_net_solar_radiation"`` variable.
 
     Notes
     -----
@@ -1918,6 +1959,13 @@ def calc_shortwave_radiation(
         sdr = geo.solar_direct_radiation(longitude, latitude, time, threshold_cos_sza=0.01)
         vector["sdr"] = sdr
 
+    # Generic contains net downward shortwave flux at TOA (SDR - RSR) in W/m2
+    generic_key = "toa_net_downward_shortwave_flux"
+    if generic_key in rad:
+        tnsr = interpolate_met(rad, vector, generic_key, **interp_kwargs)
+        vector["rsr"] = np.maximum(sdr - tnsr, 0.0)
+        return
+
     # GFS contains RSR (toa_upward_shortwave_flux) variable directly
     gfs_key = "toa_upward_shortwave_flux"
     if gfs_key in rad:
@@ -1926,10 +1974,13 @@ def calc_shortwave_radiation(
 
     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."
+        msg = (
+            f"'rad' data must contain either '{generic_key}' (generic), "
+            f"'{gfs_key}' (GFS), or '{ecmwf_key}' (ECMWF) variable."
+        )
         raise ValueError(msg)
 
-    # ECMWF contains "top_net_solar_radiation" which is SDR - RSR
+    # ECMWF also contains net downward shortwave flux at TOA, but possibly as an accumulation
     tnsr = interpolate_met(rad, vector, ecmwf_key, **interp_kwargs)
     tnsr = _rad_accumulation_to_average_instantaneous(rad, ecmwf_key, tnsr)
     vector.update({ecmwf_key: tnsr})
@@ -1958,14 +2009,20 @@ def calc_outgoing_longwave_radiation(
     Raises
     ------
     ValueError
-        If ``rad`` does not contain a ``"toa_upward_longwave_flux"``
-        or ``"top_net_thermal_radiation"`` variable.
+        If ``rad`` does not contain a ``"toa_outgoing_longwave_flux"``,
+        ``"toa_upward_longwave_flux"`` or ``"top_net_thermal_radiation"`` variable.
     """
 
     if "olr" in vector:
-        return None
+        return
+
+    # Generic contains OLR (toa_outgoing_longwave_flux) directly
+    generic_key = "toa_outgoing_longwave_flux"
+    if generic_key in rad:
+        interpolate_met(rad, vector, generic_key, "olr", **interp_kwargs)
+        return
 
-    # GFS contains OLR (toa_upward_longwave_flux) variable directly
+    # GFS contains OLR (toa_upward_longwave_flux) directly
     gfs_key = "toa_upward_longwave_flux"
     if gfs_key in rad:
         interpolate_met(rad, vector, gfs_key, "olr", **interp_kwargs)
@@ -1974,7 +2031,10 @@ def calc_outgoing_longwave_radiation(
     # ECMWF contains "top_net_thermal_radiation" which is -1 * OLR
     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."
+        msg = (
+            f"'rad' data must contain either '{generic_key}' (generic), "
+            f"'{gfs_key}' (GFS), or '{ecmwf_key}' (ECMWF) variable."
+        )
         raise ValueError(msg)
 
     tntr = interpolate_met(rad, vector, ecmwf_key, **interp_kwargs)
@@ -2300,7 +2360,7 @@ def calc_timestep_contrail_evolution(
     level_2 = geo.advect_level(level_1, vertical_velocity_1, rho_air_1, terminal_fall_speed_1, dt)
     altitude_2 = units.pl_to_m(level_2)
 
-    contrail_2 = GeoVectorDataset(
+    contrail_2 = GeoVectorDataset._from_fastpath(
         {
             "waypoint": waypoint_2,
             "flight_id": contrail_1["flight_id"],
@@ -2312,7 +2372,6 @@ def calc_timestep_contrail_evolution(
             "altitude": altitude_2,
             "level": level_2,
         },
-        copy=False,
     )
     intersection = contrail_2.coords_intersect_met(met)
     if not np.any(intersection):

pycontrails/models/cocip/output_formats.py

@@ -2259,3 +2259,4 @@ def compare_cocip_with_goes(
         plt.close()
 
         return output_path
+    return None

pycontrails/models/cocipgrid/cocip_grid.py

@@ -11,11 +11,10 @@ from typing import TYPE_CHECKING, Any, NoReturn, TypeVar, overload
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
-import xarray as xr
 
 import pycontrails
 from pycontrails.core import models
-from pycontrails.core.met import MetDataset
+from pycontrails.core.met import MetDataset, maybe_downselect_mds
 from pycontrails.core.vector import GeoVectorDataset, VectorDataset
 from pycontrails.models import humidity_scaling, sac
 from pycontrails.models.cocip import cocip, contrail_properties, wake_vortex, wind_shear
@@ -323,8 +322,8 @@ class CocipGrid(models.Model):
         If ``self.params["downselect_met"]`` is True, the :func:`_downselect_met` has
         already performed a spatial downselection of the met data.
         """
-        met = _maybe_downselect_mds(self.met, met, t0, t1)
-        rad = _maybe_downselect_mds(self.rad, rad, t0, t1)
+        met = maybe_downselect_mds(self.met, met, t0, t1)
+        rad = maybe_downselect_mds(self.rad, rad, t0, t1)
 
         return met, rad
 
@@ -615,7 +614,7 @@ class CocipGrid(models.Model):
             for idx, time in enumerate(times_in_filt):
                 # For now, sticking with the convention that every vector should
                 # have a constant time value.
-                source_slice = MetDataset(self.source.data.sel(time=[time]))
+                source_slice = MetDataset._from_fastpath(self.source.data.sel(time=[time]))
 
                 # Convert the 4D grid to a vector
                 vector = source_slice.to_vector()
@@ -1402,7 +1401,7 @@ def simulate_wake_vortex_downwash(
 
     # Experimental segment-free model
     if _is_segment_free_mode(vector):
-        return GeoVectorDataset(data, attrs=vector.attrs, copy=True)
+        return GeoVectorDataset._from_fastpath(data, attrs=vector.attrs).copy()
 
     # Stored in `_generate_new_grid_vectors`
     data["longitude_head"] = vector["longitude_head"]
@@ -1421,7 +1420,7 @@ def simulate_wake_vortex_downwash(
         # segment_length variable.
         data["segment_length"] = np.full_like(data["longitude"], segment_length)
 
-    return GeoVectorDataset(data, attrs=vector.attrs, copy=True)
+    return GeoVectorDataset._from_fastpath(data, attrs=vector.attrs).copy()
 
 
 def find_initial_persistent_contrails(
@@ -2022,7 +2021,7 @@ def advect(
         assert _is_segment_free_mode(contrail)
         assert dt_tail is None
         assert dt_head is None
-        return GeoVectorDataset(data, attrs=contrail.attrs, copy=True)
+        return GeoVectorDataset._from_fastpath(data, attrs=contrail.attrs).copy()
 
     longitude_head = contrail["longitude_head"]
     latitude_head = contrail["latitude_head"]
@@ -2064,7 +2063,7 @@ def advect(
     data["segment_length"] = segment_length_t2
     data["head_tail_dt"] = head_tail_dt_t2
 
-    return GeoVectorDataset(data, attrs=contrail.attrs, copy=True)
+    return GeoVectorDataset._from_fastpath(data, attrs=contrail.attrs).copy()
 
 
 def _aggregate_ef_summary(vector_list: list[VectorDataset]) -> VectorDataset | None:
@@ -2438,7 +2437,6 @@ def _downselect_met(
         longitude_buffer=longitude_buffer,
         level_buffer=level_buffer,
         time_buffer=(t0, t1),
-        copy=False,
     )
 
     rad = source.downselect_met(
@@ -2446,7 +2444,6 @@ def _downselect_met(
         latitude_buffer=latitude_buffer,
         longitude_buffer=longitude_buffer,
         time_buffer=(t0, t1),
-        copy=False,
     )
 
     return met, rad
@@ -2522,65 +2519,3 @@ def _check_end_time(
             f"Include additional time at the end of '{name}' or reduce 'max_age' parameter."
             f"{note}"
         )
-
-
-def _maybe_downselect_mds(
-    big_mds: MetDataset,
-    little_mds: MetDataset | None,
-    t0: np.datetime64,
-    t1: np.datetime64,
-) -> MetDataset:
-    """Possibly downselect ``big_mds`` to cover ``[t0, t1]``.
-
-    This implementation assumes ``t0 <= t1``, but this is not enforced.
-
-    If possible, ``little_mds`` is recycled to avoid re-loading data.
-
-    This function only downselects in the time domain.
-
-    If ``big_mds`` doesn't cover the time range, no error is raised.
-    """
-    if little_mds is not None:
-        little_time = little_mds.indexes["time"].to_numpy()
-        ignore_little = t0 > little_time[-1] or t1 < little_time[0]
-
-    big_time = big_mds.indexes["time"].to_numpy()
-    if little_mds is None or ignore_little:
-        i0 = np.searchsorted(big_time, t0, side="right").item()
-        i0 = max(0, i0 - 1)
-        i1 = np.searchsorted(big_time, t1, side="left").item()
-        i1 = min(i1 + 1, big_time.size)
-        return MetDataset(big_mds.data.isel(time=slice(i0, i1)), copy=False)
-
-    j0 = np.searchsorted(little_time, t0, side="right").item()
-    j0 = max(0, j0 - 1)
-    j1 = np.searchsorted(little_time, t1, side="left").item()
-    j1 = min(j1 + 1, little_time.size)
-
-    little_ds = little_mds.data.isel(time=slice(j0, j1))
-    little_time0 = little_time[j0]
-    little_time1 = little_time[j1 - 1]
-
-    if t0 >= little_time0 and t1 <= little_time1:
-        return MetDataset(little_ds, copy=False)
-
-    ds_concat = []
-    if t0 < little_time0:  # unlikely to encounter this case
-        i0 = np.searchsorted(big_time, t0, side="right").item()
-        i0 = max(0, i0 - 1)
-        i1 = np.searchsorted(big_time, little_time0, side="right").item()
-        i1 = max(i1, i0 + 1)
-        ds_concat.append(big_mds.data.isel(time=slice(i0, i1)))
-
-    ds_concat.append(little_ds)
-
-    if t1 > little_time1:
-        i0 = np.searchsorted(big_time, little_time1, side="left").item()
-        i0 = min(i0 + 1, big_time.size)
-        i1 = np.searchsorted(big_time, t1, side="left").item()
-        i1 = min(i1 + 1, big_time.size)
-        ds_concat.append(big_mds.data.isel(time=slice(i0, i1)))
-
-    # If little_mds is loaded into memory but big_mds is not,
-    # the concat operation below will load the slice of big_mds into memory.
-    return MetDataset(xr.concat(ds_concat, dim="time"), copy=False)

pycontrails/models/dry_advection.py

@@ -3,16 +3,24 @@
 from __future__ import annotations
 
 import dataclasses
+import sys
 from typing import Any, NoReturn, overload
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 import numpy as np
 import numpy.typing as npt
+import pandas as pd
 
 from pycontrails.core import models
-from pycontrails.core.met import MetDataset
+from pycontrails.core.met import MetDataset, maybe_downselect_mds
 from pycontrails.core.met_var import (
     AirTemperature,
     EastwardWind,
+    MetVariable,
     NorthwardWind,
     VerticalVelocity,
 )
@@ -55,6 +63,15 @@ class DryAdvectionParams(models.AdvectionBuffers):
     # If None, only pointwise advection is simulated without wind shear effects.
     azimuth: float | None = 0.0
 
+    #: Add additional intermediate variables to the output vector.
+    #: This includes interpolated met variables and wind-shear-derived geometry.
+    verbose_outputs: bool = False
+
+    #: Whether to include ``source`` points in the output vector. Enabling allows
+    #: the user to view additional data (e.g., interpolated met variables) for
+    #: source points as well as evolved points.
+    include_source_in_output: bool = False
+
 
 class DryAdvection(models.Model):
     """Simulate "dry advection" of an emissions plume with an elliptical cross section.
@@ -89,7 +106,12 @@ class DryAdvection(models.Model):
 
     name = "dry_advection"
     long_name = "Emission plume advection without sedimentation"
-    met_variables = AirTemperature, EastwardWind, NorthwardWind, VerticalVelocity
+    met_variables: tuple[MetVariable, ...] = (
+        AirTemperature,
+        EastwardWind,
+        NorthwardWind,
+        VerticalVelocity,
+    )
     default_params = DryAdvectionParams
 
     met: MetDataset
@@ -127,6 +149,10 @@ class DryAdvection(models.Model):
         self.update_params(params)
         self.set_source(source)
         self.source = self.require_source_type(GeoVectorDataset)
+        self.downselect_met()
+        if not self.source.coords_intersect_met(self.met).any():
+            msg = "No source coordinates intersect met data."
+            raise ValueError(msg)
 
         self.source = self._prepare_source()
 
@@ -137,36 +163,51 @@ class DryAdvection(models.Model):
         sedimentation_rate = self.params["sedimentation_rate"]
         dz_m = self.params["dz_m"]
         max_depth = self.params["max_depth"]
+        verbose_outputs = self.params["verbose_outputs"]
 
         source_time = self.source["time"]
-        t0 = source_time.min()
+        t0 = pd.Timestamp(source_time.min()).floor(pd.Timedelta(dt_integration)).to_numpy()
         t1 = source_time.max()
         timesteps = np.arange(t0 + dt_integration, t1 + dt_integration + max_age, dt_integration)
 
-        vector = GeoVectorDataset()
+        vector2 = GeoVectorDataset()
+        met = None
 
         evolved = []
         for t in timesteps:
             filt = (source_time < t) & (source_time >= t - dt_integration)
-            vector = vector + self.source.filter(filt, copy=False)
-            vector = _evolve_one_step(
-                self.met,
-                vector,
+            vector1 = vector2 + self.source.filter(filt, copy=False)
+
+            t0 = vector1["time"].min()
+            t1 = vector1["time"].max()
+            met = maybe_downselect_mds(self.met, met, t0, t1)
+
+            vector2 = _evolve_one_step(
+                met,
+                vector1,
                 t,
                 sedimentation_rate=sedimentation_rate,
                 dz_m=dz_m,
                 max_depth=max_depth,
+                verbose_outputs=verbose_outputs,
                 **interp_kwargs,
             )
+            evolved.append(vector1)
 
-            filt = (vector["age"] <= max_age) & vector.coords_intersect_met(self.met)
-            vector = vector.filter(filt)
+            filt = (vector2["age"] <= max_age) & vector2.coords_intersect_met(self.met)
+            vector2 = vector2.filter(filt)
 
-            evolved.append(vector)
-            if not vector and np.all(source_time < t):
+            if not vector2 and np.all(source_time < t):
                 break
 
-        return GeoVectorDataset.sum(evolved, fill_value=np.nan)
+        evolved.append(vector2)
+        out = GeoVectorDataset.sum(evolved, fill_value=np.nan)
+
+        if self.params["include_source_in_output"]:
+            return out
+
+        filt = out["age"] > np.timedelta64(0, "ns")
+        return out.filter(filt)
 
     def _prepare_source(self) -> GeoVectorDataset:
         r"""Prepare :attr:`source` vector for advection by wind-shear-derived variables.
@@ -202,7 +243,7 @@ class DryAdvection(models.Model):
                 raise ValueError(
                     "If 'azimuth' is None, then 'width' and 'depth' must also be None."
                 )
-            return GeoVectorDataset(self.source.select(columns, copy=False), copy=False)
+            return GeoVectorDataset._from_fastpath(self.source.select(columns, copy=False).data)
 
         if "azimuth" not in self.source:
             self.source["azimuth"] = np.full_like(self.source["longitude"], azimuth)
@@ -228,7 +269,19 @@ class DryAdvection(models.Model):
             width, depth, sigma_yz=0.0
         )
 
-        return GeoVectorDataset(self.source.select(columns, copy=False), copy=False)
+        return GeoVectorDataset._from_fastpath(self.source.select(columns, copy=False).data)
+
+    @override
+    def downselect_met(self) -> None:
+        if not self.params["downselect_met"]:
+            return
+
+        buffers = {
+            f"{coord}_buffer": self.params[f"met_{coord}_buffer"]
+            for coord in ("longitude", "latitude", "level")
+        }
+        buffers["time_buffer"] = (np.timedelta64(0, "ns"), self.params["max_age"])
+        self.met = self.source.downselect_met(self.met, **buffers)
 
 
 def _perform_interp_for_step(
@@ -329,8 +382,12 @@ def _calc_geometry(
     dz_m: float,
     dt: npt.NDArray[np.timedelta64] | np.timedelta64,
     max_depth: float | None,
+    verbose_outputs: bool,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-    """Calculate wind-shear-derived geometry of evolved plume."""
+    """Calculate wind-shear-derived geometry of evolved plume.
+
+    This method mutates the input ``vector`` in place.
+    """
 
     u_wind = vector["u_wind"]
     v_wind = vector["v_wind"]
@@ -384,6 +441,11 @@ def _calc_geometry(
         eff_heat_rate=None,
     )
 
+    if verbose_outputs:
+        vector["ds_dz"] = ds_dz
+        vector["dsn_dz"] = dsn_dz
+        vector["dT_dz"] = dT_dz
+
     sigma_yy_2, sigma_zz_2, sigma_yz_2 = contrail_properties.plume_temporal_evolution(
         width,
         depth,
@@ -442,9 +504,13 @@ def _evolve_one_step(
     sedimentation_rate: float,
     dz_m: float,
     max_depth: float | None,
+    verbose_outputs: bool,
     **interp_kwargs: Any,
 ) -> GeoVectorDataset:
-    """Evolve plume geometry by one step."""
+    """Evolve plume geometry by one step.
+
+    This method mutates the input ``vector`` in place.
+    """
 
     _perform_interp_for_step(met, vector, dz_m, **interp_kwargs)
     u_wind = vector["u_wind"]
@@ -459,20 +525,21 @@ def _evolve_one_step(
     level_2 = geo.advect_level(
         vector.level,
         vertical_velocity,
-        0.0,
-        0.0,
-        dt,  # type: ignore[arg-type]
+        rho_air=0.0,
+        terminal_fall_speed=0.0,
+        dt=dt,  # type: ignore[arg-type]
     )
 
-    out = GeoVectorDataset(
-        longitude=longitude_2,
-        latitude=latitude_2,
-        level=level_2,
-        time=np.full(longitude_2.shape, t),
-        copy=False,
+    out = GeoVectorDataset._from_fastpath(
+        {
+            "longitude": longitude_2,
+            "latitude": latitude_2,
+            "level": level_2,
+            "time": np.full(longitude_2.shape, t),
+            "age": vector["age"] + dt,
+            "waypoint": vector["waypoint"],
+        }
     )
-    out["age"] = vector["age"] + dt
-    out["waypoint"] = vector["waypoint"]
 
     azimuth = vector.get("azimuth")
     if azimuth is None:
@@ -482,9 +549,10 @@ def _evolve_one_step(
     # Attach wind-shear-derived geometry to output vector
     azimuth_2, width_2, depth_2, sigma_yz_2, area_eff_2 = _calc_geometry(
         vector,
-        dz_m,
-        dt,  # type: ignore[arg-type]
-        max_depth,  # type: ignore[arg-type]
+        dz_m=dz_m,
+        dt=dt,  # type: ignore[arg-type]
+        max_depth=max_depth,  # type: ignore[arg-type]
+        verbose_outputs=verbose_outputs,
     )
     out["azimuth"] = azimuth_2
     out["width"] = width_2