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

pycontrails/__init__.py

@@ -21,13 +21,13 @@ from __future__ import annotations
 import logging
 from importlib import metadata
 
-import dask
+import dask.config
 
 # Work around for https://github.com/pydata/xarray/issues/7259
 # Only occurs for xarray 2022.11 and above
 try:
     import netCDF4  # noqa: F401
-except ModuleNotFoundError:
+except ImportError:
     pass
 
 from pycontrails.core.cache import DiskCacheStore, GCPCacheStore

pycontrails/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.47.3'
-__version_tuple__ = version_tuple = (0, 47, 3)
+__version__ = version = '0.48.1'
+__version_tuple__ = version_tuple = (0, 48, 1)

pycontrails/core/coordinates.py

@@ -5,12 +5,13 @@ from __future__ import annotations
 import warnings
 
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 
 
 def slice_domain(
     domain: np.ndarray,
-    request: np.ndarray | tuple,
+    request: npt.ArrayLike,
     buffer: tuple[float | np.timedelta64, float | np.timedelta64] = (0.0, 0.0),
 ) -> slice:
     """Return slice of ``domain`` containing coordinates overlapping ``request``.
@@ -34,9 +35,8 @@ def slice_domain(
     ----------
     domain : np.ndarray
         Full set of domain values
-    request : np.ndarray | tuple
-        Requested values. Only the min and max values are considered. Either pass in
-        a full array-like object or a tuple of ``(min, max)``.
+    request : npt.ArrayLike
+        Requested values. Only the nanmin and nanmax values are considered.
     buffer : tuple[float | np.timedelta64, float | np.timedelta64], optional
         Extend the domain past the requested coordinates by ``buffer[0]`` on the low side
         and ``buffer[1]`` on the high side.
@@ -89,23 +89,30 @@ def slice_domain(
     if buffer == (None, None):
         return slice(None, None)
 
-    # if the request is nan, then there is nothing to slice
-    if np.isnan(request).all():
+    # Remove nans from request
+    request = np.asarray(request)
+    mask = np.isnan(request)
+    if mask.all():
         return slice(None, None)
 
+    request = request[~mask]
+
     # if the whole domain or request is nan, then there is nothing to slice
     if np.isnan(domain).all():
         raise ValueError("Domain is all nan on request")
 
     # ensure domain is sorted
-    zero: float | np.timedelta64 = 0.0
+    zero: float | np.timedelta64
     if pd.api.types.is_datetime64_dtype(domain.dtype):
         zero = np.timedelta64(0)
+    else:
+        zero = 0.0
 
     if not np.all(np.diff(domain) >= zero):
         raise ValueError("Domain must be sorted in ascending order")
 
-    if np.any(np.asarray(buffer) < zero):
+    buf0, buf1 = buffer
+    if buf0 < zero or buf1 < zero:
         warnings.warn(
             "Found buffer with negative value. This is unexpected "
             "and will reduce the size of the requested domain instead of "
@@ -116,8 +123,8 @@ def slice_domain(
     # get the index of the closest value to request min and max
     # side left returns `i`: domain[i-1] < request <= domain[i]
     # side right returns `i`: domain[i-1] <= request < domain[i]
-    idx_min = np.searchsorted(domain, np.nanmin(request) - buffer[0], side="right") - 1
-    idx_max = np.searchsorted(domain, np.nanmax(request) + buffer[1], side="left") + 1
+    idx_min = np.searchsorted(domain, np.min(request) - buf0, side="right") - 1
+    idx_max = np.searchsorted(domain, np.max(request) + buf1, side="left") + 1
 
     # clip idx_min between [0, len(domain) - 2]
     idx_min = min(len(domain) - 2, max(idx_min, 0))

pycontrails/core/datalib.py

@@ -80,16 +80,18 @@ def parse_timesteps(time: TimeInput | None, freq: str | None = "1H") -> list[dat
     elif len(time) == 1:
         time = (time[0], time[0])
     elif len(time) != 2:
-        raise ValueError("Input time bounds must have length < 2 and > 0")
+        msg = f"Input time bounds must have length < 2 and > 0, got {len(time)}"
+        raise ValueError(msg)
 
     # convert all to pandas Timestamp
     try:
         timestamps = [pd.to_datetime(t) for t in time]
     except ValueError as e:
-        raise ValueError(
-            f"Failed to parse all time inputs with error {e}. Time input "
-            "must be compatible with 'pd.to_datetime()'"
+        msg = (
+            f"Failed to parse time input {time}. "
+            "Time input must be compatible with 'pd.to_datetime()'"
         )
+        raise ValueError(msg) from e
 
     if freq is None:
         daterange = pd.DatetimeIndex([timestamps[0], timestamps[1]])
@@ -129,15 +131,15 @@ def parse_pressure_levels(
         pressure_levels = [pressure_levels]
 
     # Cast array-like to list of ints
-    # Use a new variable to satiate mypy
-    pressure_levels_ = np.asarray(pressure_levels, dtype=int).tolist()
+    out = np.asarray(pressure_levels, dtype=int).tolist()
 
     # ensure pressure levels are valid
-    for pl in pressure_levels_:
-        if supported is not None and pl not in supported:
-            raise ValueError(f"{pl} is not a valid pressure level {supported}")
+    for pl in out:
+        if supported and pl not in supported:
+            msg = f"Pressure level {pl} is not supported. Supported levels: {supported}"
+            raise ValueError(msg)
 
-    return pressure_levels_
+    return out
 
 
 def parse_variables(variables: VariableInput, supported: list[MetVariable]) -> list[MetVariable]:
@@ -175,70 +177,87 @@ def parse_variables(variables: VariableInput, supported: list[MetVariable]) -> l
     else:
         parsed_variables = variables
 
-    # unpack list of supported str values from MetVariables
-    short_names = [v.short_name for v in supported]
-    standard_names = [v.standard_name for v in supported]
-    long_names = [v.long_name for v in supported]
+    # unpack dict of supported str values from supported
+    short_names = {v.short_name: v for v in supported}
+    standard_names = {v.standard_name: v for v in supported}
+    long_names = {v.long_name: v for v in supported}
 
-    # unpack list of support int values from Met Variables
-    ecmwf_ids = [v.ecmwf_id for v in supported]
-    grib1_ids = [v.grib1_id for v in supported]
+    # unpack dict of support int values from supported
+    ecmwf_ids = {v.ecmwf_id: v for v in supported}
+    grib1_ids = {v.grib1_id: v for v in supported}
 
     for var in parsed_variables:
-        matched_variable: MetVariable | None = None
-
-        if isinstance(var, MetVariable):
-            if var in supported:
-                matched_variable = var
-
-        # list of MetVariable options
-        # here we extract the first MetVariable in var that is supported
-        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):
-                    raise ValueError("Variable options must be of type MetVariable.")
-                if v in supported:
-                    matched_variable = v
-                    break
-
-        # int code
-        elif isinstance(var, int):
-            if var in ecmwf_ids:
-                matched_variable = supported[ecmwf_ids.index(var)]
-            elif var in grib1_ids:
-                matched_variable = supported[grib1_ids.index(var)]
-
-        # string reference
-        elif isinstance(var, str):
-            if var in short_names:
-                matched_variable = supported[short_names.index(var)]
-            elif var in standard_names:
-                matched_variable = supported[standard_names.index(var)]
-            elif var in long_names:
-                matched_variable = supported[long_names.index(var)]
-
-        if matched_variable is None:
-            raise ValueError(
-                f"{var} is not in supported parameters. "
-                + f"Supported parameters include: {standard_names}"
-            )
+        matched = _find_match(
+            var,
+            supported,
+            ecmwf_ids,  # type: ignore[arg-type]
+            grib1_ids,  # type: ignore[arg-type]
+            short_names,
+            standard_names,
+            long_names,  # type: ignore[arg-type]
+        )
 
         # "replace" copies dataclass
-        met_var_list.append(dataclasses.replace(matched_variable))
+        met_var_list.append(dataclasses.replace(matched))
 
     return met_var_list
 
 
-def parse_grid(grid: float, supported: list[float]) -> float:
+def _find_match(
+    var: VariableInput,
+    supported: list[MetVariable],
+    ecmwf_ids: dict[int, MetVariable],
+    grib1_ids: dict[int, MetVariable],
+    short_names: dict[str, MetVariable],
+    standard_names: dict[str, MetVariable],
+    long_names: dict[str, MetVariable],
+) -> MetVariable:
+    """Find a match for input variable in supported."""
+
+    if isinstance(var, MetVariable):
+        if var in supported:
+            return var
+
+    # list of MetVariable options
+    # here we extract the first MetVariable in var that is supported
+    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):
+                msg = "Variable options must be of type MetVariable."
+                raise TypeError(msg)
+            if v in supported:
+                return v
+
+    # int code
+    elif isinstance(var, int):
+        if var in ecmwf_ids:
+            return ecmwf_ids[var]
+        if var in grib1_ids:
+            return grib1_ids[var]
+
+    # string reference
+    elif isinstance(var, str):
+        if var in short_names:
+            return short_names[var]
+        if var in standard_names:
+            return standard_names[var]
+        if var in long_names:
+            return long_names[var]
+
+    msg = f"{var} is not in supported parameters. Supported parameters include: {standard_names}"
+    raise ValueError(msg)
+
+
+def parse_grid(grid: float, supported: Sequence[float]) -> float:
     """Parse input grid spacing.
 
     Parameters
     ----------
     grid : float
         Input grid float
