pycontrails 0.41.0__cp310-cp310-macosx_11_0_arm64.whl → 0.42.2__cp310-cp310-macosx_11_0_arm64.whl

pycontrails/core/vector.py

@@ -97,7 +97,7 @@ class VectorDataDict(Dict[str, np.ndarray]):
         for arr in self.values():
             self._validate_array(arr)
 
-    def __setitem__(self, k: str, v: np.ndarray) -> None:
+    def __setitem__(self, k: str, v: npt.ArrayLike) -> None:
         """Set new key-value pair to instance and warn when overwriting existing key.
 
         This method casts ``v`` to a ``np.ndarray`` and ensures that the array size is
@@ -131,7 +131,7 @@ class VectorDataDict(Dict[str, np.ndarray]):
         if not len(self):
             del self._size
 
-    def setdefault(self, k: str, default: np.ndarray | None = None) -> np.ndarray:
+    def setdefault(self, k: str, default: npt.ArrayLike | None = None) -> np.ndarray:
         """Thin wrapper around ``dict.setdefault``.
 
         The main purpose of overriding is to run :meth:`_validate_array()` on set.
@@ -140,7 +140,7 @@ class VectorDataDict(Dict[str, np.ndarray]):
         ----------
         k : str
             Key
-        default : np.ndarray, optional
+        default : npt.ArrayLike, optional
             Default value for key ``k``
 
         Returns
@@ -156,10 +156,10 @@ class VectorDataDict(Dict[str, np.ndarray]):
             default = np.array([])
 
         self[k] = default
-        return default
+        return self[k]
 
     def update(  # type: ignore[override]
-        self, other: dict[str, np.ndarray] | None = None, **kwargs: np.ndarray
+        self, other: dict[str, npt.ArrayLike] | None = None, **kwargs: npt.ArrayLike
     ) -> None:
         """Update values without warning if overwriting.
 
@@ -168,24 +168,24 @@ class VectorDataDict(Dict[str, np.ndarray]):
 
         Parameters
         ----------
-        other : dict[str, np.ndarray] | None, optional
+        other : dict[str, npt.ArrayLike] | None, optional
             Fields to update as dict
-        **kwargs : np.ndarray
+        **kwargs : npt.ArrayLike
             Fields to update as kwargs
         """
         other = other or {}
-        other = {k: np.asarray(v) for k, v in other.items()}
-        for arr in other.values():
+        other_arrs = {k: np.asarray(v) for k, v in other.items()}
+        for arr in other_arrs.values():
             self._validate_array(arr)
 
-        super().update(other)
+        super().update(other_arrs)
 
         # validate any kwarg arrays
-        kwargs = {k: np.asarray(v) for k, v in kwargs.items()}
-        for arr in kwargs.values():
+        kwargs_arr = {k: np.asarray(v) for k, v in kwargs.items()}
+        for arr in kwargs_arr.values():
             self._validate_array(arr)
 
