pycontrails 0.50.1__cp39-cp39-win_amd64.whl → 0.51.0__cp39-cp39-win_amd64.whl

pycontrails/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.50.1'
-__version_tuple__ = version_tuple = (0, 50, 1)
+__version__ = version = '0.51.0'
+__version_tuple__ = version_tuple = (0, 51, 0)

pycontrails/core/datalib.py

@@ -103,6 +103,28 @@ def parse_timesteps(time: TimeInput | None, freq: str | None = "1h") -> list[dat
     return daterange.to_pydatetime().tolist()
 
 
+def validate_timestep_freq(freq: str, datasource_freq: str) -> bool:
+    """Check that input timestep frequency is compatible with the data source timestep frequency.
+
+    A data source timestep frequency of 1 hour allows input timestep frequencies of
+    1 hour, 2 hours, 3 hours, etc., but not 1.5 hours or 30 minutes.
+
+    Parameters
+    ----------
+    freq : str
+        Input timestep frequency
+    datasource_freq : str
+        Datasource timestep frequency
+
+    Returns
+    -------
+    bool
+        True if the input timestep frequency is an even multiple
+        of the data source timestep frequency.
+    """
+    return pd.Timedelta(freq) % pd.Timedelta(datasource_freq) == pd.Timedelta(0)
+
+
 def parse_pressure_levels(
     pressure_levels: PressureLevelInput, supported: list[int] | None = None
 ) -> list[int]:

pycontrails/core/flight.py

@@ -17,6 +17,7 @@ from pycontrails.core.fuel import Fuel, JetA
 from pycontrails.core.vector import AttrDict, GeoVectorDataset, VectorDataDict, VectorDataset
 from pycontrails.physics import constants, geo, units
 from pycontrails.utils import dependencies
+from pycontrails.utils.types import ArrayOrFloat
 
 logger = logging.getLogger(__name__)
 
@@ -446,6 +447,48 @@ class Flight(GeoVectorDataset):
 
         return segment_duration(self.data["time"], dtype=dtype)
 
+    def segment_haversine(self) -> npt.NDArray[np.float64]:
+        """Compute Haversine (great circle) distance between flight waypoints.
+
+        Helper function used in :meth:`resample_and_fill`.
+        `np.nan` appended so the length of the output is the same as number of waypoints.
+
+        To account for vertical displacements when computing segment lengths,
+        use :meth:`segment_length`.
+
+        Returns
+        -------
+        npt.NDArray[np.float64]
+            Array of great circle distances in [:math:`m`] between waypoints
+
+        Raises
+        ------
+        NotImplementedError
+            Raises when attr:`attrs["crs"]` is not EPSG:4326
+
+        Examples
+        --------
+        >>> from pycontrails import Flight
+        >>> fl = Flight(
+        ... longitude=np.array([1, 2, 3, 5, 8]),
+        ... latitude=np.arange(5),
+        ... altitude=np.full(shape=(5,), fill_value=11000),
+        ... time=pd.date_range('2021-01-01T12', '2021-01-01T14', periods=5),
+        ... )
+        >>> fl.segment_haversine()
+        array([157255.03346286, 157231.08336815, 248456.48781503, 351047.44358851,
+                           nan])
+
+        See Also
+        --------
+        :func:`segment_haversine`
+        :meth:`segment_length`
+        """
+        if self.attrs["crs"] != "EPSG:4326":
+            raise NotImplementedError("Only implemented for EPSG:4326 CRS.")
+
+        return geo.segment_haversine(self["longitude"], self["latitude"])
+
     def segment_length(self) -> npt.NDArray[np.float64]:
         """Compute spherical distance between flight waypoints.
 
@@ -1112,6 +1155,102 @@ class Flight(GeoVectorDataset):
         out.data.pop("level", None)  # avoid any ambiguity
         return out
 
+    def distance_to_coords(self: Flight, distance: ArrayOrFloat) -> tuple[
+        ArrayOrFloat,
+        ArrayOrFloat,
+        np.intp | npt.NDArray[np.intp],
+    ]:
+        """
+        Convert distance along flight path to geodesic coordinates.
+
+        Will return a tuple containing `(lat, lon, index)`, where index indicates which flight
+        segment contains the returned coordinate.
+
+        Parameters
+        ----------
+        distance : ArrayOrFloat
+            Distance along flight path, [:math:`m`]
+
+        Returns
+        -------
+        (ArrayOrFloat, ArrayOrFloat, int | npt.NDArray[int])
+            latitude, longitude, and segment index cooresponding to distance.
+        """
+
+        # Check if flight crosses antimeridian line
+        lon_ = self["longitude"]
+        lat_ = self["latitude"]
+        sign_ = np.sign(lon_)
+        min_pos = np.min(lon_[sign_ == 1.0], initial=np.inf)
+        max_neg = np.max(lon_[sign_ == -1.0], initial=-np.inf)
+
+        if (180.0 - min_pos) + (180.0 + max_neg) < 180.0 and min_pos < np.inf and max_neg > -np.inf:
+            # In this case, we believe the flight crosses the antimeridian
+            shift = min_pos
+            # So we shift the longitude "chart"
+            lon_ = (lon_ - shift) % 360.0
+        else:
+            shift = None
+
+        # Make a fake flight that flies at constant height so distance is just
+        # distance traveled across groud
+        flat_dataset = Flight(
+            longitude=self.coords["longitude"],
+            latitude=self.coords["latitude"],
+            time=self.coords["time"],
+            level=[self.coords["level"][0] for _ in range(self.size)],
+        )
+
+        lengths = flat_dataset.segment_length()
+        cumulative_lengths = np.nancumsum(lengths)
+        cumulative_lengths = np.insert(cumulative_lengths[:-1], 0, 0)
+        seg_idx: np.intp | npt.NDArray[np.intp]
+
+        if isinstance(distance, float):
+            seg_idx = np.argmax(cumulative_lengths > distance)
+        else:
+            seg_idx = np.argmax(cumulative_lengths > distance.reshape((distance.size, 1)), axis=1)
+
+        # If in the last segment (which has length 0), then just return the last waypoint
+        seg_idx -= 1
+
+        # linear interpolation in lat/lon - assuming the way points are within 100-200km so this
+        # should be accurate enough without needed to reproject or use spherical distance
+        lat1: ArrayOrFloat = lat_[seg_idx]
+        lon1: ArrayOrFloat = lon_[seg_idx]
+        lat2: ArrayOrFloat = lat_[seg_idx + 1]
+        lon2: ArrayOrFloat = lon_[seg_idx + 1]
+
+        dx = distance - cumulative_lengths[seg_idx]
+        fx = dx / lengths[seg_idx]
+        lat = (1 - fx) * lat1 + fx * lat2
+        lon = (1 - fx) * lon1 + fx * lon2
+
+        if isinstance(distance, float):
+            if distance < 0:
+                lat = np.nan
+                lon = np.nan
+                seg_idx = np.intp(0)
+            elif distance >= cumulative_lengths[-1]:
+                lat = lat_[-1]
+                lon = lon_[-1]
+                seg_idx = np.intp(self.size - 1)
+        else:
+            lat[distance < 0] = np.nan
+            lon[distance < 0] = np.nan
+            seg_idx[distance < 0] = 0  # type: ignore
+
+            lat[distance >= cumulative_lengths[-1]] = lat_[-1]
+            lon[distance >= cumulative_lengths[-1]] = lon_[-1]
+            seg_idx[distance >= cumulative_lengths[-1]] = self.size - 1  # type: ignore
+
+        if shift is not None:
+            # We need to translate back to the original chart here
+            lon += shift
+            lon = ((lon + 180.0) % 360.0) - 180.0
+
+        return lat, lon, seg_idx
+
     def _geodesic_interpolation(self, geodesic_threshold: float) -> pd.DataFrame | None:
         """Geodesic interpolate between large gaps between waypoints.
 
@@ -1135,7 +1274,7 @@ class Flight(GeoVectorDataset):
             raise NotImplementedError("Only implemented for EPSG:4326 CRS.")
 
         # Omit the final nan and ensure index + 1 (below) is well defined
-        segs = self.segment_length()[:-1]
+        segs = self.segment_haversine()[:-1]
 
         # For default geodesic_threshold, we expect gap_indices to be very
         # sparse (so the for loop below is cheap)

pycontrails/core/met.py

@@ -1793,10 +1793,11 @@ class MetDataArray(MetBase):
         self,
         level: float | int | None = None,
         time: np.datetime64 | datetime | None = None,
-        fill_value: float = 0.0,
+        fill_value: float = np.nan,
         iso_value: float | None = None,
         min_area: float = 0.0,
         epsilon: float = 0.0,
+        lower_bound: bool = True,
         precision: int | None = None,
         interiors: bool = True,
         convex_hull: bool = False,
@@ -1848,7 +1849,8 @@ class MetDataArray(MetBase):
             automatically.
         fill_value : float, optional
             Value used for filling missing data and for padding the underlying data array.
-            Expected to be less than the ``iso_value`` parameter. By default, 0.0.
+            Set to ``np.nan`` by default, which ensures that regions with missing data are
+            never included in polygons.
         iso_value : float, optional
             Value in field to create iso-surface.
             Defaults to the average of the min and max value of the array. (This is the
@@ -1862,6 +1864,9 @@ class MetDataArray(MetBase):
             Control the extent to which the polygon is simplified. A value of 0 does not alter
             the geometry of the polygon. The unit of this parameter is in longitude/latitude
             degrees. By default, 0.0.
+        lower_bound : bool, optional
+            Whether to use ``iso_value`` as a lower or upper bound on values in polygon interiors.
+            By default, True.
         precision : int, optional
             Number of decimal places to round coordinates to. If None, no rounding is performed.
         interiors : bool, optional
@@ -1882,6 +1887,14 @@ class MetDataArray(MetBase):
         dict[str, Any]
             Python representation of GeoJSON Feature with MultiPolygon geometry.
 
+        Notes
+        -----
+        :class:`Cocip` and :class:`CocipGrid` set some quantities to 0 and other quantities
+        to ``np.nan`` in regions where no contrails form. When computing polygons from
+        :class:`Cocip` or :class:`CocipGrid` output, take care that the choice of
+        ``fill_value`` correctly includes or excludes contrail-free regions. See the
+        :class:`Cocip` documentation for details about ``np.nan`` in model output.
+
         See Also
         --------
         :meth:`to_polyhedra`
@@ -1925,11 +1938,12 @@ class MetDataArray(MetBase):
                 "include_altitude=False."
             )
 
-        np.nan_to_num(arr, copy=False, nan=fill_value)
+        if not np.isnan(fill_value):
+            np.nan_to_num(arr, copy=False, nan=fill_value)
 
         # default iso_value
         if iso_value is None:
-            iso_value = (np.max(arr) + np.min(arr)) / 2
+            iso_value = (np.nanmax(arr) + np.nanmin(arr)) / 2
             warnings.warn(f"The 'iso_value' parameter was not specified. Using value: {iso_value}")
 
         # We'll get a nice error message if dependencies are not installed
@@ -1945,6 +1959,7 @@ class MetDataArray(MetBase):
             threshold=iso_value,
             min_area=min_area,
             epsilon=epsilon,
+            lower_bound=lower_bound,
             interiors=interiors,
             convex_hull=convex_hull,
             longitude=longitude,
@@ -1957,10 +1972,11 @@ class MetDataArray(MetBase):
     def to_polygon_feature_collection(
         self,
         time: np.datetime64 | datetime | None = None,
-        fill_value: float = 0.0,
+        fill_value: float = np.nan,
         iso_value: float | None = None,
         min_area: float = 0.0,
         epsilon: float = 0.0,
+        lower_bound: bool = True,
         precision: int | None = None,
         interiors: bool = True,
         convex_hull: bool = False,
@@ -1991,6 +2007,7 @@ class MetDataArray(MetBase):
                 iso_value=iso_value,
                 min_area=min_area,
                 epsilon=epsilon,
+                lower_bound=lower_bound,
                 precision=precision,
                 interiors=interiors,
                 convex_hull=convex_hull,
@@ -2011,6 +2028,7 @@ class MetDataArray(MetBase):
         time: datetime | None = ...,
         iso_value: float = ...,
         simplify_fraction: float = ...,
+        lower_bound: bool = ...,
         return_type: Literal["geojson"],
         path: str | None = ...,
         altitude_scale: float = ...,
@@ -2025,6 +2043,7 @@ class MetDataArray(MetBase):
         time: datetime | None = ...,
         iso_value: float = ...,
         simplify_fraction: float = ...,
+        lower_bound: bool = ...,
         return_type: Literal["mesh"],
         path: str | None = ...,
         altitude_scale: float = ...,
@@ -2038,6 +2057,7 @@ class MetDataArray(MetBase):
         time: datetime | None = None,
         iso_value: float = 0.0,
         simplify_fraction: float = 1.0,
+        lower_bound: bool = True,
         return_type: str = "geojson",
         path: str | None = None,
         altitude_scale: float = 1.0,
@@ -2058,6 +2078,9 @@ class MetDataArray(MetBase):
             Apply `open3d` `simplify_quadric_decimation` method to simplify the polyhedra geometry.
             This parameter must be in the half-open interval (0.0, 1.0].
             Defaults to 1.0, corresponding to no reduction.
+        lower_bound : bool, optional
+            Whether to use ``iso_value`` as a lower or upper bound on values in polyhedra interiors.
+            By default, True.
         return_type : str, optional
             Must be one of "geojson" or "mesh". Defaults to "geojson".
             If "geojson", this method returns a dictionary representation of a geojson MultiPolygon
@@ -2145,6 +2168,11 @@ class MetDataArray(MetBase):
         # 3d array of longitude, latitude, altitude values
         volume = self.data.sel(time=time).values
 
+        # invert if iso_value is an upper bound on interior values
+        if not lower_bound:
+            volume = -volume
+            iso_value = -iso_value
+
         # convert from array index back to coordinates
         longitude = self.indexes["longitude"].values
         latitude = self.indexes["latitude"].values

pycontrails/core/polygon.py

@@ -283,6 +283,7 @@ def find_multipolygon(
     threshold: float,
     min_area: float,
     epsilon: float,
+    lower_bound: bool = True,
     interiors: bool = True,
     convex_hull: bool = False,
     longitude: npt.NDArray[np.float64] | None = None,
@@ -304,10 +305,13 @@ def find_multipolygon(
     epsilon : float
         Epsilon value to use when simplifying the polygons. Passed into shapely's
         :meth:`shapely.geometry.Polygon.simplify` method.
+    lower_bound : bool, optional
+        Whether to treat ``threshold`` as a lower or upper bound on values in polygon interiors.
+        By default, True.
     interiors : bool, optional
-        Whether to include interior polygons.
+        Whether to include interior polygons. By default, True.
     convex_hull : bool, optional
-        Experimental. Whether to take the convex hull of each polygon.
+        Experimental. Whether to take the convex hull of each polygon. By default, False.
     longitude : npt.NDArray[np.float64] | None, optional
         If provided, the coordinates values corresponding to the longitude dimensions of ``arr``.
         The contour coordinates will be converted to longitude-latitude values by indexing
@@ -339,7 +343,10 @@ def find_multipolygon(
         buffer = 0.5
 
     arr_bin = np.empty(arr.shape, dtype=np.uint8)
-    np.greater_equal(arr, threshold, out=arr_bin)
+    if lower_bound:
+        np.greater_equal(arr, threshold, out=arr_bin)
+    else:
+        np.less_equal(arr, threshold, out=arr_bin)
 
     mode = cv2.RETR_CCOMP if interiors else cv2.RETR_EXTERNAL
     contours, hierarchy = cv2.findContours(arr_bin, mode, cv2.CHAIN_APPROX_SIMPLE)

pycontrails/datalib/ecmwf/__init__.py

@@ -4,10 +4,13 @@ from __future__ import annotations
 
 from pycontrails.datalib.ecmwf.arco_era5 import ARCOERA5
 from pycontrails.datalib.ecmwf.era5 import ERA5
+from pycontrails.datalib.ecmwf.era5_model_level import ERA5ModelLevel
 from pycontrails.datalib.ecmwf.hres import HRES
+from pycontrails.datalib.ecmwf.hres_model_level import HRESModelLevel
 from pycontrails.datalib.ecmwf.ifs import IFS
 from pycontrails.datalib.ecmwf.variables import (
     ECMWF_VARIABLES,
+    MODEL_LEVEL_VARIABLES,
     PRESSURE_LEVEL_VARIABLES,
     SURFACE_VARIABLES,
     CloudAreaFraction,
@@ -27,7 +30,9 @@ from pycontrails.datalib.ecmwf.variables import (
 __all__ = [
     "ARCOERA5",
     "ERA5",
+    "ERA5ModelLevel",
     "HRES",
+    "HRESModelLevel",
     "IFS",
     "CloudAreaFraction",
     "CloudAreaFractionInLayer",
@@ -44,4 +49,5 @@ __all__ = [
     "ECMWF_VARIABLES",
     "PRESSURE_LEVEL_VARIABLES",
     "SURFACE_VARIABLES",
+    "MODEL_LEVEL_VARIABLES",
 ]

pycontrails/datalib/ecmwf/arco_era5.py

@@ -11,7 +11,6 @@ This module supports:
 This module requires the following additional dependencies:
 
 - `metview (binaries and python bindings) <https://metview.readthedocs.io/en/latest/python.html>`_
-- `lxml <https://lxml.de/>`_
 - `gcsfs <https://gcsfs.readthedocs.io/en/latest/>`_
 - `zarr <https://zarr.readthedocs.io/en/stable/>`_
 
@@ -22,7 +21,6 @@ from __future__ import annotations
 import contextlib
 import dataclasses
 import datetime
-import functools
 import hashlib
 import multiprocessing
 import pathlib
@@ -31,7 +29,6 @@ import warnings
 from collections.abc import Iterable
 from typing import Any
 
-import pandas as pd
 import xarray as xr
 from overrides import overrides
 
@@ -39,7 +36,7 @@ from pycontrails.core import cache, datalib, met_var
 from pycontrails.core.met import MetDataset
 from pycontrails.datalib.ecmwf import common as ecmwf_common
 from pycontrails.datalib.ecmwf import variables as ecmwf_variables
-from pycontrails.physics import units
+from pycontrails.datalib.ecmwf.model_levels import pressure_levels_at_model_levels
 from pycontrails.utils import dependencies
 
 try:
@@ -76,54 +73,6 @@ MOISTURE_STORE_VARIABLES = [
 PRESSURE_LEVEL_VARIABLES = [*WIND_STORE_VARIABLES, *MOISTURE_STORE_VARIABLES, met_var.Geopotential]
 
 
-@functools.cache
-def _read_model_level_dataframe() -> pd.DataFrame:
-    """Read the ERA5 model level definitions published by ECMWF.
-
-    This requires the lxml package to be installed.
-    """
-    url = "https://confluence.ecmwf.int/display/UDOC/L137+model+level+definitions"
-    try:
-        return pd.read_html(url, na_values="-", index_col="n")[0]
-    except ImportError as exc:
-        if "lxml" in exc.msg:
-            dependencies.raise_module_not_found_error(
-                "arco_era5._read_model_level_dataframe function",
-                package_name="lxml",
-                module_not_found_error=exc,
-                extra=(
-                    "Alternatively, if instantiating an 'ARCOERA5' object, you can provide "
-                    "the 'pressure_levels' parameter directly to avoid the need to read the "
-                    "ECMWF model level definitions."
-                ),
-            )
-        raise
-
-
-def pressure_levels_at_model_levels(alt_ft_min: float, alt_ft_max: float) -> list[int]:
-    """Return the pressure levels at each model level assuming a constant surface pressure.
-
-    The pressure levels are rounded to the nearest hPa.
-
-    Parameters
-    ----------
-    alt_ft_min : float
-        Minimum altitude, [:math:`ft`].
-    alt_ft_max : float
-        Maximum altitude, [:math:`ft`].
-
-    Returns
-    -------
-    list[int]
-        List of pressure levels, [:math:`hPa`].
-    """
-    df = _read_model_level_dataframe()
-    alt_m_min = units.ft_to_m(alt_ft_min)
-    alt_m_max = units.ft_to_m(alt_ft_max)
-    filt = df["Geometric Altitude [m]"].between(alt_m_min, alt_m_max)
-    return df.loc[filt, "pf [hPa]"].round().astype(int).tolist()
-
-
 def _attribute_fix(ds: xr.Dataset | None) -> None:
     """Fix GRIB attributes.
 
@@ -444,7 +393,7 @@ class ARCOERA5(ecmwf_common.ECMWFAPI):
         self.variables = datalib.parse_variables(variables, self.supported_variables)
         self.grid = grid
         self.cachestore = cache.DiskCacheStore() if cachestore is self.__marker else cachestore
-        self.n_jobs = n_jobs
+        self.n_jobs = max(1, n_jobs)
         self.cleanup_metview_tempfiles = cleanup_metview_tempfiles
 
     @property

pycontrails/datalib/ecmwf/common.py

@@ -102,3 +102,7 @@ class ECMWFAPI(datalib.MetDataSource):
                 os.remove(cache_path)
 
             ds_t.to_netcdf(cache_path)
+
+
+class CDSCredentialsNotFound(Exception):
+    """Raise when CDS credentials are not found by :class:`cdsapi.Client` instance."""

pycontrails/datalib/ecmwf/era5.py

@@ -21,7 +21,7 @@ from overrides import overrides
 import pycontrails
 from pycontrails.core import cache, datalib
 from pycontrails.core.met import MetDataset, MetVariable
-from pycontrails.datalib.ecmwf.common import ECMWFAPI
+from pycontrails.datalib.ecmwf.common import ECMWFAPI, CDSCredentialsNotFound
 from pycontrails.datalib.ecmwf.variables import PRESSURE_LEVEL_VARIABLES, SURFACE_VARIABLES
 from pycontrails.utils import dependencies, temp
 
@@ -183,7 +183,7 @@ class ERA5(ECMWFAPI):
             grid_min = 0.25 if product_type == "reanalysis" else 0.5
             if grid < grid_min:
                 warnings.warn(
-                    f"The smallest resolution available through the CDS API is {grid_min} degrees. "
+                    f"The highest resolution available through the CDS API is {grid_min} degrees. "
                     f"Your downloaded data will have resolution {grid}, but it is a "
                     f"reinterpolation of the {grid_min} degree data. The same interpolation can be "
                     "achieved directly with xarray."
@@ -535,7 +535,3 @@ class ERA5(ECMWFAPI):
 
         ds.attrs["pycontrails_version"] = pycontrails.__version__
         return ds
-
-
-class CDSCredentialsNotFound(Exception):
-    """Raise when CDS credentials are not found by :class:`cdsapi.Client` instance."""