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

anemoi/datasets/data/rescale.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 Tuple
+from typing import Union
 
 import numpy as np
+from numpy.typing import NDArray
 
+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
@@ -24,8 +35,23 @@ from .indexing import update_tuple
 LOG = logging.getLogger(__name__)
 
 
-def make_rescale(variable, rescale):
+def make_rescale(
+    variable: str, rescale: Union[Tuple[float, float], List[str], Dict[str, float]]
+) -> Tuple[float, float]:
+    """Create rescale parameters (scale and offset) based on the input rescale specification.
 
+    Parameters
+    ----------
+    variable : str
+        The variable name.
+    rescale : Union[Tuple[float, float], List[str], Dict[str, float]]
+        The rescale specification.
+
+    Returns
+    -------
+    Tuple[float, float]
+        The scale and offset values.
+    """
     if isinstance(rescale, (tuple, list)):
 
         assert len(rescale) == 2, rescale
@@ -57,7 +83,20 @@ def make_rescale(variable, rescale):
 
 
 class Rescale(Forwards):
-    def __init__(self, dataset, rescale):
+    """A class to apply rescaling to dataset variables."""
+
+    def __init__(
+        self, dataset: Dataset, rescale: Dict[str, Union[Tuple[float, float], List[str], Dict[str, float]]]
+    ) -> None:
+        """Initialize the Rescale object.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset to be rescaled.
+        rescale : Dict[str, Union[Tuple[float, float], List[str], Dict[str, float]]]
+            The rescale specifications.
+        """
         super().__init__(dataset)
         for n in rescale:
             assert n in dataset.variables, n
@@ -80,15 +119,41 @@ class Rescale(Forwards):
         self._a = self._a.astype(self.forward.dtype)
         self._b = self._b.astype(self.forward.dtype)
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the rescale operation.
+
+        Returns
+        -------
+        Node
+            The tree representation.
+        """
         return Node(self, [self.forward.tree()], rescale=self.rescale)
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get the metadata specific to the rescale subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata dictionary.
+        """
         return dict(rescale=self.rescale)
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Get a tuple of rescaled data based on the provided index.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The index to retrieve data.
+
+        Returns
+        -------
+        NDArray[Any]
+            The rescaled data.
+        """
         index, changes = index_to_slices(index, self.shape)
         index, previous = update_tuple(index, 1, slice(None))
         result = self.forward[index]
@@ -98,13 +163,36 @@ class Rescale(Forwards):
         return result
 
     @debug_indexing
-    def __get_slice_(self, n):
+    def __get_slice_(self, n: slice) -> NDArray[Any]:
+        """Get a slice of rescaled data.
+
+        Parameters
+        ----------
+        n : slice
+            The slice to retrieve data.
+
+        Returns
+        -------
+        NDArray[Any]
+            The rescaled data.
+        """
         data = self.forward[n]
         return data * self._a + self._b
 
     @debug_indexing
-    def __getitem__(self, n):
-
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Get an item or slice of rescaled data based on the provided index.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index to retrieve data.
+
+        Returns
+        -------
+        NDArray[Any]
+            The rescaled data.
+        """
         if isinstance(n, tuple):
             return self._get_tuple(n)
 
@@ -116,7 +204,8 @@ class Rescale(Forwards):
         return data * self._a[0] + self._b[0]
 
     @cached_property
-    def statistics(self):
+    def statistics(self) -> Dict[str, NDArray[Any]]:
+        """Get the statistics of the rescaled data."""
         result = {}
         a = self._a.squeeze()
         assert np.all(a >= 0)
@@ -135,7 +224,19 @@ class Rescale(Forwards):
 
         return result
 
-    def statistics_tendencies(self, delta=None):
+    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+        """Get the tendencies of the statistics of the rescaled data.
+
+        Parameters
+        ----------
+        delta : Optional[datetime.timedelta]
+            The time delta for tendencies calculation.
+
+        Returns
+        -------
+        Dict[str, NDArray[Any]]
+            The tendencies statistics dictionary.
+        """
         result = {}
         a = self._a.squeeze()
         assert np.all(a >= 0)

anemoi/datasets/data/select.py