-        super().update(kwargs)
+        super().update(kwargs_arr)
 
     def _validate_array(self, arr: np.ndarray) -> None:
         """Ensure that `arr` is compatible with instance.
@@ -240,7 +240,7 @@ class VectorDataset:
 
     Parameters
     ----------
-    data : dict[str, np.ndarray] | pd.DataFrame | VectorDataDict | VectorDataset | None, optional
+    data : dict[str, npt.ArrayLike] | pd.DataFrame | VectorDataDict | VectorDataset | None, optional
         Initial data, by default None
     attrs : dict[str, Any] | AttrDict, optional
         Dictionary of attributes, by default None
@@ -265,7 +265,11 @@ class VectorDataset:
 
     def __init__(
         self,
-        data: dict[str, np.ndarray] | pd.DataFrame | VectorDataDict | VectorDataset | None = None,
+        data: dict[str, npt.ArrayLike]
+        | pd.DataFrame
+        | VectorDataDict
+        | VectorDataset
+        | None = None,
         attrs: dict[str, Any] | AttrDict | None = None,
         copy: bool = True,
         **attrs_kwargs: Any,
@@ -283,15 +287,17 @@ class VectorDataset:
             # Take extra caution with a time column
 
             if "time" in data:
-                if not hasattr(data["time"], "dt"):
+                time = data["time"]
+
+                if not hasattr(time, "dt"):
                     # If the time column is a string, we try to convert it to a datetime
                     # If it fails (for example, a unix integer time), we raise an error
                     # and let the user figure it out.
                     try:
-                        data["time"] = pd.to_datetime(data["time"])
+                        time = pd.to_datetime(time)
                     except ValueError:
                         raise ValueError(
-                            "Column `time` must hold datetimelike values. "
+                            "The 'time' field must hold datetime-like values. "
                             'Try data["time"] = pd.to_datetime(data["time"], unit=...) '
                             "with the appropriate unit."
                         )
@@ -301,13 +307,19 @@ class VectorDataset:
                 # we raise an error in this case. Timezone issues are complicated,
                 # and so it is better for the user to handle them rather than try
                 # to address them here.
-                if data["time"].dt.tz is not None:
+                if time.dt.tz is not None:
                     raise ValueError(
-                        "Column `time` must be timezone naive. "
+                        "The 'time' field must be timezone naive. "
                         "This can be achieved with: "
                         'data["time"] = data["time"].dt.tz_localize(None)'
                     )
-            self.data = VectorDataDict({col: ser.to_numpy(copy=copy) for col, ser in data.items()})
+
+                data = {col: ser.to_numpy(copy=copy) for col, ser in data.items() if col != "time"}
+                data["time"] = time.to_numpy(copy=copy)
+            else:
+                data = {col: ser.to_numpy(copy=copy) for col, ser in data.items()}
+
+            self.data = VectorDataDict(data)
 
         elif isinstance(data, VectorDataDict) and not copy:
             self.data = data
@@ -368,14 +380,14 @@ class VectorDataset:
         """
         return self.data.get(key, default_value)
 
-    def __setitem__(self, key: str, values: np.ndarray) -> None:
+    def __setitem__(self, key: str, values: npt.ArrayLike) -> None:
         """Set values at key `key` on :attr:`data`.
 
         Parameters
         ----------
         key : str
             Key name in :attr:`data`
-        values : np.ndarray
+        values : npt.ArrayLike
             Values to set to :attr:`data`. Array size must be compatible with existing data.
         """
         self.data[key] = values
@@ -416,27 +428,29 @@ class VectorDataset:
         return key in self.data
 
     def update(
-        self, other: VectorDataDict | dict[str, np.ndarray] | None = None, **kwargs: np.ndarray
+        self,
+        other: dict[str, npt.ArrayLike] | None = None,
+        **kwargs: npt.ArrayLike,
     ) -> None:
         """Update values in  :attr:`data` dict without warning if overwriting.
 
         Parameters
         ----------
-        other : VectorDataDict | dict[str, np.ndarray] | None, optional
+        other : dict[str, npt.ArrayLike] | None, optional
             Fields to update as dict
-        **kwargs : np.ndarray
+        **kwargs : npt.ArrayLike
             Fields to update as kwargs
         """
         self.data.update(other, **kwargs)
 
-    def setdefault(self, key: str, default: np.ndarray | None = None) -> np.ndarray:
+    def setdefault(self, key: str, default: npt.ArrayLike | None = None) -> np.ndarray:
         """Shortcut to :attr:`data.setdefault`.
 
         Parameters
         ----------
         key : str
             Key in :attr:`data` dict.
-        default : np.ndarray, optional
+        default : npt.ArrayLike, optional
             Values to use as default, if key is not defined
 
         Returns
@@ -777,14 +791,16 @@ class VectorDataset:
         data = {key: self[key] for key in keys}
         return VectorDataset(data=data, attrs=self.attrs, copy=copy)
 
-    def filter(self: VectorDatasetType, mask: np.ndarray, copy: bool = True) -> VectorDatasetType:
+    def filter(
+        self: VectorDatasetType, mask: npt.NDArray[np.bool_], copy: bool = True
+    ) -> VectorDatasetType:
         """Filter :attr:`data` according to a boolean array ``mask``.
 
         Entries corresponding to ``mask == True`` are kept.
 
         Parameters
         ----------
