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

anemoi/datasets/data/complement.py

@@ -9,9 +9,21 @@
 
 
 import logging
+from abc import abstractmethod
 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 numpy.typing import NDArray
 
 from ..grids import nearest_grid_points
+from .dataset import Dataset
+from .dataset import FullIndex
+from .dataset import Shape
+from .dataset import TupleIndex
 from .debug import Node
 from .forwards import Combined
 from .indexing import apply_index_to_slices_changes
@@ -24,101 +36,226 @@ LOG = logging.getLogger(__name__)
 
 
 class Complement(Combined):
-
-    def __init__(self, target, source, what="variables", interpolation="nearest"):
+    """A class to complement a target dataset with variables from a source dataset,
+    interpolated on the grid of the target dataset.
+
+    Attributes
+    ----------
+    target : Dataset
+        The target dataset.
+    source : Dataset
+        The source dataset.
+    variables : List[str]
+        List of variables to be added to the target dataset.
+    """
+
+    def __init__(
+        self,
+        target: Dataset,
+        source: Dataset,
+        what: str = "variables",
+        interpolation: str = "nearest",
+    ) -> None:
+        """Initializes the Complement class.
+
+        Parameters
+        ----------
+        target : Dataset
+            The target dataset.
+        source : Dataset
+            The source dataset.
+        what : str, optional
+            What to complement, default is "variables".
+        interpolation : str, optional
+            Interpolation method, default is "nearest".
+        """
         super().__init__([target, source])
 
         # We had the variables of dataset[1] to dataset[0]
         # interpoated on the grid of dataset[0]
 
-        self.target = target
-        self.source = source
+        self._target: Dataset = target
+        self._source: Dataset = source
 
         self._variables = []
 
         # Keep the same order as the original dataset
-        for v in self.source.variables:
-            if v not in self.target.variables:
+        for v in self._source.variables:
+            if v not in self._target.variables:
                 self._variables.append(v)
 
         if not self._variables:
             raise ValueError("Augment: no missing variables")
 
     @property
-    def variables(self):
+    def variables(self) -> List[str]:
+        """Returns the list of variables to be added to the target dataset."""
         return self._variables
 
     @property
-    def name_to_index(self):
+    def name_to_index(self) -> Dict[str, int]:
+        """Returns a dictionary mapping variable names to their indices."""
         return {v: i for i, v in enumerate(self.variables)}
 
     @property
-    def shape(self):
-        shape = self.target.shape
+    def shape(self) -> Shape:
+        """Returns the shape of the complemented dataset."""
+        shape = self._target.shape
         return (shape[0], len(self._variables)) + shape[2:]
 
     @property
-    def variables_metadata(self):
-        return {k: v for k, v in self.source.variables_metadata.items() if k in self._variables}
-
-    def check_same_variables(self, d1, d2):
+    def variables_metadata(self) -> Dict[str, Any]:
+        """Returns the metadata of the variables to be added to the target dataset."""
+        return {k: v for k, v in self._source.variables_metadata.items() if k in self._variables}
+
+    def check_same_variables(self, d1: Dataset, d2: Dataset) -> None:
+        """Checks if the variables in two datasets are the same.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            The first dataset.
+        d2 : Dataset
+            The second dataset.
+        """
         pass
 
     @cached_property
-    def missing(self):
-        missing = self.source.missing.copy()
-        missing = missing | self.target.missing
+    def missing(self) -> Set[int]:
+        """Returns the set of missing indices in the source and target datasets."""
+        missing = self._source.missing.copy()
+        missing = missing | self._target.missing
         return set(missing)
 
-    def tree(self):
-        """Generates a hierarchical tree structure for the `Cutout` instance and
-        its associated datasets.
+    def tree(self) -> Node:
+        """Generates a hierarchical tree structure for the Complement instance and its associated datasets.
 
-        Returns:
-            Node: A `Node` object representing the `Cutout` instance as the root
-            node, with each dataset in `self.datasets` represented as a child
-            node.
+        Returns
+        -------
+        Node
+            A Node object representing the Complement instance as the root node, with each dataset in self.datasets represented as a child node.
         """
-        return Node(self, [d.tree() for d in (self.target, self.source)])
+        return Node(self, [d.tree() for d in (self._target, self._source)])
+
+    def __getitem__(self, index: FullIndex) -> NDArray[Any]:
+        """Gets the data at the specified index.
+
+        Parameters
+        ----------
+        index : FullIndex
+            The index to retrieve data from.
 
-    def __getitem__(self, index):
+        Returns
+        -------
+        NDArray[Any]
+            The data at the specified index.
+        """
         if isinstance(index, (int, slice)):
             index = (index, slice(None), slice(None), slice(None))
         return self._get_tuple(index)
 
+    @abstractmethod
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Abstract method to get the data at the specified tuple index.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The tuple index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data at the specified tuple index.
+        """
+        pass
+
 
 class ComplementNone(Complement):
