anemoi-datasets 0.5.26__py3-none-any.whl → 0.5.28__py3-none-any.whl

anemoi/datasets/create/statistics/__init__.py

@@ -17,9 +17,6 @@ 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
@@ -77,7 +74,7 @@ def default_statistics_dates(dates: list[datetime.datetime]) -> tuple[datetime.d
     return dates[0], end
 
 
-def to_datetime(date: Union[str, datetime.datetime]) -> np.datetime64:
+def to_datetime(date: str | datetime.datetime) -> np.datetime64:
     """Convert a date to numpy datetime64 format.
 
     Parameters
@@ -97,7 +94,7 @@ def to_datetime(date: Union[str, datetime.datetime]) -> np.datetime64:
     return date
 
 
-def to_datetimes(dates: list[Union[str, datetime.datetime]]) -> list[np.datetime64]:
+def to_datetimes(dates: list[str | datetime.datetime]) -> list[np.datetime64]:
     """Convert a list of dates to numpy datetime64 format.
 
     Parameters
@@ -221,7 +218,7 @@ def check_variance(
 
 
 def compute_statistics(
-    array: NDArray[Any], check_variables_names: Optional[List[str]] = None, allow_nans: bool = False
+    array: NDArray[Any], check_variables_names: list[str] | None = 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.
 

anemoi/datasets/create/testing.py

@@ -0,0 +1,4 @@
+class TestingContext:
+    """A context for testing plugins."""
+
+    pass

anemoi/datasets/create/typing.py

@@ -8,8 +8,7 @@
 # nor does it submit to any jurisdiction.
 
 import datetime
-from typing import List
 
 Date = datetime.datetime
 
-DateList = List[Date]
+DateList = list[Date]

anemoi/datasets/create/utils.py

@@ -13,7 +13,6 @@ import os
 import warnings
 from contextlib import contextmanager
 from typing import Any
-from typing import Union
 
 import numpy as np
 from earthkit.data import settings
@@ -97,48 +96,6 @@ def to_datetime(*args: Any, **kwargs: Any) -> datetime.datetime:
     return to_datetime_(*args, **kwargs)
 
 
-def make_list_int(value: Union[str, list, tuple, int]) -> list[int]:
-    """Convert a string, list, tuple, or integer to a list of integers.
-
-    Parameters
-    ----------
-    value : str or list or tuple or int
-        The value to convert.
-
-    Returns
-    -------
-    list[int]
-        A list of integers.
-
-    Raises
-    ------
-    ValueError
-        If the value cannot be converted to a list of integers.
-    """
-    # Convert a string like "1/2/3" or "1/to/3" or "1/to/10/by/2" to a list of integers.
-    # Moved to anemoi.utils.humanize
-    # replace with from anemoi.utils.humanize import make_list_int
-    # when anemoi-utils is released and pyproject.toml is updated
-    if isinstance(value, str):
-        if "/" not in value:
-            return [value]
-        bits = value.split("/")
-        if len(bits) == 3 and bits[1].lower() == "to":
-            value = list(range(int(bits[0]), int(bits[2]) + 1, 1))
-
-        elif len(bits) == 5 and bits[1].lower() == "to" and bits[3].lower() == "by":
-            value = list(range(int(bits[0]), int(bits[2]) + int(bits[4]), int(bits[4])))
-
-    if isinstance(value, list):
-        return value
-    if isinstance(value, tuple):
-        return value
-    if isinstance(value, int):
-        return [value]
-
-    raise ValueError(f"Cannot make list from {value}")
-
-
 def normalize_and_check_dates(
     dates: list[datetime.datetime],
     start: datetime.datetime,

anemoi/datasets/create/zarr.py

@@ -11,7 +11,6 @@ import datetime
 import logging
 import shutil
 from typing import Any
-from typing import Optional
 
 import numpy as np
 import zarr
@@ -120,7 +119,7 @@ class ZarrBuiltRegistry:
     flags = None
     z = None
 
-    def __init__(self, path: str, synchronizer_path: Optional[str] = None, use_threads: bool = False):
+    def __init__(self, path: str, synchronizer_path: str | None = None, use_threads: bool = False):
         """Initialize the ZarrBuiltRegistry.
 
         Parameters
@@ -154,6 +153,12 @@ class ZarrBuiltRegistry:
             except FileNotFoundError:
                 pass
 
+        _build = self.zarr_path + "/_build"
+        try:
+            shutil.rmtree(_build)
+        except FileNotFoundError:
+            pass
+
     def _open_write(self) -> zarr.Group:
         """Open the Zarr store in write mode."""
         import zarr

anemoi/datasets/data/__init__.py

@@ -8,9 +8,9 @@
 # nor does it submit to any jurisdiction.
 
 import logging
+import os
 from typing import TYPE_CHECKING
 from typing import Any
-from typing import Set
 
 # from .dataset import FullIndex
 # from .dataset import Shape
@@ -82,6 +82,9 @@ def open_dataset(*args: Any, **kwargs: Any) -> "Dataset":
     Dataset
         The opened dataset.
     """
+
+    trace = int(os.environ.get("ANEMOI_DATASETS_TRACE", 0))
+
     # That will get rid of OmegaConf objects
 
     args, kwargs = _convert(args), _convert(kwargs)
@@ -90,22 +93,28 @@ def open_dataset(*args: Any, **kwargs: Any) -> "Dataset":
     ds = ds.mutate()
     ds.arguments = {"args": args, "kwargs": kwargs}
     ds._check()
+
+    if trace:
+        from anemoi.datasets.testing import Trace
+
+        ds = Trace(ds)
+
     return ds
 
 
-def save_dataset(recipe: dict, zarr_path: str, n_workers: int = 1) -> None:
+def save_dataset(dataset: "Dataset", zarr_path: str, n_workers: int = 1) -> None:
     """Open a dataset and save it to disk.
 
     Parameters
     ----------
-    recipe : dict
-        Recipe used with open_dataset (not a dataset creation recipe).
+    dataset : Dataset
+        anemoi-dataset opened from python to save to Zarr store
     zarr_path : str
         Path to store the obtained anemoi dataset to disk.
     n_workers : int
         Number of workers to use for parallel processing. If none, sequential processing will be performed.
     """
-    _save_dataset(recipe, zarr_path, n_workers)
+    _save_dataset(dataset, zarr_path, n_workers)
 
 
 def list_dataset_names(*args: Any, **kwargs: Any) -> list[str]:
@@ -124,6 +133,6 @@ def list_dataset_names(*args: Any, **kwargs: Any) -> list[str]:
         The list of dataset names.
     """
     ds = _open_dataset(*args, **kwargs)
-    names: Set[str] = set()
+    names: set[str] = set()
     ds.get_dataset_names(names)
     return sorted(names)

anemoi/datasets/data/complement.py

@@ -12,11 +12,6 @@ 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 Optional
-from typing import Set
-from typing import Tuple
 
 import numpy as np
 from numpy.typing import NDArray
@@ -91,26 +86,26 @@ class Complement(Combined):
             raise ValueError("Augment: no missing variables")
 
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Returns the list of variables to be added to the target dataset."""
         return self._variables
 
     @property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         datasets = [self._source, self._target]
         return {
             k: [d.statistics[k][d.name_to_index[i]] for d in datasets for i in d.variables if i in self.variables]
             for k in datasets[0].statistics
         }
 
-    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+    def statistics_tendencies(self, delta: datetime.timedelta | None = None) -> dict[str, NDArray[Any]]:
         index = [self._source.name_to_index[v] for v in self._variables]
         if delta is None:
             delta = self.frequency
         return {k: v[index] for k, v in self._source.statistics_tendencies(delta).items()}
 
     @property
-    def name_to_index(self) -> Dict[str, int]:
+    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)}
 
@@ -121,7 +116,7 @@ class Complement(Combined):
         return (shape[0], len(self._variables)) + shape[2:]
 
     @property
-    def variables_metadata(self) -> Dict[str, Any]:
+    def variables_metadata(self) -> dict[str, Any]:
         """Returns the metadata of the variables to be added to the target dataset."""
         # Merge the two dicts first
         all_meta = {**self._source.variables_metadata, **self._target.variables_metadata}
@@ -142,7 +137,7 @@ class Complement(Combined):
         pass
 
     @cached_property
-    def missing(self) -> Set[int]:
+    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
@@ -317,7 +312,7 @@ class ComplementNearest(Complement):
         return apply_index_to_slices_changes(result, changes)
 
 
-def complement_factory(args: Tuple, kwargs: dict) -> Dataset:
+def complement_factory(args: tuple, kwargs: dict) -> Dataset:
     """Factory function to create a Complement instance based on the provided arguments.
 
     Parameters

anemoi/datasets/data/concat.py

@@ -11,9 +11,6 @@
 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
@@ -123,12 +120,12 @@ class ConcatMixin:
         return np.concatenate(result)
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Returns the set of missing indices in the concatenated datasets."""
-        result: Set[int] = set()
+        result: set[int] = set()
         offset = 0
         for d in self.datasets:
-            result = result | set(m + offset for m in d.missing)
+            result = result | {m + offset for m in d.missing}
             offset += len(d)
         return result
 
@@ -195,7 +192,7 @@ class Concat(ConcatMixin, Combined):
         return Node(self, [d.tree() for d in self.datasets])
 
     @classmethod
-    def check_dataset_compatibility(cls, datasets: List[Any], fill_missing_gaps: bool = False) -> List[Any]:
+    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
@@ -259,7 +256,7 @@ class Concat(ConcatMixin, Combined):
         return {}
 
 
-def concat_factory(args: Tuple[Any, ...], kwargs: dict) -> Concat:
+def concat_factory(args: tuple[Any, ...], kwargs: dict) -> Concat:
     """Factory function to create a Concat object.
 
     Parameters

anemoi/datasets/data/dataset.py

@@ -22,15 +22,10 @@ try:
 except ImportError:
     # Python 3.9
     EllipsisType = type(Ellipsis)
+from collections.abc import Sequence
+from collections.abc import Sized
 from typing import TYPE_CHECKING
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Sequence
-from typing import Set
-from typing import Sized
-from typing import Tuple
 from typing import Union
 
 import numpy as np
@@ -48,8 +43,8 @@ if TYPE_CHECKING:
 LOG = logging.getLogger(__name__)
 
 
-Shape = Tuple[int, ...]
-TupleIndex = Tuple[Union[int, slice, EllipsisType], ...]
+Shape = tuple[int, ...]
+TupleIndex = tuple[Union[int, slice, EllipsisType], ...]
 FullIndex = Union[int, slice, TupleIndex]
 
 
@@ -92,8 +87,8 @@ def _tidy(v: Any) -> Any:
 
 
 class Dataset(ABC, Sized):
-    arguments: Dict[str, Any] = {}
-    _name: Union[str, None] = None
+    arguments: dict[str, Any] = {}
+    _name: str | None = None
 
     def mutate(self) -> "Dataset":
         """Give an opportunity to a subclass to return a new Dataset object of a different class, if needed.
@@ -148,7 +143,7 @@ class Dataset(ABC, Sized):
         return result
 
     @property
-    def name(self) -> Union[str, None]:
+    def name(self) -> str | None:
         """Return the name of the dataset."""
         return self._name
 
@@ -298,6 +293,12 @@ class Dataset(ABC, Sized):
             if skip_missing_dates:
                 return SkipMissingDates(self, expected_access)._subset(**kwargs).mutate()
 
+        if "rolling_average" in kwargs:
+            from .rolling_average import RollingAverage
+
+            rolling_average = kwargs.pop("rolling_average")
+            return RollingAverage(self, rolling_average)._subset(**kwargs).mutate()
+
         if "interpolate_frequency" in kwargs:
             from .interpolate import InterpolateFrequency
 
@@ -360,9 +361,9 @@ class Dataset(ABC, Sized):
 
     def _dates_to_indices(
         self,
-        start: Union[None, str, datetime.datetime],
-        end: Union[None, str, datetime.datetime],
-    ) -> List[int]:
+        start: None | str | datetime.datetime,
+        end: None | str | datetime.datetime,
+    ) -> list[int]:
         """Convert date range to a list of indices.
 
         Parameters
@@ -387,7 +388,7 @@ class Dataset(ABC, Sized):
 
         return [i for i, date in enumerate(self.dates) if start <= date <= end]
 
-    def _select_to_columns(self, vars: Union[str, List[str], Tuple[str], set]) -> List[int]:
+    def _select_to_columns(self, vars: str | list[str] | tuple[str] | set) -> list[int]:
         """Convert variable names to a list of column indices.
 
         Parameters
@@ -411,7 +412,7 @@ class Dataset(ABC, Sized):
 
         return [self.name_to_index[v] for v in vars]
 
-    def _drop_to_columns(self, vars: Union[str, Sequence[str]]) -> List[int]:
+    def _drop_to_columns(self, vars: str | Sequence[str]) -> list[int]:
         """Convert variable names to a list of column indices to drop.
 
         Parameters
@@ -432,7 +433,7 @@ class Dataset(ABC, Sized):
 
         return sorted([v for k, v in self.name_to_index.items() if k not in vars])
 
-    def _reorder_to_columns(self, vars: Union[str, List[str], Tuple[str], Dict[str, int]]) -> List[int]:
+    def _reorder_to_columns(self, vars: str | list[str] | tuple[str] | dict[str, int]) -> list[int]:
         """Convert variable names to a list of reordered column indices.
 
         Parameters
@@ -465,8 +466,8 @@ class Dataset(ABC, Sized):
         return indices
 
     def dates_interval_to_indices(
-        self, start: Union[None, str, datetime.datetime], end: Union[None, str, datetime.datetime]
-    ) -> List[int]:
+        self, start: None | str | datetime.datetime, end: None | str | datetime.datetime
+    ) -> list[int]:
         """Convert date interval to a list of indices.
 
         Parameters
@@ -483,7 +484,7 @@ class Dataset(ABC, Sized):
         """
         return self._dates_to_indices(start, end)
 
