pycontrails 0.42.0__cp311-cp311-macosx_11_0_arm64.whl → 0.42.2__cp311-cp311-macosx_11_0_arm64.whl

pycontrails/core/flightplan.py

@@ -0,0 +1,238 @@
+"""ATC Flight Plan Parser."""
+
+import re
+from typing import Any
+
+
+def to_atc_plan(plan: dict[str, Any]) -> str:
+    """Write dictionary from :func:`parse_atc_plan` as ATC flight plan string.
+
+    Parameters
+    ----------
+    plan: dict[str, Any]
+        Dictionary representation of ATC flight plan returned from :func:`parse_atc_plan`.
+
+    Returns
+    -------
+    str
+        ATC flight plan string conforming to ICAO Doc 4444-ATM/501
+
+    See Also
+    --------
+    :func:`parse_atc_plan`
+    """
+    ret = f'(FPL-{plan["callsign"]}-{plan["flight_rules"]}'
+    ret += f'{plan["type_of_flight"]}\n'
+    ret += "-"
+    if "number_aircraft" in plan and plan["number_aircraft"] <= 10:
+        ret += plan["number_aircraft"]
+    ret += f'{plan["type_of_aircraft"]}/{plan["wake_category"]}-'
+    ret += f'{plan["equipment"]}/{plan["transponder"]}\n'
+    ret += f'-{plan["departure_icao"]}{plan["time"]}\n'
+    ret += f'-{plan["speed_type"]}{plan["speed"]}{plan["level_type"]}'
+    ret += f'{plan["level"]} {plan["route"]}\n'
+    if "destination_icao" in plan and "duration" in plan:
+        ret += f'-{plan["destination_icao"]}{plan["duration"]}'
+    if "alt_icao" in plan:
+        ret += f' {plan["alt_icao"]}'
+    if "second_alt_icao" in plan:
+        ret += f' {plan["second_alt_icao"]}'
+    ret += "\n"
+    ret += f'-{plan["other_info"]})\n'
+    if "supplementary_info" in plan:
+        ret += " ".join([f"{i[0]}/{i[1]}" for i in plan["supplementary_info"].items()])
+
+    if ret[-1] == "\n":
+        ret = ret[:-1]
+
+    return ret
+
+
+def parse_atc_plan(atc_plan: str) -> dict[str, str]:
+    """Parse an ATC flight plan string into a dictionary.
+
+    The route string is not converted to lat/lon in this process.
+
+    Parameters
+    ----------
+    atc_plan : str
+        An ATC flight plan string conforming to ICAO Doc 4444-ATM/501 (Appendix 2)
+
+    Returns
+    -------
+    dict[str, str]
+        A dictionary consisting of parsed components of the ATC flight plan.
+        A full ATC plan will contain the keys:
+
+        - ``callsign``: ICAO flight callsign
+        - ``flight_rules``: Flight rules ("I", "V", "Y", "Z")
+        - ``type_of_flight``: Type of flight ("S", "N", "G", "M", "X")
+        - ``number_aircraft``: The number of aircraft, if more than one
+        - ``type_of_aircraft``: ICAO aircraft type
+        - ``wake_category``: Wake turbulence category
+        - ``equipment``: Radiocommunication, navigation and approach aid equipment and capabilities
+        - ``transponder``: Surveillance equipment and capabilities
+        - ``departure_icao``: ICAO departure airport
+        - ``time``: Estimated off-block (departure) time (UTC)
+        - ``speed_type``: Speed units ("K": km / hr, "N": knots)
+        - ``speed``: Cruise true airspeed in ``speed_type`` units
+        - ``level_type``: Level units ("F", "S", "A", "M")
+        - ``level``: Cruise level
+        - ``route``: Route string
+        - ``destination_icao``: ICAO destination airport
+        - ``duration``: The total estimated elapsed time for the flight plan
+        - ``alt_icao``: ICAO alternate destination airport
+        - ``second_alt_icao``: ICAO second alternate destination airport
+        - ``other_info``: Other information
+        - ``supplementary_info``: Supplementary information
+
+    References
+    ----------
+    - https://applications.icao.int/tools/ATMiKIT/story_content/external_files/story_content/external_files/DOC%204444_PANS%20ATM_en.pdf
+
+    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()
+    atc_plan = atc_plan.strip()
+
+    if len(atc_plan) == 0:
+        raise ValueError("Empty or invalid flight plan")
+
+    atc_plan = atc_plan.replace("(FPL", "")
+    atc_plan = atc_plan.replace(")", "")
+    atc_plan = atc_plan.replace("--", "-")
+
+    basic = atc_plan.split("-")
+
+    flightplan: dict[str, Any] = {}
+
+    # Callsign
+    if len(basic) > 1:
+        flightplan["callsign"] = basic[1]
+
+    # Flight Rules
+    if len(basic) > 2:
+        flightplan["flight_rules"] = basic[2][0]
+        flightplan["type_of_flight"] = basic[2][1]
+
+    # Aircraft
+    if len(basic) > 3:
+        aircraft = basic[3].split("/")
+        matches = re.match("(\d{1})(\S{3,4})", aircraft[0])
+        if matches:
+            groups = matches.groups()
+        else:
+            groups = ()
+
+        if matches and len(groups) > 2:
+            flightplan["number_aircraft"] = groups[1]
+            flightplan["type_of_aircraft"] = groups[2]
+        else:
+            flightplan["type_of_aircraft"] = aircraft[0]
+
+        if len(aircraft) > 1:
+            flightplan["wake_category"] = aircraft[1]
+
+    # Equipment
+    if len(basic) > 4:
+        equip = basic[4].split("/")
+        flightplan["equipment"] = equip[0]
+        if len(equip) > 1:
+            flightplan["transponder"] = equip[1]
+
+    # Dep. airport info
+    if len(basic) > 5:
+        matches = re.match("(\D*)(\d*)", basic[5])
+        if matches:
+            groups = matches.groups()
+        else:
+            groups = ()
+
+        if len(groups) > 0:
+            flightplan["departure_icao"] = groups[0]
+        if len(groups) > 1:
+            flightplan["time"] = groups[1]
+
+    # Speed and route info
+    if len(basic) > 6:
+        matches = re.match("(\D*)(\d*)(\D*)(\d*)", basic[6])
+        if matches:
+            groups = matches.groups()
+        else:
+            groups = ()
+
+        # match speed and level
+        if len(groups) > 0:
+            flightplan["speed_type"] = groups[0]
+            if len(groups) > 1:
+                flightplan["speed"] = groups[1]
+            if len(groups) > 2:
+                flightplan["level_type"] = groups[2]
+            if len(groups) > 3:
+                flightplan["level"] = groups[3]
+
+            flightplan["route"] = basic[6][len("".join(groups)) :].strip()
+        else:
+            flightplan["route"] = basic[6].strip()
+
+    # Dest. airport info
+    if len(basic) > 7:
+        matches = re.match("(\D{4})(\d{4})", basic[7])
+        if matches:
+            groups = matches.groups()
+        else:
+            groups = ()
+
+        if len(groups) > 0:
+            flightplan["destination_icao"] = groups[0]
+        if len(groups) > 1:
+            flightplan["duration"] = groups[1]
+
+        matches = re.match("(\D{4})(\d{4})(\s{1})(\D{4})", basic[7])
+        if matches:
+            groups = matches.groups()
+        else:
+            groups = ()
+
+        if len(groups) > 3:
+            flightplan["alt_icao"] = groups[3]
+
+        matches = re.match("(\D{4})(\d{4})(\s{1})(\D{4})(\s{1})(\D{4})", basic[7])
+        if matches:
+            groups = matches.groups()
+        else:
+            groups = ()
+
+        if len(groups) > 5:
+            flightplan["second_alt_icao"] = groups[5]
+
+    # Other info
+    if len(basic) > 8:
+        flightplan["other_info"] = basic[8]
+
+    # Supl. Info
+    if len(basic) > 9:
+        sup_match = re.findall("(\D{1}[\/]{1})", basic[9])
+        if len(sup_match) > 0:
+            suplInfo = {}
+            for i in range(len(sup_match) - 1):
+                this_key = sup_match[i]
+                this_idx = basic[9].find(this_key)
+
+                next_key = sup_match[i + 1]
+                next_idx = basic[9].find(next_key)
+
+                val = basic[9][this_idx + 2 : next_idx - 1]
+                suplInfo[this_key[0]] = val
+
+            last_key = sup_match[-1]
+            last_idx = basic[9].find(last_key)
+            suplInfo[last_key[0]] = basic[9][last_idx + 2 :]
+
+            flightplan["supplementary_info"] = suplInfo
+
+    return flightplan

pycontrails/core/interpolation.py

@@ -74,7 +74,7 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         self.bounds_error = bounds_error
         self.fill_value = fill_value
 
-    def _prepare_xi_simple(self, xi: npt.NDArray[np.float64]) -> npt.NDArray[np.bool_] | None:
+    def _prepare_xi_simple(self, xi: npt.NDArray[np.float64]) -> npt.NDArray[np.bool_]:
         """Run looser version of :meth:`_prepare_xi`.
 
         Parameters
