pycontrails 0.49.4__cp310-cp310-macosx_11_0_arm64.whl → 0.50.0__cp310-cp310-macosx_11_0_arm64.whl

pycontrails/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.49.4'
-__version_tuple__ = version_tuple = (0, 49, 4)
+__version__ = version = '0.50.0'
+__version_tuple__ = version_tuple = (0, 50, 0)

pycontrails/core/datalib.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import abc
-import dataclasses
 import hashlib
 import logging
 import pathlib
@@ -109,6 +108,11 @@ def parse_pressure_levels(
 ) -> list[int]:
     """Check input pressure levels are consistent type and ensure levels exist in ECMWF data source.
 
+    .. versionchanged:: 0.50.0
+
+        The returned pressure levels are now sorted. Pressure levels must be unique.
+        Raises ValueError if pressure levels have mixed signs.
+
     Parameters
     ----------
     pressure_levels : PressureLevelInput
@@ -127,18 +131,31 @@ def parse_pressure_levels(
     ValueError
         Raises ValueError if pressure level is not supported by ECMWF data source
     """
-    # ensure pressure_levels is list-like
+    # Ensure pressure_levels is array-like
     if isinstance(pressure_levels, (int, float)):
         pressure_levels = [pressure_levels]
 
-    # Cast array-like to list of ints
-    out = np.asarray(pressure_levels, dtype=int).tolist()
+    # Cast array-like to int dtype and sort
+    arr = np.asarray(pressure_levels, dtype=int)
+    arr.sort()
 
-    # ensure pressure levels are valid
-    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)
+    # If any values are non-positive, the entire array should be [-1]
+    if np.any(arr <= 0) and not np.array_equal(arr, [-1]):
+        msg = f"Pressure levels must be all positive or all -1, got {arr}"
+        raise ValueError(msg)
+
+    # Ensure pressure levels are unique
+    if np.any(np.diff(arr) == 0):
+        msg = f"Pressure levels must be unique, got {arr}"
+        raise ValueError(msg)
+
+    out = arr.tolist()
+    if supported is None:
+        return out
+
+    if missing := set(out).difference(supported):
+        msg = f"Pressure levels {sorted(missing)} are not supported. Supported levels: {supported}"
+        raise ValueError(msg)
 
     return out
 