+    """A class to complement a target dataset with variables from a source dataset without interpolation."""
+
+    def __init__(self, target: Any, source: Any) -> None:
+        """Initializes the ComplementNone class.
 
-    def __init__(self, target, source):
+        Parameters
+        ----------
+        target : Any
+            The target dataset.
+        source : Any
+            The source dataset.
+        """
         super().__init__(target, source)
 
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Gets the data at the specified tuple index without interpolation.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The tuple index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data at the specified tuple index.
+        """
         index, changes = index_to_slices(index, self.shape)
-        result = self.source[index]
+        result = self._source[index]
         return apply_index_to_slices_changes(result, changes)
 
 
 class ComplementNearest(Complement):
+    """A class to complement a target dataset with variables from a source dataset using nearest neighbor interpolation."""
 
-    def __init__(self, target, source):
+    def __init__(self, target: Any, source: Any) -> None:
+        """Initializes the ComplementNearest class.
+
+        Parameters
+        ----------
+        target : Any
+            The target dataset.
+        source : Any
+            The source dataset.
+        """
         super().__init__(target, source)
 
         self._nearest_grid_points = nearest_grid_points(
-            self.source.latitudes,
-            self.source.longitudes,
-            self.target.latitudes,
-            self.target.longitudes,
+            self._source.latitudes,
+            self._source.longitudes,
+            self._target.latitudes,
+            self._target.longitudes,
         )
 
-    def check_compatibility(self, d1, d2):
+    def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
+        """Checks the compatibility of two datasets for nearest neighbor interpolation.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            The first dataset.
+        d2 : Dataset
+            The second dataset.
+        """
         pass
 
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Gets the data at the specified tuple index using nearest neighbor interpolation.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The tuple index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data at the specified tuple index.
+        """
         variable_index = 1
         index, changes = index_to_slices(index, self.shape)
         index, previous = update_tuple(index, variable_index, slice(None))
-        source_index = [self.source.name_to_index[x] for x in self.variables[previous]]
-        source_data = self.source[index[0], source_index, index[2], ...]
+        source_index = [self._source.name_to_index[x] for x in self.variables[previous]]
+        source_data = self._source[index[0], source_index, index[2], ...]
         target_data = source_data[..., self._nearest_grid_points]
 
         result = target_data[..., index[3]]
@@ -126,7 +263,21 @@ class ComplementNearest(Complement):
         return apply_index_to_slices_changes(result, changes)
 
 
-def complement_factory(args, kwargs):
+def complement_factory(args: Tuple, kwargs: dict) -> Dataset:
+    """Factory function to create a Complement instance based on the provided arguments.
+
+    Parameters
+    ----------
+    args : Tuple
+        Positional arguments.
+    kwargs : dict
+        Keyword arguments.
+
+    Returns
+    -------
+    Dataset
+        The complemented dataset.
+    """
     from .select import Select
 
     assert len(args) == 0, args
@@ -159,6 +310,6 @@ def complement_factory(args, kwargs):
     reorder = source.variables
     complemented = _open([target, complement])
     ordered = (
-        Select(complemented, complemented._reorder_to_columns(reorder), {"reoder": reorder})._subset(**kwargs).mutate()
+        Select(complemented, complemented._reorder_to_columns(reorder), {"reoder": reorder})._subset(**kwargs).mutate(),
     )
     return ordered

anemoi/datasets/data/concat.py

@@ -10,10 +10,19 @@
 
 import logging
 from functools import cached_property
+from typing import Any
+from typing import List
+from typing import Set
+from typing import Tuple
 
 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 Combined
@@ -29,13 +38,31 @@ LOG = logging.getLogger(__name__)
 
 
 class ConcatMixin:
-
-    def __len__(self):
+    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):
+    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]
@@ -46,7 +73,19 @@ class ConcatMixin:
         return apply_index_to_slices_changes(result, changes)
 
     @debug_indexing
-    def __getitem__(self, n):
+    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)
 
@@ -61,7 +100,19 @@ class ConcatMixin:
         return self.datasets[k][n]
 
     @debug_indexing
-    def _get_slice(self, s):
+    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]
@@ -72,8 +123,9 @@ class ConcatMixin:
         return np.concatenate(result)
 
     @cached_property
-    def missing(self):
-        result = set()
+    def missing(self) -> Set[int]:
+        """Returns the set of missing indices in the concatenated datasets."""
+        result: Set[int] = set()
         offset = 0
         for d in self.datasets:
             result = result | set(m + offset for m in d.missing)
@@ -82,32 +134,82 @@ class ConcatMixin:
 
 
 class Concat(ConcatMixin, Combined):
-
-    def check_compatibility(self, d1, d2):
+    def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
+        """Checks 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, d2):
+    def check_same_lengths(self, d1: Dataset, d2: Dataset) -> None:
+        """Checks 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, d2):
+    def check_same_dates(self, d1: Dataset, d2: Dataset) -> None:
+        """Checks 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):
+    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):
+    def shape(self) -> Shape:
+        """Returns the shape of the concatenated datasets."""
         return (len(self),) + self.datasets[0].shape[1:]
 
-    def tree(self):
+    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])
 
     @classmethod
-    def check_dataset_compatibility(cls, datasets, fill_missing_gaps=False):
+    def check_dataset_compatibility(cls, datasets: List[Any], fill_missing_gaps: bool = False) -> List[Any]:
+        """Checks the compatibility of the datasets for concatenation and fills missing gaps if required.
+
+        Parameters
+        ----------
+        datasets : List[Any]
+            List of datasets to check.
+        fill_missing_gaps : bool, optional
+            Whether to fill missing gaps between datasets, by default False.
+
+        Returns
+        -------
+        List[Any]
+            List of compatible datasets.
+        """
         # Study the dates
         ranges = [(d.dates[0].astype(object), d.dates[-1].astype(object)) for d in datasets]
 
@@ -146,9 +248,32 @@ class Concat(ConcatMixin, Combined):
 
         return result
 
+    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 concat_factory(args: Tuple[Any, ...], kwargs: dict) -> Concat:
+    """Factory function to create a Concat object.
 
-def concat_factory(args, kwargs):
+    Parameters
+    ----------
+    args : Tuple[Any, ...]
+        Positional arguments.
+    kwargs : dict
+        Keyword arguments.
 
+    Returns
+    -------
+    Concat
+        A Concat object.
+    """
     datasets = kwargs.pop("concat")
     fill_missing_gaps = kwargs.pop("fill_missing_gaps", False)