-    supported : list[float]
-        List of support grid values
+    supported : Sequence[float]
+        Sequence of support grid values
 
     Returns
     -------
@@ -251,7 +270,8 @@ def parse_grid(grid: float, supported: list[float]) -> float:
         Raises ValueError when ``grid`` is not in supported
     """
     if grid not in supported:
-        raise ValueError(f"Grid input must be one of {supported}")
+        msg = f"Grid input {grid} must be one of {supported}"
+        raise ValueError(msg)
 
     return grid
 
@@ -277,7 +297,8 @@ def round_hour(time: datetime, hour: int) -> datetime:
         Description
     """
     if hour not in range(1, 24):
-        raise ValueError("hour must be between [1, 23]")
+        msg = f"hour must be between [1, 23], got {hour}"
+        raise ValueError(msg)
 
     hour = (time.hour // hour) * hour
     return datetime(time.year, time.month, time.day, hour, 0, 0)
@@ -505,6 +526,19 @@ class MetDataSource(abc.ABC):
         :func:`xarray.open_mfdataset`
         """
 
+    @abc.abstractmethod
+    def set_metadata(self, ds: xr.Dataset | MetDataset) -> None:
+        """Set met source metadata on ``ds.attrs``.
+
+        This is called within the :meth:`open_metdataset` method to set metadata
+        on the returned :class:`MetDataset` instance.
+
+        Parameters
+        ----------
+        ds : xr.Dataset | MetDataset
+            Dataset to set metadata on. Mutated in place.
+        """
+
     # ----------------------
     # Common utility methods
     # ----------------------
@@ -518,7 +552,7 @@ class MetDataSource(abc.ABC):
         """
         if times_to_download := self.list_timesteps_not_cached(**xr_kwargs):
             logger.debug(
-                f"Not all files found in cachestore. Downloading times {times_to_download}"
+                "Not all files found in cachestore. Downloading times %s", times_to_download
             )
             self.download_dataset(times_to_download)
         else:
@@ -547,7 +581,7 @@ class MetDataSource(abc.ABC):
     def is_datafile_cached(self, t: datetime, **xr_kwargs: Any) -> bool:
         """Check datafile defined by datetime for variables and pressure levels in class.
 
-        If using a cloud cache store (i.e. :class:`cache.GCPCacheStore`) this is where the datafile
+        If using a cloud cache store (i.e. :class:`cache.GCPCacheStore`), this is where the datafile
         will be mirrored to a local file for access.
 
         Parameters
@@ -575,56 +609,64 @@ class MetDataSource(abc.ABC):
 
         # see if cache data file exists, and if so, get the file + path
         cache_path = self.create_cachepath(t)
-        if self.cachestore.exists(cache_path):
-            logger.debug(f"Cachepath {cache_path} exists, getting from cache.")
-            # If GCP cache is used, this will download file and return the local mirrored path
-            # If the local file already exists, this will return the local path
-            disk_path = self.cachestore.get(cache_path)
-
-            # check if all variables and pressure levels are in that path
-            try:
-                with self.open_dataset(disk_path, **xr_kwargs) as ds:
-                    for var in self.variable_shortnames:
-                        if var not in ds:
-                            logger.warning(
-                                f"Variable {var} not in downloaded dataset. "
-                                f"Found variables: {ds.data_vars}"
-                            )
-                            return False
-
-                    for pl in self.pressure_levels:
-                        if pl not in ds["level"].values:
-                            logger.warning(
-                                f"Pressure Level {pl} not in downloaded dataset. "
-                                f"Found pressure levels: {ds['level'].values}"
-                            )
-                            return False
-
-                    logger.debug(f"All variables and pressure levels found in {cache_path}")
-                    return True
-
-            except OSError as err:
-                if isinstance(self.cachestore, cache.GCPCacheStore):
-                    # If a GCPCacheStore is used, remove the corrupt file and
-                    # try again.
-                    # If the file is corrupt in the GCP bucket, we'll
-                    # get stuck in an infinite loop here
-                    logger.warning(
-                        "Found corrupt file %s on local disk. Try again to download from %s.",
-                        disk_path,
-                        self.cachestore,
-                    )
-                    self.cachestore.clear_disk(disk_path)
-                    return self.is_datafile_cached(t, **xr_kwargs)
-
-                raise OSError(
-                    f"Unable to open NETCDF file at {disk_path} "
-                    "This may be due to a incomplete download. "
-                    f"Consider manually removing {disk_path} and retrying."
-                ) from err
-
-        logger.debug(f"Cachepath {cache_path} does not exist in cache")
-        return False
+        if not self.cachestore.exists(cache_path):
+            logger.debug("Cachepath %s does not exist in cache", cache_path)
+            return False
+
+        logger.debug("Cachepath %s exists, getting from cache.", cache_path)
+
+        # If GCP cache is used, this will download file and return the local mirrored path
+        # If the local file already exists, this will return the local path
+        disk_path = self.cachestore.get(cache_path)
+
+        # check if all variables and pressure levels are in that path
+        try:
+            with self.open_dataset(disk_path, **xr_kwargs) as ds:
+                return self._check_is_ds_complete(ds, cache_path)
+
+        except OSError as err:
+            if isinstance(self.cachestore, cache.GCPCacheStore):
+                # If a GCPCacheStore is used, remove the corrupt file and try again.
+                # If the file is corrupt in the bucket, we'll get stuck in an infinite loop here.
+                logger.warning(
+                    "Found corrupt file %s on local disk. Try again to download from %s.",
+                    disk_path,
+                    self.cachestore,
+                    exc_info=err,
+                )
+                self.cachestore.clear_disk(disk_path)
+                return self.is_datafile_cached(t, **xr_kwargs)
+
+            msg = (
+                f"Unable to open NETCDF file at '{disk_path}'. "
+                "This may be due to a incomplete download. "
+                f"Consider manually removing '{disk_path}' and retrying."
+            )
+            raise OSError(msg) from err
+
+    def _check_is_ds_complete(self, ds: xr.Dataset, cache_path: str) -> bool:
+        """Check if ``ds`` has all variables and pressure levels defined by the instance."""
+        for var in self.variable_shortnames:
+            if var not in ds:
+                logger.warning(
+                    "Variable %s not in downloaded dataset. Found variables: %s",
+                    var,
+                    ds.data_vars,
+                )
+                return False
+
+        pl = np.asarray(self.pressure_levels)
+        cond = np.isin(pl, ds["level"].values)
+        if not np.all(cond):
+            logger.warning(
+                "Pressure Levels %s not in downloaded dataset. Found pressure levels: %s",
+                pl[~cond].tolist(),
+                ds["level"].values.tolist(),
+            )
+            return False
+
+        logger.debug("All variables and pressure levels found in %s", cache_path)
+        return True
 
     def open_dataset(
         self,

pycontrails/core/flight.py

@@ -27,12 +27,26 @@ if TYPE_CHECKING:
 
 
 class FlightPhase(enum.IntEnum):
-    """Flight phase enumeration."""
+    """Flight phase enumeration.
 
+    Use :func:`segment_phase` or :meth:`Flight.segment_phase` to determine flight phase.
+    """
+
+    #: Waypoints at which the flight is in a climb phase
     CLIMB = enum.auto()
+
+    #: Waypoints at which the flight is in a cruise phase
     CRUISE = enum.auto()
+
+    #: Waypoints at which the flight is in a descent phase
     DESCENT = enum.auto()
+
+    #: Waypoints at which the flight is not in a climb, cruise, or descent phase.
+    #: In practice, this category is used for waypoints at which the ROCD resembles
+    #: that of a cruise phase, but the altitude is below the minimum cruise altitude.
     LEVEL_FLIGHT = enum.auto()
+
+    #: Waypoints at which the ROCD is not defined.
     NAN = enum.auto()
 
 
@@ -890,11 +904,10 @@ class Flight(GeoVectorDataset):
             df["longitude"] = ((df["longitude"] + 180.0) % 360.0) - 180.0
 
         # STEP 6: Interpolate nan values in altitude
-        if df["altitude"].isna().any():
+        altitude = df["altitude"].to_numpy()
+        if np.any(np.isnan(altitude)):
             df_freq = pd.Timedelta(freq).to_numpy()
-            new_alt = _altitude_interpolation(
-                df["altitude"].to_numpy(), nominal_rocd, df_freq, climb_descend_at_end
-            )
+            new_alt = _altitude_interpolation(altitude, nominal_rocd, df_freq, climb_descend_at_end)
             _verify_altitude(new_alt, nominal_rocd, df_freq)
             df["altitude"] = new_alt
 
@@ -954,17 +967,18 @@ class Flight(GeoVectorDataset):
         elapsed_time = np.nancumsum(np.roll(seg_dur, 1))
         alt_ft = fit_altitude(
             elapsed_time,
-            np.copy(self.altitude_ft),
-            max_segments,
-            pop,
-            r2_target,
-            max_cruise_rocd,
-            sg_window,
+            self.altitude_ft,
+            max_segments=max_segments,
+            pop=pop,
+            r2_target=r2_target,
+            max_cruise_rocd=max_cruise_rocd,
+            sg_window=sg_window,
+            sg_polyorder=sg_polyorder,
         )
 
-        flight = self.copy()
-        flight.update(altitude_ft=alt_ft)
-        return flight
+        out = self.copy()
+        out.update(altitude_ft=alt_ft)
+        return out
 
     def _geodesic_interpolation(self, geodesic_threshold: float) -> pd.DataFrame | None:
         """Geodesic interpolate between large gaps between waypoints.
@@ -1014,17 +1028,18 @@ class Flight(GeoVectorDataset):
         latitudes: list[float] = []
         times: list[np.ndarray] = []
 
+        longitude = self["longitude"]
+        latitude = self["latitude"]
+        time = self["time"]
+
         for index in gap_indices:
-            lon0, lat0, t0 = (
-                self["longitude"][index],
-                self["latitude"][index],
-                self["time"][index],
-            )
-            lon1, lat1, t1 = (
-                self["longitude"][index + 1],
-                self["latitude"][index + 1],
-                self["time"][index + 1],
-            )
+            lon0 = longitude[index]
+            lat0 = latitude[index]
+            t0 = time[index]
+            lon1 = longitude[index + 1]
+            lat1 = latitude[index + 1]
+            t1 = time[index + 1]
+
             distance = segs[index]
             n_steps = distance // geodesic_threshold  # number of new waypoints to generate
 
@@ -1037,8 +1052,9 @@ class Flight(GeoVectorDataset):
             latitudes.extend(lats)
 
             # + 1 to denominator to stay consistent with geod.npts (only interior points)
-            t_step = (t1 - t0) / (n_steps + 1)
-            # substract 0.5 * t_step to ensure round-off error doesn't put final arange point
+            t_step = (t1 - t0) / (n_steps + 1.0)
+
+            # subtract 0.5 * t_step to ensure round-off error doesn't put final arange point
             # very close to t1
             t_range = np.arange(t0 + t_step, t1 - 0.5 * t_step, t_step)
             times.append(t_range)
@@ -1808,8 +1824,8 @@ def segment_phase(
     cruise = (
         (rocd < threshold_rocd) & (rocd > -threshold_rocd) & (altitude_ft > min_cruise_altitude_ft)
     )
-    climb = ~cruise & (rocd > 0)
-    descent = ~cruise & (rocd < 0)
+    climb = ~cruise & (rocd > 0.0)
+    descent = ~cruise & (rocd < 0.0)
     level_flight = ~(nan | cruise | climb | descent)
 
     phase = np.empty(rocd.shape, dtype=np.uint8)
@@ -1857,6 +1873,7 @@ def segment_rocd(
 def fit_altitude(
     elapsed_time: npt.NDArray[np.float_],
     altitude_ft: npt.NDArray[np.float_],
+    *,
     max_segments: int = 30,
     pop: int = 3,
     r2_target: float = 0.999,