anemoi-datasets 0.5.26__py3-none-any.whl → 0.5.27__py3-none-any.whl

anemoi/datasets/data/stores.py

@@ -15,11 +15,6 @@ import tempfile
 import warnings
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Set
-from typing import Union
 from urllib.parse import urlparse
 
 import numpy as np
@@ -90,7 +85,7 @@ class S3Store(ReadOnlyStore):
     options using the anemoi configs.
     """
 
-    def __init__(self, url: str, region: Optional[str] = None) -> None:
+    def __init__(self, url: str, region: str | None = None) -> None:
         """Initialize the S3Store with a URL and optional region."""
         from anemoi.utils.remote.s3 import s3_client
 
@@ -199,7 +194,7 @@ def open_zarr(path: str, dont_fail: bool = False, cache: int = None) -> zarr.hie
 class Zarr(Dataset):
     """A zarr dataset."""
 
-    def __init__(self, path: Union[str, zarr.hierarchy.Group]) -> None:
+    def __init__(self, path: str | zarr.hierarchy.Group) -> None:
         """Initialize the Zarr dataset with a path or zarr group."""
         if isinstance(path, zarr.hierarchy.Group):
             self.was_zarr = True
@@ -215,7 +210,7 @@ class Zarr(Dataset):
         self._missing = set()
 
     @property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Return the missing dates of the dataset."""
         return self._missing
 
@@ -236,7 +231,7 @@ class Zarr(Dataset):
         """Retrieve an item from the dataset."""
         return self.data[n]
 
-    def _unwind(self, index: Union[int, slice, list, tuple], rest: list, shape: tuple, axis: int, axes: list) -> iter:
+    def _unwind(self, index: int | slice | list | tuple, rest: list, shape: tuple, axis: int, axes: list) -> iter:
         """Unwind the index for multi-dimensional indexing."""
         if not isinstance(index, (int, slice, list, tuple)):
             try:
@@ -298,7 +293,7 @@ class Zarr(Dataset):
             return self.z.longitude[:]
 
     @property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Return the statistics of the dataset."""
         return dict(
             mean=self.z.mean[:],
@@ -307,7 +302,7 @@ class Zarr(Dataset):
             minimum=self.z.minimum[:],
         )
 
-    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+    def statistics_tendencies(self, delta: datetime.timedelta | None = None) -> dict[str, NDArray[Any]]:
         """Return the statistical tendencies of the dataset."""
         if delta is None:
             delta = self.frequency
@@ -354,14 +349,14 @@ class Zarr(Dataset):
         return dates[1].astype(object) - dates[0].astype(object)
 
     @property
-    def name_to_index(self) -> Dict[str, int]:
+    def name_to_index(self) -> dict[str, int]:
         """Return the name to index mapping of the dataset."""
         if "variables" in self.z.attrs:
             return {n: i for i, n in enumerate(self.z.attrs["variables"])}
         return self.z.attrs["name_to_index"]
 
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Return the variables of the dataset."""
         return [
             k
@@ -372,7 +367,7 @@ class Zarr(Dataset):
         ]
 
     @cached_property
-    def constant_fields(self) -> List[str]:
+    def constant_fields(self) -> list[str]:
         """Return the constant fields of the dataset."""
         result = self.z.attrs.get("constant_fields")
         if result is None:
@@ -380,7 +375,7 @@ class Zarr(Dataset):
         return self.computed_constant_fields()
 
     @property
-    def variables_metadata(self) -> Dict[str, Any]:
+    def variables_metadata(self) -> dict[str, Any]:
         """Return the metadata of the variables."""
         return self.z.attrs.get("variables_metadata", {})
 
@@ -392,7 +387,7 @@ class Zarr(Dataset):
         """Return the end date of the statistics."""
         return self.dates[-1]
 
