pycontrails 0.52.2__cp310-cp310-win_amd64.whl → 0.53.0__cp310-cp310-win_amd64.whl

pycontrails/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.52.2'
-__version_tuple__ = version_tuple = (0, 52, 2)
+__version__ = version = '0.53.0'
+__version_tuple__ = version_tuple = (0, 53, 0)

pycontrails/core/cache.py

@@ -146,7 +146,7 @@ class CacheStore(ABC):
         """
 
         # TODO: run in parallel?
-        return [self.put(d, cp) for d, cp in zip(data_path, cache_path)]
+        return [self.put(d, cp) for d, cp in zip(data_path, cache_path, strict=True)]
 
     # In the three methods below, child classes have a complete docstring.
 

pycontrails/core/flight.py

@@ -1356,7 +1356,7 @@ class Flight(GeoVectorDataset):
             # NOTE: geod.npts does not return the initial or terminal points
             lonlats: list[tuple[float, float]] = geod.npts(lon0, lat0, lon1, lat1, n_steps)
 
-            lons, lats = zip(*lonlats)
+            lons, lats = zip(*lonlats, strict=True)
             longitudes.extend(lons)
             latitudes.extend(lats)
 
@@ -1657,10 +1657,11 @@ def _return_linestring(data: dict[str, npt.NDArray[np.float64]]) -> list[list[fl
         The list of coordinates
     """
     # rounding to reduce the size of resultant json arrays