-        mask : np.ndarray
+        mask : npt.NDArray[np.bool_]
             Boolean array with compatible shape.
         copy : bool, optional
             Copy data on filter. Defaults to True. See
@@ -1035,26 +1051,26 @@ class GeoVectorDataset(VectorDataset):
 
     Parameters
     ----------
-    data : dict[str, np.ndarray] | pd.DataFrame | VectorDataDict | VectorDataset | None, optional
+    data : dict[str, npt.ArrayLike] | pd.DataFrame | VectorDataDict | VectorDataset | None, optional
         Data dictionary or :class:`pandas.DataFrame` .
         Must include keys/columns ``time``, ``latitude``, ``longitude``, ``altitude`` or ``level``.
         Keyword arguments for ``time``, ``latitude``, ``longitude``, ``altitude`` or ``level``
         override ``data`` inputs. Expects ``altitude`` in meters and ``time``
         as a DatetimeLike (or array that can processed with :meth:`pd.to_datetime`).
         Additional waypoint-specific data can be included as additional keys/columns.
-    longitude : np.ndarray, optional
+    longitude : npt.ArrayLike, optional
         Longitude data.
         Defaults to None.
-    latitude : np.ndarray, optional
+    latitude : npt.ArrayLike, optional
         Latitude data.
         Defaults to None.
-    altitude : np.ndarray, optional
+    altitude : npt.ArrayLike, optional
         Altitude data, [:math:`m`].
         Defaults to None.
-    level : np.ndarray, optional
+    level : npt.ArrayLike, optional
         Level data, [:math:`hPa`].
         Defaults to None.
-    time : np.ndarray, optional
+    time : npt.ArrayLike, optional
         Time data.
         Expects an array of DatetimeLike values,
         or array that can proccessed with :meth:`pd.to_datetime`.
@@ -1084,7 +1100,11 @@ class GeoVectorDataset(VectorDataset):
 
     def __init__(
         self,
-        data: dict[str, np.ndarray] | pd.DataFrame | VectorDataDict | VectorDataset | None = None,
+        data: dict[str, npt.ArrayLike]
+        | pd.DataFrame
+        | VectorDataDict
+        | VectorDataset
+        | None = None,
         longitude: npt.ArrayLike | None = None,
         latitude: npt.ArrayLike | None = None,
         altitude: npt.ArrayLike | None = None,
@@ -1201,7 +1221,7 @@ class GeoVectorDataset(VectorDataset):
         return attrs
 
     @property
-    def level(self) -> np.ndarray:
+    def level(self) -> npt.NDArray[np.float_]:
         """Get pressure ``level`` values for points.
 
         Automatically calculates pressure level using :func:`units.m_to_pl` using ``altitude`` key.
@@ -1212,7 +1232,7 @@ class GeoVectorDataset(VectorDataset):
 
         Returns
         -------
-        np.ndarray
+        npt.NDArray[np.float_]
             Point pressure level values, [:math:`hPa`]
         """
         try:
@@ -1221,7 +1241,7 @@ class GeoVectorDataset(VectorDataset):
             return units.m_to_pl(self.altitude)
 
     @property
-    def altitude(self) -> np.ndarray:
+    def altitude(self) -> npt.NDArray[np.float_]:
         """Get altitude.
 
         Automatically calculates altitude using :func:`units.pl_to_m` using ``level`` key.
@@ -1232,7 +1252,7 @@ class GeoVectorDataset(VectorDataset):
 
         Returns
         -------
-        np.ndarray
+        npt.NDArray[np.float_]
             Altitude, [:math:`m`]
         """
         try:
@@ -1246,12 +1266,12 @@ class GeoVectorDataset(VectorDataset):
             return units.ft_to_m(self["altitude_ft"])
 
     @property
-    def air_pressure(self) -> np.ndarray:
+    def air_pressure(self) -> npt.NDArray[np.float_]:
         """Get ``air_pressure`` values for points.
 
         Returns
         -------
-        np.ndarray
+        npt.NDArray[np.float_]
             Point air pressure values, [:math:`Pa`]
         """
         try:
@@ -1260,12 +1280,12 @@ class GeoVectorDataset(VectorDataset):
             return 100 * self.level
 
     @property
-    def altitude_ft(self) -> np.ndarray:
+    def altitude_ft(self) -> npt.NDArray[np.float_]:
         """Get altitude in feet.
 
         Returns
         -------
-        np.ndarray
+        npt.NDArray[np.float_]
             Altitude, [:math:`ft`]
         """
         try:
@@ -1368,7 +1388,7 @@ class GeoVectorDataset(VectorDataset):
 
     def coords_intersect_met(
         self, met: met_module.MetDataset | met_module.MetDataArray
-    ) -> np.ndarray:
+    ) -> npt.NDArray[np.bool_]:
         """Return boolean mask of data inside the bounding box defined by ``met``.
 
         Parameters
@@ -1378,7 +1398,7 @@ class GeoVectorDataset(VectorDataset):
 
         Returns
         -------
-        np.ndarray
+        npt.NDArray[np.bool_]
             True if point is inside the bounding box defined by ``met``.
         """
 
