pycontrails 0.42.0__cp39-cp39-macosx_11_0_arm64.whl → 0.42.2__cp39-cp39-macosx_11_0_arm64.whl

pycontrails/_version.py

@@ -1,4 +1,4 @@
 # file generated by setuptools_scm
 # don't change, don't track in version control
-__version__ = version = '0.42.0'
-__version_tuple__ = version_tuple = (0, 42, 0)
+__version__ = version = '0.42.2'
+__version_tuple__ = version_tuple = (0, 42, 2)

pycontrails/core/cache.py

@@ -370,18 +370,16 @@ class DiskCacheStore(CacheStore):
         if disk_path.is_file():
             logger.debug("Remove file at path %s", disk_path)
             disk_path.unlink()
+            return
 
         # Assume anything else is a directory
-        elif disk_path.exists():
+        if disk_path.exists():
             # rm directory recursively
             logger.debug("Remove directory at path %s", disk_path)
             shutil.rmtree(disk_path, ignore_errors=True)
+            return
 
-        else:
-            warnings.warn(f"No cache path found at {disk_path}")
-
-        # make sure local cache directory exists
-        # pathlib.Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
+        warnings.warn(f"No cache path found at {disk_path}")
 
 
 class GCPCacheStore(CacheStore):

pycontrails/core/datalib.py

@@ -32,6 +32,9 @@ NETCDF_ENGINE: str = "netcdf4"
 #: Default chunking strategy when opening datasets with xarray
 DEFAULT_CHUNKS: dict[str, int] = {"time": 1}
 
+#: Whether to open multi-file datasets in parallel
+OPEN_IN_PARALLEL: bool = False
+
 
 def parse_timesteps(time: TimeInput | None, freq: str | None = "1H") -> list[datetime]:
     """Parse time input into set of time steps.
@@ -645,7 +648,7 @@ class MetDataSource(abc.ABC):
 
                 - chunks: {"time": 1}
                 - engine: "netcdf4"
-                - parallel: True
+                - parallel: False
 
         Returns
         -------
@@ -654,5 +657,5 @@ class MetDataSource(abc.ABC):
         """
         xr_kwargs.setdefault("engine", NETCDF_ENGINE)
         xr_kwargs.setdefault("chunks", DEFAULT_CHUNKS)
-        xr_kwargs.setdefault("parallel", True)
+        xr_kwargs.setdefault("parallel", OPEN_IN_PARALLEL)
         return xr.open_mfdataset(disk_paths, **xr_kwargs)

pycontrails/core/fleet.py

@@ -93,6 +93,49 @@ class Fleet(Flight):
         final_waypoints[final_waypoint_indices] = True
         return final_waypoints
 
