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

anemoi/datasets/data/masked.py

@@ -10,11 +10,21 @@
 
 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
 
 from ..grids import cropping_mask
 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
@@ -27,7 +37,18 @@ LOG = logging.getLogger(__name__)
 
 
 class Masked(Forwards):
-    def __init__(self, forward, mask):
+    """A class to represent a masked dataset."""
+
+    def __init__(self, forward: Dataset, mask: NDArray[np.bool_]) -> None:
+        """Initialize the Masked class.
+
+        Parameters
+        ----------
+        forward : Dataset
+            The dataset to be masked.
+        mask : NDArray[np.bool_]
+            The mask array.
+        """
         super().__init__(forward)
         assert len(forward.shape) == 4, "Grids must be 1D for now"
         self.mask = mask
@@ -36,19 +57,34 @@ class Masked(Forwards):
         self.mask_name = f"{self.__class__.__name__.lower()}_mask"
 
     @cached_property
-    def shape(self):
+    def shape(self) -> Shape:
+        """Get the shape of the masked dataset."""
         return self.forward.shape[:-1] + (np.count_nonzero(self.mask),)
 
     @cached_property
-    def latitudes(self):
+    def latitudes(self) -> NDArray[Any]:
+        """Get the masked latitudes."""
         return self.forward.latitudes[self.mask]
 
     @cached_property
-    def longitudes(self):
+    def longitudes(self) -> NDArray[Any]:
+        """Get the masked longitudes."""
         return self.forward.longitudes[self.mask]
 
     @debug_indexing
-    def __getitem__(self, index):
+    def __getitem__(self, index: FullIndex) -> NDArray[Any]:
+        """Get the masked data at the specified index.
+
+        Parameters
+        ----------
+        index : FullIndex
+            The index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The masked data at the specified index.
+        """
         if isinstance(index, tuple):
             return self._get_tuple(index)
 
@@ -60,7 +96,19 @@ class Masked(Forwards):
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Get the masked data for a tuple index.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The tuple index to retrieve data from.
+
+        Returns
+        -------
+        NDArray[Any]
+            The masked data for the tuple index.
+        """
         index, changes = index_to_slices(index, self.shape)
         index, previous = update_tuple(index, self.axis, slice(None))
         result = self.forward[index]
@@ -69,14 +117,35 @@ class Masked(Forwards):
         result = apply_index_to_slices_changes(result, changes)
         return result
 
-    def collect_supporting_arrays(self, collected, *path):
+    def collect_supporting_arrays(self, collected: List[Tuple], *path: Any) -> None:
+        """Collect supporting arrays.
+
+        Parameters
+        ----------
+        collected : List[Tuple]
+            The list to collect supporting arrays into.
+        path : Any
+            Additional path arguments.
+        """
         super().collect_supporting_arrays(collected, *path)
         collected.append((path, self.mask_name, self.mask))
 
 
 class Thinning(Masked):
-
-    def __init__(self, forward, thinning, method):
+    """A class to represent a thinned dataset."""
+
+    def __init__(self, forward: Dataset, thinning: Optional[int], method: str) -> None:
+        """Initialize the Thinning class.
+
+        Parameters
+        ----------
+        forward : Dataset
+            The dataset to be thinned.
+        thinning : Optional[int]
+            The thinning factor.
+        method : str
+            The thinning method.
+        """
         self.thinning = thinning
         self.method = method
 
@@ -105,20 +174,51 @@ class Thinning(Masked):
         super().__init__(forward, mask)
 
     def mutate(self) -> Dataset:
+        """Mutate the dataset.
+
+        Returns
+        -------
+        Dataset
+            The mutated dataset.
+        """
         if self.thinning is None:
             return self.forward.mutate()
         return super().mutate()
 
-    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()], thinning=self.thinning, method=self.method)
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get the metadata specific to the Thinning subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata specific to the Thinning subclass.
+        """
         return dict(thinning=self.thinning, method=self.method)
 
 
 class Cropping(Masked):
-
-    def __init__(self, forward, area):
+    """A class to represent a cropped dataset."""
+
+    def __init__(self, forward: Dataset, area: Union[Dataset, Tuple[float, float, float, float]]) -> None:
+        """Initialize the Cropping class.
+
+        Parameters
+        ----------
+        forward : Dataset
+            The dataset to be cropped.
+        area : Union[Dataset, Tuple[float, float, float, float]]
+            The cropping area.
+        """
         from ..data import open_dataset
 
         area = area if isinstance(area, (list, tuple)) else open_dataset(area)
@@ -135,8 +235,22 @@ class Cropping(Masked):
 
         super().__init__(forward, mask)
 