@@ -1405,26 +1425,26 @@ class GeoVectorDataset(VectorDataset):
         self,
         mda: met_module.MetDataArray,
         *,
-        longitude: np.ndarray | None = None,
-        latitude: np.ndarray | None = None,
-        level: np.ndarray | None = None,
-        time: np.ndarray | None = None,
+        longitude: npt.NDArray[np.float_] | None = None,
+        latitude: npt.NDArray[np.float_] | None = None,
+        level: npt.NDArray[np.float_] | None = None,
+        time: npt.NDArray[np.datetime64] | None = None,
         use_indices: bool = False,
         **interp_kwargs: Any,
-    ) -> np.ndarray:
+    ) -> npt.NDArray[np.float_]:
         """Intersect waypoints with MetDataArray.
 
         Parameters
         ----------
         mda : MetDataArray
             MetDataArray containing a meteorological variable at spatio-temporal coordinates.
-        longitude : np.ndarray, optional
+        longitude : npt.NDArray[np.float_], optional
             Override existing coordinates for met interpolation
-        latitude : np.ndarray, optional
+        latitude : npt.NDArray[np.float_], optional
             Override existing coordinates for met interpolation
-        level : np.ndarray, optional
+        level : npt.NDArray[np.float_], optional
             Override existing coordinates for met interpolation
-        time : np.ndarray, optional
+        time : npt.NDArray[np.datetime64], optional
             Override existing coordinates for met interpolation
         use_indices : bool, optional
             Experimental.
@@ -1437,7 +1457,7 @@ class GeoVectorDataset(VectorDataset):
 
         Returns
         -------
-        np.ndarray
+        npt.NDArray[np.float_]
             Interpolated values
 
         Examples
@@ -1453,7 +1473,7 @@ class GeoVectorDataset(VectorDataset):
         >>> 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()
 
         >>> # Example flight
         >>> df = pd.DataFrame()
@@ -1550,9 +1570,7 @@ class GeoVectorDataset(VectorDataset):
         self["_distances_y"] = distances_y
         self["_distances_z"] = distances_z
         self["_distances_t"] = distances_t
-
-        if out_of_bounds is not None:
-            self["_out_of_bounds"] = out_of_bounds
+        self["_out_of_bounds"] = out_of_bounds
 
     def _get_indices(self) -> interpolation.RGIArtifacts | None:
         """Get entries from call to :meth:`_put_indices`.