+    def fit_altitude(
+        self,
+        max_segments: int = 30,
+        pop: int = 3,
+        r2_target: float = 0.999,
+        max_cruise_rocd: float = 10,
+        sg_window: int = 7,
+        sg_polyorder: int = 1,
+    ) -> Fleet:
+        """Use piecewise linear fitting to smooth a flight profile.
+
+        Fit a flight profile to a series of line segments. Segments that have a
+        small rocd will be set to have a slope of zero and snapped to the
+        nearest thousand foot level.  A Savitzky-Golay filter will then be
+        applied to the profile to smooth the climbs and descents.  This filter
+        works best for high frequency flight data, sampled at a 1-3 second
+        sampling period.
+
+        Parameters
+        ----------
+        max_segments : int, optional
+            The maximum number of line segements to fit to the flight profile.
+        pop: int, optional
+            Population parameter used for the stocastic optimization routine
+            used to fit the flight profile.
+        r2_target: float, optional
+            Target r^2 value for solver. Solver will continue to add line
+            segments until the resulting r^2 value is greater than this.
+        max_cruise_rocd: float, optional
+            The maximum ROCD for a segment that will be forced to a slope of
+            zero, [:math:`ft s^{-1}`]
+        sg_window: int, optional
+            Parameter for :func:`scipy.signal.savgol_filter`
+        sg_polyorder: int, optional
+            Parameter for :func:`scipy.signal.savgol_filter`
+
+        Returns
+        -------
+        Fleet
+            Smoothed flight
+        """
+        raise NotImplementedError("Only implemented for Flight instances")
+
     @classmethod
     def from_seq(
         cls,
@@ -257,8 +300,8 @@ class Fleet(Flight):
 
     def segment_true_airspeed(
         self,
-        u_wind: npt.NDArray[np.float_] | None = None,
-        v_wind: npt.NDArray[np.float_] | None = None,
+        u_wind: npt.NDArray[np.float_] | float = 0.0,
+        v_wind: npt.NDArray[np.float_] | float = 0.0,
         smooth: bool = True,
         window_length: int = 7,
         polyorder: int = 1,
@@ -277,13 +320,14 @@ class Fleet(Flight):
         RuntimeError
             Unexpected key `__u_wind` or `__v_wind` found in :attr:`data`.
         """
-        if u_wind is not None:
+        if isinstance(u_wind, np.ndarray):
             # Choosing a key we don't think exists
             key = "__u_wind"
             if key in self:
                 raise RuntimeError(f"Unexpected key {key} found")
             self[key] = u_wind
-        if v_wind is not None:
+
+        if isinstance(v_wind, np.ndarray):
             # Choosing a key we don't think exists
             key = "__v_wind"
             if key in self:
@@ -292,11 +336,11 @@ class Fleet(Flight):
 
         # Calculate TAS on each flight individually
         def calc_tas(fl: Flight) -> npt.NDArray[np.float_]:
-            u_wind = fl.get("__u_wind", None)
-            v_wind = fl.get("__v_wind", None)
+            u = fl.get("__u_wind", u_wind)
+            v = fl.get("__v_wind", v_wind)
 
             return fl.segment_true_airspeed(
-                u_wind, v_wind, smooth=smooth, window_length=window_length, polyorder=polyorder
+                u, v, smooth=smooth, window_length=window_length, polyorder=polyorder
             )
 
         fls = self.to_flight_list(copy=False)
@@ -330,6 +374,14 @@ class Fleet(Flight):
     def segment_length(self) -> npt.NDArray[np.float_]:
         return np.where(self.final_waypoints, np.nan, super().segment_length())
 
+    @property
+    @overrides
+    def max_distance_gap(self) -> float:
+        if self.attrs["crs"] != "EPSG:4326":
+            raise NotImplementedError("Only implemented for EPSG:4326 CRS.")
+
+        return np.nanmax(self.segment_length()).item()
+
     @overrides
     def segment_azimuth(self) -> npt.NDArray[np.float_]:
         return np.where(self.final_waypoints, np.nan, super().segment_azimuth())

pycontrails/core/flight.py

@@ -242,7 +242,7 @@ class Flight(GeoVectorDataset):
         self.fuel = fuel or JetA()
 
         # Check flight data for possible errors
-        if np.any(self.altitude > 16000):
+        if np.any(self.altitude > 16000.0):
             flight_id = self.attrs.get("flight_id", "")
             flight_id = flight_id and f" for flight {flight_id}"
             warnings.warn(
@@ -257,35 +257,36 @@ class Flight(GeoVectorDataset):
                 "with segment-based methods (e.g. 'segment_true_airspeed')."
             )
 
-        diff_ = np.diff(self["time"])
+        time_diff = np.diff(self["time"])
 
         # Ensure that time is sorted
-        if self and np.any(diff_ < np.timedelta64(0)):
+        if self and np.any(time_diff < np.timedelta64(0)):
             if not copy:
                 raise ValueError(
-                    "`time` data must be sorted if `copy` is False on creation. "
+                    "The 'time' array must be sorted if 'copy=False' on creation. "
                     "Set copy=False, or sort data before creating Flight."
                 )
-            warnings.warn("Sorting Flight data by `time`.")
+            warnings.warn("Sorting Flight data by time.")
 
             sorted_flight = self.sort("time")
             self.data = sorted_flight.data
-            # Update diff_ ... we use it again below
-            diff_ = np.diff(self["time"])
+
+            # Update time_diff ... we use it again below
+            time_diff = np.diff(self["time"])
 
         # Check for duplicate times. If dropping duplicates,
-        # keep the *first* occurrence of each time. This is achieved with
-        # the np.insert (rather than np.append) function.
-        filt = np.insert(diff_ > np.timedelta64(0), 0, True)
-        if not np.all(filt):
+        # keep the *first* occurrence of each time.
+        duplicated_times = time_diff == np.timedelta64(0)
+        if self and np.any(duplicated_times):
             if drop_duplicated_times:
-                filtered_flight = self.filter(filt, copy=False)
+                mask = np.insert(duplicated_times, 0, False)
+                filtered_flight = self.filter(~mask, copy=False)
                 self.data = filtered_flight.data
             else:
                 warnings.warn(
-                    "Flight contains duplicate times. This will cause errors "
-                    "with segment-based methods (e.g. 'segment_true_airspeed'). Set "
-                    "'drop_duplicated=True' or call the 'resample_and_fill' method."
+                    f"Flight contains {duplicated_times.sum()} duplicate times. "
+                    "This will cause errors with segment-based methods. Set "
+                    "'drop_duplicated_times=True' or call the 'resample_and_fill' method."
                 )
 
     @overrides
@@ -417,7 +418,7 @@ class Flight(GeoVectorDataset):
     # Segment Properties
     # ------------
 
-    def segment_duration(self, dtype: np.dtype = np.dtype(np.float32)) -> npt.NDArray[np.float_]:
+    def segment_duration(self, dtype: npt.DTypeLike = np.float32) -> npt.NDArray[np.float_]:
         r"""Compute time elapsed between waypoints in seconds.
 
         ``np.nan`` appended so the length of the output is the same as number of waypoints.
@@ -581,13 +582,11 @@ class Flight(GeoVectorDataset):
         npt.NDArray[np.float_]
             Groundspeed of the segment, [:math:`m s^{-1}`]
         """
-        # get horizontal distance - set altitude to 0
+        # get horizontal distance (altitude is ignored)
         horizontal_segment_length = geo.segment_haversine(self["longitude"], self["latitude"])
 
         # time between waypoints, in seconds
-        dt_sec = np.empty_like(horizontal_segment_length)
-        dt_sec[:-1] = np.diff(self["time"]) / np.timedelta64(1, "s")
-        dt_sec[-1] = np.nan
+        dt_sec = self.segment_duration(dtype=horizontal_segment_length.dtype)
 
         # calculate groundspeed
         groundspeed = horizontal_segment_length / dt_sec
@@ -602,8 +601,8 @@ class Flight(GeoVectorDataset):
 
     def segment_true_airspeed(
         self,
-        u_wind: npt.NDArray[np.float_] | None = None,
-        v_wind: npt.NDArray[np.float_] | None = None,
+        u_wind: npt.NDArray[np.float_] | float = 0.0,
+        v_wind: npt.NDArray[np.float_] | float = 0.0,
         smooth: bool = True,
         window_length: int = 7,
         polyorder: int = 1,
@@ -614,10 +613,10 @@ class Flight(GeoVectorDataset):
 
         Parameters
         ----------
-        u_wind : npt.NDArray[np.float_], optional
+        u_wind : npt.NDArray[np.float_] | float
             U wind speed, [:math:`m \ s^{-1}`].
             Defaults to 0 for all waypoints.
-        v_wind : npt.NDArray[np.float_], optional
+        v_wind : npt.NDArray[np.float_] | float
             V wind speed, [:math:`m \ s^{-1}`].
             Defaults to 0 for all waypoints.
         smooth : bool, optional
@@ -635,12 +634,6 @@ class Flight(GeoVectorDataset):
         """
         groundspeed = self.segment_groundspeed(smooth, window_length, polyorder)
 
-        if u_wind is None:
-            u_wind = np.zeros_like(groundspeed)
-
-        if v_wind is None:
-            v_wind = np.zeros_like(groundspeed)
-
         sin_a, cos_a = self.segment_angle()
         gs_x = groundspeed * cos_a
         gs_y = groundspeed * sin_a
@@ -820,20 +813,20 @@ class Flight(GeoVectorDataset):
                    2       50.0       0.0       0.0 2020-01-01 02:00:00
 
         >>> fl.resample_and_fill('10T').dataframe  # resample with 10 minute frequency
-                          time  longitude  latitude  altitude
-        0  2020-01-01 00:00:00   0.000000       0.0       0.0
-        1  2020-01-01 00:10:00   0.000000       0.0       0.0
-        2  2020-01-01 00:20:00   0.000000       0.0       0.0
-        3  2020-01-01 00:30:00   0.000000       0.0       0.0
-        4  2020-01-01 00:40:00   0.000000       0.0       0.0
-        5  2020-01-01 00:50:00   0.000000       0.0       0.0
-        6  2020-01-01 01:00:00   0.000000       0.0       0.0
-        7  2020-01-01 01:10:00   8.928571       0.0       0.0
-        8  2020-01-01 01:20:00  16.964286       0.0       0.0
-        9  2020-01-01 01:30:00  25.892857       0.0       0.0
-        10 2020-01-01 01:40:00  33.928571       0.0       0.0
-        11 2020-01-01 01:50:00  41.964286       0.0       0.0
-        12 2020-01-01 02:00:00  50.000000       0.0       0.0
+            longitude  latitude  altitude                time
+        0    0.000000       0.0       0.0 2020-01-01 00:00:00
+        1    0.000000       0.0       0.0 2020-01-01 00:10:00
+        2    0.000000       0.0       0.0 2020-01-01 00:20:00
+        3    0.000000       0.0       0.0 2020-01-01 00:30:00
+        4    0.000000       0.0       0.0 2020-01-01 00:40:00
+        5    0.000000       0.0       0.0 2020-01-01 00:50:00
+        6    0.000000       0.0       0.0 2020-01-01 01:00:00
+        7    8.928571       0.0       0.0 2020-01-01 01:10:00
+        8   16.964286       0.0       0.0 2020-01-01 01:20:00
+        9   25.892857       0.0       0.0 2020-01-01 01:30:00
+        10  33.928571       0.0       0.0 2020-01-01 01:40:00
+        11  41.964286       0.0       0.0 2020-01-01 01:50:00
+        12  50.000000       0.0       0.0 2020-01-01 02:00:00
         """
         methods = "geodesic", "linear"
         if fill_method not in methods:
@@ -908,6 +901,65 @@ class Flight(GeoVectorDataset):
         df = df.reset_index()
         return Flight(data=df, attrs=self.attrs)
 
+    def fit_altitude(
+        self,
+        max_segments: int = 30,
+        pop: int = 3,
+        r2_target: float = 0.999,
+        max_cruise_rocd: float = 10.0,
+        sg_window: int = 7,
+        sg_polyorder: int = 1,
+    ) -> Flight:
+        """Use piecewise linear fitting to smooth a flight profile.
+
+        Fit a flight profile to a series of line segments. Segments that have a
+        small rocd will be set to have a slope of zero and snapped to the
+        nearest thousand foot level.  A Savitzky-Golay filter will then be
+        applied to the profile to smooth the climbs and descents.  This filter
+        works best for high frequency flight data, sampled at a 1-3 second
+        sampling period.
+
+        Parameters
+        ----------
+        max_segments : int, optional
+            The maximum number of line segements to fit to the flight profile.
+        pop: int, optional
+            Population parameter used for the stocastic optimization routine
+            used to fit the flight profile.
+        r2_target: float, optional
+            Target r^2 value for solver. Solver will continue to add line
+            segments until the resulting r^2 value is greater than this.
+        max_cruise_rocd: float, optional
+            The maximum ROCD for a segment that will be forced to a slope of
+            zero, [:math:`ft s^{-1}`]
+        sg_window: int, optional
+            Parameter for :func:`scipy.signal.savgol_filter`
+        sg_polyorder: int, optional
+            Parameter for :func:`scipy.signal.savgol_filter`
+
+        Returns
+        -------
+        Flight
+            Smoothed flight
+        """
+        # np.roll pushes the last NaN value from `segment_duration` to the front
+        # so the elapsed time at the first waypoint will be 0
+        seg_dur = self.segment_duration(dtype=np.float64)
+        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,
+        )
+
+        flight = self.copy()
+        flight.update(altitude_ft=alt_ft)
+        return flight
+
     def _geodesic_interpolation(self, geodesic_threshold: float) -> pd.DataFrame | None:
         """Geodesic interpolate between large gaps between waypoints.
 
@@ -1123,7 +1175,7 @@ class Flight(GeoVectorDataset):
         >>> variables = ["air_temperature", "specific_humidity"]
         >>> levels = [300, 250, 200]
         >>> era5 = ERA5(time=times, variables=variables, pressure_levels=levels)
-        >>> met = era5.open_metdataset(xr_kwargs=dict(parallel=False))
+        >>> met = era5.open_metdataset()
 
         >>> # Build flight
         >>> df = pd.DataFrame()
@@ -1232,9 +1284,9 @@ def _return_linestring(data: dict[str, npt.NDArray[np.float_]]) -> list[list[flo
     """
     # rounding to reduce the size of resultant json arrays
     points = zip(  # pylint: disable=zip-builtin-not-iterating
-        np.around(data["longitude"], decimals=4),
-        np.around(data["latitude"], decimals=4),
-        np.around(data["altitude"], decimals=4),
+        np.round(data["longitude"], decimals=4),
+        np.round(data["latitude"], decimals=4),
+        np.round(data["altitude"], decimals=4),
     )
     return [list(p) for p in points]
 
@@ -1515,7 +1567,7 @@ def filter_altitude(
 
 
 def segment_duration(
-    time: npt.NDArray[np.datetime64], dtype: np.dtype = np.dtype(np.float32)
+    time: npt.NDArray[np.datetime64], dtype: npt.DTypeLike = np.float32
 ) -> npt.NDArray[np.float_]:
     """Calculate the time difference between waypoints.
 
@@ -1634,3 +1686,77 @@ def segment_rocd(
     out[-1] = np.nan
 
     return out
+
+
+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,
+    max_cruise_rocd: float = 10.0,
+    sg_window: int = 7,
+    sg_polyorder: int = 1,
+) -> npt.NDArray[np.float_]:
+    """Use piecewise linear fitting to smooth a flight profile.
+
+    Fit a flight profile to a series of line segments. Segments that have a
+    small rocd will be set to have a slope of zero and snapped to the
+    nearest thousand foot level.  A Savitzky-Golay filter will then be
+    applied to the profile to smooth the climbs and descents.  This filter
+    works best for high frequency flight data, sampled at a 1-3 second
+    sampling period.
+
+    Parameters
+    ----------
+    elapsed_time: npt.NDArray[np.float_]
+        Cumulative time of flight between waypoints, [:math:`s`]
+    altitude_ft: npt.NDArray[np.float_]
+        Altitude of each waypoint, [:math:`ft`
+    max_segments: int, optional
+        The maximum number of line segements to fit to the flight profile.
+    pop: int, optional
+        Population parameter used for the stocastic optimization routine
+        used to fit the flight profile.
+    r2_target: float, optional
+        Target r^2 value for solver. Solver will continue to add line
+        segments until the resulting r^2 value is greater than this.
+    max_cruise_rocd: float, optional
+        The maximum ROCD for a segment that will be forced to a slope of
+        zero, [:math:`ft s^{-1}`]
+    sg_window: int, optional
+        Parameter for :func:`scipy.signal.savgol_filter`
+    sg_polyorder: int, optional
+        Parameter for :func:`scipy.signal.savgol_filter`
+
+    Returns
+    -------
+    npt.NDArray[np.float_]
+        Smoothed flight altitudes
+    """
+    try:
+        import pwlf
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "The 'fit_altitude' function requires the 'pwlf' package."
+            "This can be installed with 'pip install pwlf'."
+        )
+    for i in range(1, max_segments):
+        m2 = pwlf.PiecewiseLinFit(elapsed_time, altitude_ft)
+        r = m2.fitfast(i, pop)
+        r2 = m2.r_squared()
+        if r2 > r2_target:
+            break
+
+    mask = abs(m2.slopes) < max_cruise_rocd / 60.0
+    bounds = r[:-1][mask], r[1:][mask]
+    lvl = np.round(m2.intercepts[mask], -3)
+    time_stack = np.repeat(elapsed_time[:, np.newaxis], lvl.size, axis=1)
+    filt = (time_stack >= bounds[0]) & (time_stack <= bounds[1])
+    altitude_ft = np.copy(altitude_ft)
+    for i in range(lvl.size):
+        altitude_ft[filt[:, i]] = lvl[i]
+
+    altitude_ft = scipy.signal.savgol_filter(altitude_ft, sg_window, sg_polyorder)
+
+    return altitude_ft