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

anemoi/datasets/create/statistics/__init__.py

@@ -16,10 +16,15 @@ import os
 import pickle
 import shutil
 import socket
+from typing import Any
+from typing import List
+from typing import Optional
+from typing import Union
 
 import numpy as np
 import tqdm
 from anemoi.utils.provenance import gather_provenance_info
+from numpy.typing import NDArray
 
 from ..check import check_data_values
 from .summary import Summary
@@ -27,15 +32,18 @@ from .summary import Summary
 LOG = logging.getLogger(__name__)
 
 
-def default_statistics_dates(dates):
+def default_statistics_dates(dates: list[datetime.datetime]) -> tuple[datetime.datetime, datetime.datetime]:
     """Calculate default statistics dates based on the given list of dates.
 
-    Args:
-        dates (list): List of datetime objects representing dates.
-
-    Returns:
-        tuple: A tuple containing the default start and end dates.
+    Parameters
+    ----------
+    dates : list of datetime.datetime
+        List of datetime objects representing dates.
 
+    Returns
+    -------
+    tuple of datetime.datetime
+        A tuple containing the default start and end dates.
     """
 
     def to_datetime(d):
@@ -69,7 +77,19 @@ def default_statistics_dates(dates):
     return dates[0], end
 
 
-def to_datetime(date):
+def to_datetime(date: Union[str, datetime.datetime]) -> np.datetime64:
+    """Convert a date to numpy datetime64 format.
+
+    Parameters
+    ----------
+    date : str or datetime.datetime
+        The date to convert.
+
+    Returns
+    -------
+    numpy.datetime64
+        The converted date.
+    """
     if isinstance(date, str):
         return np.datetime64(date)
     if isinstance(date, datetime.datetime):
@@ -77,11 +97,43 @@ def to_datetime(date):
     return date
 
 
-def to_datetimes(dates):
+def to_datetimes(dates: list[Union[str, datetime.datetime]]) -> list[np.datetime64]:
+    """Convert a list of dates to numpy datetime64 format.
+
+    Parameters
+    ----------
+    dates : list of str or datetime.datetime
+        List of dates to convert.
+
+    Returns
+    -------
+    list of numpy.datetime64
+        List of converted dates.
+    """
     return [to_datetime(d) for d in dates]
 
 
-def fix_variance(x, name, count, sums, squares):
+def fix_variance(x: float, name: str, count: NDArray[Any], sums: NDArray[Any], squares: NDArray[Any]) -> float:
+    """Fix negative variance values due to numerical errors.
+
+    Parameters
+    ----------
+    x : float
+        The variance value.
+    name : str
+        The variable name.
+    count : numpy.ndarray
+        The count array.
+    sums : numpy.ndarray
+        The sums array.
+    squares : numpy.ndarray
+        The squares array.
+
+    Returns
+    -------
+    float
+        The fixed variance value.
+    """
     assert count.shape == sums.shape == squares.shape
     assert isinstance(x, float)
 
@@ -112,7 +164,42 @@ def fix_variance(x, name, count, sums, squares):
     return 0
 
 
-def check_variance(x, variables_names, minimum, maximum, mean, count, sums, squares):
+def check_variance(
+    x: NDArray[Any],
+    variables_names: list[str],
+    minimum: NDArray[Any],
+    maximum: NDArray[Any],
+    mean: NDArray[Any],
+    count: NDArray[Any],
+    sums: NDArray[Any],
+    squares: NDArray[Any],
+) -> None:
+    """Check for negative variance values and raise an error if found.
+
+    Parameters
+    ----------
+    x : numpy.ndarray
+        The variance array.
+    variables_names : list of str
+        List of variable names.
+    minimum : numpy.ndarray
+        The minimum values array.
+    maximum : numpy.ndarray
+        The maximum values array.
+    mean : numpy.ndarray
+        The mean values array.
+    count : numpy.ndarray
+        The count array.
+    sums : numpy.ndarray
+        The sums array.
+    squares : numpy.ndarray
+        The squares array.
+
+    Raises
+    ------
+    ValueError
+        If negative variance is found.
+    """
     if (x >= 0).all():
         return
     print(x)