-    def provenance(self) -> Dict[str, Any]:
+    def provenance(self) -> dict[str, Any]:
         """Return the provenance information of the dataset.
 
         Returns
@@ -511,7 +512,7 @@ class Dataset(ABC, Sized):
         return tuple(shape)
 
     @property
-    def typed_variables(self) -> Dict[str, Any]:
+    def typed_variables(self) -> dict[str, Any]:
         """Return the variables with their types."""
         from anemoi.transform.variables import Variable
 
@@ -532,7 +533,7 @@ class Dataset(ABC, Sized):
 
         return result
 
-    def _input_sources(self) -> List[Any]:
+    def _input_sources(self) -> list[Any]:
         """Return the input sources of the dataset.
 
         Returns
@@ -544,7 +545,7 @@ class Dataset(ABC, Sized):
         self.collect_input_sources(sources)
         return sources
 
-    def metadata(self) -> Dict[str, Any]:
+    def metadata(self) -> dict[str, Any]:
         """Return the metadata of the dataset.
 
         Returns
@@ -588,7 +589,7 @@ class Dataset(ABC, Sized):
         """Return the end date of the dataset."""
         return self.dates[-1]
 
-    def dataset_metadata(self) -> Dict[str, Any]:
+    def dataset_metadata(self) -> dict[str, Any]:
         """Return the metadata of the dataset.
 
         Returns
