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

anemoi/datasets/data/forwards.py

@@ -14,10 +14,6 @@ import warnings
 from abc import abstractmethod
 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
 
 import numpy as np
 from numpy.typing import NDArray
@@ -75,7 +71,7 @@ class Forwards(Dataset):
         return self.forward[n]
 
     @property
-    def name(self) -> Optional[str]:
+    def name(self) -> str | None:
         """Returns the name of the forward dataset."""
         if self._name is not None:
             return self._name
@@ -112,26 +108,26 @@ class Forwards(Dataset):
         return self.forward.longitudes
 
     @property
-    def name_to_index(self) -> Dict[str, int]:
+    def name_to_index(self) -> dict[str, int]:
         """Returns a dictionary mapping variable names to their indices."""
         return self.forward.name_to_index
 
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Returns the variables of the forward dataset."""
         return self.forward.variables
 
     @property
-    def variables_metadata(self) -> Dict[str, Any]:
+    def variables_metadata(self) -> dict[str, Any]:
         """Returns the metadata of the variables in the forward dataset."""
         return self.forward.variables_metadata
 
     @property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Returns the statistics of the forward dataset."""
         return self.forward.statistics
 
-    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]]:
         """Returns the statistics tendencies of the forward dataset.
 
         Parameters
@@ -159,7 +155,7 @@ class Forwards(Dataset):
         return self.forward.dtype
 
     @property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Returns the missing data information of the forward dataset."""
         return self.forward.missing
 
@@ -168,7 +164,7 @@ class Forwards(Dataset):
         """Returns the grids of the forward dataset."""
         return self.forward.grids
 
-    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+    def metadata_specific(self, **kwargs: Any) -> dict[str, Any]:
         """Returns metadata specific to the forward dataset.
 
         Parameters
@@ -187,7 +183,7 @@ class Forwards(Dataset):
             **kwargs,
         )
 
-    def collect_supporting_arrays(self, collected: List[Any], *path: Any) -> None:
+    def collect_supporting_arrays(self, collected: list[Any], *path: Any) -> None:
         """Collects supporting arrays from the forward dataset.
 
         Parameters
@@ -199,7 +195,7 @@ class Forwards(Dataset):
         """
         self.forward.collect_supporting_arrays(collected, *path)
 
-    def collect_input_sources(self, collected: List[Any]) -> None:
+    def collect_input_sources(self, collected: list[Any]) -> None:
         """Collects input sources from the forward dataset.
 
         Parameters
@@ -225,11 +221,11 @@ class Forwards(Dataset):
         return self.forward.source(index)
 
     @abstractmethod
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Returns metadata specific to the subclass."""
         pass
 
-    def get_dataset_names(self, names: Set[str]) -> None:
+    def get_dataset_names(self, names: set[str]) -> None:
         """Collects the names of the datasets.
 
         Parameters
@@ -240,7 +236,7 @@ class Forwards(Dataset):
         self.forward.get_dataset_names(names)
 
     @property
-    def constant_fields(self) -> List[str]:
+    def constant_fields(self) -> list[str]:
         """Returns the constant fields of the forward dataset."""
         return self.forward.constant_fields
 
@@ -248,7 +244,7 @@ class Forwards(Dataset):
 class Combined(Forwards):
     """A class to combine multiple datasets into a single dataset."""
 
-    def __init__(self, datasets: List[Dataset]) -> None:
+    def __init__(self, datasets: list[Dataset]) -> None:
         """Initializes a Combined object.
 
         Parameters
@@ -466,7 +462,7 @@ class Combined(Forwards):
         self.check_same_variables(d1, d2)
         self.check_same_dates(d1, d2)
 
-    def provenance(self) -> List[Any]:
+    def provenance(self) -> list[Any]:
         """Returns the provenance of the combined datasets.
 
         Returns
@@ -487,7 +483,7 @@ class Combined(Forwards):
         lst = ", ".join(repr(d) for d in self.datasets)
         return f"{self.__class__.__name__}({lst})"
 