-    points = zip(  # pylint: disable=zip-builtin-not-iterating
+    points = zip(
         np.round(data["longitude"], decimals=4),
         np.round(data["latitude"], decimals=4),
         np.round(data["altitude"], decimals=4),
+        strict=True,
     )
     return [list(p) for p in points]
 
@@ -1949,7 +1950,9 @@ def _altitude_interpolation_climb_descend_middle(
     # Form array of cumulative altitude values if the flight were to climb
     # at nominal_rocd over each group of nan
     cumalt_list = []
-    for start_na_idx, end_na_idx, size in zip(start_na_idxs, end_na_idxs, na_group_size):
+    for start_na_idx, end_na_idx, size in zip(
+        start_na_idxs, end_na_idxs, na_group_size, strict=True
+    ):
         if s[start_na_idx] <= s[end_na_idx]:
             cumalt_list.append(np.arange(1, size, dtype=float))
         else:
@@ -2053,7 +2056,7 @@ def filter_altitude(
     --------
     :meth:`traffic.core.flight.Flight.filter`
     :func:`scipy.signal.medfilt`
-    """  # noqa: E501
+    """
     if not len(altitude_ft):
         raise ValueError("Altitude must have non-zero length to filter")
 
@@ -2114,7 +2117,7 @@ def filter_altitude(
 
     result = np.copy(altitude_ft)
     if np.any(start_idxs):
-        for i0, i1 in zip(start_idxs, end_idxs):
+        for i0, i1 in zip(start_idxs, end_idxs, strict=True):
             result[i0:i1] = altitude_filt[i0:i1]
 
     # reapply Savitzky-Golay filter to smooth climb and descent

pycontrails/core/flightplan.py

@@ -93,7 +93,7 @@ def parse_atc_plan(atc_plan: str) -> dict[str, str]:
     See Also
     --------
     :func:`to_atc_plan`
-    """  # noqa: E501
+    """
     atc_plan = atc_plan.replace("\r", "")
     atc_plan = atc_plan.replace("\n", "")
     atc_plan = atc_plan.upper()

pycontrails/core/interpolation.py

@@ -76,6 +76,7 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         self.method = _pick_method(scipy.__version__, method)
         self.bounds_error = bounds_error
         self.fill_value = fill_value
+        self._spline = None
 
     def _prepare_xi_simple(self, xi: npt.NDArray[np.float64]) -> npt.NDArray[np.bool_]:
         """Run looser version of :meth:`_prepare_xi`.
@@ -215,7 +216,8 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
 
         if ndim == 1:
             # np.interp could be better ... although that may also promote the dtype
-            return rgi_cython.evaluate_linear_1d(values, indices, norm_distances, out)
+            # 1-d view is required for evaluate_linear_1d
+            return rgi_cython.evaluate_linear_1d(values, indices[0, :], norm_distances[0, :], out)
 
         msg = f"Invalid number of dimensions: {ndim}"
         raise ValueError(msg)

pycontrails/core/met.py

@@ -9,7 +9,15 @@ import pathlib
 import typing
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Hashable, Iterable, Iterator, Mapping, MutableMapping, Sequence
+from collections.abc import (
+    Generator,
+    Hashable,
+    Iterable,
+    Iterator,
+    Mapping,
+    MutableMapping,
+    Sequence,
+)
 from contextlib import ExitStack
 from datetime import datetime
 from typing import (
@@ -62,12 +70,12 @@ class MetBase(ABC, Generic[XArrayType]):
     cachestore: CacheStore | None
 
     #: Default dimension order for DataArray or Dataset (x, y, z, t)
-    dim_order: list[Hashable] = [
+    dim_order: tuple[Hashable, Hashable, Hashable, Hashable] = (
         "longitude",
         "latitude",
         "level",
         "time",
-    ]
+    )
 
     def __repr__(self) -> str:
         data = getattr(self, "data", None)
@@ -192,10 +200,8 @@ class MetBase(ABC, Generic[XArrayType]):
     def _validate_transpose(self) -> None:
         """Check that data is transposed according to :attr:`dim_order`."""
 
-        dims_tuple = tuple(self.dim_order)
-
         def _check_da(da: xr.DataArray, key: Hashable | None = None) -> None:
-            if da.dims != dims_tuple:
+            if da.dims != self.dim_order:
                 if key is not None:
                     msg = (
                         f"Data dimension not transposed on variable '{key}'. Initiate with"
@@ -263,7 +269,7 @@ class MetBase(ABC, Generic[XArrayType]):
         self.data["time"] = self.data["time"].astype("datetime64[ns]", copy=False)
 
         # sortby to ensure each coordinate has ascending order
-        self.data = self.data.sortby(self.dim_order, ascending=True)
+        self.data = self.data.sortby(list(self.dim_order), ascending=True)
 
         if not self.is_wrapped:
             # Ensure longitude is contained in interval [-180, 180)
@@ -285,7 +291,7 @@ class MetBase(ABC, Generic[XArrayType]):
         self._validate_latitude()
 
         # transpose to have ordering (x, y, z, t, ...)
-        dim_order = self.dim_order + [d for d in self.data.dims if d not in self.dim_order]
+        dim_order = [*self.dim_order, *(d for d in self.data.dims if d not in self.dim_order)]
         self.data = self.data.transpose(*dim_order)
 
         # single level data
@@ -481,7 +487,7 @@ class MetBase(ABC, Generic[XArrayType]):
         self.cachestore = self.cachestore or DiskCacheStore()
 
         # group by hour and save one dataset for each hour to temp file
-        times, datasets = zip(*dataset.groupby("time", squeeze=False))
+        times, datasets = zip(*dataset.groupby("time", squeeze=False), strict=True)
 
         # Open ExitStack to control temp_file context manager
         with ExitStack() as stack:
@@ -912,7 +918,7 @@ class MetDataset(MetBase):
         KeyError
             Raises when dataset does not contain variable in ``vars``
         """
-        if isinstance(vars, (MetVariable, str)):
+        if isinstance(vars, MetVariable | str):
             vars = (vars,)
 
         met_keys: list[str] = []
@@ -1014,7 +1020,7 @@ class MetDataset(MetBase):
 
     @overrides
     def broadcast_coords(self, name: str) -> xr.DataArray:
-        da = xr.ones_like(self.data[list(self.data.keys())[0]]) * self.data[name]
+        da = xr.ones_like(self.data[next(iter(self.data.keys()))]) * self.data[name]
         da.name = name
 
         return da
@@ -1066,7 +1072,7 @@ class MetDataset(MetBase):
         coords_vals = [indexes[key].values for key in coords_keys]
         coords_meshes = np.meshgrid(*coords_vals, indexing="ij")
         raveled_coords = (mesh.ravel() for mesh in coords_meshes)
-        data = dict(zip(coords_keys, raveled_coords))
+        data = dict(zip(coords_keys, raveled_coords, strict=True))
 
         out = vector_module.GeoVectorDataset(data, copy=False)
         for key, da in self.data.items():
@@ -1502,6 +1508,7 @@ class MetDataArray(MetBase):
         bounds_error: bool = ...,
         fill_value: float | np.float64 | None = ...,
         localize: bool = ...,
+        lowmem: bool = ...,
         indices: interpolation.RGIArtifacts | None = ...,
         return_indices: Literal[False] = ...,
     ) -> npt.NDArray[np.float64]: ...
@@ -1518,6 +1525,7 @@ class MetDataArray(MetBase):
         bounds_error: bool = ...,
         fill_value: float | np.float64 | None = ...,
         localize: bool = ...,
+        lowmem: bool = ...,
         indices: interpolation.RGIArtifacts | None = ...,
         return_indices: Literal[True],
     ) -> tuple[npt.NDArray[np.float64], interpolation.RGIArtifacts]: ...
@@ -1533,6 +1541,7 @@ class MetDataArray(MetBase):
         bounds_error: bool = False,
         fill_value: float | np.float64 | None = np.nan,
         localize: bool = False,
+        lowmem: bool = False,
         indices: interpolation.RGIArtifacts | None = None,
         return_indices: bool = False,
     ) -> npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], interpolation.RGIArtifacts]:
@@ -1540,7 +1549,9 @@ class MetDataArray(MetBase):
 
         Zero dimensional coordinates are reshaped to 1D arrays.
 
-        Method automatically loads underlying :attr:`data` into memory.
+        If ``lowmem == False``, method automatically loads underlying :attr:`data` into
+        memory. Otherwise, method iterates through smaller subsets of :attr:`data` and releases
+        subsets from memory once interpolation against each subset is finished.
 
         If ``method == "nearest"``, the out array will have the same ``dtype`` as
         the underlying :attr:`data`.
@@ -1586,10 +1597,18 @@ class MetDataArray(MetBase):
         localize: bool, optional
             Experimental. If True, downselect gridded data to smallest bounding box containing
             all points.  By default False.
+        lowmem: bool, optional
+            Experimental. If True, iterate through points binned by the time coordinate of the
+            grided data, and downselect gridded data to the smallest bounding box containing
+            each binned set of point *before loading into memory*. This can significantly reduce
+            memory consumption with large numbers of points at the cost of increased runtime.
+            By default False.
         indices: tuple | None, optional
             Experimental. See :func:`interpolation.interp`. None by default.
         return_indices: bool, optional
             Experimental. See :func:`interpolation.interp`. False by default.
+            Note that values returned differ when ``lowmem=True`` and ``lowmem=False``,
+            so output should only be re-used in calls with the same ``lowmem`` value.
 
         Returns
         -------
@@ -1632,10 +1651,29 @@ class MetDataArray(MetBase):
         >>> level = np.linspace(200, 300, 10)
         >>> time = pd.date_range("2022-03-01T14", periods=10, freq="5min")
         >>> mda.interpolate(longitude, latitude, level, time)