@@ -608,7 +609,7 @@ class Dataset(ABC, Sized):
             name=self.name,
         )
 
-    def _supporting_arrays(self, *path: str) -> Dict[str, NDArray[Any]]:
+    def _supporting_arrays(self, *path: str) -> dict[str, NDArray[Any]]:
         """Return the supporting arrays of the dataset.
 
         Parameters
@@ -646,7 +647,7 @@ class Dataset(ABC, Sized):
 
         return result
 
-    def supporting_arrays(self) -> Dict[str, NDArray[Any]]:
+    def supporting_arrays(self) -> dict[str, NDArray[Any]]:
         """Return the supporting arrays to be saved in the checkpoints.
 
         Returns
@@ -657,7 +658,7 @@ class Dataset(ABC, Sized):
         arrays, _ = self._supporting_arrays_and_sources()
         return arrays
 
-    def _supporting_arrays_and_sources(self) -> Tuple[Dict[str, NDArray], Dict[int, List[str]]]:
+    def _supporting_arrays_and_sources(self) -> tuple[dict[str, NDArray], dict[int, list[str]]]:
         """Return the supporting arrays and their sources.
 
         Returns
@@ -684,7 +685,7 @@ class Dataset(ABC, Sized):
 
         return result, source_to_arrays
 
-    def collect_supporting_arrays(self, collected: List[Tuple[Tuple[str, ...], str, NDArray[Any]]], *path: str) -> None:
+    def collect_supporting_arrays(self, collected: list[tuple[tuple[str, ...], str, NDArray[Any]]], *path: str) -> None:
         """Collect supporting arrays.
 
         Parameters
