anemoi-datasets 0.5.16__py3-none-any.whl → 0.5.17__py3-none-any.whl

anemoi/datasets/data/indexing.py

@@ -9,13 +9,30 @@
 
 
 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
 
+from .dataset import FullIndex
+from .dataset import Shape
+from .dataset import TupleIndex
 
-def _tuple_with_slices(t, shape):
-    """Replace all integers in a tuple with slices, so we preserve the dimensionality."""
 
+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:
+    t (TupleIndex): The tuple index to process.
+    shape (Shape): The shape of the array.
+
+    Returns:
+    Tuple[TupleIndex, Tuple[int, ...]]: A tuple containing the modified index and the changes.
+    """
     result = tuple(slice(i, i + 1) if isinstance(i, int) else i for i in t)
     changes = tuple(j for (j, i) in enumerate(t) if isinstance(i, int))
     result = tuple(slice(*s.indices(shape[i])) for (i, s) in enumerate(result))
@@ -23,7 +40,16 @@ def _tuple_with_slices(t, shape):
     return result, changes
 
 
-def _extend_shape(index, shape):
+def _extend_shape(index: TupleIndex, shape: Shape) -> TupleIndex:
+    """Extend the shape of the index to match the shape of the array.
+
+    Parameters:
+    index (TupleIndex): The index to extend.
+    shape (Shape): The shape of the array.
+
+    Returns:
+    TupleIndex: The extended index.
+    """
     if Ellipsis in index:
         if index.count(Ellipsis) > 1:
             raise IndexError("Only one Ellipsis is allowed")
@@ -40,7 +66,16 @@ def _extend_shape(index, shape):
     return index
 
 
-def _index_to_tuple(index, shape):
+def _index_to_tuple(index: FullIndex, shape: Shape) -> TupleIndex:
+    """Convert an index to a tuple index.
+
+    Parameters:
+    index (FullIndex): The index to convert.
+    shape (Shape): The shape of the array.
+
+    Returns:
+    TupleIndex: The converted tuple index.
+    """
     if isinstance(index, int):
         return _extend_shape((index,), shape)
     if isinstance(index, slice):
@@ -52,12 +87,29 @@ def _index_to_tuple(index, shape):
     raise ValueError(f"Invalid index: {index}")
 
 
-def index_to_slices(index, shape):
-    """Convert an index to a tuple of slices, with the same dimensionality as the shape."""
+def index_to_slices(index: Union[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:
+    index (Union[int, slice, Tuple]): The index to convert.
+    shape (Shape): The shape of the array.
+
+    Returns:
+    Tuple[TupleIndex, Tuple[int, ...]]: A tuple containing the slices and the changes.
+    """
     return _tuple_with_slices(_index_to_tuple(index, shape), shape)
 
 
-def apply_index_to_slices_changes(result, changes):
+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:
+    result (NDArray[Any]): The result array.
+    changes (Tuple[int, ...]): The changes to apply.
+
+    Returns:
+    NDArray[Any]: The modified result array.
+    """
     if changes:
         shape = result.shape
         for i in changes:
@@ -66,16 +118,33 @@ def apply_index_to_slices_changes(result, changes):
     return result
 
 
-def update_tuple(t, index, value):
-    """Replace the elements of a tuple at the given index with a new value."""
+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:
+    tp (Tuple): The original tuple.
+    index (int): The index to update.
+    value (Any): The new value.
+
+    Returns:
+    Tuple[Tuple, Any]: The updated tuple and the previous value.
+    """
     t = list(t)
     prev = t[index]
     t[index] = value
     return tuple(t), prev
 
 
-def length_to_slices(index, lengths):
-    """Convert an index to a list of slices, given the lengths of the dimensions."""
+def length_to_slices(index: slice, lengths: List[int]) -> List[Union[slice, None]]:
+    """Convert an index to a list of slices, given the lengths of the dimensions.
+
+    Parameters:
+    index (slice): The index to convert.
+    lengths (List[int]): The lengths of the dimensions.
+
+    Returns:
+    List[Union[slice, None]]: A list of slices.
+    """
     total = sum(lengths)
     start, stop, step = index.indices(total)
 