@@ -84,12 +84,9 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
 
         Returns
         -------
-        npt.NDArray[np.bool_] | None
+        npt.NDArray[np.bool_]
             A 1-dimensional Boolean array indicating which points are out of bounds.
-            If ``bounds_error`` is ``True``, this will be ``None``, indicating that
-            no points are out of bounds. (This is the same convention as
-            :meth:`scipy.interpolate.RegularGridInterpolator._prepare_xi`). If
-            every point is in bounds, this is set to ``None``.
+            If ``bounds_error`` is ``True``, this will be all ``False``.
         """
 
         if self.bounds_error:
@@ -99,12 +96,9 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
                 if not (np.all(p >= g0) and np.all(p <= g1)):
                     raise ValueError(f"One of the requested xi is out of bounds in dimension {i}")
 
-            return None
+            return np.zeros(xi.shape[0], dtype=bool)
 
-        out_of_bounds = self._find_out_of_bounds(xi.T)
-        if not np.any(out_of_bounds):
-            return None
-        return out_of_bounds
+        return self._find_out_of_bounds(xi.T)
 
     def __call__(
         self, xi: npt.NDArray[np.float64], method: str | None = None
@@ -137,7 +131,9 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         return self._set_out_of_bounds(out, out_of_bounds)
 
     def _set_out_of_bounds(
-        self, out: npt.NDArray[np.float_], out_of_bounds: npt.NDArray[np.bool_] | None
+        self,
+        out: npt.NDArray[np.float_],
+        out_of_bounds: npt.NDArray[np.bool_],
     ) -> npt.NDArray[np.float_]:
         """Set out-of-bounds values to the fill value.
 
