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

pycontrails/models/cocipgrid/cocip_grid.py

@@ -74,11 +74,11 @@ class CocipGrid(models.Model):
     """
 
     __slots__ = (
-        "rad",
-        "timesteps",
+        "_target_dtype",
         "contrail",
         "contrail_list",
-        "_target_dtype",
+        "rad",
+        "timesteps",
     )
 
     name = "contrail_grid"
@@ -241,12 +241,12 @@ class CocipGrid(models.Model):
         existing_vectors: Iterator[GeoVectorDataset] = iter(())
 
         for time_idx, time_end in enumerate(self.timesteps):
-            met, rad = self._maybe_downselect_met_rad(met, rad, time_end)
-
             evolved_this_step = []
             ef_summary_this_step = []
             downwash_vectors_this_step = []
             for vector in self._generate_new_vectors(time_idx):
+                t0 = vector["time"].min()
+                met, rad = self._maybe_downselect_met_rad(met, rad, t0, time_end)
                 downwash, verbose_dict = _run_downwash(vector, met, rad, self.params)
 
                 if downwash:
@@ -264,6 +264,8 @@ class CocipGrid(models.Model):
                     pbar.update()
 
             for vector in itertools.chain(existing_vectors, downwash_vectors_this_step):
+                t0 = vector["time"].min()
+                met, rad = self._maybe_downselect_met_rad(met, rad, t0, time_end)
                 contrail, ef = _evolve_vector(
                     vector,
                     met=met,
@@ -304,83 +306,25 @@ class CocipGrid(models.Model):
         self,
         met: MetDataset | None,
         rad: MetDataset | None,
-        time_end: np.datetime64,
+        t0: np.datetime64,
+        t1: np.datetime64,
     ) -> tuple[MetDataset, MetDataset]:
-        """Downselect ``self.met`` and ``self.rad`` if necessary to cover ``time_end``.
+        """Downselect ``self.met`` and ``self.rad`` if necessary to cover ``[t0, t1]``.
+
+        This implementation assumes ``t0 <= t1``, but does not enforce this.
 
         If the currently used ``met`` and ``rad`` slices do not include the time
-        ``time_end``, new slices are selected from the larger ``self.met`` and
-        ``self.rad`` data. The slicing only occurs in the time domain.
+        interval ``[t0, t1]``, new slices are selected from the larger ``self.met``
+        and ``self.rad`` data. The slicing only occurs in the time domain.
 
-        The end of currently-used ``met`` and ``rad`` will be used as the start
-        of newly-selected met slices when possible to avoid losing and re-loading
-        already-loaded met data.
+        Existing slices from ``met`` and ``rad`` will be used when possible to avoid
+        losing and re-loading already-loaded met data.
 
-        If ``self.params["downselect_met"]`` is True, :func:`_downselect_met` has
+        If ``self.params["downselect_met"]`` is True, the :func:`_downselect_met` has
         already performed a spatial downselection of the met data.
         """