-    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+    def metadata_specific(self, **kwargs: Any) -> dict[str, Any]:
         """Return the specific metadata of the dataset."""
         return super().metadata_specific(
             attrs=dict(self.z.attrs),
@@ -416,7 +411,7 @@ class Zarr(Dataset):
         """Return the tree representation of the dataset."""
         return Node(self, [], path=self.path)
 
-    def get_dataset_names(self, names: Set[str]) -> None:
+    def get_dataset_names(self, names: set[str]) -> None:
         """Get the names of the datasets."""
         name, _ = os.path.splitext(os.path.basename(self.path))
         names.add(name)
@@ -433,17 +428,17 @@ class Zarr(Dataset):
 class ZarrWithMissingDates(Zarr):
     """A zarr dataset with missing dates."""
 
-    def __init__(self, path: Union[str, zarr.hierarchy.Group]) -> None:
+    def __init__(self, path: str | zarr.hierarchy.Group) -> None:
         """Initialize the ZarrWithMissingDates dataset with a path or zarr group."""
         super().__init__(path)
 
         missing_dates = self.z.attrs.get("missing_dates", [])
-        missing_dates = set([np.datetime64(x, "s") for x in missing_dates])
+        missing_dates = {np.datetime64(x, "s") for x in missing_dates}
         self.missing_to_dates = {i: d for i, d in enumerate(self.dates) if d in missing_dates}
         self._missing = set(self.missing_to_dates)
 
     @property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Return the missing dates of the dataset."""
         return self._missing
 
@@ -506,7 +501,7 @@ class ZarrWithMissingDates(Zarr):
 QUIET = set()
 
 
-def zarr_lookup(name: str, fail: bool = True) -> Optional[str]:
+def zarr_lookup(name: str, fail: bool = True) -> str | None:
     """Look up a zarr dataset by name."""
 
     config = load_config()["datasets"]

anemoi/datasets/data/subset.py

@@ -10,12 +10,9 @@
 
 import datetime
 import logging
+from collections.abc import Sequence
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Sequence
-from typing import Set
 from typing import Union
 
 import numpy as np
@@ -95,7 +92,7 @@ def _end(a: int, b: int, dates: NDArray[np.datetime64]) -> int:
         return b
 
 
-def _combine_reasons(reason1: Dict[str, Any], reason2: Dict[str, Any], dates: NDArray[np.datetime64]) -> Dict[str, Any]:
+def _combine_reasons(reason1: dict[str, Any], reason2: dict[str, Any], dates: NDArray[np.datetime64]) -> dict[str, Any]:
     """Combine two reason dictionaries.
 
     Parameters:
@@ -126,7 +123,7 @@ class Subset(Forwards):
     reason (Dict[str, Any]): Dictionary of reasons.
     """
 
-    def __init__(self, dataset: Union[Dataset, "Subset"], indices: Sequence[int], reason: Dict[str, Any]) -> None:
+    def __init__(self, dataset: Union[Dataset, "Subset"], indices: Sequence[int], reason: dict[str, Any]) -> None:
         """Initialize the Subset.
 
         Parameters:
@@ -140,8 +137,8 @@ class Subset(Forwards):
             dataset = dataset.dataset
 
         self.dataset: Dataset = dataset
-        self.indices: List[int] = list(indices)
-        self.reason: Dict[str, Any] = {k: v for k, v in reason.items() if v is not None}
+        self.indices: list[int] = list(indices)
+        self.reason: dict[str, Any] = {k: v for k, v in reason.items() if v is not None}
 
         # Forward other properties to the super dataset
         super().__init__(dataset)
@@ -274,10 +271,10 @@ class Subset(Forwards):
         return f"Subset({self.dataset},{self.dates[0]}...{self.dates[-1]}/{self.frequency})"
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the missing indices of the subset."""
         missing = self.dataset.missing
-        result: Set[int] = set()
+        result: set[int] = set()
         for j, i in enumerate(self.indices):
             if i in missing:
                 result.add(j)
@@ -291,7 +288,7 @@ class Subset(Forwards):
         """
         return Node(self, [self.dataset.tree()], **self.reason)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the forwards subclass.
 
         Returns:

anemoi/datasets/data/unchecked.py

@@ -10,14 +10,10 @@
 
 import datetime
 import logging
+from collections.abc import Callable
 from functools import cached_property
 from functools import wraps
 from typing import Any
-from typing import Callable
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Set
 
 import numpy as np
 from numpy.typing import NDArray
@@ -163,13 +159,13 @@ class Unchecked(Combined):
 
     @check("check_same_variables")
     @property
-    def name_to_index(self) -> Dict[str, int]:
+    def name_to_index(self) -> dict[str, int]:
         """Get the mapping of variable names to their indices."""
         raise NotImplementedError()
 
     @check("check_same_variables")
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Get the list of variables in the dataset."""
         raise NotImplementedError()
 
@@ -181,12 +177,12 @@ class Unchecked(Combined):
 
     @check("check_same_variables")
     @property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Get the statistics of the dataset."""
         raise NotImplementedError()
 
     @check("check_same_variables")
-    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+    def statistics_tendencies(self, delta: datetime.timedelta | None = None) -> dict[str, NDArray[Any]]:
         """Get the statistics tendencies of the dataset.
 
         Parameters
@@ -207,9 +203,9 @@ class Unchecked(Combined):
         raise NotImplementedError()
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the missing data indices."""
-        result: Set[int] = set()
+        result: set[int] = set()
         for d in self.datasets:
             result = result | d.missing
         return result

anemoi/datasets/data/xy.py

@@ -11,10 +11,6 @@
 import logging
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Set
-from typing import Tuple
 
 from .dataset import Dataset
 from .dataset import FullIndex
@@ -29,7 +25,7 @@ LOG = logging.getLogger(__name__)
 class ZipBase(Combined):
     """Base class for handling zipped datasets."""
 
-    def __init__(self, datasets: List[Any], check_compatibility: bool = True) -> None:
+    def __init__(self, datasets: list[Any], check_compatibility: bool = True) -> None:
         """Initialize ZipBase with a list of datasets.
 
         Parameters
@@ -58,7 +54,7 @@ class ZipBase(Combined):
         new_parents = [parent.clone(ds) for ds in self.datasets]
         return self.clone(new_parents)
 
-    def clone(self, datasets: List[Any]) -> "ZipBase":
+    def clone(self, datasets: list[Any]) -> "ZipBase":
         """Clone the ZipBase with new datasets.
 
         Parameters
@@ -81,7 +77,11 @@ class ZipBase(Combined):
         Node
             Tree representation of the datasets.
         """
-        return Node(self, [d.tree() for d in self.datasets], check_compatibility=self._check_compatibility)
+        return Node(
+            self,
+            [d.tree() for d in self.datasets],
+            check_compatibility=self._check_compatibility,
+        )
 
     def __len__(self) -> int:
         """Get the length of the smallest dataset.
@@ -93,7 +93,7 @@ class ZipBase(Combined):
         """
         return min(len(d) for d in self.datasets)
 
-    def __getitem__(self, n: FullIndex) -> Tuple[Any, ...]:
+    def __getitem__(self, n: FullIndex) -> tuple[Any, ...]:
         """Get the item at the specified index from all datasets.
 
         Parameters
@@ -145,55 +145,55 @@ class ZipBase(Combined):
         pass
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the set of missing indices from all datasets."""
-        result: Set[int] = set()
+        result: set[int] = set()
         for d in self.datasets:
             result = result | d.missing
         return result
 
     @property
-    def shape(self) -> Tuple[Any, ...]:
+    def shape(self) -> tuple[Any, ...]:
         """Get the shape of all datasets."""
         return tuple(d.shape for d in self.datasets)
 
     @property
-    def field_shape(self) -> Tuple[Any, ...]:
+    def field_shape(self) -> tuple[Any, ...]:
         """Get the field shape of all datasets."""
         return tuple(d.shape for d in self.datasets)
 
     @property
-    def latitudes(self) -> Tuple[Any, ...]:
+    def latitudes(self) -> tuple[Any, ...]:
         """Get the latitudes of all datasets."""
         return tuple(d.latitudes for d in self.datasets)
 
     @property
-    def longitudes(self) -> Tuple[Any, ...]:
+    def longitudes(self) -> tuple[Any, ...]:
         """Get the longitudes of all datasets."""
         return tuple(d.longitudes for d in self.datasets)
 
     @property
-    def dtype(self) -> Tuple[Any, ...]:
+    def dtype(self) -> tuple[Any, ...]:
         """Get the data types of all datasets."""
         return tuple(d.dtype for d in self.datasets)
 
     @property
-    def grids(self) -> Tuple[Any, ...]:
+    def grids(self) -> tuple[Any, ...]:
         """Get the grids of all datasets."""
         return tuple(d.grids for d in self.datasets)
 
     @property
-    def statistics(self) -> Tuple[Any, ...]:
+    def statistics(self) -> tuple[Any, ...]:
         """Get the statistics of all datasets."""
         return tuple(d.statistics for d in self.datasets)
 
     @property
-    def resolution(self) -> Tuple[Any, ...]:
+    def resolution(self) -> tuple[Any, ...]:
         """Get the resolution of all datasets."""
         return tuple(d.resolution for d in self.datasets)
 
     @property
-    def name_to_index(self) -> Tuple[Any, ...]:
+    def name_to_index(self) -> tuple[Any, ...]:
         """Get the name to index mapping of all datasets."""
         return tuple(d.name_to_index for d in self.datasets)
 
@@ -210,6 +210,10 @@ class ZipBase(Combined):
         if self._check_compatibility:
             super().check_compatibility(d1, d2)
 
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
+        """Get the metadata specific to the subclass."""
+        return {}
+
 
 class Zip(ZipBase):
     """Class for handling zipped datasets."""
@@ -223,7 +227,7 @@ class XY(ZipBase):
     pass
 
 
-def xy_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> XY:
+def xy_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> XY:
     """Factory function to create an XY instance.
 
     Parameters
@@ -256,7 +260,7 @@ def xy_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> XY:
     return XY(datasets, check_compatibility=check_compatibility)._subset(**kwargs)
 
 
-def zip_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Zip:
+def zip_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Zip:
     """Factory function to create a Zip instance.
 
     Parameters

anemoi/datasets/dates/__init__.py

@@ -10,15 +10,10 @@
 
 import datetime
 import warnings
+from collections.abc import Iterator
 from functools import reduce
 from math import gcd
 from typing import Any
-from typing import Dict
-from typing import Iterator
-from typing import List
-from typing import Optional
-from typing import Tuple
-from typing import Union
 
 # from anemoi.utils.dates import as_datetime
 from anemoi.utils.dates import DateTimes
@@ -29,7 +24,7 @@ from anemoi.utils.hindcasts import HindcastDatesTimes
 from anemoi.utils.humanize import print_dates
 
 
-def extend(x: Union[str, List[Any], Tuple[Any, ...]]) -> Iterator[datetime.datetime]:
+def extend(x: str | list[Any] | tuple[Any, ...]) -> Iterator[datetime.datetime]:
     """Extend a date range or list of dates into individual datetime objects.
 
     Args:
@@ -86,7 +81,7 @@ class DatesProvider:
     3
     """
 
-    def __init__(self, missing: Optional[List[Union[str, datetime.datetime]]] = None) -> None:
+    def __init__(self, missing: list[str | datetime.datetime] | None = None) -> None:
         """Initialize the DatesProvider with optional missing dates.
 
         Parameters
@@ -168,7 +163,7 @@ class ValuesDates(DatesProvider):
         **kwargs (Any): Additional arguments.
     """
 
-    def __init__(self, values: List[Union[str, datetime.datetime]], **kwargs: Any) -> None:
+    def __init__(self, values: list[str | datetime.datetime], **kwargs: Any) -> None:
         """Initialize ValuesDates with a list of values.
 
         Args:
@@ -188,7 +183,7 @@ class ValuesDates(DatesProvider):
         """
         return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the ValuesDates instance to a dictionary.
 
         Returns
@@ -215,9 +210,9 @@ class StartEndDates(DatesProvider):
 
     def __init__(
         self,
-        start: Union[str, datetime.datetime],
-        end: Union[str, datetime.datetime],
-        frequency: Union[int, str] = 1,
+        start: str | datetime.datetime,
+        end: str | datetime.datetime,
+        frequency: int | str = 1,
         **kwargs: Any,
     ) -> None:
         """Initialize StartEndDates with start, end, and frequency.
@@ -259,7 +254,7 @@ class StartEndDates(DatesProvider):
 
         super().__init__(missing=missing)
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the StartEndDates instance to a dictionary.
 
         Returns
@@ -314,9 +309,9 @@ class HindcastsDates(DatesProvider):
 
     def __init__(
         self,
-        start: Union[str, List[str]],
-        end: Union[str, List[str]],
-        steps: List[int] = [0],
+        start: str | list[str],
+        end: str | list[str],
+        steps: list[int] = [0],
         years: int = 20,
         **kwargs: Any,
     ) -> None:
@@ -403,7 +398,7 @@ class HindcastsDates(DatesProvider):
         """
         return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
 
-    def as_dict(self) -> Dict[str, Any]:
+    def as_dict(self) -> dict[str, Any]:
         """Convert the HindcastsDates instance to a dictionary.
 
         Returns

anemoi/datasets/dates/groups.py

@@ -12,19 +12,16 @@ import datetime
 import itertools
 from abc import ABC
 from abc import abstractmethod
+from collections.abc import Callable
+from collections.abc import Iterator
 from functools import cached_property
 from typing import Any
-from typing import Callable
-from typing import Iterator
-from typing import List
-from typing import Tuple
-from typing import Union
 
 from anemoi.datasets.dates import DatesProvider
 from anemoi.datasets.dates import as_datetime
 
 
-def _shorten(dates: Union[List[datetime.datetime], Tuple[datetime.datetime, ...]]) -> Union[str, List[str]]:
+def _shorten(dates: list[datetime.datetime] | tuple[datetime.datetime, ...]) -> str | list[str]:
     """Shorten the list of dates for display.
 
     Args:
@@ -43,7 +40,7 @@ def _shorten(dates: Union[List[datetime.datetime], Tuple[datetime.datetime, ...]
 class GroupOfDates:
     """A class to represent a group of dates."""
 
-    def __init__(self, dates: List[datetime.datetime], provider: DatesProvider, partial_ok: bool = False) -> None:
+    def __init__(self, dates: list[datetime.datetime], provider: DatesProvider, partial_ok: bool = False) -> None:
         assert isinstance(provider, DatesProvider), type(provider)
         assert isinstance(dates, list)
 
@@ -197,10 +194,10 @@ class Groups:
 class Filter:
     """A class to filter out missing dates."""
 
-    def __init__(self, missing: List[datetime.datetime]) -> None:
-        self.missing = set(as_datetime(m) for m in missing)
+    def __init__(self, missing: list[datetime.datetime]) -> None:
+        self.missing = {as_datetime(m) for m in missing}
 
-    def __call__(self, dates: List[datetime.datetime]) -> List[datetime.datetime]:
+    def __call__(self, dates: list[datetime.datetime]) -> list[datetime.datetime]:
         """Filter out missing dates from the list of dates.
 
         Args:

anemoi/datasets/grids.py

@@ -11,10 +11,6 @@
 import base64
 import logging
 from typing import Any
-from typing import List
-from typing import Optional
-from typing import Tuple
-from typing import Union
 
 import numpy as np
 from numpy.typing import NDArray
@@ -94,7 +90,7 @@ def plot_mask(
 
 # TODO: Use the one from anemoi.utils.grids instead
 # from anemoi.utils.grids import ...
-def xyz_to_latlon(x: NDArray[Any], y: NDArray[Any], z: NDArray[Any]) -> Tuple[NDArray[Any], NDArray[Any]]:
+def xyz_to_latlon(x: NDArray[Any], y: NDArray[Any], z: NDArray[Any]) -> tuple[NDArray[Any], NDArray[Any]]:
     """Convert Cartesian coordinates to latitude and longitude.
 
     Parameters
@@ -121,7 +117,7 @@ def xyz_to_latlon(x: NDArray[Any], y: NDArray[Any], z: NDArray[Any]) -> Tuple[ND
 # from anemoi.utils.grids import ...
 def latlon_to_xyz(
     lat: NDArray[Any], lon: NDArray[Any], radius: float = 1.0
-) -> Tuple[NDArray[Any], NDArray[Any], NDArray[Any]]:
+) -> tuple[NDArray[Any], NDArray[Any], NDArray[Any]]:
     """Convert latitude and longitude to Cartesian coordinates.
 
     Parameters
@@ -272,8 +268,8 @@ def cutout_mask(
     global_lons: NDArray[Any],
     cropping_distance: float = 2.0,
     neighbours: int = 5,
-    min_distance_km: Optional[Union[int, float]] = None,
-    plot: Optional[str] = None,
+    min_distance_km: int | float | None = None,
+    plot: str | None = None,
 ) -> NDArray[Any]:
     """Return a mask for the points in [global_lats, global_lons] that are inside of [lats, lons].
 
@@ -465,7 +461,7 @@ def thinning_mask(
     return np.array([i for i in indices])
 
 
-def outline(lats: NDArray[Any], lons: NDArray[Any], neighbours: int = 5) -> List[int]:
+def outline(lats: NDArray[Any], lons: NDArray[Any], neighbours: int = 5) -> list[int]:
     """Find the outline of the grid points.
 
     Parameters