@@ -8,9 +8,20 @@
 # 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 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
@@ -24,10 +35,20 @@ LOG = logging.getLogger(__name__)
 
 
 class Select(Forwards):
-    """Select a subset of the variables."""
-
-    def __init__(self, dataset, indices, reason):
-
+    """Class to select a subset of variables from a dataset."""
+
+    def __init__(self, dataset: Dataset, indices: List[int], reason: Dict[str, Any]) -> None:
+        """Initialize the Select class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset to select from.
+        indices : List[int]
+            The indices of the variables to select.
+        reason : Dict[str, Any]
+            The reason for the selection.
+        """
         reason = reason.copy()
 
         while isinstance(dataset, Select):
@@ -35,23 +56,53 @@ class Select(Forwards):
             reason.update(dataset.reason)
             dataset = dataset.dataset
 
-        self.dataset = dataset
+        self.dataset: Dataset = dataset
         self.indices = list(indices)
         assert len(self.indices) > 0
         self.reason = reason or {"indices": self.indices}
 
-        # Forward other properties to the main dataset
         super().__init__(dataset)
 
-    def clone(self, dataset):
+    def clone(self, dataset: Dataset) -> Dataset:
+        """Clone the Select object with a new dataset.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The new dataset.
+
+        Returns
+        -------
+        Select
+            The cloned Select object.
+        """
         return self.__class__(dataset, self.indices, self.reason).mutate()
 
-    def mutate(self):
+    def mutate(self) -> Dataset:
+        """Mutate the dataset.
+
+        Returns
+        -------
+        Dataset
+            The mutated dataset.
+        """
         return self.forward.swap_with_parent(parent=self)
 
     @debug_indexing
     @expand_list_indexing
-    def _get_tuple(self, index):
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        """Get a tuple of data.
+
+        Parameters
+        ----------
+        index : TupleIndex
+            The index to retrieve.
+
+        Returns
+        -------
+        NDArray[Any]
+            The retrieved data.
+        """
         index, changes = index_to_slices(index, self.shape)
         index, previous = update_tuple(index, 1, slice(None))
         result = self.dataset[index]
@@ -61,7 +112,19 @@ class Select(Forwards):
         return result
 
     @debug_indexing
-    def __getitem__(self, n):
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        """Get an item from the dataset.
+
+        Parameters
+        ----------
+        n : FullIndex
+            The index to retrieve.
+
+        Returns
+        -------
+        NDArray[Any]
+            The retrieved data.
+        """
         if isinstance(n, tuple):
             return self._get_tuple(n)
 
@@ -72,46 +135,112 @@ class Select(Forwards):
         return row[self.indices]
 
     @cached_property
-    def shape(self):
+    def shape(self) -> Shape:
+        """Get the shape of the dataset."""
         return (len(self), len(self.indices)) + self.dataset.shape[2:]
 
     @cached_property
-    def variables(self):
+    def variables(self) -> List[str]:
+        """Get the variables of the dataset."""
         return [self.dataset.variables[i] for i in self.indices]
 
     @cached_property
-    def variables_metadata(self):
+    def variables_metadata(self) -> Dict[str, Any]:
+        """Get the metadata of the variables."""
         return {k: v for k, v in self.dataset.variables_metadata.items() if k in self.variables}
 
     @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)}
 
     @cached_property
-    def statistics(self):
+    def statistics(self) -> Dict[str, NDArray[Any]]:
+        """Get the statistics of the dataset."""
         return {k: v[self.indices] for k, v in self.dataset.statistics.items()}
 
-    def statistics_tendencies(self, delta=None):
+    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+        """Get the statistical tendencies of the dataset.
+
+        Parameters
+        ----------
+        delta : Optional[datetime.timedelta]
+            The time delta for the tendencies.
+
+        Returns
+        -------
+        Dict[str, NDArray[Any]]
+            The statistical tendencies.
+        """
         if delta is None:
             delta = self.frequency
         return {k: v[self.indices] for k, v in self.dataset.statistics_tendencies(delta).items()}
 
-    def metadata_specific(self, **kwargs):
+    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+        """Get the specific metadata of the dataset.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The specific metadata.
+        """
         return super().metadata_specific(indices=self.indices, **kwargs)
 