-
-        if met is None:
-            # idx is the first index at which self.met.variables["time"].to_numpy() >= time_end
-            idx = np.searchsorted(self.met.indexes["time"].to_numpy(), time_end).item()
-            sl = slice(max(0, idx - 1), idx + 1)
-            logger.debug("Select met slice %s", sl)
-            met = MetDataset(self.met.data.isel(time=sl), copy=False)
-
-        elif time_end > met.indexes["time"].to_numpy()[-1]:
-            current_times = met.indexes["time"].to_numpy()
-            all_times = self.met.indexes["time"].to_numpy()
-            # idx is the first index at which all_times >= time_end
-            idx = np.searchsorted(all_times, time_end).item()
-            sl = slice(max(0, idx - 1), idx + 1)
-
-            # case 1: cannot re-use end of current met as start of new met
-            if current_times[-1] != all_times[sl.start]:
-                logger.debug("Select met slice %s", sl)
-                met = MetDataset(self.met.data.isel(time=sl), copy=False)
-            # case 2: can re-use end of current met plus one step of new met
-            elif sl.start < all_times.size - 1:
-                sl = slice(sl.start + 1, sl.stop)
-                logger.debug("Reuse end of met and select met slice %s", sl)
-                met = MetDataset(
-                    xr.concat((met.data.isel(time=[-1]), self.met.data.isel(time=sl)), dim="time"),
-                    copy=False,
-                )
-            # case 3: can re-use end of current met and nothing else
-            else:
-                logger.debug("Reuse end of met")
-                met = MetDataset(met.data.isel(time=[-1]), copy=False)
-
-        if rad is None:
-            # idx is the first index at which self.rad.variables["time"].to_numpy() >= time_end
-            idx = np.searchsorted(self.rad.indexes["time"].to_numpy(), time_end).item()
-            sl = slice(max(0, idx - 1), idx + 1)
-            logger.debug("Select rad slice %s", sl)
-            rad = MetDataset(self.rad.data.isel(time=sl), copy=False)
-
-        elif time_end > rad.indexes["time"].to_numpy()[-1]:
-            current_times = rad.indexes["time"].to_numpy()
-            all_times = self.rad.indexes["time"].to_numpy()
-            # idx is the first index at which all_times >= time_end
-            idx = np.searchsorted(all_times, time_end).item()
-            sl = slice(max(0, idx - 1), idx + 1)
-
-            # case 1: cannot re-use end of current rad as start of new rad
-            if current_times[-1] != all_times[sl.start]:
-                logger.debug("Select rad slice %s", sl)
-                rad = MetDataset(self.rad.data.isel(time=sl), copy=False)
-            # case 2: can re-use end of current rad plus one step of new rad
-            elif sl.start < all_times.size - 1:
-                sl = slice(sl.start + 1, sl.stop)
-                logger.debug("Reuse end of rad and select rad slice %s", sl)
-                rad = MetDataset(
-                    xr.concat((rad.data.isel(time=[-1]), self.rad.data.isel(time=sl)), dim="time"),
-                    copy=False,
-                )
-            # case 3: can re-use end of current rad and nothing else
-            else:
-                logger.debug("Reuse end of rad")
-                rad = MetDataset(rad.data.isel(time=[-1]), copy=False)
+        met = _maybe_downselect_mds(self.met, met, t0, t1)
+        rad = _maybe_downselect_mds(self.rad, rad, t0, t1)
 
         return met, rad
 
@@ -806,10 +750,10 @@ class CocipGrid(models.Model):
 
     @staticmethod
     def create_source(
-        level: npt.NDArray[np.float64] | list[float] | float,
+        level: npt.NDArray[np.floating] | list[float] | float,
         time: npt.NDArray[np.datetime64] | list[np.datetime64] | np.datetime64,
-        longitude: npt.NDArray[np.float64] | list[float] | None = None,
-        latitude: npt.NDArray[np.float64] | list[float] | None = None,
+        longitude: npt.NDArray[np.floating] | list[float] | None = None,
+        latitude: npt.NDArray[np.floating] | list[float] | None = None,
         lon_step: float = 1.0,
         lat_step: float = 1.0,
     ) -> MetDataset:
@@ -821,7 +765,7 @@ class CocipGrid(models.Model):
 
         Parameters
         ----------
-        level : level: npt.NDArray[np.float64] | list[float] | float
+        level : level: npt.NDArray[np.floating] | list[float] | float
             Pressure levels for gridded cocip.
             To avoid interpolating outside of the passed ``met`` and ``rad`` data, this
             parameter should avoid the extreme values of the ``met`` and `rad` levels.
