pycontrails 0.54.1__cp313-cp313-win_amd64.whl → 0.54.3__cp313-cp313-win_amd64.whl

pycontrails/core/vector.py

@@ -5,15 +5,25 @@ from __future__ import annotations
 import hashlib
 import json
 import logging
+import sys
 import warnings
 from collections.abc import Generator, Iterable, Iterator, Sequence
-from typing import Any, TypeVar, overload
+from typing import Any, overload
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
 
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
 import xarray as xr
-from overrides import overrides
 
 from pycontrails.core import coordinates, interpolation
 from pycontrails.core import met as met_module
@@ -23,10 +33,6 @@ from pycontrails.utils import json as json_utils
 
 logger = logging.getLogger(__name__)
 
-#: Vector types
-VectorDatasetType = TypeVar("VectorDatasetType", bound="VectorDataset")
-GeoVectorDatasetType = TypeVar("GeoVectorDatasetType", bound="GeoVectorDataset")
-
 
 class AttrDict(dict[str, Any]):
     """Thin wrapper around dict to warn when setting a key that already exists."""
@@ -574,7 +580,7 @@ class VectorDataset:
         """
         return self.size > 0
 
-    def __add__(self: VectorDatasetType, other: VectorDatasetType | None) -> VectorDatasetType:
+    def __add__(self, other: Self | None) -> Self:
         """Concatenate two compatible instances of VectorDataset.
 
         In this context, compatibility means that both have identical :attr:`data` keys.
@@ -586,12 +592,12 @@ class VectorDataset:
 
         Parameters
         ----------
-        other : VectorDatasetType
+        other : Self | None
             Other values to concatenate
 
         Returns
         -------
-        VectorDatasetType
+        Self
             Concatenated values.
 
         Raises
@@ -610,11 +616,11 @@ class VectorDataset:
 
     @classmethod
     def sum(
-        cls: type[VectorDatasetType],
+        cls,
         vectors: Sequence[VectorDataset],
         infer_attrs: bool = True,
         fill_value: float | None = None,
-    ) -> VectorDatasetType:
+    ) -> Self:
         """Sum a list of :class:`VectorDataset` instances.
 
         Parameters
@@ -692,7 +698,7 @@ class VectorDataset:
             return cls(data, attrs=vectors[0].attrs, copy=False)
         return cls(data, copy=False)
 
-    def __eq__(self: VectorDatasetType, other: object) -> bool:
+    def __eq__(self, other: object) -> bool:
         """Determine if two instances are equal.
 
         NaN values are considered equal in this comparison.
@@ -700,7 +706,7 @@ class VectorDataset:
         Parameters
         ----------
         other : object
-            VectorDatasetType to compare with
+            VectorDataset to compare with
 
         Returns
         -------
@@ -784,8 +790,8 @@ class VectorDataset:
     # Utilities
     # ------------
 
-    def copy(self: VectorDatasetType, **kwargs: Any) -> VectorDatasetType:
-        """Return a copy of this VectorDatasetType class.
+    def copy(self, **kwargs: Any) -> Self:
+        """Return a copy of this instance.
 
         Parameters
         ----------
@@ -794,7 +800,7 @@ class VectorDataset:
 
         Returns
         -------
-        VectorDatasetType
+        Self
             Copy of class
         """
         return type(self)(data=self.data, attrs=self.attrs, copy=True, **kwargs)
@@ -820,9 +826,7 @@ class VectorDataset:
         data = {key: self[key] for key in keys}
         return VectorDataset(data=data, attrs=self.attrs, copy=copy)
 
-    def filter(
-        self: VectorDatasetType, mask: npt.NDArray[np.bool_], copy: bool = True, **kwargs: Any
-    ) -> VectorDatasetType:
+    def filter(self, mask: npt.NDArray[np.bool_], copy: bool = True, **kwargs: Any) -> Self:
         """Filter :attr:`data` according to a boolean array ``mask``.
 
         Entries corresponding to ``mask == True`` are kept.
@@ -840,7 +844,7 @@ class VectorDataset:
 
         Returns
         -------
-        VectorDatasetType
+        Self
             Containing filtered data
 
         Raises
@@ -855,7 +859,7 @@ class VectorDataset:
         data = {key: value[mask] for key, value in self.data.items()}
         return type(self)(data=data, attrs=self.attrs, copy=copy, **kwargs)
 
-    def sort(self: VectorDatasetType, by: str | list[str]) -> VectorDatasetType:
+    def sort(self, by: str | list[str]) -> Self:
         """Sort data by key(s).
 
         This method always creates a copy of the data by calling
@@ -868,7 +872,7 @@ class VectorDataset:
 
         Returns
         -------
-        VectorDatasetType
+        Self
             Instance with sorted data.
         """
         return type(self)(data=self.dataframe.sort_values(by=by), attrs=self.attrs, copy=False)