@@ -105,8 +174,17 @@ def length_to_slices(index, lengths):
     return result
 
 
-def _as_tuples(index):
-    def _(i):
+def _as_tuples(index: Tuple) -> Tuple:
+    """Convert elements of the index to tuples if they are lists or arrays.
+
+    Parameters:
+    index (Tuple): The index to convert.
+
+    Returns:
+    Tuple: The converted index.
+    """
+
+    def _(i: Any) -> Any:
         if hasattr(i, "tolist"):
             # NumPy arrays, TensorFlow tensors, etc.
             i = i.tolist()
@@ -121,18 +199,27 @@ def _as_tuples(index):
     return tuple(_(i) for i in index)
 
 
-def expand_list_indexing(method):
-    """Allows to use slices, lists, and tuples to select data from the dataset. Zarr does not support indexing with lists/arrays directly, so we need to implement it ourselves."""
+def expand_list_indexing(method: Callable[..., NDArray[Any]]) -> Callable[..., NDArray[Any]]:
+    """Allows to use slices, lists, and tuples to select data from the dataset.
+    Zarr does not support indexing with lists/arrays directly,
+    so we need to implement it ourselves.
+
+    Parameters:
+    method (Callable[..., NDArray[Any]]): The method to wrap.
+
+    Returns:
+    Callable[..., NDArray[Any]]: The wrapped method.
+    """
 
     @wraps(method)
-    def wrapper(self, index):
+    def wrapper(self: Any, index: FullIndex) -> NDArray[Any]:
         if not isinstance(index, tuple):
             return method(self, index)
 
         if not any(isinstance(i, (list, tuple)) for i in index):
             return method(self, index)
 
-        which = []
+        which: List[int] = []
         for i, idx in enumerate(index):
             if isinstance(idx, (list, tuple)):
                 which.append(i)
@@ -154,9 +241,15 @@ def expand_list_indexing(method):
     return wrapper
 
 
-def make_slice_or_index_from_list_or_tuple(indices):
-    """Convert a list or tuple of indices to a slice or an index, if possible."""
+def make_slice_or_index_from_list_or_tuple(indices: List[int]) -> Union[List[int], slice]:
+    """Convert a list or tuple of indices to a slice or an index, if possible.
+
+    Parameters:
+    indices (List[int]): The list or tuple of indices.
 
+    Returns:
+    Union[List[int], slice]: The slice or index.
+    """
     if len(indices) < 2:
         return indices
 

anemoi/datasets/data/interpolate.py

@@ -8,12 +8,21 @@
 # nor does it submit to any jurisdiction.
 
 
+import datetime
 import logging
 from functools import cached_property
+from typing import Any
+from typing import Dict
+from typing import Set
 
 import numpy as np
 from anemoi.utils.dates import frequency_to_timedelta
+from numpy.typing import NDArray
 
+from .dataset import Dataset
+from .dataset import FullIndex
+from .dataset import Shape
+from .dataset import TupleIndex
 from .debug import Node
 from .debug import debug_indexing
 from .forwards import Forwards
@@ -26,8 +35,18 @@ LOG = logging.getLogger(__name__)
 
 
 class InterpolateFrequency(Forwards):
-
-    def __init__(self, dataset, frequency):
+    """A class to represent a dataset with interpolated frequency."""
+
+    def __init__(self, dataset: Dataset, frequency: str) -> None:
+        """Initialize the InterpolateFrequency class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset to be interpolated.
+        frequency : str
+            The interpolation frequency.
+        """
         super().__init__(dataset)
         self._frequency = frequency_to_timedelta(frequency)
 