@@ -697,7 +698,7 @@ class Dataset(ABC, Sized):
         # Override this method to add more arrays
         pass
 
-    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+    def metadata_specific(self, **kwargs: Any) -> dict[str, Any]:
         """Return specific metadata of the dataset.
 
         Parameters
@@ -764,7 +765,7 @@ class Dataset(ABC, Sized):
         """Return the label of the dataset."""
         return self.__class__.__name__.lower()
 
-    def computed_constant_fields(self) -> List[str]:
+    def computed_constant_fields(self) -> list[str]:
         """Return the computed constant fields of the dataset.
 
         Returns
@@ -781,7 +782,7 @@ class Dataset(ABC, Sized):
 
         return sorted(self._compute_constant_fields_from_a_few_samples())
 
-    def _compute_constant_fields_from_a_few_samples(self) -> List[str]:
+    def _compute_constant_fields_from_a_few_samples(self) -> list[str]:
         """Compute constant fields from a few samples.
 
         Returns
@@ -822,7 +823,7 @@ class Dataset(ABC, Sized):
 
         return [v for i, v in enumerate(self.variables) if constants[i]]
 
-    def _compute_constant_fields_from_statistics(self) -> List[str]:
+    def _compute_constant_fields_from_statistics(self) -> list[str]:
         """Compute constant fields from statistics.
 
         Returns