+        array([220.44347694, 223.08900738, 225.74338924, 228.41642088,
+               231.10858599, 233.54857391, 235.71504913, 237.86478872,
+               239.99274623, 242.10792167])
+
+        >>> # Can easily switch to alternative low-memory implementation
+        >>> mda.interpolate(longitude, latitude, level, time, lowmem=True)
         array([220.44347694, 223.08900738, 225.74338924, 228.41642088,
                231.10858599, 233.54857391, 235.71504913, 237.86478872,
                239.99274623, 242.10792167])
         """
+        if lowmem:
+            return self._interp_lowmem(
+                longitude,
+                latitude,
+                level,
+                time,
+                method=method,
+                bounds_error=bounds_error,
+                fill_value=fill_value,
+                indices=indices,
+                return_indices=return_indices,
+            )
+
         # Load if necessary
         if not self.in_memory:
             self._check_memory("Interpolation over")
@@ -1660,6 +1698,100 @@ class MetDataArray(MetBase):
             return_indices=return_indices,
         )
 
+    def _interp_lowmem(
+        self,
+        longitude: float | npt.NDArray[np.float64],
+        latitude: float | npt.NDArray[np.float64],
+        level: float | npt.NDArray[np.float64],
+        time: np.datetime64 | npt.NDArray[np.datetime64],
+        *,
+        method: str = "linear",
+        bounds_error: bool = False,
+        fill_value: float | np.float64 | None = np.nan,
+        minimize_memory: bool = False,
+        indices: interpolation.RGIArtifacts | None = None,
+        return_indices: bool = False,
+    ) -> npt.NDArray[np.float64] | tuple[npt.NDArray[np.float64], interpolation.RGIArtifacts]:
+        """Interpolate values against underlying DataArray.
+
+        This method is used by :meth:`interpolate` when ``lowmem=True``.
+        Parameters and return types are identical to :meth:`interpolate`, except
+        that the ``localize`` keyword argument is omitted.
+        """
+        # Convert all inputs to 1d arrays
+        # Not validating against ndim >= 2
+        longitude, latitude, level, time = np.atleast_1d(longitude, latitude, level, time)
+
+        if bounds_error:
+            _lowmem_boundscheck(time, self.data)
+
+        # Create buffers for holding interpolation output
+        # Use np.full rather than np.empty so points not covered
+        # by masks are filled with correct out-of-bounds values.
+        out = np.full(longitude.shape, fill_value, dtype=self.data.dtype)
+        if return_indices:
+            rgi_artifacts = interpolation.RGIArtifacts(
+                xi_indices=np.full((4, longitude.size), -1, dtype=np.int64),
+                norm_distances=np.full((4, longitude.size), np.nan, dtype=np.float64),
+                out_of_bounds=np.full((longitude.size,), True, dtype=np.bool_),
+            )
+
+        # Iterate over portions of points between adjacent time steps in gridded data
+        for mask in _lowmem_masks(time, self.data["time"].values):
+            if mask is None or not np.any(mask):
+                continue
+
+            lon_sl = longitude[mask]
+            lat_sl = latitude[mask]
+            lev_sl = level[mask]
+            t_sl = time[mask]
+            if indices is not None:
+                indices_sl = interpolation.RGIArtifacts(
+                    xi_indices=indices.xi_indices[:, mask],
+                    norm_distances=indices.norm_distances[:, mask],
+                    out_of_bounds=indices.out_of_bounds[mask],
+                )
+            else:
+                indices_sl = None
+
+            coords = {"longitude": lon_sl, "latitude": lat_sl, "level": lev_sl, "time": t_sl}
+            if any(np.all(np.isnan(coord)) for coord in coords.values()):
+                continue
+            da = interpolation._localize(self.data, coords)
+            if not da._in_memory:
+                logger.debug(
+                    "Loading %s MB subset of %s into memory.",
+                    round(da.nbytes / 1_000_000, 2),
+                    da.name,
+                )
+                da.load()
+
+            tmp = interpolation.interp(
+                longitude=lon_sl,
+                latitude=lat_sl,
+                level=lev_sl,
+                time=t_sl,
+                da=da,
+                method=method,
+                bounds_error=bounds_error,
+                fill_value=fill_value,
+                localize=False,  # would be no-op; da is localized already
+                indices=indices_sl,
+                return_indices=return_indices,
+            )
+
+            if return_indices:
+                out[mask], rgi_sl = tmp
+                rgi_artifacts.xi_indices[:, mask] = rgi_sl.xi_indices
+                rgi_artifacts.norm_distances[:, mask] = rgi_sl.norm_distances
+                rgi_artifacts.out_of_bounds[mask] = rgi_sl.out_of_bounds
+            else:
+                out[mask] = tmp
+
+        if return_indices:
+            return out, rgi_artifacts
+        return out
+
     def _check_memory(self, msg_start: str) -> None:
         """Check the memory usage of the underlying data.
 