@@ -1109,12 +1113,12 @@ class VectorDataset:
 
     @classmethod
     def create_empty(
-        cls: type[VectorDatasetType],
+        cls,
         keys: Iterable[str],
         attrs: dict[str, Any] | None = None,
         **attrs_kwargs: Any,
-    ) -> VectorDatasetType:
-        """Create instance with variables defined by `keys` and size 0.
+    ) -> Self:
+        """Create instance with variables defined by ``keys`` and size 0.
 
         If instance requires additional variables to be defined, these keys will automatically
         be attached to returned instance.
@@ -1130,15 +1134,13 @@ class VectorDataset:
 
         Returns
         -------
-        VectorDatasetType
+        Self
             Empty VectorDataset instance.
         """
         return cls(data=_empty_vector_dict(keys or set()), attrs=attrs, copy=False, **attrs_kwargs)
 
     @classmethod
-    def from_dict(
-        cls: type[VectorDatasetType], obj: dict[str, Any], copy: bool = True, **obj_kwargs: Any
-    ) -> VectorDatasetType:
+    def from_dict(cls, obj: dict[str, Any], copy: bool = True, **obj_kwargs: Any) -> Self:
         """Create instance from dict representation containing data and attrs.
 
         Parameters
@@ -1153,7 +1155,7 @@ class VectorDataset:
 
         Returns
         -------
-        VectorDatasetType
+        Self
             VectorDataset instance.
 
         See Also
@@ -1171,9 +1173,7 @@ class VectorDataset:
 
         return cls(data=data, attrs=attrs, copy=copy)
 
-    def generate_splits(
-        self: VectorDatasetType, n_splits: int, copy: bool = True
-    ) -> Generator[VectorDatasetType, None, None]:
+    def generate_splits(self, n_splits: int, copy: bool = True) -> Generator[Self, None, None]:
         """Split instance into ``n_split`` sub-vectors.
 
         Parameters
@@ -1186,7 +1186,7 @@ class VectorDataset:
 
         Returns
         -------
-        Generator[VectorDatasetType, None, None]
+        Generator[Self, None, None]
             Generator of split vectors.
 
         See Also
@@ -1367,7 +1367,7 @@ class GeoVectorDataset(VectorDataset):
         if np.any(latitude > 90.0) or np.any(latitude < -90.0):
             raise ValueError("EPSG:4326 latitude coordinates should lie between [-90, 90].")
 
-    @overrides
+    @override
     def _display_attrs(self) -> dict[str, str]:
         try:
             time0 = pd.Timestamp(np.nanmin(self["time"]))
@@ -1922,13 +1922,13 @@ class GeoVectorDataset(VectorDataset):
     # ------------
 
     @classmethod
-    @overrides
+    @override
     def create_empty(
-        cls: type[GeoVectorDatasetType],
+        cls,
         keys: Iterable[str] | None = None,
         attrs: dict[str, Any] | None = None,
         **attrs_kwargs: Any,
-    ) -> GeoVectorDatasetType:
+    ) -> Self:
         keys = *cls.required_keys, "altitude", *(keys or ())
         return super().create_empty(keys, attrs, **attrs_kwargs)
 

pycontrails/datalib/ecmwf/arco_era5.py

@@ -19,11 +19,16 @@ from __future__ import annotations
 
 import datetime
 import hashlib
+import sys
 from typing import Any
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 import numpy.typing as npt
 import xarray as xr
-from overrides import overrides
 
 from pycontrails.core import cache, met_var
 from pycontrails.core.met import MetDataset
@@ -274,7 +279,7 @@ class ERA5ARCO(ecmwf_common.ECMWFAPI):
         """
         return ecmwf_variables.SURFACE_VARIABLES
 
-    @overrides
+    @override
     def download_dataset(self, times: list[datetime.datetime]) -> None:
         if not times:
             return
@@ -286,7 +291,7 @@ class ERA5ARCO(ecmwf_common.ECMWFAPI):
 
         self.cache_dataset(ds)
 
-    @overrides
+    @override
     def create_cachepath(self, t: datetime.datetime) -> str:
         if self.cachestore is None:
             msg = "Attribute self.cachestore must be defined to create cache path"
@@ -302,7 +307,7 @@ class ERA5ARCO(ecmwf_common.ECMWFAPI):
 
         return self.cachestore.path(cache_path)
 
