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

anemoi/datasets/data/subset.py

@@ -8,12 +8,24 @@
 # 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 Sequence
+from typing import Set
+from typing import Union
 
 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 Source
 from .debug import debug_indexing
@@ -27,11 +39,31 @@ from .indexing import update_tuple
 LOG = logging.getLogger(__name__)
 
 
-def _default(a, b, dates):
+def _default(a: int, b: int, dates: NDArray[np.datetime64]) -> list[int]:
+    """Default combination function for reasons.
+
+    Parameters:
+    a (int): First integer value.
+    b (int): Second integer value.
+    dates (NDArray[np.datetime64]): Array of datetime64 dates.
+
+    Returns:
+    list[int]: List containing the two input integers.
+    """
     return [a, b]
 
 
-def _start(a, b, dates):
+def _start(a: int, b: int, dates: NDArray[np.datetime64]) -> int:
+    """Determine the start date between two dates.
+
+    Parameters:
+    a (int): First integer value.
+    b (int): Second integer value.
+    dates (NDArray[np.datetime64]): Array of datetime64 dates.
+
+    Returns:
+    int: The index of the start date.
+    """
     from .misc import as_first_date
 
     c = as_first_date(a, dates)
@@ -42,7 +74,17 @@ def _start(a, b, dates):
         return a
 
 
-def _end(a, b, dates):
+def _end(a: int, b: int, dates: NDArray[np.datetime64]) -> int:
+    """Determine the end date between two dates.
+
+    Parameters:
+    a (int): First integer value.
+    b (int): Second integer value.
+    dates (NDArray[np.datetime64]): Array of datetime64 dates.
+
+    Returns:
+    int: The index of the end date.
+    """
     from .misc import as_last_date
 
     c = as_last_date(a, dates)
@@ -53,7 +95,17 @@ def _end(a, b, dates):
         return b
 
 
-def _combine_reasons(reason1, reason2, dates):
+def _combine_reasons(reason1: Dict[str, Any], reason2: Dict[str, Any], dates: NDArray[np.datetime64]) -> Dict[str, Any]:
+    """Combine two reason dictionaries.
+
+    Parameters:
+    reason1 (Dict[str, Any]): First reason dictionary.
+    reason2 (Dict[str, Any]): Second reason dictionary.
+    dates (NDArray[np.datetime64]): Array of datetime64 dates.
+
+    Returns:
+    Dict[str, Any]: Combined reason dictionary.
+    """
 
     reason = reason1.copy()
     for k, v in reason2.items():
@@ -66,29 +118,63 @@ def _combine_reasons(reason1, reason2, dates):
 
 
 class Subset(Forwards):
-    """Select a subset of the dates."""
-
-    def __init__(self, dataset, indices, reason):
+    """Select a subset of the dates.
+
+    Attributes:
+    dataset (Dataset): The dataset.
+    indices (List[int]): List of indices.
+    reason (Dict[str, Any]): Dictionary of reasons.
+    """
+
+    def __init__(self, dataset: Union[Dataset, "Subset"], indices: Sequence[int], reason: Dict[str, Any]) -> None:
+        """Initialize the Subset.
+
+        Parameters:
+        dataset (Dataset | Subset): The dataset or subset.
+        indices (Sequence[int]): Sequence of indices.
+        reason (Dict[str, Any]): Dictionary of reasons.
+        """
         while isinstance(dataset, Subset):
             indices = [dataset.indices[i] for i in indices]
             reason = _combine_reasons(reason, dataset.reason, dataset.dates)
             dataset = dataset.dataset
 
-        self.dataset = dataset
-        self.indices = list(indices)
-        self.reason = {k: v for k, v in reason.items() if v is not None}
+        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}
 
         # Forward other properties to the super dataset
         super().__init__(dataset)
 
-    def clone(self, dataset):
+    def clone(self, dataset: Dataset) -> Dataset:
+        """Clone the subset with a new dataset.
+
+        Parameters:
+        dataset (Dataset): The new dataset.
+
+        Returns:
+        Dataset: The cloned subset.
+        """
         return self.__class__(dataset, self.indices, self.reason).mutate()
 