@@ -56,17 +75,53 @@ class InterpolateFrequency(Forwards):
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Get the interpolated data for a tuple index.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The tuple index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The interpolated data for the tuple index.
+        """
         index, changes = index_to_slices(index, self.shape)
         index, previous = update_tuple(index, 0, slice(None))
         result = self._get_slice(previous)
         return apply_index_to_slices_changes(result[index], changes)
 
-    def _get_slice(self, s):
+    def _get_slice(self, s: slice) -> NDArray[Any]:
+        """Get the interpolated data for a slice.
+
+        Parameters
+        ----------
+        s : slice
+            The slice to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The interpolated data for the slice.
+        """
         return np.stack([self[i] for i in range(*s.indices(self._len))])
 
     @debug_indexing
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Get the interpolated data at the specified index.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The interpolated data at the specified index.
+        """
         if isinstance(n, tuple):
             return self._get_tuple(n)
 
@@ -92,15 +147,24 @@ class InterpolateFrequency(Forwards):
         assert 0 < alpha < 1, alpha
         return self.forward[i] * (1 - alpha) + self.forward[i + 1] * alpha
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Get the length of the interpolated dataset.
+
+        Returns
+        -------
+        int
+            The length of the interpolated dataset.
+        """
         return (self.other_len - 1) * self.ratio + 1
 
     @property
-    def frequency(self):
+    def frequency(self) -> datetime.timedelta:
+        """Get the interpolation frequency."""
         return self._frequency
 
     @cached_property
-    def dates(self):
+    def dates(self) -> NDArray[np.datetime64]:
+        """Get the interpolated dates."""
         result = []
         deltas = [np.timedelta64(self.seconds * i, "s") for i in range(self.ratio)]
         for d in self.forward.dates[:-1]:
@@ -110,14 +174,23 @@ class InterpolateFrequency(Forwards):
         return np.array(result)
 
     @property
-    def shape(self):
+    def shape(self) -> Shape:
+        """Get the shape of the interpolated dataset."""
         return (self._len,) + self.forward.shape[1:]
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the dataset.
+
+        Returns
+        -------
+        Node
+            The tree representation of the dataset.
+        """
         return Node(self, [self.forward.tree()], frequency=self.frequency)
 
     @cached_property
-    def missing(self):
+    def missing(self) -> Set[int]:
+        """Get the missing data indices."""
         result = []
         j = 0
         for i in range(self.other_len):
@@ -130,7 +203,14 @@ class InterpolateFrequency(Forwards):
         result = set(x for x in result if x < self._len)
         return result
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get the metadata specific to the InterpolateFrequency subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata specific to the InterpolateFrequency subclass.
+        """
         return {
             # "frequency": frequency_to_string(self._frequency),
         }

anemoi/datasets/data/join.py

@@ -8,11 +8,22 @@
 # nor does it submit to any jurisdiction.
 
 
+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
 
+from .dataset import Dataset
+from .dataset import FullIndex
+from .dataset import Shape
+from .dataset import TupleIndex
 from .debug import Node
 from .debug import Source
 from .debug import debug_indexing
@@ -30,20 +41,57 @@ LOG = logging.getLogger(__name__)
 class Join(Combined):
     """Join the datasets along the variables axis."""
 
-    def check_compatibility(self, d1, d2):
+    def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
+        """Check the compatibility of two datasets.
+
+        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=1)
 
-    def check_same_variables(self, d1, d2):
+    def check_same_variables(self, d1: Dataset, d2: Dataset) -> None:
+        """Check if the datasets have the same variables.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            The first dataset.
+        d2 : Dataset
+            The second dataset.
+        """
         # Turned off because we are joining along the variables axis
         pass
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Get the length of the joined dataset.
+
+        Returns
+        -------
+        int
+            The length of the joined dataset.
+        """
         return len(self.datasets[0])
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Get the data for a tuple index.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The tuple index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data for the tuple index.
+        """
         index, changes = index_to_slices(index, self.shape)
         index, previous = update_tuple(index, 1, slice(None))
 