-    def source(self, index):
+    def source(self, index: int) -> Source:
+        """Get the source of the dataset.
+
+        Parameters
+        ----------
+        index : int
+            The index of the source.
+
+        Returns
+        -------
+        Source
+            The source of the dataset.
+        """
         return Source(self, index, self.dataset.source(self.indices[index]))
 
-    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.dataset.tree()], **self.reason)
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get the metadata specific to the subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata specific to the subclass.
+        """
         # return dict(indices=self.indices)
         return dict(reason=self.reason)
 
 
 class Rename(Forwards):
-    def __init__(self, dataset, rename):
+    """Class to rename variables in a dataset."""
+
+    def __init__(self, dataset: Dataset, rename: Dict[str, str]) -> None:
+        """Initialize the Rename class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset to rename.
+        rename : Dict[str, str]
+            The mapping of old names to new names.
+        """
         super().__init__(dataset)
         for n in rename:
             assert n in dataset.variables, n
@@ -122,19 +251,32 @@ class Rename(Forwards):
         self.rename = rename
 
     @property
-    def variables(self):
+    def variables(self) -> List[str]:
+        """Get the renamed variables."""
         return self._variables
 
     @property
-    def variables_metadata(self):
+    def variables_metadata(self) -> Dict[str, Any]:
+        """Get the renamed variables metadata."""
         return self._variables_metadata
 
     @cached_property
-    def name_to_index(self):
+    def name_to_index(self) -> Dict[str, int]:
+        """Get the mapping of renamed variable names to indices."""
         return {k: i for i, k in enumerate(self.variables)}
 
-    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()], rename=self.rename)
 
-    def subclass_metadata_specific(self):
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get the metadata specific to the subclass.
+
+        Returns:
+            Dict[str, Any]: The metadata specific to the subclass.
+        """
         return dict(rename=self.rename)

anemoi/datasets/data/statistics.py

@@ -8,10 +8,18 @@
 # 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 Optional
+from typing import Set
+
+from numpy.typing import NDArray
 
 from . import open_dataset
+from .dataset import Dataset
 from .debug import Node
 from .forwards import Forwards
 
@@ -19,7 +27,26 @@ LOG = logging.getLogger(__name__)
 
 
 class Statistics(Forwards):
-    def __init__(self, dataset, statistic):
+    """A class to represent statistics for a dataset.
+
+    Attributes
+    ----------
+    dataset : Dataset
+        The dataset object.
+    statistic : Any
+        The statistic data.
+    """
+
+    def __init__(self, dataset: Dataset, statistic: Any) -> None:
+        """Initialize the Statistics object.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset object.
+        statistic : Any
+            The statistic data.
+        """
         super().__init__(dataset)
         self._statistic = open_dataset(statistic, select=dataset.variables)
         # TODO: relax that check to allow for a subset of variables
@@ -29,20 +56,54 @@ class Statistics(Forwards):
             )
 
     @cached_property
-    def statistics(self):
+    def statistics(self) -> Dict[str, NDArray[Any]]:
+        """Get the statistics."""
         return self._statistic.statistics
 
-    def statistics_tendencies(self, delta=None):
+    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+        """Get the statistics tendencies.
+
+        Parameters
+        ----------
+        delta : Optional[datetime.timedelta]
+            The time delta.
+
+        Returns
+        -------
+        Dict[str, NDArray[Any]]
+            The statistics tendencies.
+        """
         if delta is None:
             delta = self.frequency
         return self._statistic.statistics_tendencies(delta)
 
-    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 dict(statistics=self._statistic.metadata_specific())
 
-    def tree(self):
+    def tree(self) -> Node:
+        """Get the tree representation of the statistics.
+
+        Returns
+        -------
+        Node
+            The tree representation of the statistics.
+        """
         return Node(self, [self.forward.tree()])
 
-    def get_dataset_names(self, names):
+    def get_dataset_names(self, names: Set[str]) -> None:
+        """Get the dataset names.
+
+        Parameters
+        ----------
+        names : Set[str]
+            The set of dataset names.
+        """
         super().get_dataset_names(names)
         self._statistic.get_dataset_names(names)