@@ -829,7 +773,7 @@ class CocipGrid(models.Model):
             ``met.data['level'].values[1: -1]``.
         time: npt.NDArray[np.datetime64 | list[np.datetime64] | np.datetime64,
             One or more time values for gridded cocip.
-        longitude, latitude : npt.NDArray[np.float64] | list[float], optional
+        longitude, latitude : npt.NDArray[np.floating] | list[float], optional
             Longitude and latitude arrays, by default None. If not specified, values of
             ``lon_step`` and ``lat_step`` are used to define ``longitude`` and ``latitude``.
         lon_step, lat_step : float, optional
@@ -2173,7 +2117,7 @@ def _aggregate_ef_summary(vector_list: list[VectorDataset]) -> VectorDataset | N
 
 def result_to_metdataset(
     result: VectorDataset | None,
-    verbose_dict: dict[str, npt.NDArray[np.float64]],
+    verbose_dict: dict[str, npt.NDArray[np.floating]],
     source: MetDataset,
     nominal_segment_length: float,
     attrs: dict[str, str],
@@ -2185,7 +2129,7 @@ def result_to_metdataset(
     result : VectorDataset | None
         Aggregated data arising from contrail evolution. Expected to contain keys:
         ``index``, ``age``, ``ef``.
-    verbose_dict : dict[str, npt.NDArray[np.float64]]:
+    verbose_dict : dict[str, npt.NDArray[np.floating]]:
         Verbose outputs to attach to results.
     source : MetDataset
         :attr:`CocipGrid.`source` data on which to attach results.
@@ -2244,9 +2188,9 @@ def result_to_metdataset(
 
 def result_merge_source(
     result: VectorDataset | None,
-    verbose_dict: dict[str, npt.NDArray[np.float64]],
+    verbose_dict: dict[str, npt.NDArray[np.floating]],
     source: GeoVectorDataset,
-    nominal_segment_length: float | npt.NDArray[np.float64],
+    nominal_segment_length: float | npt.NDArray[np.floating],
     attrs: dict[str, str],
 ) -> GeoVectorDataset:
     """Merge ``results`` and ``verbose_dict`` onto ``source``."""
@@ -2282,7 +2226,7 @@ def _concat_verbose_dicts(
     verbose_dicts: list[dict[str, pd.Series]],
     source_size: int,
     verbose_outputs_formation: set[str],
-) -> dict[str, npt.NDArray[np.float64]]:
+) -> dict[str, npt.NDArray[np.floating]]:
     # Concatenate the values and return
     ret: dict[str, np.ndarray] = {}
     for key in verbose_outputs_formation:
@@ -2375,7 +2319,7 @@ def _warn_not_wrap(met: MetDataset) -> None:
         )
 
 
-def _get_uncertainty_params(contrail: VectorDataset) -> dict[str, npt.NDArray[np.float64]]:
+def _get_uncertainty_params(contrail: VectorDataset) -> dict[str, npt.NDArray[np.floating]]:
     """Return uncertainty parameters in ``contrail``.
 
     This function assumes the underlying humidity scaling model is
@@ -2398,7 +2342,7 @@ def _get_uncertainty_params(contrail: VectorDataset) -> dict[str, npt.NDArray[np
 
     Returns
     -------
-    dict[str, npt.NDArray[np.float64]]
+    dict[str, npt.NDArray[np.floating]]
         Dictionary of uncertainty parameters.
     """
     keys = (
@@ -2578,3 +2522,65 @@ 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

@@ -10,14 +10,19 @@ import numpy.typing as npt
 
 from pycontrails.core import models
 from pycontrails.core.met import MetDataset
-from pycontrails.core.met_var import AirTemperature, EastwardWind, NorthwardWind, VerticalVelocity
+from pycontrails.core.met_var import (
+    AirTemperature,
+    EastwardWind,
+    NorthwardWind,
+    VerticalVelocity,
+)
 from pycontrails.core.vector import GeoVectorDataset
 from pycontrails.models.cocip import contrail_properties, wind_shear
 from pycontrails.physics import geo, thermo
 
 
 @dataclasses.dataclass
-class DryAdvectionParams(models.ModelParams):
+class DryAdvectionParams(models.AdvectionBuffers):
     """Parameters for the :class:`DryAdvection` model."""
 
     #: Apply Euler's method with a fixed step size of ``dt_integration``. Advected waypoints
@@ -138,12 +143,12 @@ class DryAdvection(models.Model):
         t1 = source_time.max()
         timesteps = np.arange(t0 + dt_integration, t1 + dt_integration + max_age, dt_integration)
 
-        vector = None
+        vector = GeoVectorDataset()
 
         evolved = []
         for t in timesteps:
             filt = (source_time < t) & (source_time >= t - dt_integration)
-            vector = self.source.filter(filt, copy=False) + vector
+            vector = vector + self.source.filter(filt, copy=False)
             vector = _evolve_one_step(
                 self.met,
                 vector,
@@ -237,7 +242,6 @@ def _perform_interp_for_step(
     vector.setdefault("level", vector.level)
     air_pressure = vector.setdefault("air_pressure", vector.air_pressure)
 
-    air_temperature = models.interpolate_met(met, vector, "air_temperature", **interp_kwargs)
     models.interpolate_met(met, vector, "northward_wind", "v_wind", **interp_kwargs)
     models.interpolate_met(met, vector, "eastward_wind", "u_wind", **interp_kwargs)
     models.interpolate_met(
@@ -253,6 +257,7 @@ def _perform_interp_for_step(
         # Early exit for pointwise only simulation
         return
 
+    air_temperature = models.interpolate_met(met, vector, "air_temperature", **interp_kwargs)
     air_pressure_lower = thermo.pressure_dz(air_temperature, air_pressure, dz_m)
     vector["air_pressure_lower"] = air_pressure_lower
     level_lower = air_pressure_lower / 100.0

pycontrails/models/emissions/__init__.py

@@ -11,10 +11,10 @@ from pycontrails.models.emissions.emissions import (
 )
 
 __all__ = [
-    "Emissions",
-    "EmissionsParams",
     "EDBGaseous",
     "EDBnvpm",
+    "Emissions",
+    "EmissionsParams",
     "load_default_aircraft_engine_mapping",
     "load_engine_nvpm_profile_from_edb",
     "load_engine_params_from_edb",