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

anemoi/datasets/data/debug.py

@@ -11,12 +11,10 @@
 import logging
 import os
 import textwrap
+from collections.abc import Callable
 from functools import wraps
 from typing import TYPE_CHECKING
 from typing import Any
-from typing import Callable
-from typing import List
-from typing import Optional
 
 from anemoi.utils.text import Tree
 from numpy.typing import NDArray
@@ -56,7 +54,7 @@ def css(name: str) -> str:
 class Node:
     """A class to represent a node in a dataset tree."""
 
-    def __init__(self, dataset: "Dataset", kids: List[Any], **kwargs: Any) -> None:
+    def __init__(self, dataset: "Dataset", kids: list[Any], **kwargs: Any) -> None:
         """Initializes a Node object.
 
         Parameters
@@ -72,7 +70,7 @@ class Node:
         self.kids = kids
         self.kwargs = kwargs
 
-    def _put(self, indent: int, result: List[str]) -> None:
+    def _put(self, indent: int, result: list[str]) -> None:
         """Helper method to add the node representation to the result list.
 
         Parameters
@@ -103,11 +101,11 @@ class Node:
         str
             String representation of the node.
         """
-        result: List[str] = []
+        result: list[str] = []
         self._put(0, result)
         return "\n".join(result)
 
-    def graph(self, digraph: List[str], nodes: dict) -> None:
+    def graph(self, digraph: list[str], nodes: dict) -> None:
         """Generates a graph representation of the node.
 
         Parameters
@@ -170,7 +168,7 @@ class Node:
         digraph.append("}")
         return "\n".join(digraph)
 
-    def _html(self, indent: str, rows: List[List[str]]) -> None:
+    def _html(self, indent: str, rows: list[list[str]]) -> None:
         """Helper method to add the node representation to the HTML rows.
 
         Parameters
@@ -273,7 +271,7 @@ class Node:
 class Source:
     """A class used to follow the provenance of a data point."""
 
-    def __init__(self, dataset: Any, index: int, source: Optional[Any] = None, info: Optional[Any] = None) -> None:
+    def __init__(self, dataset: Any, index: int, source: Any | None = None, info: Any | None = None) -> None:
         """Initializes a Source object.
 
         Parameters

anemoi/datasets/data/ensemble.py

@@ -10,8 +10,6 @@
 
 import logging
 from typing import Any
-from typing import Dict
-from typing import Tuple
 
 import numpy as np
 from numpy.typing import NDArray
@@ -105,7 +103,7 @@ class Number(Forwards):
         """
         return Node(self, [self.forward.tree()], numbers=[n + 1 for n in self.members])
 
-    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+    def metadata_specific(self, **kwargs: Any) -> dict[str, Any]:
         """Returns metadata specific to the Number object.
 
         Parameters
@@ -122,7 +120,7 @@ class Number(Forwards):
             "numbers": [n + 1 for n in self.members],
         }
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Returns metadata specific to the Number object."""
         return {}
 
@@ -140,7 +138,7 @@ class Ensemble(GivenAxis):
         """
         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]:
         """Get the metadata specific to the forwards subclass.
 
         Returns:
@@ -149,7 +147,7 @@ class Ensemble(GivenAxis):
         return {}
 
 
-def ensemble_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Ensemble:
+def ensemble_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Ensemble:
     """Factory function to create an Ensemble object.
 
     Parameters

anemoi/datasets/data/fill_missing.py

@@ -10,9 +10,6 @@
 
 import logging
 from typing import Any
-from typing import Dict
-from typing import Optional
-from typing import Set
 
 import numpy as np
 from numpy.typing import NDArray
@@ -46,7 +43,7 @@ class MissingDatesFill(Forwards):
         """
         super().__init__(dataset)
         self._missing = set(dataset.missing)
-        self._warnings: Set[int] = set()
+        self._warnings: set[int] = set()
 
     @debug_indexing
     @expand_list_indexing
@@ -84,7 +81,7 @@ class MissingDatesFill(Forwards):
         return np.stack([self[i] for i in range(*s.indices(self._len))])
 
     @property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Get the set of missing dates."""
         return set()
 
@@ -153,7 +150,7 @@ class MissingDatesClosest(MissingDatesFill):
         self.closest = closest
         self._closest = {}
 
-    def _fill_missing(self, n: int, a: Optional[int], b: Optional[int]) -> NDArray[Any]:
+    def _fill_missing(self, n: int, a: int | None, b: int | None) -> NDArray[Any]:
         """Fill the missing date at the given index.
 
         Parameters
@@ -189,7 +186,7 @@ class MissingDatesClosest(MissingDatesFill):
 
         return self.forward[self._closest[n]]
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get metadata specific to the subclass.
 
         Returns
@@ -224,7 +221,7 @@ class MissingDatesInterpolate(MissingDatesFill):
         super().__init__(dataset)
         self._alpha = {}
 
-    def _fill_missing(self, n: int, a: Optional[int], b: Optional[int]) -> NDArray[Any]:
+    def _fill_missing(self, n: int, a: int | None, b: int | None) -> NDArray[Any]:
         """Fill the missing date at the given index using interpolation.
 
         Parameters
@@ -264,7 +261,7 @@ class MissingDatesInterpolate(MissingDatesFill):
         alpha = self._alpha[n]
         return self.forward[a] * (1 - alpha) + self.forward[b] * alpha
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get metadata specific to the subclass.
 
         Returns