-    @overrides
+    @override
     def open_metdataset(
         self,
         dataset: xr.Dataset | None = None,
@@ -331,7 +336,7 @@ class ERA5ARCO(ecmwf_common.ECMWFAPI):
         self.set_metadata(mds)
         return mds
 
-    @overrides
+    @override
     def set_metadata(self, ds: xr.Dataset | MetDataset) -> None:
         ds.attrs.update(
             provider="ECMWF",

pycontrails/datalib/ecmwf/common.py

@@ -4,14 +4,19 @@ from __future__ import annotations
 
 import logging
 import os
+import sys
 from typing import Any
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 LOG = logging.getLogger(__name__)
 
 import numpy as np
 import pandas as pd
 import xarray as xr
-from overrides import overrides
 
 from pycontrails.core import met
 from pycontrails.datalib._met_utils import metsource
@@ -88,7 +93,7 @@ class ECMWFAPI(metsource.MetDataSource):
         kwargs.setdefault("cachestore", self.cachestore)
         return met.MetDataset(ds, **kwargs)
 
-    @overrides
+    @override
     def cache_dataset(self, dataset: xr.Dataset) -> None:
         if self.cachestore is None:
             LOG.debug("Cache is turned off, skipping")

pycontrails/datalib/ecmwf/era5.py

@@ -7,16 +7,21 @@ import hashlib
 import logging
 import os
 import pathlib
+import sys
 import warnings
 from contextlib import ExitStack
 from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 LOG = logging.getLogger(__name__)
 
 import pandas as pd
 import xarray as xr
-from overrides import overrides
 
 import pycontrails
 from pycontrails.core import cache
@@ -347,7 +352,7 @@ class ERA5(ECMWFAPI):
         # return cache path
         return self.cachestore.path(f"{datestr}-{suffix}.nc")
 
-    @overrides
+    @override
     def download_dataset(self, times: list[datetime]) -> None:
         download_times: dict[datetime, list[datetime]] = collections.defaultdict(list)
         for t in times:
@@ -359,7 +364,7 @@ class ERA5(ECMWFAPI):
         for times_for_day in download_times.values():
             self._download_file(times_for_day)
 
-    @overrides
+    @override
     def open_metdataset(
         self,
         dataset: xr.Dataset | None = None,
@@ -399,7 +404,7 @@ class ERA5(ECMWFAPI):
         self.set_metadata(mds)
         return mds
 
-    @overrides
+    @override
     def set_metadata(self, ds: xr.Dataset | MetDataset) -> None:
         if self.product_type == "reanalysis":
             product = "reanalysis"

pycontrails/datalib/ecmwf/era5_model_level.py

@@ -25,12 +25,16 @@ import contextlib
 import hashlib
 import logging
 import os
+import sys
 import threading
 import warnings
 from datetime import datetime
 from typing import Any
 
-from overrides import overrides
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
 
 LOG = logging.getLogger(__name__)
 
@@ -245,7 +249,7 @@ class ERA5ModelLevel(ECMWFAPI):
         """
         return "reanalysis-era5-complete"
 
-    @overrides
+    @override
     def create_cachepath(self, t: datetime | pd.Timestamp) -> str:
         """Return cachepath to local ERA5 data file based on datetime.
 
@@ -277,7 +281,7 @@ class ERA5ModelLevel(ECMWFAPI):
 
         return self.cachestore.path(cache_path)
 
-    @overrides
+    @override
     def download_dataset(self, times: list[datetime]) -> None:
         # group data to request by month (nominal) or by day (ensemble)
         requests: dict[datetime, list[datetime]] = collections.defaultdict(list)
@@ -294,7 +298,7 @@ class ERA5ModelLevel(ECMWFAPI):
         for times_in_request in requests.values():
             self._download_convert_cache_handler(times_in_request)
 
-    @overrides
+    @override
     def open_metdataset(
         self,
         dataset: xr.Dataset | None = None,
@@ -320,7 +324,7 @@ class ERA5ModelLevel(ECMWFAPI):
         self.set_metadata(mds)
         return mds
 
-    @overrides
+    @override
     def set_metadata(self, ds: xr.Dataset | MetDataset) -> None:
         if self.product_type == "reanalysis":
             product = "reanalysis"

pycontrails/datalib/ecmwf/hres.py

@@ -5,16 +5,21 @@ from __future__ import annotations
 import hashlib
 import logging
 import pathlib
+import sys
 from contextlib import ExitStack
 from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 LOG = logging.getLogger(__name__)
 
 import numpy as np
 import pandas as pd
 import xarray as xr
-from overrides import overrides
 
 import pycontrails
 from pycontrails.core import cache
@@ -326,9 +331,9 @@ class HRES(ECMWFAPI):
             f" {getattr(self, 'steps', '')}"
         )
 
-    @classmethod
+    @staticmethod
     def create_synoptic_time_ranges(
-        self, timesteps: list[pd.Timestamp]
+        timesteps: list[pd.Timestamp],
     ) -> list[tuple[pd.Timestamp, pd.Timestamp]]:
         """Create synoptic time bounds encompassing date range.
 
@@ -556,7 +561,7 @@ class HRES(ECMWFAPI):
             f"\n\tgrid={request['grid']},\n\tlevtype={request['levtype']}{levelist}"
         )
 
-    @overrides
+    @override
     def create_cachepath(self, t: datetime) -> str:
         if self.cachestore is None:
             raise ValueError("self.cachestore attribute must be defined to create cache path")
@@ -574,7 +579,7 @@ class HRES(ECMWFAPI):
         # return cache path
         return self.cachestore.path(f"{datestr}-{step}-{suffix}.nc")
 
-    @overrides
+    @override
     def download_dataset(self, times: list[datetime]) -> None:
         """Download data from data source for input times.
 
@@ -595,7 +600,7 @@ class HRES(ECMWFAPI):
         elif len(steps) > 0:
             self._download_file(steps)
 
-    @overrides
+    @override
     def open_metdataset(
         self,
         dataset: xr.Dataset | None = None,
@@ -635,7 +640,7 @@ class HRES(ECMWFAPI):
         self.set_metadata(mds)
         return mds
 
-    @overrides
+    @override
     def set_metadata(self, ds: xr.Dataset | MetDataset) -> None:
         if self.stream == "oper":
             product = "forecast"

pycontrails/datalib/ecmwf/hres_model_level.py

@@ -13,15 +13,20 @@ from __future__ import annotations
 import contextlib
 import hashlib
 import logging
+import sys
 import warnings
 from datetime import datetime, timedelta
 from typing import Any
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 LOG = logging.getLogger(__name__)
 
 import pandas as pd
 import xarray as xr
-from overrides import overrides
 
 import pycontrails
 from pycontrails.core import cache
@@ -283,7 +288,7 @@ class HRESModelLevel(ECMWFAPI):
         """
         return []
 
-    @overrides
+    @override
     def create_cachepath(self, t: datetime | pd.Timestamp) -> str:
         """Return cachepath to local HRES data file based on datetime.
 
@@ -316,13 +321,13 @@ class HRESModelLevel(ECMWFAPI):
 
         return self.cachestore.path(cache_path)
 
-    @overrides
+    @override
     def download_dataset(self, times: list[datetime]) -> None:
         # will always submit a single MARS request since each forecast is a separate file on tape
         LOG.debug(f"Retrieving ERA5 data for times {times} from forecast {self.forecast_time}")
         self._download_convert_cache_handler(times)
 
-    @overrides
+    @override
     def open_metdataset(
         self,
         dataset: xr.Dataset | None = None,
@@ -348,7 +353,7 @@ class HRESModelLevel(ECMWFAPI):
         self.set_metadata(mds)
         return mds
 
-    @overrides
+    @override
     def set_metadata(self, ds: xr.Dataset | MetDataset) -> None:
         ds.attrs.update(
             provider="ECMWF", dataset="HRES", product="forecast", radiation_accumulated=True

pycontrails/datalib/ecmwf/ifs.py

@@ -4,16 +4,21 @@ from __future__ import annotations
 
 import logging
 import pathlib
+import sys
 import warnings
 from datetime import datetime
 from typing import Any
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 LOG = logging.getLogger(__name__)
 
 import numpy as np
 import pandas as pd
 import xarray as xr
-from overrides import overrides
 
 from pycontrails.core import met
 from pycontrails.datalib._met_utils import metsource
@@ -119,7 +124,7 @@ class IFS(metsource.MetDataSource):
         """
         return None
 
-    @overrides
+    @override
     def open_metdataset(
         self,
         dataset: xr.Dataset | None = None,
@@ -190,7 +195,7 @@ class IFS(metsource.MetDataSource):
         self.set_metadata(ds)
         return met.MetDataset(ds, **kwargs)
 
-    @overrides
+    @override
     def set_metadata(self, ds: xr.Dataset | met.MetDataset) -> None:
         ds.attrs.update(
             provider="ECMWF",
@@ -198,15 +203,15 @@ class IFS(metsource.MetDataSource):
             product="forecast",
         )
 
-    @overrides
+    @override
     def download_dataset(self, times: list[datetime]) -> None:
         raise NotImplementedError("IFS download is not supported")
 
-    @overrides
+    @override
     def cache_dataset(self, dataset: xr.Dataset) -> None:
         raise NotImplementedError("IFS dataset caching not supported")
 
-    @overrides
+    @override
     def create_cachepath(self, t: datetime) -> str:
         raise NotImplementedError("IFS download is not supported")
 

pycontrails/datalib/ecmwf/variables.py

@@ -107,6 +107,7 @@ RelativeHumidity = MetVariable(
     long_name=met_var.RelativeHumidity.long_name,
     units="%",
     level_type=met_var.RelativeHumidity.level_type,
+    grib1_id=met_var.RelativeHumidity.grib1_id,
     ecmwf_id=met_var.RelativeHumidity.ecmwf_id,
     grib2_id=met_var.RelativeHumidity.grib2_id,
     description=(