pycontrails 0.54.2__cp312-cp312-macosx_11_0_arm64.whl → 0.54.4__cp312-cp312-macosx_11_0_arm64.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,19 +750,22 @@ 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:
         """
         Shortcut to create a :class:`MetDataset` source from coordinate arrays.
 
+        .. versionchanged:: 0.54.3
+            By default, the returned latitude values now extend to the poles.
+
         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.
@@ -826,11 +773,9 @@ 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``.
-            To avoid model degradation at the poles, latitude values are expected to be
-            between -80 and 80 degrees.
         lon_step, lat_step : float, optional
             Longitude and latitude resolution, by default 1.0.
             Only used if parameter ``longitude`` (respective ``latitude``) not specified.
@@ -847,15 +792,11 @@ class CocipGrid(models.Model):
         if longitude is None:
             longitude = np.arange(-180, 180, lon_step, dtype=float)
         if latitude is None:
-            latitude = np.arange(-80, 80.000001, lat_step, dtype=float)
-
-        out = MetDataset.from_coords(longitude=longitude, latitude=latitude, level=level, time=time)
+            latitude = np.arange(-90, 90.000001, lat_step, dtype=float)
 
-        if np.any(out.data.latitude > 80.0001) or np.any(out.data.latitude < -80.0001):
-            msg = "Model only supports latitude between -80 and 80."
-            raise ValueError(msg)
-
-        return out
+        return MetDataset.from_coords(
+            longitude=longitude, latitude=latitude, level=level, time=time
+        )
 
 
 ################################
@@ -2054,10 +1995,13 @@ def advect(
     time_t2 = time + dt
     age_t2 = age + dt
 
-    longitude_t2 = geo.advect_longitude(
-        longitude=longitude, latitude=latitude, u_wind=u_wind, dt=dt
+    longitude_t2, latitude_t2 = geo.advect_horizontal(
+        longitude=longitude,
+        latitude=latitude,
+        u_wind=u_wind,
+        v_wind=v_wind,
+        dt=dt,
     )
-    latitude_t2 = geo.advect_latitude(latitude=latitude, v_wind=v_wind, dt=dt)
     level_t2 = geo.advect_level(level, vertical_velocity, rho_air, terminal_fall_speed, dt)
     altitude_t2 = units.pl_to_m(level_t2)
 
@@ -2089,15 +2033,20 @@ def advect(
     u_wind_tail = contrail["eastward_wind_tail"]
     v_wind_tail = contrail["northward_wind_tail"]
 
-    longitude_head_t2 = geo.advect_longitude(
-        longitude=longitude_head, latitude=latitude_head, u_wind=u_wind_head, dt=dt_head
+    longitude_head_t2, latitude_head_t2 = geo.advect_horizontal(
+        longitude=longitude_head,
+        latitude=latitude_head,
+        u_wind=u_wind_head,
+        v_wind=v_wind_head,
+        dt=dt_head,
     )
-    latitude_head_t2 = geo.advect_latitude(latitude=latitude_head, v_wind=v_wind_head, dt=dt_head)
-
-    longitude_tail_t2 = geo.advect_longitude(
-        longitude=longitude_tail, latitude=latitude_tail, u_wind=u_wind_tail, dt=dt_tail
+    longitude_tail_t2, latitude_tail_t2 = geo.advect_horizontal(
+        longitude=longitude_tail,
+        latitude=latitude_tail,
+        u_wind=u_wind_tail,
+        v_wind=v_wind_tail,
+        dt=dt_tail,
     )
-    latitude_tail_t2 = geo.advect_latitude(latitude=latitude_tail, v_wind=v_wind_tail, dt=dt_tail)
 
     segment_length_t2 = geo.haversine(
         lons0=longitude_head_t2,
@@ -2168,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],
@@ -2180,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.
@@ -2239,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``."""
@@ -2277,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:
@@ -2370,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
@@ -2393,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 = (
@@ -2573,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

@@ -9,16 +9,20 @@ import numpy as np
 import numpy.typing as npt
 
 from pycontrails.core import models
-from pycontrails.core.flight import Flight
 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
@@ -92,9 +96,6 @@ class DryAdvection(models.Model):
     met_required = True
     source: GeoVectorDataset
 
-    @overload
-    def eval(self, source: Flight, **params: Any) -> Flight: ...
-
     @overload
     def eval(self, source: GeoVectorDataset, **params: Any) -> GeoVectorDataset: ...
 
@@ -109,7 +110,12 @@ class DryAdvection(models.Model):
         Parameters
         ----------
         source : GeoVectorDataset
-            Arbitrary points to advect.
+            Arbitrary points to advect. A :class:`Flight` instance is not treated any
+            differently than a :class:`GeoVectorDataset`. In particular, the user must
+            explicitly set ``flight["azimuth"] = flight.segment_azimuth()`` if they
+            want to use wind shear effects for a flight.
+            In the current implementation, any existing meteorological variables in the ``source``
+            are ignored. The ``source`` will be interpolated against the :attr:`met` dataset.
         params : Any
             Overwrite model parameters defined in ``__init__``.
 
@@ -122,7 +128,7 @@ class DryAdvection(models.Model):
         self.set_source(source)
         self.source = self.require_source_type(GeoVectorDataset)
 
-        self._prepare_source()
+        self.source = self._prepare_source()
 
         interp_kwargs = self.interp_kwargs
 
@@ -137,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) + vector
+            vector = vector + self.source.filter(filt, copy=False)
             vector = _evolve_one_step(
                 self.met,
                 vector,
@@ -162,49 +168,44 @@ class DryAdvection(models.Model):
 
         return GeoVectorDataset.sum(evolved, fill_value=np.nan)
 
-    def _prepare_source(self) -> None:
+    def _prepare_source(self) -> GeoVectorDataset:
         r"""Prepare :attr:`source` vector for advection by wind-shear-derived variables.
 
-        This method adds the following variables to :attr:`source` if the `"azimuth"`
-        parameter is not None:
+        The following variables are always guaranteed to be present in :attr:`source`:
 
         - ``age``: Age of plume.
+        - ``waypoint``: Identifier for each waypoint.
+
+        If `"azimuth"` is present in :attr:`source`, `source.attrs`, or :attr:`params`,
+        the following variables will also be added:
+
         - ``azimuth``: Initial plume direction, measured in clockwise direction from
-          true north, [:math:`\deg`].
+            true north, [:math:`\deg`].
         - ``width``: Initial plume width, [:math:`m`].
         - ``depth``: Initial plume depth, [:math:`m`].
         - ``sigma_yz``: All zeros for cross-term term in covariance matrix of plume.
-        """
 
+        Returns
+        -------
+        GeoVectorDataset
+            A filtered version of the source with only the required columns.
+        """
         self.source.setdefault("level", self.source.level)
-
-        columns: tuple[str, ...] = ("longitude", "latitude", "level", "time")
-        if "azimuth" in self.source:
-            columns += ("azimuth",)
-        self.source = GeoVectorDataset(self.source.select(columns, copy=False))
-
-        # Get waypoint index if not already set
+        self.source["age"] = np.full(self.source.size, np.timedelta64(0, "ns"))
         self.source.setdefault("waypoint", np.arange(self.source.size))
 
-        self.source["age"] = np.full(self.source.size, np.timedelta64(0, "ns"))
+        columns = ["longitude", "latitude", "level", "time", "age", "waypoint"]
+        azimuth = self.get_source_param("azimuth", set_attr=False)
+        if azimuth is None:
+            # Early exit for pointwise only simulation
+            if self.params["width"] is not None or self.params["depth"] is not None:
+                raise ValueError(
+                    "If 'azimuth' is None, then 'width' and 'depth' must also be None."
+                )
+            return GeoVectorDataset(self.source.select(columns, copy=False), copy=False)
 
         if "azimuth" not in self.source:
-            if isinstance(self.source, Flight):
-                pointwise_only = False
-                self.source["azimuth"] = self.source.segment_azimuth()
-            else:
-                try:
-                    self.source.broadcast_attrs("azimuth")
-                except KeyError:
-                    if (azimuth := self.params["azimuth"]) is not None:
-                        pointwise_only = False
-                        self.source["azimuth"] = np.full_like(self.source["longitude"], azimuth)
-                    else:
-                        pointwise_only = True
-                else:
-                    pointwise_only = False
-        else:
-            pointwise_only = False
+            self.source["azimuth"] = np.full_like(self.source["longitude"], azimuth)
 
         for key in ("width", "depth"):
             if key in self.source:
@@ -214,18 +215,12 @@ class DryAdvection(models.Model):
                 continue
 
             val = self.params[key]
-            if val is None and not pointwise_only:
+            if val is None:
                 raise ValueError(f"If '{key}' is None, then 'azimuth' must also be None.")
 
-            if val is not None and pointwise_only:
-                raise ValueError(f"Cannot specify '{key}' without specifying 'azimuth'.")
-
-            if not pointwise_only:
-                self.source[key] = np.full_like(self.source["longitude"], val)
-
-        if pointwise_only:
-            return
+            self.source[key] = np.full_like(self.source["longitude"], val)
 
+        columns.extend(["azimuth", "width", "depth", "sigma_yz", "area_eff"])
         self.source["sigma_yz"] = np.zeros_like(self.source["longitude"])
         width = self.source["width"]
         depth = self.source["depth"]
@@ -233,6 +228,8 @@ class DryAdvection(models.Model):
             width, depth, sigma_yz=0.0
         )
 
+        return GeoVectorDataset(self.source.select(columns, copy=False), copy=False)
+
 
 def _perform_interp_for_step(
     met: MetDataset,
@@ -245,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(
@@ -261,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
@@ -412,15 +409,20 @@ def _calc_geometry(
     u_wind_tail = vector.data.pop("eastward_wind_tail")
     v_wind_tail = vector.data.pop("northward_wind_tail")
 
-    longitude_head_t2 = geo.advect_longitude(
-        longitude=longitude_head, latitude=latitude_head, u_wind=u_wind_head, dt=dt
+    longitude_head_t2, latitude_head_t2 = geo.advect_horizontal(
+        longitude=longitude_head,
+        latitude=latitude_head,
+        u_wind=u_wind_head,
+        v_wind=v_wind_head,
+        dt=dt,
     )
-    latitude_head_t2 = geo.advect_latitude(latitude=latitude_head, v_wind=v_wind_head, dt=dt)
-
-    longitude_tail_t2 = geo.advect_longitude(
-        longitude=longitude_tail, latitude=latitude_tail, u_wind=u_wind_tail, dt=dt
+    longitude_tail_t2, latitude_tail_t2 = geo.advect_horizontal(
+        longitude=longitude_tail,
+        latitude=latitude_tail,
+        u_wind=u_wind_tail,
+        v_wind=v_wind_tail,
+        dt=dt,
     )
-    latitude_tail_t2 = geo.advect_latitude(latitude=latitude_tail, v_wind=v_wind_tail, dt=dt)
 
     azimuth_2 = geo.azimuth(
         lons0=longitude_tail_t2,
@@ -453,8 +455,7 @@ def _evolve_one_step(
     longitude = vector["longitude"]
 
     dt = t - vector["time"]
-    longitude_2 = geo.advect_longitude(longitude, latitude, u_wind, dt)  # type: ignore[arg-type]
-    latitude_2 = geo.advect_latitude(latitude, v_wind, dt)  # type: ignore[arg-type]
+    longitude_2, latitude_2 = geo.advect_horizontal(longitude, latitude, u_wind, v_wind, dt)  # type: ignore[arg-type]
     level_2 = geo.advect_level(
         vector.level,
         vertical_velocity,

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",