@@ -285,7 +282,7 @@ class MissingDatesInterpolate(MissingDatesFill):
         return Node(self, [self.forward.tree()])
 
 
-def fill_missing_dates_factory(dataset: Any, method: str, kwargs: Dict[str, Any]) -> Dataset:
+def fill_missing_dates_factory(dataset: Any, method: str, kwargs: dict[str, Any]) -> Dataset:
     """Factory function to create an instance of a class to fill missing dates.
 
     Parameters

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
@@ -25,171 +21,19 @@ from .dataset import FullIndex
 from .dataset import Shape
 from .dataset import TupleIndex
 from .debug import Node
-from .debug import debug_indexing
-from .forwards import Combined
 from .forwards import GivenAxis
 from .indexing import apply_index_to_slices_changes
-from .indexing import expand_list_indexing
 from .indexing import index_to_slices
-from .indexing import length_to_slices
-from .indexing import update_tuple
 from .misc import _auto_adjust
 from .misc import _open
 
 LOG = logging.getLogger(__name__)
 
 
-class Concat(Combined):
-    """A class to represent concatenated datasets."""
-
-    def __len__(self) -> int:
-        """Returns the total length of the concatenated datasets.
-
-        Returns
-        -------
-        int
-            Total length of the concatenated datasets.
-        """
-        return sum(len(i) for i in self.datasets)
-
-    @debug_indexing
-    @expand_list_indexing
-    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
-        """Retrieves a tuple of data from the concatenated datasets based on the given index.
-
-        Parameters
-        ----------
-        index : TupleIndex
-            Index specifying the data to retrieve.
-
-        Returns
-        -------
-        NDArray[Any]
-            Concatenated data array from the specified index.
-        """
-        index, changes = index_to_slices(index, self.shape)
-        # print(index, changes)
-        lengths = [d.shape[0] for d in self.datasets]
-        slices = length_to_slices(index[0], lengths)
-        # print("slies", slices)
-        result = [d[update_tuple(index, 0, i)[0]] for (d, i) in zip(self.datasets, slices) if i is not None]
-        result = np.concatenate(result, axis=0)
-        return apply_index_to_slices_changes(result, changes)
-
-    @debug_indexing
-    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
-        """Retrieves data from the concatenated datasets based on the given index.
-
-        Parameters
-        ----------
-        n : FullIndex
-            Index specifying the data to retrieve.
-
-        Returns
-        -------
-        NDArray[Any]
-            Data array from the concatenated datasets based on the index.
-        """
-        if isinstance(n, tuple):
-            return self._get_tuple(n)
-
-        if isinstance(n, slice):
-            return self._get_slice(n)
-
-        # TODO: optimize
-        k = 0
-        while n >= self.datasets[k]._len:
-            n -= self.datasets[k]._len
-            k += 1
-        return self.datasets[k][n]
-
-    @debug_indexing
-    def _get_slice(self, s: slice) -> NDArray[Any]:
-        """Retrieves a slice of data from the concatenated datasets.
-
-        Parameters
-        ----------
-        s : slice
-            Slice object specifying the range of data to retrieve.
-
-        Returns
-        -------
-        NDArray[Any]
-            Concatenated data array from the specified slice.
-        """
-        result = []
-
-        lengths = [d.shape[0] for d in self.datasets]
-        slices = length_to_slices(s, lengths)
-
-        result = [d[i] for (d, i) in zip(self.datasets, slices) if i is not None]
-
-        return np.concatenate(result)
-
-    def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
-        """Check the compatibility of two datasets for concatenation.
-
-        Parameters
-        ----------
-        d1 : Dataset
-            The first dataset.
-        d2 : Dataset
-            The second dataset.
-        """
-        super().check_compatibility(d1, d2)
-        self.check_same_sub_shapes(d1, d2, drop_axis=0)
-
-    def check_same_lengths(self, d1: Dataset, d2: Dataset) -> None:
-        """Check if the lengths of two datasets are the same.
-
-        Parameters
-        ----------
-        d1 : Dataset
-            The first dataset.
-        d2 : Dataset
-            The second dataset.
-        """
-        # Turned off because we are concatenating along the first axis
-        pass
-
-    def check_same_dates(self, d1: Dataset, d2: Dataset) -> None:
-        """Check if the dates of two datasets are the same.
-
-        Parameters
-        ----------
-        d1 : Dataset
-            The first dataset.
-        d2 : Dataset
-            The second dataset.
-        """
-        # Turned off because we are concatenating along the dates axis
-        pass
-
-    @property
-    def dates(self) -> NDArray[np.datetime64]:
-        """Returns the concatenated dates of all datasets."""
-        return np.concatenate([d.dates for d in self.datasets])
-
-    @property
-    def shape(self) -> Shape:
-        """Returns the shape of the concatenated datasets."""
-        return (len(self),) + self.datasets[0].shape[1:]
-
-    def tree(self) -> Node:
-        """Generates a hierarchical tree structure for the concatenated datasets.
-
-        Returns
-        -------
-        Node
-            A Node object representing the concatenated datasets.
-        """
-        return Node(self, [d.tree() for d in self.datasets])
-
-
 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 +73,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 +90,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 +119,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 +136,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 +150,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 +331,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 +421,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 +432,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 +462,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