@@ -842,8 +843,8 @@ class Dataset(ABC, Sized):
 
     def plot(
         self,
-        date: Union[int, datetime.datetime, np.datetime64, str],
-        variable: Union[int, str],
+        date: int | datetime.datetime | np.datetime64 | str,
+        variable: int | str,
         member: int = 0,
         **kwargs: Any,
     ) -> "matplotlib.pyplot.Axes":
@@ -873,10 +874,10 @@ class Dataset(ABC, Sized):
 
     def to_index(
         self,
-        date: Union[int, datetime.datetime, np.datetime64, str],
-        variable: Union[int, str],
+        date: int | datetime.datetime | np.datetime64 | str,
+        variable: int | str,
         member: int = 0,
-    ) -> Tuple[int, int, int]:
+    ) -> tuple[int, int, int]:
         """Convert date, variable, and member to indices.
 
         Parameters
@@ -945,7 +946,7 @@ class Dataset(ABC, Sized):
 
     @property
     @abstractmethod
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Return the list of variables in the dataset."""
         pass
 
@@ -969,7 +970,7 @@ class Dataset(ABC, Sized):
 
     @property
     @abstractmethod
-    def name_to_index(self) -> Dict[str, int]:
+    def name_to_index(self) -> dict[str, int]:
         """Return the mapping of variable names to indices."""
         pass
 
@@ -1005,30 +1006,30 @@ class Dataset(ABC, Sized):
 
     @property
     @abstractmethod
-    def variables_metadata(self) -> Dict[str, Any]:
+    def variables_metadata(self) -> dict[str, Any]:
         """Return the metadata of the variables in the dataset."""
         pass
 
     @abstractmethod
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Return the set of missing indices in the dataset."""
         pass
 
     @abstractmethod
     @cached_property
-    def constant_fields(self) -> List[str]:
+    def constant_fields(self) -> list[str]:
         """Return the list of constant fields in the dataset."""
         pass
 
     @abstractmethod
     @cached_property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Return the statistics of the dataset."""
         pass
 
     @abstractmethod
-    def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+    def statistics_tendencies(self, delta: datetime.timedelta | None = None) -> dict[str, NDArray[Any]]:
         """Return the tendencies of the statistics in the dataset.
 
         Parameters
@@ -1071,7 +1072,7 @@ class Dataset(ABC, Sized):
         pass
 
     @abstractmethod
-    def collect_input_sources(self, sources: List[Any]) -> None:
+    def collect_input_sources(self, sources: list[Any]) -> None:
         """Collect the input sources of the dataset.
 
         Parameters
@@ -1082,7 +1083,7 @@ class Dataset(ABC, Sized):
         pass
 
     @abstractmethod
-    def get_dataset_names(self, names: Set[str]) -> None:
+    def get_dataset_names(self, names: set[str]) -> None:
         """Get the names of the datasets.
 
         Parameters