@@ -133,8 +220,25 @@ def check_variance(x, variables_names, minimum, maximum, mean, count, sums, squa
     raise ValueError("Negative variance")
 
 
-def compute_statistics(array, check_variables_names=None, allow_nans=False):
-    """Compute statistics for a given array, provides minimum, maximum, sum, squares, count and has_nans as a dictionary."""
+def compute_statistics(
+    array: NDArray[Any], check_variables_names: Optional[List[str]] = None, allow_nans: bool = False
+) -> dict[str, np.ndarray]:
+    """Compute statistics for a given array, provides minimum, maximum, sum, squares, count and has_nans as a dictionary.
+
+    Parameters
+    ----------
+    array : numpy.ndarray
+        The array to compute statistics for.
+    check_variables_names : list of str, optional
+        List of variable names to check. Defaults to None.
+    allow_nans : bool, optional
+        Whether to allow NaN values. Defaults to False.
+
+    Returns
+    -------
+    dict of str to numpy.ndarray
+        A dictionary containing the computed statistics.
+    """
     LOG.info(f"Computing statistics for {array.shape} array")
     nvars = array.shape[1]
 
@@ -180,16 +284,34 @@ def compute_statistics(array, check_variables_names=None, allow_nans=False):
 
 
 class TmpStatistics:
+    """Temporary statistics storage class."""
+
     version = 3
     # Used in parrallel, during data loading,
     # to write statistics in pickled npz files.
     # can provide statistics for a subset of dates.
 
-    def __init__(self, dirname, overwrite=False):
+    def __init__(self, dirname: str, overwrite: bool = False) -> None:
+        """Initialize TmpStatistics.
+
+        Parameters
+        ----------
+        dirname : str
+            Directory name for storing statistics.
+        overwrite : bool, optional
+            Whether to overwrite existing files. Defaults to False.
+        """
         self.dirname = dirname
         self.overwrite = overwrite
 
-    def add_provenance(self, **kwargs):
+    def add_provenance(self, **kwargs: dict) -> None:
+        """Add provenance information.
+
+        Parameters
+        ----------
+        **kwargs : dict
+            Additional provenance information.
+        """
         self.create(exist_ok=True)
         path = os.path.join(self.dirname, "provenance.json")
         if os.path.exists(path):
@@ -198,16 +320,35 @@ class TmpStatistics:
         with open(path, "w") as f:
             json.dump(out, f)
 
-    def create(self, exist_ok):
+    def create(self, exist_ok: bool) -> None:
+        """Create the directory for storing statistics.
+
+        Parameters
+        ----------
+        exist_ok : bool
+            Whether to ignore if the directory already exists.
+        """
         os.makedirs(self.dirname, exist_ok=exist_ok)
 
-    def delete(self):
+    def delete(self) -> None:
+        """Delete the directory for storing statistics."""
         try:
             shutil.rmtree(self.dirname)
         except FileNotFoundError:
             pass
 
-    def write(self, key, data, dates):
+    def write(self, key: str, data: any, dates: list[datetime.datetime]) -> None:
+        """Write statistics data to a file.
+
+        Parameters
+        ----------
+        key : str
+            The key for the data.
+        data : any
+            The data to write.
+        dates : list of datetime.datetime
+            List of dates associated with the data.
+        """
         self.create(exist_ok=True)
         h = hashlib.sha256(str(dates).encode("utf-8")).hexdigest()
         path = os.path.join(self.dirname, f"{h}.npz")
@@ -222,7 +363,14 @@ class TmpStatistics:
 
         LOG.debug(f"Written statistics data for {len(dates)} dates in {path} ({dates})")
 
-    def _gather_data(self):
+    def _gather_data(self) -> tuple[str, list[datetime.datetime], dict]:
+        """Gather data from stored files.
+
+        Yields
+        ------
+        tuple of str, list of datetime.datetime, dict
+            A tuple containing key, dates, and data.
+        """
         # use glob to read all pickles
         files = glob.glob(self.dirname + "/*.npz")
         LOG.debug(f"Reading stats data, found {len(files)} files in {self.dirname}")
@@ -231,37 +379,67 @@ class TmpStatistics:
             with open(f, "rb") as f:
                 yield pickle.load(f)
 
-    def get_aggregated(self, *args, **kwargs):
+    def get_aggregated(self, *args: Any, **kwargs: Any) -> Summary:
+        """Get aggregated statistics.
+
+        Parameters
+        ----------
+        *args : Any
+            Additional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        Summary
+            The aggregated statistics summary.
+        """
         aggregator = StatAggregator(self, *args, **kwargs)
         return aggregator.aggregate()
 
-    def __str__(self):
-        return f"TmpStatistics({self.dirname})"
-
-
-def normalise_date(d):
-    if isinstance(d, str):
-        d = np.datetime64(d)
-    return d
+    def __str__(self) -> str:
+        """String representation of TmpStatistics.
 
-
-def normalise_dates(dates):
-    return [normalise_date(d) for d in dates]
+        Returns
+        -------
+        str
+            The string representation.
+        """
+        return f"TmpStatistics({self.dirname})"
 
 
 class StatAggregator:
+    """Statistics aggregator class."""
+
     NAMES = ["minimum", "maximum", "sums", "squares", "count", "has_nans"]
 
-    def __init__(self, owner, dates, variables_names, allow_nans):
+    def __init__(
+        self, owner: TmpStatistics, dates: list[datetime.datetime], variables_names: list[str], allow_nans: bool
+    ) -> None:
+        """Initialize StatAggregator.
+
+        Parameters
+        ----------
+        owner : TmpStatistics
+            The owner TmpStatistics instance.
+        dates : list of datetime.datetime
+            List of dates to aggregate.
+        variables_names : list of str
+            List of variable names.
+        allow_nans : bool
+            Whether to allow NaN values.
+        """
         dates = sorted(dates)
         dates = to_datetimes(dates)
         assert dates, "No dates selected"
         self.owner = owner
         self.dates = dates
+        self._number_of_dates = len(dates)
+        self._set_of_dates = set(dates)
         self.variables_names = variables_names
         self.allow_nans = allow_nans
 
-        self.shape = (len(self.dates), len(self.variables_names))
+        self.shape = (self._number_of_dates, len(self.variables_names))
         LOG.debug(f"Aggregating statistics on shape={self.shape}. Variables : {self.variables_names}")
 
         self.minimum = np.full(self.shape, np.nan, dtype=np.float64)
@@ -273,12 +451,16 @@ class StatAggregator:
 
         self._read()
 
-    def _read(self):
+    def _read(self) -> None:
+        """Read and aggregate statistics data from files."""
+
         def check_type(a, b):
-            a = list(a)
-            b = list(b)
-            a = a[0] if a else None
-            b = b[0] if b else None
+            if not isinstance(a, set):
+                a = set(list(a))
+            if not isinstance(b, set):
+                b = set(list(b))
+            a = next(iter(a)) if a else None
+            b = next(iter(b)) if b else None
             assert type(a) is type(b), (type(a), type(b))
 
         found = set()
@@ -294,20 +476,20 @@ class StatAggregator:
             for n in self.NAMES:
                 assert n in stats, (n, list(stats.keys()))
             _dates = to_datetimes(_dates)
-            check_type(_dates, self.dates)
+            check_type(_dates, self._set_of_dates)
             if found:
-                check_type(found, self.dates)
+                check_type(found, self._set_of_dates)
                 assert found.isdisjoint(_dates), "Duplicate dates found in precomputed statistics"
 
             # filter dates
-            dates = set(_dates) & set(self.dates)
+            dates = set(_dates) & self._set_of_dates
 
             if not dates:
                 # dates have been completely filtered for this chunk
                 continue
 
             # filter data
-            bitmap = np.isin(_dates, self.dates)
+            bitmap = np.array([d in self._set_of_dates for d in _dates])
             for k in self.NAMES:
                 stats[k] = stats[k][bitmap]
 
@@ -323,11 +505,18 @@ class StatAggregator:
 
         for d in self.dates:
             assert d in found, f"Statistics for date {d} not precomputed."
-        assert len(self.dates) == len(found), "Not all dates found in precomputed statistics"
-        assert len(self.dates) == offset, "Not all dates found in precomputed statistics."
+        assert self._number_of_dates == len(found), "Not all dates found in precomputed statistics"
+        assert self._number_of_dates == offset, "Not all dates found in precomputed statistics."
         LOG.debug(f"Statistics for {len(found)} dates found.")
 
-    def aggregate(self):
+    def aggregate(self) -> Summary:
+        """Aggregate the statistics data.
+
+        Returns
+        -------
+        Summary
+            The aggregated statistics summary.
+        """
         minimum = np.nanmin(self.minimum, axis=0)
         maximum = np.nanmax(self.maximum, axis=0)
 

anemoi/datasets/create/statistics/summary.py

@@ -9,6 +9,7 @@
 
 import json
 from collections import defaultdict
+from typing import Any
 
 import numpy as np
 
@@ -28,15 +29,32 @@ class Summary(dict):
         "has_nans",
     ]  # order matter for __str__.
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the Summary object with given keyword arguments.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Arbitrary keyword arguments representing summary statistics.
+        """
         super().__init__(**kwargs)
         self.check()
 
     @property
-    def size(self):
+    def size(self) -> int:
+        """Get the size of the summary, which is the number of variables."""
         return len(self["variables_names"])
 
-    def check(self):
+    def check(self) -> None:
+        """Perform checks on the summary statistics to ensure they are valid.
+
+        Raises
+        ------
+        AssertionError
+            If any of the checks fail.
+        StatisticsValueError
+            If any of the statistical checks fail.
+        """
         for k, v in self.items():
             if k == "variables_names":
                 assert len(v) == self.size
@@ -63,7 +81,14 @@ class Summary(dict):
                 e.args += (i, name)
                 raise
 
-    def __str__(self):
+    def __str__(self) -> str:
+        """Return a string representation of the summary statistics.
+
+        Returns
+        -------
+        str
+            A formatted string of the summary statistics.
+        """
         header = ["Variables"] + self.STATS_NAMES
         out = [" ".join(header)]
 
@@ -73,7 +98,16 @@ class Summary(dict):
         ]
         return "\n".join(out)
 
-    def save(self, filename, **metadata):
+    def save(self, filename: str, **metadata: Any) -> None:
+        """Save the summary statistics to a JSON file.
+
+        Parameters
+        ----------
+        filename : str
+            The name of the file to save the summary statistics.
+        **metadata : Any
+            Additional metadata to include in the JSON file.
+        """
         assert filename.endswith(".json"), filename
         dic = {}
         for k in self.STATS_NAMES:
@@ -89,7 +123,19 @@ class Summary(dict):
         with open(filename, "w") as f:
             json.dump(out, f, indent=2)
 
-    def load(self, filename):
+    def load(self, filename: str) -> "Summary":
+        """Load the summary statistics from a JSON file.
+
+        Parameters
+        ----------
+        filename : str
+            The name of the file to load the summary statistics from.
+
+        Returns
+        -------
+        Summary
+            The loaded Summary object.
+        """
         assert filename.endswith(".json"), filename
         with open(filename) as f:
             dic = json.load(f)

anemoi/datasets/create/testing.py

@@ -0,0 +1,76 @@
+# (C) Copyright 2025- Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import tempfile
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+import yaml
+
+from anemoi.datasets.create import creator_factory
+
+
+class TestingContext:
+    pass
+
+
+def create_dataset(
+    *,
+    config: Union[str, Dict[str, Any]],
+    output: Optional[str],
+    delta: Optional[List[str]] = None,
+    is_test: bool = False,
+) -> str:
+    """Create a dataset based on the provided configuration.
+
+    Parameters
+    ----------
+    config : Union[str, Dict[str, Any]]
+        The configuration for the dataset. Can be a path to a YAML file or a dictionary.
+    output : Optional[str]
+        The output path for the dataset. If None, a temporary directory will be created.
+    delta : Optional[List[str]], optional
+        List of delta for secondary statistics, by default None.
+    is_test : bool, optional
+        Flag indicating if the dataset creation is for testing purposes, by default False.
+
+    Returns
+    -------
+    str
+        The path to the created dataset.
+    """
+    if isinstance(config, dict):
+        temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml")
+        yaml.dump(config, temp_file)
+        config = temp_file.name
+
+    if output is None:
+        output = tempfile.mkdtemp(suffix=".zarr")
+
+    creator_factory("init", config=config, path=output, overwrite=True, test=is_test).run()
+    creator_factory("load", path=output).run()
+    creator_factory("finalise", path=output).run()
+    creator_factory("patch", path=output).run()
+
+    if delta is not None:
+        creator_factory("init_additions", path=output, delta=delta).run()
+        creator_factory("run_additions", path=output, delta=delta).run()
+        creator_factory("finalise_additions", path=output, delta=delta).run()
+
+    creator_factory("cleanup", path=output).run()
+
+    if delta is not None:
+        creator_factory("cleanup", path=output, delta=delta).run()
+
+    creator_factory("verify", path=output).run()
+
+    return output

anemoi/datasets/create/{functions/filters/noop.py → typing.py}

@@ -1,4 +1,4 @@
-# (C) Copyright 2024 Anemoi contributors.
+# (C) Copyright 2025- Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
@@ -7,6 +7,9 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+import datetime
+from typing import List
 
-def execute(context, input, *args, **kwargs):
-    return input
+Date = datetime.datetime
+
+DateList = List[Date]