@@ -54,11 +102,35 @@ class Join(Combined):
         return apply_index_to_slices_changes(result[:, previous], changes)
 
     @debug_indexing
-    def _get_slice(self, s):
+    def _get_slice(self, s: slice) -> NDArray[Any]:
+        """Get the data for a slice.
+
+        Parameters
+        ----------
+        s : slice
+            The slice to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data for the slice.
+        """
         return np.stack([self[i] for i in range(*s.indices(self._len))])
 
     @debug_indexing
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Get the data at the specified index.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data at the specified index.
+        """
         if isinstance(n, tuple):
             return self._get_tuple(n)
 
@@ -68,11 +140,19 @@ class Join(Combined):
         return np.concatenate([d[n] for d in self.datasets])
 
     @cached_property
-    def shape(self):
+    def shape(self) -> Shape:
+        """Get the shape of the joined dataset."""
         cols = sum(d.shape[1] for d in self.datasets)
         return (len(self), cols) + self.datasets[0].shape[2:]
 
-    def _overlay(self):
+    def _overlay(self) -> Dataset:
+        """Overlay the datasets.
+
+        Returns
+        -------
+        Dataset
+            The overlaid dataset.
+        """
         indices = {}
         i = 0
         for d in self.datasets:
@@ -102,9 +182,10 @@ class Join(Combined):
         return Select(self, indices, {"overlay": variables})
 
     @cached_property
-    def variables(self):
+    def variables(self) -> List[str]:
+        """Get the variables of the joined dataset."""
         seen = set()
-        result = []
+        result: List[str] = []
         for d in reversed(self.datasets):
             for v in reversed(d.variables):
                 while v in seen:
@@ -115,7 +196,8 @@ class Join(Combined):
         return result
 
     @property
-    def variables_metadata(self):
+    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(")"))]
 
@@ -134,16 +216,30 @@ class Join(Combined):
         return result
 
     @cached_property
-    def name_to_index(self):
+    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):
+    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=None):
+    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+        """Get the statistics tendencies of the joined dataset.
+
+        Parameters
+        ----------
+        delta : Optional[datetime.timedelta]
+            The time delta for the tendencies.
+
+        Returns
+        -------
+        Dict[str, NDArray[Any]]
+            The statistics tendencies of the joined dataset.
+        """
         if delta is None:
             delta = self.frequency
         return {
@@ -151,7 +247,19 @@ class Join(Combined):
             for k in self.datasets[0].statistics_tendencies(delta)
         }
 
-    def source(self, index):
+    def source(self, index: int) -> Source:
+        """Get the source of the data at the specified index.
+
+        Parameters
+        ----------
+        index : int
+            The index to retrieve the source from.
+
+        Returns
+        -------
+        Source
+            The source of the data at the specified index.
+        """
         i = index
         for dataset in self.datasets:
             if i < dataset.shape[1]:
@@ -160,18 +268,49 @@ class Join(Combined):
         assert False
 
     @cached_property
-    def missing(self):
-        result = set()
+    def missing(self) -> Set[int]:
+        """Get the missing data indices."""
+        result: Set[int] = set()
         for d in self.datasets:
             result = result | d.missing
         return result
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the dataset.
+
+        Returns
+        -------
+        Node
+            The tree representation of the dataset.
+        """
         return Node(self, [d.tree() for d in self.datasets])
 
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
+        """Get the metadata specific to the forwards subclass.
+
+        Returns
+        -------
+        dict[str, Any]
+            The metadata specific to the forwards subclass.
+        """
+        return {}
+
+
+def join_factory(args: tuple, kwargs: dict) -> Dataset:
+    """Create a joined dataset.
 
-def join_factory(args, kwargs):
+    Parameters
+    ----------
+    args : tuple
+        The positional arguments.
+    kwargs : dict
+        The keyword arguments.
 
+    Returns
+    -------
+    Dataset
+        The joined dataset.
+    """
     datasets = kwargs.pop("join")
     assert isinstance(datasets, (list, tuple))
     assert len(args) == 0