@@ -1731,7 +1863,7 @@ class MetDataArray(MetBase):
         cachestore = cachestore or DiskCacheStore()
         chunks = chunks or {}
         data = _load(hash, cachestore, chunks)
-        return cls(data[list(data.data_vars)[0]])
+        return cls(data[next(iter(data.data_vars))])
 
     @property
     def proportion(self) -> float:
@@ -2124,7 +2256,7 @@ class MetDataArray(MetBase):
         -----
         Uses the `scikit-image Marching Cubes  <https://scikit-image.org/docs/dev/auto_examples/edges/plot_marching_cubes.html>`_
         algorithm to reconstruct a surface from the point-cloud like arrays.
-        """  # noqa: E501
+        """
         try:
             from skimage import measure
         except ModuleNotFoundError as e:
@@ -2656,3 +2788,46 @@ def _add_vertical_coords(data: XArrayType) -> XArrayType:
         data.coords["altitude"] = data.coords["altitude"].astype(dtype, copy=False)
 
     return data
+
+
+def _lowmem_boundscheck(time: npt.NDArray[np.datetime64], da: xr.DataArray) -> None:
+    """Extra bounds check required with low-memory interpolation strategy.
+
+    Because the main loop in `_interp_lowmem` processes points between time steps
+    in gridded data, it will never encounter points that are out-of-bounds in time
+    and may fail to produce requested out-of-bounds errors.
+    """
+    da_time = da["time"].to_numpy()
+    if not np.all((time >= da_time.min()) & (time <= da_time.max())):
+        axis = da.get_axis_num("time")
+        msg = f"One of the requested xi is out of bounds in dimension {axis}"
+        raise ValueError(msg)
+
+
+def _lowmem_masks(
+    time: npt.NDArray[np.datetime64], t_met: npt.NDArray[np.datetime64]
+) -> Generator[npt.NDArray[np.bool_], None, None]:
+    """Generate sequence of masks for low-memory interpolation."""
+    t_met_max = t_met.max()
+    t_met_min = t_met.min()
+    inbounds = (time >= t_met_min) & (time <= t_met_max)
+    if not np.any(inbounds):
+        return
+
+    earliest = np.nanmin(time)
+    istart = 0 if earliest < t_met_min else np.flatnonzero(t_met <= earliest).max()
+    latest = np.nanmax(time)
+    iend = t_met.size - 1 if latest > t_met_max else np.flatnonzero(t_met >= latest).min()
+    if istart == iend:
+        yield inbounds
+        return
+
+    # Sequence of masks covers elements in time in the interval [t_met[istart], t_met[iend]].
+    # The first iteration masks elements in the interval [t_met[istart], t_met[istart+1]]
+    # (inclusive of both endpoints).
+    # Subsequent iterations mask elements in the interval (t_met[i], t_met[i+1]]
+    # (inclusive of right endpoint only).
+    for i in range(istart, iend):
+        mask = ((time >= t_met[i]) if i == istart else (time > t_met[i])) & (time <= t_met[i + 1])
+        if np.any(mask):
+            yield mask

pycontrails/core/met_var.py

@@ -22,7 +22,7 @@ class MetVariable:
     - `NCEP Grib v2 Code Table <https://www.nco.ncep.noaa.gov/pmb/docs/grib2/grib2_doc/grib2_table4-2.shtml>`_
 
     Used for defining support parameters in a grib-like fashion.
-    """  # noqa: E501
+    """
 
     #: Short variable name.
     #: Chosen for greatest consistency between data sources.

pycontrails/core/models.py

@@ -11,7 +11,7 @@ import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from dataclasses import dataclass, fields
-from typing import Any, NoReturn, TypeVar, Union, overload
+from typing import Any, NoReturn, TypeVar, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -30,13 +30,13 @@ from pycontrails.utils.types import type_guard
 logger = logging.getLogger(__name__)
 
 #: Model input source types
-ModelInput = Union[MetDataset, GeoVectorDataset, Flight, Sequence[Flight], None]
+ModelInput = MetDataset | GeoVectorDataset | Flight | Sequence[Flight] | None
 
 #: Model output source types