@@ -1577,14 +1595,13 @@ class GeoVectorDataset(VectorDataset):
             distances_y = self["_distances_y"]
             distances_z = self["_distances_z"]
             distances_t = self["_distances_t"]
+            out_of_bounds = self["_out_of_bounds"]
         except KeyError:
             return None
 
         indices = np.asarray([indices_x, indices_y, indices_z, indices_t])
         distances = np.asarray([distances_x, distances_y, distances_z, distances_t])
 
-        out_of_bounds = self.get("_out_of_bounds", None)
-
         return interpolation.RGIArtifacts(indices, distances, out_of_bounds)
 
     def _invalidate_indices(self) -> None:

pycontrails/datalib/__init__.py

@@ -5,5 +5,5 @@ See individual modules for met variables and additional exports.
 
 - :module:`pycontrails.datalib.ecmwf`
 - :module:`pycontrails.datalib.gfs`
-
+- :module:`pycontrails.datalib.spire`
 """

pycontrails/datalib/ecmwf/common.py

@@ -77,21 +77,18 @@ class ECMWFAPI(datalib.MetDataSource):
             raise KeyError(f"Input dataset is missing variables {e}")
 
         # downselect times
-        try:
-            if self.timesteps:
+        if not self.timesteps:
+            self.timesteps = ds["time"].values.astype("datetime64[ns]").tolist()
+        else:
+            try:
                 ds = ds.sel(time=self.timesteps)
-            else:
-                # set timesteps from dataset "time" coordinates
-                # np.datetime64 doesn't covert to list[datetime] unless its unit is us
-                self.timesteps = ds["time"].values.astype("datetime64[us]").tolist()
-        except KeyError:
-            # this snippet shows the missing times for convenience
-            np_timesteps = [np.datetime64(t, "ns") for t in self.timesteps]
-            missing_times = list(set(np_timesteps) - set(ds["time"].values))
-            missing_times.sort()
-            raise KeyError(
-                f"Input dataset is missing time coordinates {[str(t) for t in missing_times]}"
-            )
+            except KeyError:
+                # this snippet shows the missing times for convenience
+                np_timesteps = [np.datetime64(t, "ns") for t in self.timesteps]
+                missing_times = sorted(set(np_timesteps) - set(ds["time"].values))
+                raise KeyError(
+                    f"Input dataset is missing time coordinates {[str(t) for t in missing_times]}"
+                )
 
         # downselect pressure level
         # if "level" is not in dims and
@@ -104,16 +101,12 @@ class ECMWFAPI(datalib.MetDataSource):
             ds = ds.sel(level=self.pressure_levels)
         except KeyError:
             # this snippet shows the missing levels for convenience
-            missing_levels = list(set(self.pressure_levels) - set(ds["level"].values))
-            missing_levels.sort()
+            missing_levels = sorted(set(self.pressure_levels) - set(ds["level"].values))
             raise KeyError(f"Input dataset is missing level coordinates {missing_levels}")
 
         # harmonize variable names
         ds = met.standardize_variables(ds, self.variables)
 
-        if "cachestore" not in kwargs:
-            kwargs["cachestore"] = self.cachestore
-
         # modify values
 
         # rescale relative humidity from % -> dimensionless if its in dataset
@@ -129,4 +122,6 @@ class ECMWFAPI(datalib.MetDataSource):
             ] = "Relative humidity rescaled to [0 - 1] instead of %"
 
         ds.attrs["met_source"] = type(self).__name__
+
+        kwargs.setdefault("cachestore", self.cachestore)
         return met.MetDataset(ds, **kwargs)

pycontrails/datalib/spire/__init__.py

@@ -0,0 +1,19 @@
+"""ECMWF Data Access."""
+
+from __future__ import annotations
+
+from pycontrails.datalib.spire.spire import (
+    clean,
+    generate_flight_id,
+    identify_flights,
+    is_valid_trajectory,
+    validate_flights,
+)
+
+__all__ = [
+    "clean",
+    "generate_flight_id",
+    "identify_flights",
+    "is_valid_trajectory",
+    "validate_flights",
+]