-    def mutate(self):
+    def mutate(self) -> Dataset:
+        """Mutate the subset.
+
+        Returns:
+        Dataset: The mutated subset.
+        """
         return self.forward.swap_with_parent(parent=self)
 
     @debug_indexing
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Get item by index.
+
+        Parameters:
+        n (FullIndex): The index.
+
+        Returns:
+        NDArray[Any]: The indexed data.
+        """
         if isinstance(n, tuple):
             return self._get_tuple(n)
 
@@ -100,7 +186,15 @@ class Subset(Forwards):
         return self.dataset[n]
 
     @debug_indexing
-    def _get_slice(self, s):
+    def _get_slice(self, s: slice) -> NDArray[Any]:
+        """Get slice of data.
+
+        Parameters:
+        s (slice): The slice.
+
+        Returns:
+        NDArray[Any]: The sliced data.
+        """
         # TODO: check if the indices can be simplified to a slice
         # the time checking maybe be longer than the time saved
         # using a slice
@@ -112,7 +206,15 @@ class Subset(Forwards):
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, n):
+    def _get_tuple(self, n: TupleIndex) -> NDArray[Any]:
+        """Get tuple of data.
+
+        Parameters:
+        n (TupleIndex): The tuple index.
+
+        Returns:
+        NDArray[Any]: The indexed data.
+        """
         index, changes = index_to_slices(n, self.shape)
         indices = [self.indices[i] for i in range(*index[0].indices(self._len))]
         indices = make_slice_or_index_from_list_or_tuple(indices)
@@ -121,43 +223,75 @@ class Subset(Forwards):
         result = apply_index_to_slices_changes(result, changes)
         return result
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Get the length of the subset.
+
+        Returns:
+        int: The length of the subset.
+        """
         return len(self.indices)
 
     @cached_property
-    def shape(self):
+    def shape(self) -> Shape:
+        """Get the shape of the subset."""
         return (len(self),) + self.dataset.shape[1:]
 
     @cached_property
-    def dates(self):
+    def dates(self) -> NDArray[np.datetime64]:
+        """Get the dates of the subset."""
         return self.dataset.dates[self.indices]
 
     @cached_property
-    def frequency(self):
+    def frequency(self) -> datetime.timedelta:
+        """Get the frequency of the subset."""
         dates = self.dates
         if len(dates) < 2:
             raise ValueError(f"Cannot determine frequency of a subset with less than two dates ({self.dates}).")
         return frequency_to_timedelta(dates[1].astype(object) - dates[0].astype(object))
 
-    def source(self, index):
+    def source(self, index: int) -> Source:
+        """Get the source of the subset.
+
+        Parameters:
+        index (int): The index.
+
+        Returns:
+        Source: The source of the subset.
+        """
         return Source(self, index, self.forward.source(index))
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Get the string representation of the subset.
+
+        Returns:
+        str: The string representation of the subset.
+        """
         return f"Subset({self.dataset},{self.dates[0]}...{self.dates[-1]}/{self.frequency})"
 
     @cached_property
-    def missing(self):
+    def missing(self) -> Set[int]:
+        """Get the missing indices of the subset."""
         missing = self.dataset.missing
-        result = set()
+        result: Set[int] = set()
         for j, i in enumerate(self.indices):
             if i in missing:
                 result.add(j)
         return result
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the subset.
+
+        Returns:
+        Node: The tree representation of the subset.
+        """
         return Node(self, [self.dataset.tree()], **self.reason)
 