-ModelOutput = Union[MetDataArray, MetDataset, GeoVectorDataset, Flight, list[Flight], NoReturn]
+ModelOutput = MetDataArray | MetDataset | GeoVectorDataset | Flight | list[Flight]
 
 #: Model attribute source types
-SourceType = Union[MetDataset, GeoVectorDataset, Flight, Fleet]
+SourceType = MetDataset | GeoVectorDataset | Flight | Fleet
 
 _Source = TypeVar("_Source")
 
@@ -453,7 +453,7 @@ class Model(ABC):
             return Fleet.from_seq(source)
 
         # Raise error if source is not a MetDataset or GeoVectorDataset
-        if not isinstance(source, (MetDataset, GeoVectorDataset)):
+        if not isinstance(source, MetDataset | GeoVectorDataset):
             msg = f"Unknown source type: {type(source)}"
             raise TypeError(msg)
 

pycontrails/core/vector.py

@@ -847,7 +847,7 @@ class VectorDataset:
         ------
         TypeError
             If ``mask`` is not a boolean array.
-        """  # noqa: E501
+        """
         self.data._validate_array(mask)
         if mask.dtype != bool:
             raise TypeError("Parameter `mask` must be a boolean array.")
@@ -983,7 +983,7 @@ class VectorDataset:
         numeric_attrs = (
             attr
             for attr, val in self.attrs.items()
-            if (isinstance(val, (int, float, np.number)) and attr not in ignore_keys)
+            if (isinstance(val, int | float | np.number) and attr not in ignore_keys)
         )
         self.broadcast_attrs(numeric_attrs, overwrite)
 
@@ -1072,7 +1072,7 @@ class VectorDataset:
                 obj = obj.to_numpy()
 
             # Convert numpy objects to python objects
-            if isinstance(obj, (np.ndarray, np.generic)):
+            if isinstance(obj, np.ndarray | np.generic):
 
                 # round time to unix seconds
                 if key == "time":
@@ -1166,7 +1166,7 @@ class VectorDataset:
         attrs = {}
 
         for k, v in {**obj, **obj_kwargs}.items():
-            if isinstance(v, (list, np.ndarray)):
+            if isinstance(v, list | np.ndarray):
                 data[k] = v
             else:
                 attrs[k] = v
@@ -1194,7 +1194,7 @@ class VectorDataset:
         See Also
         --------
         :func:`numpy.array_split`
-        """  # noqa: E501
+        """
         full_index = np.arange(self.size)
         index_splits = np.array_split(full_index, n_splits)
         for index in index_splits:

pycontrails/datalib/_leo_utils/vis.py

@@ -6,16 +6,6 @@ import numpy as np
 
 from pycontrails.utils import dependencies
 
-try:
-    import skimage as ski
-except ModuleNotFoundError as exc:
-    dependencies.raise_module_not_found_error(
-        name="landsat module",
-        package_name="scikit-image",
-        module_not_found_error=exc,
-        pycontrails_optional_package="sat",
-    )
-
 
 def normalize(channel: np.ndarray) -> np.ndarray:
     """Normalize channel values to range [0, 1], preserving ``np.nan`` in output.
@@ -53,8 +43,17 @@ def equalize(channel: np.ndarray, **equalize_kwargs: Any) -> np.ndarray:
     NaN values are converted to 0 before passing to :py:func:`ski.exposure.equalize_adapthist`
     and may affect equalized values in the neighborhood where they occur.
     """
+    try:
+        import skimage.exposure
+    except ModuleNotFoundError as exc:
+        dependencies.raise_module_not_found_error(
+            name="landsat module",
+            package_name="scikit-image",
+            module_not_found_error=exc,
+            pycontrails_optional_package="sat",
+        )
     return np.where(
         np.isnan(channel),
         np.nan,
-        ski.exposure.equalize_adapthist(np.nan_to_num(channel, nan=0), **equalize_kwargs),
+        skimage.exposure.equalize_adapthist(np.nan_to_num(channel, nan=0.0), **equalize_kwargs),
     )