-    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+    def metadata_specific(self, **kwargs: Any) -> dict[str, Any]:
         """Returns metadata specific to the combined datasets.
 
         Parameters
@@ -508,7 +504,7 @@ class Combined(Forwards):
             **kwargs,
         )
 
-    def collect_supporting_arrays(self, collected: List[Any], *path: Any) -> None:
+    def collect_supporting_arrays(self, collected: list[Any], *path: Any) -> None:
         """Collects supporting arrays from the combined datasets.
 
         Parameters
@@ -524,7 +520,7 @@ class Combined(Forwards):
             d.collect_supporting_arrays(collected, *path, name)
 
     @property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Returns the missing data information of the combined datasets.
 
         Raises
@@ -534,7 +530,7 @@ class Combined(Forwards):
         """
         raise NotImplementedError("missing() not implemented for Combined")
 
-    def get_dataset_names(self, names: Set[str]) -> None:
+    def get_dataset_names(self, names: set[str]) -> None:
         """Collects the names of the combined datasets.
 
         Parameters
@@ -549,7 +545,7 @@ class Combined(Forwards):
 class GivenAxis(Combined):
     """A class to combine datasets along a given axis."""
 
-    def __init__(self, datasets: List[Any], axis: int) -> None:
+    def __init__(self, datasets: list[Any], axis: int) -> None:
         """Initializes a GivenAxis object.
 
         Parameters
@@ -656,10 +652,10 @@ class GivenAxis(Combined):
         return np.concatenate([d[n] for d in self.datasets], axis=self.axis - 1)
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Returns the missing data information of the combined dataset along the given axis."""
         offset = 0
-        result: Set[int] = set()
+        result: set[int] = set()
         for d in self.datasets:
             result.update(offset + m for m in d.missing)
             if self.axis == 0:  # Advance if axis is time

anemoi/datasets/data/grids.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 Optional
-from typing import Tuple
 
 import numpy as np
 from numpy.typing import NDArray
@@ -189,7 +185,7 @@ class Concat(Combined):
 class GridsBase(GivenAxis):
     """A base class for handling grids in datasets."""
 
-    def __init__(self, datasets: List[Any], axis: int) -> None:
+    def __init__(self, datasets: list[Any], axis: int) -> None:
         """Initializes a GridsBase object.
 
         Parameters
@@ -229,7 +225,7 @@ class GridsBase(GivenAxis):
         # We don't check the resolution, because we want to be able to combine
         pass
 
-    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+    def metadata_specific(self, **kwargs: Any) -> dict[str, Any]:
         """Returns metadata specific to the GridsBase object.
 
         Parameters
@@ -246,7 +242,7 @@ class GridsBase(GivenAxis):
             multi_grids=True,
         )
 
-    def collect_input_sources(self, collected: List[Any]) -> None:
+    def collect_input_sources(self, collected: list[Any]) -> None:
         """Collects input sources from the datasets.
 
         Parameters
@@ -275,7 +271,7 @@ class Grids(GridsBase):
         return np.concatenate([d.longitudes for d in self.datasets])
 
     @property
-    def grids(self) -> Tuple[Any, ...]:
+    def grids(self) -> tuple[Any, ...]:
         """Returns the grids of all datasets."""
         result = []
         for d in self.datasets:
@@ -292,7 +288,7 @@ class Grids(GridsBase):
         """
         return Node(self, [d.tree() for d in self.datasets], mode="concat")
 
-    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:
@@ -306,12 +302,12 @@ class Cutout(GridsBase):
 
     def __init__(
         self,
-        datasets: List[Any],
+        datasets: list[Any],
         axis: int = 3,
         cropping_distance: float = 2.0,
         neighbours: int = 5,
-        min_distance_km: Optional[float] = None,
-        plot: Optional[bool] = None,
+        min_distance_km: float | None = None,
+        plot: bool | None = None,
     ) -> None:
         """Initializes a Cutout object for hierarchical management of Limited Area
         Models (LAMs) and a global dataset, handling overlapping regions.
@@ -487,7 +483,7 @@ class Cutout(GridsBase):
 
         return apply_index_to_slices_changes(result, changes)
 
-    def collect_supporting_arrays(self, collected: List[Any], *path: Any) -> None:
+    def collect_supporting_arrays(self, collected: list[Any], *path: Any) -> None:
         """Collect supporting arrays, including masks for each LAM and the global dataset.
 
         Parameters
