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

anemoi/datasets/data/fill_missing.py

@@ -9,11 +9,19 @@
 
 
 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
 
 from anemoi.datasets.data 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 Forwards
@@ -26,29 +34,74 @@ LOG = logging.getLogger(__name__)
 
 
 class MissingDatesFill(Forwards):
-    def __init__(self, dataset):
+    """Class to handle filling missing dates in a dataset."""
+
+    def __init__(self, dataset: Dataset) -> None:
+        """Initialize the MissingDatesFill class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset with missing dates.
+        """
         super().__init__(dataset)
         self._missing = set(dataset.missing)
-        self._warnings = set()
+        self._warnings: Set[int] = set()
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Get a tuple of data for the given index.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The index to retrieve data for.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data for the given 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 a slice of data.
+
+        Parameters
+        ----------
+        s : slice
+            The slice to retrieve data for.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data for the given slice.
+        """
         return np.stack([self[i] for i in range(*s.indices(self._len))])
 
     @property
-    def missing(self):
+    def missing(self) -> Set[int]:
+        """Get the set of missing dates."""
         return set()
 
     @debug_indexing
-    def __getitem__(self, n):
-
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Get the data for the given index.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index to retrieve data for.
+
+        Returns
+        -------
+        NDArray[Any]
+            The data for the given index.
+        """
         try:
             return self.forward[n]
         except MissingDateError:
@@ -84,14 +137,39 @@ class MissingDatesFill(Forwards):
 
 
 class MissingDatesClosest(MissingDatesFill):
-
-    def __init__(self, dataset, closest):
+    """Class to handle filling missing dates using the closest available date."""
+
+    def __init__(self, dataset: Any, closest: str) -> None:
+        """Initialize the MissingDatesClosest class.
+
+        Parameters
+        ----------
+        dataset : Any
+            The dataset with missing dates.
+        closest : str
+            The strategy to use for filling missing dates ('up' or 'down').
+        """
         super().__init__(dataset)
         self.closest = closest
         self._closest = {}
 
-    def _fill_missing(self, n, a, b):
-
+    def _fill_missing(self, n: int, a: Optional[int], b: Optional[int]) -> NDArray[Any]:
+        """Fill the missing date at the given index.
+
+        Parameters
+        ----------
+        n : int
+            The index of the missing date.
+        a : Optional[int]
+            The previous available date index.
+        b : Optional[int]
+            The next available date index.
+
+        Returns
+        -------
+        NDArray[Any]
+            The filled data for the missing date.
+        """
         if n not in self._warnings:
             LOG.warning(f"Missing date at index {n} ({self.dates[n]})")
             if abs(n - a) == abs(b - n):
@@ -111,19 +189,58 @@ class MissingDatesClosest(MissingDatesFill):
 
         return self.forward[self._closest[n]]
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get metadata specific to the subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata specific to the subclass.
+        """
         return {"closest": self.closest}
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the object.
+
+        Returns
+        -------
+        Node
+            The tree representation of the object.
+        """
         return Node(self, [self.forward.tree()], closest=self.closest)
 
 
 class MissingDatesInterpolate(MissingDatesFill):
-    def __init__(self, dataset):
+    """Class to handle filling missing dates using interpolation."""
+
+    def __init__(self, dataset: Any) -> None:
+        """Initialize the MissingDatesInterpolate class.
+
+        Parameters
+        ----------
+        dataset : Any
+            The dataset with missing dates.
+        """
         super().__init__(dataset)
         self._alpha = {}
 
-    def _fill_missing(self, n, a, b):
+    def _fill_missing(self, n: int, a: Optional[int], b: Optional[int]) -> NDArray[Any]:
+        """Fill the missing date at the given index using interpolation.
+
+        Parameters
+        ----------
+        n : int
+            The index of the missing date.
+        a : Optional[int]
+            The previous available date index.
+        b : Optional[int]
+            The next available date index.
+
+        Returns
+        -------
+        NDArray[Any]
+            The filled data for the missing date.
+        """
         if n not in self._warnings:
             LOG.warning(f"Missing date at index {n} ({self.dates[n]})")
 
@@ -147,14 +264,44 @@ class MissingDatesInterpolate(MissingDatesFill):
         alpha = self._alpha[n]
         return self.forward[a] * (1 - alpha) + self.forward[b] * alpha
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get metadata specific to the subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata specific to the subclass.
+        """
         return {}
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the object.
+
+        Returns
+        -------
+        Node
+            The tree representation of the object.
+        """
         return Node(self, [self.forward.tree()])
 
 
-def fill_missing_dates_factory(dataset, method, kwargs):
+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
+    ----------
+    dataset : Any
+        The dataset with missing dates.
+    method : str
+        The method to use for filling missing dates ('closest' or 'interpolate').
+    kwargs : Dict[str, Any]
+        Additional arguments for the method.
+
+    Returns
+    -------
+    Dataset
+        An instance of a class to fill missing dates.
+    """
     if method == "closest":
         closest = kwargs.get("closest", "up")
         return MissingDatesClosest(dataset, closest=closest)