pycontrails/datalib/_met_utils/metsource.py

@@ -8,7 +8,7 @@ import logging
 import pathlib
 from collections.abc import Sequence
 from datetime import datetime
-from typing import Any, Union
+from typing import Any, TypeAlias
 
 import numpy as np
 import pandas as pd
@@ -20,11 +20,13 @@ from pycontrails.utils.types import DatetimeLike
 
 logger = logging.getLogger(__name__)
 
-TimeInput = Union[str, DatetimeLike, Sequence[Union[str, DatetimeLike]]]
-VariableInput = Union[
-    str, int, MetVariable, np.ndarray, Sequence[Union[str, int, MetVariable, Sequence[MetVariable]]]
-]
-PressureLevelInput = Union[int, float, np.ndarray, Sequence[Union[int, float]]]
+# https://github.com/python/mypy/issues/14824
+TimeInput: TypeAlias = str | DatetimeLike | Sequence[str | DatetimeLike]
+VariableInput = (
+    str | int | MetVariable | np.ndarray | Sequence[str | int | MetVariable | Sequence[MetVariable]]
+)
+
+PressureLevelInput = int | float | np.ndarray | Sequence[int | float]
 
 #: NetCDF engine to use for parsing netcdf files
 NETCDF_ENGINE: str = "netcdf4"
@@ -66,13 +68,13 @@ def parse_timesteps(time: TimeInput | None, freq: str | None = "1h") -> list[dat
     ------
     ValueError
         Raises when the time has len > 2 or when time elements fail to be parsed with pd.to_datetime
-    """  # noqa: E501
+    """
 
     if time is None:
         return []
 
     # confirm input is tuple or list-like of length 2
-    if isinstance(time, (str, datetime, pd.Timestamp, np.datetime64)):
+    if isinstance(time, str | datetime | pd.Timestamp | np.datetime64):
         time = (time, time)
     elif len(time) == 1:
         time = (time[0], time[0])
@@ -151,7 +153,7 @@ def parse_pressure_levels(
         Raises ValueError if pressure level is not supported by ECMWF data source
     """
     # Ensure pressure_levels is array-like
-    if isinstance(pressure_levels, (int, float)):
+    if isinstance(pressure_levels, int | float):
         pressure_levels = [pressure_levels]
 
     # Cast array-like to int dtype and sort
@@ -212,7 +214,7 @@ def parse_variables(variables: VariableInput, supported: list[MetVariable]) -> l
     met_var_list: list[MetVariable] = []
 
     # ensure input variables are a list of str
-    if isinstance(variables, (str, int, MetVariable)):
+    if isinstance(variables, str | int | MetVariable):
         parsed_variables = [variables]
     elif isinstance(variables, np.ndarray):
         parsed_variables = variables.tolist()
@@ -257,7 +259,7 @@ def _find_match(
 
     # list of MetVariable options
     # here we extract the first MetVariable in var that is supported
-    elif isinstance(var, (list, tuple)):
+    elif isinstance(var, list | tuple):
         for v in var:
             # sanity check since we don't support other types as lists
             if not isinstance(v, MetVariable):

pycontrails/datalib/ecmwf/era5.py

@@ -121,7 +121,7 @@ class ERA5(ECMWFAPI):
     ...     pressure_levels=[350, 300],
     ...     cachestore=gcp_cache
     ... )
-    """  # noqa: E501
+    """
 
     __slots__ = (
         "product_type",

pycontrails/datalib/ecmwf/era5_model_level.py

@@ -349,7 +349,7 @@ class ERA5ModelLevel(ECMWFAPI):
         unique_dates = set(t.strftime("%Y-%m-%d") for t in times)
         unique_times = set(t.strftime("%H:%M:%S") for t in times)
         # param 152 = log surface pressure, needed for metview level conversion
-        grib_params = set(self.variable_ecmwfids + [152])
+        grib_params = set((*self.variable_ecmwfids, 152))
         common = {
             "class": "ea",
             "date": "/".join(sorted(unique_dates)),