@@ -577,7 +573,7 @@ class Cutout(GridsBase):
         """
         return Node(self, [d.tree() for d in self.datasets])
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Returns metadata specific to the Cutout object.
 
         Returns
@@ -588,7 +584,7 @@ class Cutout(GridsBase):
         return {}
 
 
-def grids_factory(args: Tuple[Any, ...], kwargs: dict) -> Dataset:
+def grids_factory(args: tuple[Any, ...], kwargs: dict) -> Dataset:
     """Factory function to create a Grids object.
 
     Parameters
@@ -618,7 +614,7 @@ def grids_factory(args: Tuple[Any, ...], kwargs: dict) -> Dataset:
     return Grids(datasets, axis=axis)._subset(**kwargs)
 
 
-def cutout_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Dataset:
+def cutout_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Dataset:
     """Factory function to create a Cutout object.
 
     Parameters

anemoi/datasets/data/indexing.py

@@ -8,12 +8,9 @@
 # nor does it submit to any jurisdiction.
 
 
+from collections.abc import Callable
 from functools import wraps
 from typing import Any
-from typing import Callable
-from typing import List
-from typing import Tuple
-from typing import Union
 
 import numpy as np
 from numpy.typing import NDArray
@@ -23,7 +20,7 @@ from .dataset import Shape
 from .dataset import TupleIndex
 
 
-def _tuple_with_slices(t: TupleIndex, shape: Shape) -> Tuple[TupleIndex, Tuple[int, ...]]:
+def _tuple_with_slices(t: TupleIndex, shape: Shape) -> tuple[TupleIndex, tuple[int, ...]]:
     """Replace all integers in a tuple with slices, so we preserve the dimensionality.
 
     Parameters:
@@ -87,7 +84,7 @@ def _index_to_tuple(index: FullIndex, shape: Shape) -> TupleIndex:
     raise ValueError(f"Invalid index: {index}")
 
 
-def index_to_slices(index: Union[int, slice, Tuple], shape: Shape) -> Tuple[TupleIndex, Tuple[int, ...]]:
+def index_to_slices(index: int | slice | tuple, shape: Shape) -> tuple[TupleIndex, tuple[int, ...]]:
     """Convert an index to a tuple of slices, with the same dimensionality as the shape.
 
     Parameters:
@@ -100,7 +97,7 @@ def index_to_slices(index: Union[int, slice, Tuple], shape: Shape) -> Tuple[Tupl
     return _tuple_with_slices(_index_to_tuple(index, shape), shape)
 
 
-def apply_index_to_slices_changes(result: NDArray[Any], changes: Tuple[int, ...]) -> NDArray[Any]:
+def apply_index_to_slices_changes(result: NDArray[Any], changes: tuple[int, ...]) -> NDArray[Any]:
     """Apply changes to the result array based on the slices.
 
     Parameters:
@@ -118,7 +115,7 @@ def apply_index_to_slices_changes(result: NDArray[Any], changes: Tuple[int, ...]
     return result
 
 
-def update_tuple(t: Tuple, index: int, value: Any) -> Tuple[Tuple, Any]:
+def update_tuple(t: tuple, index: int, value: Any) -> tuple[tuple, Any]:
     """Replace the elements of a tuple at the given index with a new value.
 
     Parameters:
@@ -135,7 +132,7 @@ def update_tuple(t: Tuple, index: int, value: Any) -> Tuple[Tuple, Any]:
     return tuple(t), prev
 
 
-def length_to_slices(index: slice, lengths: List[int]) -> List[Union[slice, None]]:
+def length_to_slices(index: slice, lengths: list[int]) -> list[slice | None]:
     """Convert an index to a list of slices, given the lengths of the dimensions.
 
     Parameters:
@@ -174,7 +171,7 @@ def length_to_slices(index: slice, lengths: List[int]) -> List[Union[slice, None
     return result
 
 
-def _as_tuples(index: Tuple) -> Tuple:
+def _as_tuples(index: tuple) -> tuple:
     """Convert elements of the index to tuples if they are lists or arrays.
 
     Parameters:
@@ -219,7 +216,7 @@ def expand_list_indexing(method: Callable[..., NDArray[Any]]) -> Callable[..., N
         if not any(isinstance(i, (list, tuple)) for i in index):
             return method(self, index)
 
-        which: List[int] = []
+        which: list[int] = []
         for i, idx in enumerate(index):
             if isinstance(idx, (list, tuple)):
                 which.append(i)
@@ -241,7 +238,7 @@ def expand_list_indexing(method: Callable[..., NDArray[Any]]) -> Callable[..., N
     return wrapper
 
 
-def make_slice_or_index_from_list_or_tuple(indices: List[int]) -> Union[List[int], slice]:
+def make_slice_or_index_from_list_or_tuple(indices: list[int]) -> list[int] | slice:
     """Convert a list or tuple of indices to a slice or an index, if possible.
 
     Parameters:

anemoi/datasets/data/interpolate.py

@@ -12,12 +12,6 @@ import datetime
 import logging
 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 Tuple
-from typing import Union
 
 import numpy as np
 from anemoi.utils.dates import frequency_to_timedelta
@@ -193,7 +187,7 @@ class InterpolateFrequency(Forwards):
         return Node(self, [self.forward.tree()], frequency=self.frequency)
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the missing data indices."""
         result = []
         j = 0
@@ -204,10 +198,10 @@ class InterpolateFrequency(Forwards):
                     result.append(j)
                 j += 1
 
-        result = set(x for x in result if x < self._len)
+        result = {x for x in result if x < self._len}
         return result
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the InterpolateFrequency subclass.
 
         Returns
@@ -221,9 +215,7 @@ class InterpolateFrequency(Forwards):
 
 
 class InterpolateNearest(Forwards):
-    def __init__(
-        self, dataset: Dataset, interpolate_variables: List[str], max_distance: Optional[float] = None
-    ) -> None:
+    def __init__(self, dataset: Dataset, interpolate_variables: list[str], max_distance: float | None = None) -> None:
         """Initialize the InterpolateNearest class.
 
         Parameters
@@ -262,7 +254,7 @@ class InterpolateNearest(Forwards):
         return self.forward.shape
 
     @property
-    def metadata(self) -> Dict[str, Any]:
+    def metadata(self) -> dict[str, Any]:
         return self.forward.metadata()
 
     @staticmethod
@@ -281,12 +273,12 @@ class InterpolateNearest(Forwards):
         result = target_data[(slice(None),) + index[1:]]
         return apply_index_to_slices_changes(result, changes)
 
-    def __getitem__(self, index: Union[int, slice, Tuple[Union[int, slice], ...]]) -> NDArray[Any]:
+    def __getitem__(self, index: int | slice | tuple[int | slice, ...]) -> NDArray[Any]:
         if isinstance(index, (int, slice)):
             index = (index, slice(None), slice(None), slice(None))
         return self._get_tuple(index)
 
-    def subclass_metadata_specific(self) -> Dict[str, Any]:
+    def subclass_metadata_specific(self) -> dict[str, Any]:
         return {
             "interpolate_variables": self.vars,
         }

anemoi/datasets/data/join.py

@@ -12,10 +12,6 @@ import datetime
 import logging
 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
 
 import numpy as np
 from numpy.typing import NDArray
@@ -182,10 +178,10 @@ class Join(Combined):
         return Select(self, indices, {"overlay": variables})
 
     @cached_property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Get the variables of the joined dataset."""
         seen = set()
-        result: List[str] = []
+        result: list[str] = []
         for d in reversed(self.datasets):
             for v in reversed(d.variables):
                 while v in seen:
@@ -196,7 +192,7 @@ class Join(Combined):
         return result
 
     @property
-    def variables_metadata(self) -> Dict[str, Any]:
+    def variables_metadata(self) -> dict[str, Any]:
         """Get the metadata of the variables."""
         result = {}
         variables = [v for v in self.variables if not (v.startswith("(") and v.endswith(")"))]
@@ -216,18 +212,18 @@ class Join(Combined):
         return result
 
     @cached_property
-    def name_to_index(self) -> Dict[str, int]:
+    def name_to_index(self) -> dict[str, int]:
         """Get the mapping of variable names to indices."""
         return {k: i for i, k in enumerate(self.variables)}
 
     @property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Get the statistics of the joined dataset."""
         return {
             k: np.concatenate([d.statistics[k] for d in self.datasets], axis=0) for k in self.datasets[0].statistics
         }
 
-    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 joined dataset.
 
         Parameters
@@ -268,9 +264,9 @@ class Join(Combined):
         assert False
 
     @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/masked.py

@@ -11,11 +11,6 @@
 import logging
 from functools import cached_property
 from typing import Any
-from typing import Dict
-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
@@ -117,7 +112,7 @@ class Masked(Forwards):
         result = apply_index_to_slices_changes(result, changes)
         return result
 
-    def collect_supporting_arrays(self, collected: List[Tuple], *path: Any) -> None:
+    def collect_supporting_arrays(self, collected: list[tuple], *path: Any) -> None:
         """Collect supporting arrays.
 
         Parameters
@@ -134,7 +129,7 @@ class Masked(Forwards):
 class Thinning(Masked):
     """A class to represent a thinned dataset."""
 
-    def __init__(self, forward: Dataset, thinning: Optional[int], method: str) -> None:
+    def __init__(self, forward: Dataset, thinning: int | None, method: str) -> None:
         """Initialize the Thinning class.
 
         Parameters
@@ -195,7 +190,7 @@ class Thinning(Masked):
         """
         return Node(self, [self.forward.tree()], thinning=self.thinning, method=self.method)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the Thinning subclass.
 
         Returns
@@ -209,7 +204,7 @@ class Thinning(Masked):
 class Cropping(Masked):
     """A class to represent a cropped dataset."""
 
-    def __init__(self, forward: Dataset, area: Union[Dataset, Tuple[float, float, float, float]]) -> None:
+    def __init__(self, forward: Dataset, area: Dataset | tuple[float, float, float, float]) -> None:
         """Initialize the Cropping class.
 
         Parameters
@@ -245,7 +240,7 @@ class Cropping(Masked):
         """
         return Node(self, [self.forward.tree()], area=self.area)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the Cropping subclass.
 
         Returns
@@ -314,7 +309,7 @@ class TrimEdge(Masked):
         """
         return Node(self, [self.forward.tree()], edge=self.edge)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the TrimEdge subclass.
 
         Returns

anemoi/datasets/data/merge.py

@@ -12,10 +12,6 @@ import datetime
 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
 
 import numpy as np
 from numpy.typing import NDArray
@@ -40,7 +36,7 @@ LOG = logging.getLogger(__name__)
 class Merge(Combined):
     """A class to merge multiple datasets along the dates axis, handling gaps in dates if allowed."""
 
-    def __init__(self, datasets: List[Dataset], allow_gaps_in_dates: bool = False) -> None:
+    def __init__(self, datasets: list[Dataset], allow_gaps_in_dates: bool = False) -> None:
         """Initialize the Merge object.
 
         Parameters
@@ -128,10 +124,10 @@ class Merge(Combined):
         return self._frequency
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the indices of missing dates in the merged dataset."""
         # TODO: optimize
-        result: Set[int] = set()
+        result: set[int] = set()
 
         for i, (dataset, row) in enumerate(self._indices):
             if dataset == self._missing_index:
@@ -192,7 +188,7 @@ class Merge(Combined):
         """
         return Node(self, [d.tree() for d in self.datasets], allow_gaps_in_dates=self.allow_gaps_in_dates)
 
-    def metadata_specific(self) -> Dict[str, Any]:
+    def metadata_specific(self) -> dict[str, Any]:
         """Get the specific metadata for the merged dataset.
 
         Returns
@@ -265,7 +261,7 @@ class Merge(Combined):
         return np.stack([self[i] for i in range(*s.indices(self._len))])
 
 
-def merge_factory(args: Tuple, kwargs: Dict[str, Any]) -> Dataset:
+def merge_factory(args: tuple, kwargs: dict[str, Any]) -> Dataset:
     """Factory function to create a merged dataset.
 
     Parameters