-    def subclass_metadata_specific(self):
+    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 {
             # "indices": self.indices,
             "reason": self.reason,

anemoi/datasets/data/unchecked.py

@@ -8,11 +8,24 @@
 # nor does it submit to any jurisdiction.
 
 
+import datetime
 import logging
 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
 
 from .concat import ConcatMixin
+from .dataset import Dataset
+from .dataset import FullIndex
+from .dataset import Shape
 from .debug import Node
 from .forwards import Combined
 from .misc import _auto_adjust
@@ -22,18 +35,48 @@ LOG = logging.getLogger(__name__)
 
 
 class check:
+    """A decorator class to perform checks before calling a method."""
+
+    def __init__(self, check: str) -> None:
+        """Initialize the check decorator.
 
-    def __init__(self, check):
+        Parameters
+        ----------
+        check : str
+            The name of the check method.
+        """
         self.check = check
 
-    def __call__(self, method):
+    def __call__(self, method: Callable) -> Callable:
+        """Call the check decorator.
+
+        Parameters
+        ----------
+        method : Callable
+            The method to decorate.
+
+        Returns
+        -------
+        Callable
+            The decorated method.
+        """
         name = method.__name__
         check = self.check
 
         @wraps(method)
-        def wrapper(obj):
-            """This is a decorator that checks the compatibility of the datasets before calling the method. If the datasets are compatible, it will return the result of the method, otherwise it will raise an exception."""
-
+        def wrapper(obj: "Unchecked") -> Any:
+            """Wrapper function to check compatibility before calling the method.
+
+            Parameters
+            ----------
+            obj : Unchecked
+                The Unchecked object.
+
+            Returns
+            -------
+            Any
+                The result of the method.
+            """
             for d in obj.datasets[1:]:
                 getattr(obj, check)(obj.datasets[0], d)
 
@@ -43,126 +86,194 @@ class check:
 
 
 class Unchecked(Combined):
+    """A class representing a dataset without compatibility checks."""
+
+    def tree(self) -> Node:
+        """Get the tree representation of the dataset.
 
-    def tree(self):
+        Returns
+        -------
+        Node
+            The tree representation.
+        """
         return Node(self, [d.tree() for d in self.datasets])
 
-    def _subset(self, **kwargs):
+    def _subset(self, **kwargs: dict) -> "Unchecked":
+        """Get a subset of the dataset.
+
+        Parameters
+        ----------
+        **kwargs : dict
+            Subset parameters.
+
+        Returns
+        -------
+        Unchecked
+            The subset of the dataset.
+        """
         assert not kwargs
         return self
 
-    def check_compatibility(self, d1, d2):
+    def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
+        """Check compatibility between two datasets.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            The first dataset.
+        d2 : Dataset
+            The second dataset.
+        """
         pass
 
-    ###########################################
     @property
     @check("check_same_dates")
-    def dates(self):
+    def dates(self) -> NDArray[np.datetime64]:
+        """Get the dates of the dataset."""
         pass
 
     @property
     @check("check_same_resolution")
-    def resolution(self):
+    def resolution(self) -> Any:
+        """Get the resolution of the dataset."""
         pass
 
     @property
-    def field_shape(self):
+    def field_shape(self) -> tuple:
+        """Get the field shape of the dataset."""
         raise NotImplementedError()
 
     @property
     @check("check_same_frequency")
-    def frequency(self):
+    def frequency(self) -> datetime.timedelta:
+        """Get the frequency of the dataset."""
         raise NotImplementedError()
 
     @property
     @check("check_same_grid")
-    def latitudes(self):
+    def latitudes(self) -> NDArray[Any]:
+        """Get the latitudes of the dataset."""
         raise NotImplementedError()
 
     @property
     @check("check_same_grid")
-    def longitudes(self):
+    def longitudes(self) -> NDArray[Any]:
+        """Get the longitudes of the dataset."""
         raise NotImplementedError()
 
-    @property
     @check("check_same_variables")
-    def name_to_index(self):
+    @property
+    def name_to_index(self) -> Dict[str, int]:
+        """Get the mapping of variable names to their indices."""
         raise NotImplementedError()
 
-    @property
     @check("check_same_variables")
-    def variables(self):
+    @property
+    def variables(self) -> List[str]:
+        """Get the list of variables in the dataset."""
         raise NotImplementedError()
 
-    @property
     @check("check_same_variables")
-    def variables_metadata(self):
+    @property
+    def variables_metadata(self) -> dict:
+        """Get the metadata for the variables."""
         raise NotImplementedError()
 
-    @property
     @check("check_same_variables")
-    def statistics(self):
+    @property
+    def statistics(self) -> Dict[str, NDArray[Any]]:
+        """Get the statistics of the dataset."""
         raise NotImplementedError()
 
     @check("check_same_variables")
-    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 dataset.
+
+        Parameters
+        ----------
+        delta : Optional[datetime.timedelta]
+            The time delta for tendencies.
+
+        Returns
+        -------
+        Dict[str, NDArray[Any]]
+            The statistics tendencies.
+        """
         raise NotImplementedError()
 
     @property
-    def shape(self):
+    def shape(self) -> Shape:
+        """Get the shape of the dataset."""
         raise NotImplementedError()
 
-    # @property
-    # def field_shape(self):
-    #     return tuple(d.shape for d in self.datasets)
-
-    # @property
-    # def latitudes(self):
-    #     return tuple(d.latitudes for d in self.datasets)
-
-    # @property
-    # def longitudes(self):
-    #     return tuple(d.longitudes for d in self.datasets)
-
-    # @property
-    # def statistics(self):
-    #     return tuple(d.statistics for d in self.datasets)
-
-    # @property
-    # def resolution(self):
-    #     return tuple(d.resolution for d in self.datasets)
-
-    # @property
-    # def name_to_index(self):
-    #     return tuple(d.name_to_index for d in self.datasets)
-
     @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
 
 
 class Chain(ConcatMixin, Unchecked):
-    """Same as Concat, but with no checks"""
+    """A class representing a chain of datasets without compatibility checks."""
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Get the length of the dataset.
+
+        Returns
+        -------
+        int
+            The length of the dataset.
+        """
         return sum(len(d) for d in self.datasets)
 
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> tuple:
+        """Get an item from the dataset.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index of the item.
+
+        Returns
+        -------
+        tuple
+            The item at the specified index.
+        """
         return tuple(d[n] for d in self.datasets)
 
     @property
-    def dates(self):
+    def dates(self) -> NDArray[np.datetime64]:
+        """Get the dates of the dataset."""
         raise NotImplementedError()
 
-    def dataset_metadata(self):
+    def dataset_metadata(self) -> dict:
+        """Get the metadata of the dataset.
+
+        Returns
+        -------
+        dict
+            The metadata of the dataset.
+        """
         return {"multiple": [d.dataset_metadata() for d in self.datasets]}
 
 
-def chain_factory(args, kwargs):
+def chain_factory(args: tuple, kwargs: dict) -> Dataset:
+    """Factory function to create a Chain dataset.
+
+    Parameters
+    ----------
+    args : tuple
+        Positional arguments.
+    kwargs : dict
+        Keyword arguments.
 
+    Returns
+    -------
+    Dataset
+        The Chain dataset.
+    """
     chain = kwargs.pop("chain")
     assert len(args) == 0
     assert isinstance(chain, (list, tuple))