@@ -145,7 +141,7 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         ----------
         out : npt.NDArray[np.float_]
             Values from interpolation. This is modified in-place.
-        out_of_bounds : npt.NDArray[np.bool_] | None
+        out_of_bounds : npt.NDArray[np.bool_]
             A 1-dimensional Boolean array indicating which points are out of bounds.
 
         Returns
@@ -153,7 +149,7 @@ class PycontrailsRegularGridInterpolator(scipy.interpolate.RegularGridInterpolat
         out : npt.NDArray[np.float_]
             A reference to the ``out`` array.
         """
-        if out_of_bounds is not None and self.fill_value is not None:
+        if self.fill_value is not None and np.any(out_of_bounds):
             out[out_of_bounds] = self.fill_value
 
         return out
@@ -557,7 +553,7 @@ class RGIArtifacts:
 
     xi_indices: npt.NDArray[np.int64]
     norm_distances: npt.NDArray[np.float64]
-    out_of_bounds: npt.NDArray[np.bool_] | None
+    out_of_bounds: npt.NDArray[np.bool_]
 
 
 # ------------------------------------------------------------------------------

pycontrails/core/met.py

@@ -416,7 +416,7 @@ class MetBase(ABC, Generic[XArrayType]):
         >>> variables = "air_temperature", "specific_humidity"
         >>> levels = [200, 300]
         >>> era5 = ERA5(times, variables, levels)
-        >>> mds = era5.open_metdataset(xr_kwargs=dict(parallel=False))
+        >>> mds = era5.open_metdataset()
         >>> mds.variables["level"].values  # faster access than mds.data["level"]
         array([200., 300.])
 
@@ -606,7 +606,7 @@ class MetDataset(MetBase):
     >>> era5 = ERA5(time, variables, pressure_levels)
 
     >>> # Open directly as `MetDataset`
-    >>> met = era5.open_metdataset(xr_kwargs=dict(parallel=False))
+    >>> met = era5.open_metdataset()
     >>> # Use `data` attribute to access `xarray` object
     >>> assert isinstance(met.data, xr.Dataset)
 
@@ -992,7 +992,7 @@ class MetDataset(MetBase):
         >>> variables = ["air_temperature", "specific_humidity"]
         >>> levels = [250, 200]
         >>> era5 = ERA5(time=times, variables=variables, pressure_levels=levels)
-        >>> met = era5.open_metdataset(xr_kwargs=dict(parallel=False))
+        >>> met = era5.open_metdataset()
         >>> met.to_vector(transfer_attrs=False)
         GeoVectorDataset [6 keys x 4152960 length, 1 attributes]
             Keys: longitude, latitude, level, time, air_temperature, ..., specific_humidity
@@ -1507,7 +1507,7 @@ class MetDataArray(MetBase):
         >>> variables = "air_temperature"
         >>> levels = [200, 250, 300]
         >>> era5 = ERA5(times, variables, levels)
-        >>> met = era5.open_metdataset(xr_kwargs=dict(parallel=False))
+        >>> met = era5.open_metdataset()
         >>> mda = met["air_temperature"]
 
         >>> # Interpolation at a grid point agrees with value
@@ -1779,7 +1779,7 @@ class MetDataArray(MetBase):
         >>> from pprint import pprint
         >>> from pycontrails.datalib.ecmwf import ERA5
         >>> era5 = ERA5("2022-03-01", variables="air_temperature", pressure_levels=250)
-        >>> mda = era5.open_metdataset(xr_kwargs=dict(parallel=False))["air_temperature"]
+        >>> mda = era5.open_metdataset()["air_temperature"]
         >>> mda.shape
         (1440, 721, 1, 1)
 

pycontrails/core/models.py

@@ -566,6 +566,10 @@ class Model(ABC):
             try:
                 # This case is when self.source is a subgrid of self.met
                 # The call to .sel will raise a KeyError if this is not the case
+
+                # XXX: Sometimes this hangs when using dask!
+                # This issue is somewhat similar to
+                # https://github.com/pydata/xarray/issues/4406
                 self.source[met_key] = da.sel(self.source.coords)
 
             except KeyError:

pycontrails/core/vector.py

@@ -287,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."
                         )
@@ -305,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
@@ -1465,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()
@@ -1562,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`.
@@ -1589,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/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/ext/bada/__init__.py

@@ -20,6 +20,12 @@ try:
         BADAParams,
     )
 
+except ImportError as e:
+    raise ImportError(
+        'Failed to import `pycontrails-bada` extension. Install with `pip install "pycontrails-bada'
+        ' @ git+ssh://git@github.com/contrailcirrus/pycontrails-bada.git"`'
+    ) from e
+else:
     __all__ = [
         "BADA",
         "BADA3",
@@ -33,9 +39,3 @@ try:
         "bada4",
         "bada_model",
     ]
-
-except ImportError as e:
-    raise ImportError(
-        'Failed to import `pycontrails-bada` extension. Install with `pip install "pycontrails-bada'
-        ' @ git+ssh://git@github.com/contrailcirrus/pycontrails-bada.git"`'
-    ) from e

pycontrails/ext/cirium/__init__.py

@@ -5,10 +5,10 @@ from __future__ import annotations
 try:
     from pycontrails_cirium import Cirium
 
-    __all__ = ["Cirium"]
-
 except ImportError as e:
     raise ImportError(
         "Failed to import `pycontrails-cirium` extension. Install with `pip install"
         ' "pycontrails-cirium @ git+ssh://git@github.com/contrailcirrus/pycontrails-cirium.git"`'
     ) from e
+else:
+    __all__ = ["Cirium"]