-    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()], area=self.area)
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get the metadata specific to the Cropping subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata specific to the Cropping subclass.
+        """
         return dict(area=self.area)

anemoi/datasets/data/merge.py

@@ -8,12 +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 Set
+from typing import Tuple
 
 import numpy as np
+from numpy.typing import NDArray
 
 from . import MissingDateError
+from .dataset import Dataset
+from .dataset import FullIndex
+from .dataset import TupleIndex
 from .debug import Node
 from .debug import debug_indexing
 from .forwards import Combined
@@ -28,14 +38,18 @@ LOG = logging.getLogger(__name__)
 
 
 class Merge(Combined):
-
-    # d0 d2 d4 d6 ...
-    # d1 d3 d5 d7 ...
-
-    # gives
-    # d0 d1 d2 d3 ...
-
-    def __init__(self, datasets, allow_gaps_in_dates=False):
+    """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:
+        """Initialize the Merge object.
+
+        Parameters
+        ----------
+        datasets : List[Dataset]
+            List of datasets to merge.
+        allow_gaps_in_dates : bool, optional
+            Whether to allow gaps in dates. Defaults to False.
+        """
         super().__init__(datasets)
 
         self.allow_gaps_in_dates = allow_gaps_in_dates
@@ -91,23 +105,33 @@ class Merge(Combined):
 
         self._dates = np.array(_dates, dtype="datetime64[s]")
         self._indices = np.array(indices)
-        self._frequency = frequency  # .astype(object)
+        self._frequency = frequency.astype(object)
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Get the number of dates in the merged dataset.
+
+        Returns
+        -------
+        int
+            Number of dates.
+        """
         return len(self._dates)
 
     @property
-    def dates(self):
+    def dates(self) -> NDArray[np.datetime64]:
+        """Get the dates of the merged dataset."""
         return self._dates
 
     @property
-    def frequency(self):
+    def frequency(self) -> datetime.timedelta:
+        """Get the frequency of the dates in the merged dataset."""
         return self._frequency
 
     @cached_property
-    def missing(self):
+    def missing(self) -> Set[int]:
+        """Get the indices of missing dates in the merged dataset."""
         # TODO: optimize
-        result = set()
+        result: Set[int] = set()
 
         for i, (dataset, row) in enumerate(self._indices):
             if dataset == self._missing_index:
@@ -119,26 +143,79 @@ class Merge(Combined):
 
         return result
 
-    def check_same_lengths(self, d1, d2):
+    def check_same_lengths(self, d1: Dataset, d2: Dataset) -> None:
+        """Check if the lengths of two datasets are the same. (Disabled for merging).
+
+        Parameters
+        ----------
+        d1 : Dataset
+            First dataset.
+        d2 : Dataset
+            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:
+        """Check if the dates of two datasets are the same. (Disabled for merging).
+
+        Parameters
+        ----------
+        d1 : Dataset
+            First dataset.
+        d2 : Dataset
+            Second dataset.
+        """
         # Turned off because we are concatenating along the dates axis
         pass
 
-    def check_compatibility(self, d1, d2):
+    def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
+        """Check if two datasets are compatible for merging.
+
+        Parameters
+        ----------
+        d1 : Dataset
+            First dataset.
+        d2 : Dataset
+            Second dataset.
+        """
         super().check_compatibility(d1, d2)
         self.check_same_sub_shapes(d1, d2, drop_axis=0)
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the merged dataset.
+
+        Returns
+        -------
+        Node
+            Tree representation of the merged dataset.
+        """
         return Node(self, [d.tree() for d in self.datasets], allow_gaps_in_dates=self.allow_gaps_in_dates)
 
-    def metadata_specific(self):
+    def metadata_specific(self) -> Dict[str, Any]:
+        """Get the specific metadata for the merged dataset.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Specific metadata.
+        """
         return {"allow_gaps_in_dates": self.allow_gaps_in_dates}
 
     @debug_indexing
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Get the item at the specified index.
+
+        Parameters
+        ----------
+        n : FullIndex
+            Index to retrieve.
+
+        Returns
+        -------
+        NDArray[Any]
+            Retrieved item.
+        """
         if isinstance(n, tuple):
             return self._get_tuple(n)
 
@@ -154,18 +231,55 @@ class Merge(Combined):
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Get the item at the specified tuple index.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            Tuple index to retrieve.
+
+        Returns
+        -------
+        NDArray[Any]
+            Retrieved item.
+        """
         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 items in the specified slice.
+
+        Parameters
+        ----------
+        s : slice
+            Slice to retrieve.
+
+        Returns
+        -------
+        NDArray[Any]
+            Retrieved items.
+        """
         return np.stack([self[i] for i in range(*s.indices(self._len))])
 
 
-def merge_factory(args, kwargs):
+def merge_factory(args: Tuple, kwargs: Dict[str, Any]) -> Dataset:
+    """Factory function to create a merged dataset.
+
+    Parameters
+    ----------
+    args : Tuple
+        Positional arguments.
+    kwargs : Dict[str, Any]
+        Keyword arguments.
 
+    Returns
+    -------
+    Dataset
+        Merged dataset.
+    """
     datasets = kwargs.pop("merge")
 
     assert isinstance(datasets, (list, tuple))