@@ -146,6 +163,11 @@ def parse_pressure_levels(
 def parse_variables(variables: VariableInput, supported: list[MetVariable]) -> list[MetVariable]:
     """Parse input variables.
 
+    .. versionchanged:: 0.50.0
+
+        The output is no longer copied. Each :class:`MetVariable` is a frozen dataclass,
+        so copying is unnecessary.
+
     Parameters
     ----------
     variables : VariableInput
@@ -178,35 +200,31 @@ def parse_variables(variables: VariableInput, supported: list[MetVariable]) -> l
     else:
         parsed_variables = variables
 
-    # 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 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}
+    supported_set = set(supported)
 
     for var in parsed_variables:
         matched = _find_match(
             var,
-            supported,
+            supported_set,
             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))
+        met_var_list.append(matched)
 
     return met_var_list
 
 
 def _find_match(
     var: VariableInput,
-    supported: list[MetVariable],
+    supported: set[MetVariable],
     ecmwf_ids: dict[int, MetVariable],
     grib1_ids: dict[int, MetVariable],
     short_names: dict[str, MetVariable],
@@ -215,9 +233,8 @@ def _find_match(
 ) -> MetVariable:
     """Find a match for input variable in supported."""
 
-    if isinstance(var, MetVariable):
-        if var in supported:
-            return var
+    if isinstance(var, MetVariable) and var in supported:
+        return var
 
     # list of MetVariable options
     # here we extract the first MetVariable in var that is supported
@@ -230,21 +247,19 @@ def _find_match(
             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]
+        if ret := ecmwf_ids.get(var):
+            return ret
+        if ret := grib1_ids.get(var):
+            return ret
 
-    # 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]
+        if ret := short_names.get(var):
+            return ret
+        if ret := standard_names.get(var):
+            return ret
+        if ret := long_names.get(var):
+            return ret
 
     msg = f"{var} is not in supported parameters. Supported parameters include: {standard_names}"
     raise ValueError(msg)
@@ -395,6 +410,14 @@ class MetDataSource(abc.ABC):
         """
         return [v.standard_name for v in self.variables]
 
+    @property
+    def is_single_level(self) -> bool:
+        """Return True if the datasource is single level data.
+
+        .. versionadded:: 0.50.0
+        """
+        return self.pressure_levels == [-1]
+
     @property
     def pressure_level_variables(self) -> list[MetVariable]:
         """Parameters available from data source.
@@ -426,10 +449,9 @@ class MetDataSource(abc.ABC):
         list[MetVariable] | None
             List of MetVariable available in datasource
         """
-        if self.pressure_levels != [-1]:
-            return self.pressure_level_variables
-
-        return self.single_level_variables
+        return (
+            self.single_level_variables if self.is_single_level else self.pressure_level_variables
+        )
 
     @property
     def supported_pressure_levels(self) -> list[int] | None:
@@ -497,7 +519,7 @@ class MetDataSource(abc.ABC):
     def open_metdataset(
         self,
         dataset: xr.Dataset | None = None,
-        xr_kwargs: dict[str, int] | None = None,
+        xr_kwargs: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> MetDataset:
         """Open MetDataset from data source.
@@ -510,7 +532,7 @@ class MetDataSource(abc.ABC):
         dataset : xr.Dataset | None, optional
             Input :class:`xr.Dataset` loaded manually.
             The dataset must have the same format as the original data source API or files.
-        xr_kwargs : dict[str, int] | None, optional
+        xr_kwargs : dict[str, Any] | None, optional
             Dictionary of keyword arguments passed into :func:`xarray.open_mfdataset`
             when opening files. Examples include "chunks", "engine", "parallel", etc.
             Ignored if ``dataset`` is input.

pycontrails/core/met.py

@@ -1065,8 +1065,8 @@ class MetDataset(MetBase):
 
         """
         coords_keys = self.data.dims
-        variables = self.indexes
-        coords_vals = [variables[key].values for key in coords_keys]
+        indexes = self.indexes
+        coords_vals = [indexes[key].values for key in coords_keys]
         coords_meshes = np.meshgrid(*coords_vals, indexing="ij")
         raveled_coords = (mesh.ravel() for mesh in coords_meshes)
         data = dict(zip(coords_keys, raveled_coords))
@@ -1181,12 +1181,12 @@ class MetDataset(MetBase):
         level: npt.ArrayLike | float,
         time: npt.ArrayLike | np.datetime64,
     ) -> MetDataset:
-        """Create a :class:`MetDataset` containing a coordinate skeleton from coordinate arrays.
+        r"""Create a :class:`MetDataset` containing a coordinate skeleton from coordinate arrays.
 
         Parameters
         ----------
         longitude, latitude : npt.ArrayLike | float
-            Horizontal coordinates, in [:math:`degrees`]
+            Horizontal coordinates, in [:math:`\deg`]
         level : npt.ArrayLike | float
             Vertical coordinate, in [:math:`hPa`]
         time: npt.ArrayLike | np.datetime64,

pycontrails/core/met_var.py

@@ -76,12 +76,12 @@ class MetVariable:
 
     @property
     def ecmwf_link(self) -> str | None:
-        """Database link in the ECMWF Paramter Database if :attr:`ecmwf_id` is defined.
+        """Database link in the ECMWF Parameter Database if :attr:`ecmwf_id` is defined.
 
         Returns
         -------
         str
-            Database link in the ECMWF Paramter Database
+            Database link in the ECMWF Parameter Database
         """
         return (
             f"https://apps.ecmwf.int/codes/grib/param-db?id={self.ecmwf_id}"

pycontrails/core/models.py

@@ -440,17 +440,21 @@ class Model(ABC):
             # Return dataset with the same coords as self.met, but empty data_vars
             return MetDataset(xr.Dataset(coords=self.met.data.coords))
 
+        copy_source = self.params["copy_source"]
+
         # Turn Sequence into Fleet
         if isinstance(source, Sequence):
-            # TODO: fix type guard here
-            return Fleet.from_seq(source, copy=self.params["copy_source"])
+            if not copy_source:
+                msg = "Parameter copy_source=False is not supported for Sequence[Flight] source"
+                raise ValueError(msg)
+            return Fleet.from_seq(source)
 
         # Raise error if source is not a MetDataset or GeoVectorDataset
         if not isinstance(source, (MetDataset, GeoVectorDataset)):
             msg = f"Unknown source type: {type(source)}"
             raise TypeError(msg)
 
-        if self.params["copy_source"]:
+        if copy_source:
             source = source.copy()
 
         if not isinstance(source, Flight):

pycontrails/core/vector.py

@@ -1899,19 +1899,19 @@ class GeoVectorDataset(VectorDataset):
         MetDataset | MetDataArray
             Copy of downselected MetDataset or MetDataArray.
         """
-        variables = met.indexes
+        indexes = met.indexes
         lon_slice = coordinates.slice_domain(
-            variables["longitude"].to_numpy(),
+            indexes["longitude"].to_numpy(),
             self["longitude"],
             buffer=longitude_buffer,
         )
         lat_slice = coordinates.slice_domain(
-            variables["latitude"].to_numpy(),
+            indexes["latitude"].to_numpy(),
             self["latitude"],
             buffer=latitude_buffer,
         )
         time_slice = coordinates.slice_domain(
-            variables["time"].to_numpy(),
+            indexes["time"].to_numpy(),
             self["time"],
             buffer=time_buffer,
         )
@@ -1921,7 +1921,7 @@ class GeoVectorDataset(VectorDataset):
             level_slice = slice(None)
         else:
             level_slice = coordinates.slice_domain(
-                variables["level"].to_numpy(),
+                indexes["level"].to_numpy(),
                 self.level,
                 buffer=level_buffer,
             )

pycontrails/datalib/ecmwf/__init__.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from pycontrails.datalib.ecmwf.arco_era5 import ARCOERA5
 from pycontrails.datalib.ecmwf.era5 import ERA5
 from pycontrails.datalib.ecmwf.hres import HRES
 from pycontrails.datalib.ecmwf.ifs import IFS
@@ -11,6 +12,7 @@ from pycontrails.datalib.ecmwf.variables import (
     SURFACE_VARIABLES,
     CloudAreaFraction,
     CloudAreaFractionInLayer,
+    Divergence,
     PotentialVorticity,
     RelativeHumidity,
     RelativeVorticity,
@@ -23,11 +25,13 @@ from pycontrails.datalib.ecmwf.variables import (
 )
 
 __all__ = [
+    "ARCOERA5",
     "ERA5",
     "HRES",
     "IFS",
     "CloudAreaFraction",
     "CloudAreaFractionInLayer",
+    "Divergence",
     "PotentialVorticity",
     "RelativeHumidity",
     "RelativeVorticity",