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

anemoi/datasets/data/observations/__init__.py

@@ -10,8 +10,6 @@ import logging
 import os
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import Tuple
 
 import numpy as np
 from anemoi.utils.dates import frequency_to_timedelta
@@ -82,10 +80,8 @@ class ObservationsBase(Dataset):
         #    return [self.getitem(j) for j in i]
 
         raise ValueError(
-            (
-                f"Expected int, got {i} of type {type(i)}. Only int is supported to index "
-                "observations datasets. Please use a second [] to select part of the data [i][a,b,c]"
-            )
+            f"Expected int, got {i} of type {type(i)}. Only int is supported to index "
+            "observations datasets. Please use a second [] to select part of the data [i][a,b,c]"
         )
 
     @property
@@ -195,13 +191,11 @@ class ObservationsZarr(ObservationsBase):
 
         if len(self.forward) != len(self.dates):
             raise ValueError(
-                (
-                    f"Dates are not consistent with the number of items in the dataset. "
-                    f"The dataset contains {len(self.forward)} time windows. "
-                    f"This is not compatible with the "
-                    f"{len(self.dates)} requested dates with frequency={frequency_hours}"
-                    f"{self.dates[0]}, {self.dates[1]}, ..., {self.dates[-2]}, {self.dates[-1]} "
-                )
+                f"Dates are not consistent with the number of items in the dataset. "
+                f"The dataset contains {len(self.forward)} time windows. "
+                f"This is not compatible with the "
+                f"{len(self.dates)} requested dates with frequency={frequency_hours}"
+                f"{self.dates[0]}, {self.dates[1]}, ..., {self.dates[-2]}, {self.dates[-1]} "
             )
 
     @property
@@ -307,7 +301,7 @@ class ObservationsZarr(ObservationsBase):
         return f"Observations({os.path.basename(self.path)}, {self.dates[0]};{self.dates[-1]}, {len(self)})"
 
 
-def observations_factory(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> ObservationsBase:
+def observations_factory(args: tuple[Any, ...], kwargs: dict[str, Any]) -> ObservationsBase:
     observations = kwargs.pop("observations")
 
     if not isinstance(observations, dict):

anemoi/datasets/data/padded.py

@@ -12,8 +12,6 @@ import datetime
 import logging
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import Set
 
 import numpy as np
 from anemoi.utils.dates import frequency_to_timedelta
@@ -38,7 +36,7 @@ class Padded(Forwards):
     _after: int = 0
     _inside: int = 0
 
-    def __init__(self, dataset: Dataset, start: str, end: str, frequency: str, reason: Dict[str, Any]) -> None:
+    def __init__(self, dataset: Dataset, start: str, end: str, frequency: str, reason: dict[str, Any]) -> None:
         """Create a padded subset of a dataset.
 
         Attributes:
@@ -195,7 +193,7 @@ class Padded(Forwards):
         return (len(self.dates),) + self.dataset.shape[1:]
 
     @cached_property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         raise NotImplementedError("Need to decide whether to include the added dates as missing or not")
         # return self.forward.missing
 
@@ -207,7 +205,7 @@ class Padded(Forwards):
         """
         return Node(self, [self.dataset.tree()], **self.reason)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the forwards subclass.
 
         Returns:

anemoi/datasets/data/records/backends/__init__.py

@@ -35,7 +35,7 @@ class Npz1Backend(Backend):
             return dict(np.load(f))
 
     def read_metadata(self):
-        with open(os.path.join(self.path, "metadata.json"), "r") as f:
+        with open(os.path.join(self.path, "metadata.json")) as f:
             return json.load(f)
 
     def read_statistics(self):
@@ -56,7 +56,7 @@ class Npz2Backend(Backend):
             return dict(np.load(f))
 
     def read_metadata(self):
-        with open(os.path.join(self.path, "metadata.json"), "r") as f:
+        with open(os.path.join(self.path, "metadata.json")) as f:
             return json.load(f)
 
     def read_statistics(self):

anemoi/datasets/data/rescale.py

@@ -12,11 +12,6 @@ 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
@@ -35,9 +30,7 @@ from .indexing import update_tuple
 LOG = logging.getLogger(__name__)
 
 
-def make_rescale(
-    variable: str, rescale: Union[Tuple[float, float], List[str], Dict[str, float]]
-) -> Tuple[float, float]:
+def make_rescale(variable: str, rescale: tuple[float, float] | list[str] | dict[str, float]) -> tuple[float, float]:
     """Create rescale parameters (scale and offset) based on the input rescale specification.
 
     Parameters
@@ -86,7 +79,7 @@ class Rescale(Forwards):
     """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]]]
+        self, dataset: Dataset, rescale: dict[str, tuple[float, float] | list[str] | dict[str, float]]
     ) -> None:
         """Initialize the Rescale object.
 
@@ -129,7 +122,7 @@ class Rescale(Forwards):
         """
         return Node(self, [self.forward.tree()], rescale=self.rescale)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the rescale subclass.
 
         Returns
@@ -204,7 +197,7 @@ class Rescale(Forwards):
         return data * self._a[0] + self._b[0]
 
     @cached_property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Get the statistics of the rescaled data."""
         result = {}
         a = self._a.squeeze()
@@ -224,7 +217,7 @@ class Rescale(Forwards):
 
         return result
 
-    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]]:
         """Get the tendencies of the statistics of the rescaled data.
 
         Parameters

anemoi/datasets/data/rolling_average.py

@@ -0,0 +1,141 @@
+# (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 logging
+from functools import cached_property
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from anemoi.datasets.data.indexing import expand_list_indexing
+
+from .dataset import Dataset
+from .dataset import FullIndex
+from .debug import Node
+from .debug import debug_indexing
+from .forwards import Forwards
+
+LOG = logging.getLogger(__name__)
+
+
+class RollingAverage(Forwards):
+    """A class to represent a dataset with interpolated frequency."""
+
+    def __init__(self, dataset: Dataset, window: str | tuple[int, int, str]) -> None:
+        """Initialize the RollingAverage class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset to be averaged with a rolling window.
+        window : (int, int, str)
+            The rolling average window (start, end, 'freq').
+            'freq' means the window is in number of time steps in the dataset.
+            Both start and end are inclusive, i.e. window = (-2, 2, 'freq') means a window of 5 time steps.
+            For now, only 'freq' is supported, in the future other units may be supported.
+            Windows such as "[-2h, +2h]" are not supported yet.
+        """
+        super().__init__(dataset)
+        if not (isinstance(window, (list, tuple)) and len(window) == 3):
+            raise ValueError(f"Window must be (int, int, str), got {window}")
+        if not isinstance(window[0], int) or not isinstance(window[1], int) or not isinstance(window[2], str):
+            raise ValueError(f"Window must be (int, int, str), got {window}")
+        if window[2] not in ["freq", "frequency"]:
+            raise NotImplementedError(f"Window must be (int, int, 'freq'), got {window}")
+
+        # window = (0, 0, 'freq') means no change
+        self.i_start = -window[0]
+        self.i_end = window[1] + 1
+        if self.i_start <= 0:
+            raise ValueError(f"Window start must be negative, got {window}")
+        if self.i_end <= 0:
+            raise ValueError(f"Window end must be positive, got {window}")
+
+        self.window_str = f"-{self.i_start}-to-{self.i_end}"
+
+    @property
+    def shape(self):
+        shape = list(self.forward.shape)
+        shape[0] = len(self)
+        return tuple(shape)
+
+    @debug_indexing
+    @expand_list_indexing
+    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
+        def f(array):
+            return np.nanmean(array, axis=0)
+
+        if isinstance(n, slice):
+            n = (n,)
+
+        if isinstance(n, tuple):
+            first = n[0]
+            if len(n) > 1:
+                rest = n[1:]
+            else:
+                rest = ()
+
+            if isinstance(first, int):
+                slice_ = slice(first, first + self.i_start + self.i_end)
+                data = self.forward[(slice_,) + rest]
+                return f(data)
+
+            if isinstance(first, slice):
+                first = list(range(first.start or 0, first.stop or len(self), first.step or 1))
+
+            if isinstance(first, (list, tuple)):
+                first = [i if i >= 0 else len(self) + i for i in first]
+                if any(i >= len(self) for i in first):
+                    raise IndexError(f"Index out of range: {first}")
+                slices = [slice(i, i + self.i_start + self.i_end) for i in first]
+                data = [self.forward[(slice_,) + rest] for slice_ in slices]
+                res = [f(d) for d in data]
+                return np.array(res)
+
+            assert False, f"Expected int, slice, list or tuple as first element of tuple, got {type(first)}"
+
+        assert isinstance(n, int), f"Expected int, slice, tuple, got {type(n)}"
+
+        if n < 0:
+            n = len(self) + n
+        if n >= len(self):
+            raise IndexError(f"Index out of range: {n}")
+
+        slice_ = slice(n, n + self.i_start + self.i_end)
+        data = self.forward[slice_]
+        return f(data)
+
+    def __len__(self) -> int:
+        return len(self.forward) - (self.i_end + self.i_start - 1)
+
+    @cached_property
+    def dates(self) -> NDArray[np.datetime64]:
+        """Get the interpolated dates."""
+        dates = self.forward.dates
+        return dates[self.i_start : len(dates) - self.i_end + 1]
+
+    def tree(self) -> Node:
+        return Node(self, [self.forward.tree()], window=self.window_str)
+
+    @cached_property
+    def missing(self) -> set[int]:
+        """Get the missing data indices."""
+        result = []
+
+        for i in self.forward.missing:
+            for j in range(0, self.i_end + self.i_start):
+                result.append(i + j)
+
+        result = {x for x in result if x < self._len}
+        return result
+
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
+        return {}

anemoi/datasets/data/select.py

@@ -12,9 +12,6 @@ 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
 
@@ -37,7 +34,7 @@ LOG = logging.getLogger(__name__)
 class Select(Forwards):
     """Class to select a subset of variables from a dataset."""
 
-    def __init__(self, dataset: Dataset, indices: List[int], reason: Dict[str, Any]) -> None:
+    def __init__(self, dataset: Dataset, indices: list[int], reason: dict[str, Any]) -> None:
         """Initialize the Select class.
 
         Parameters
@@ -140,26 +137,26 @@ class Select(Forwards):
         return (len(self), len(self.indices)) + self.dataset.shape[2:]
 
     @cached_property
-    def variables(self) -> List[str]:
+    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) -> Dict[str, Any]:
+    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) -> Dict[str, int]:
+    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) -> Dict[str, NDArray[Any]]:
+    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: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
+    def statistics_tendencies(self, delta: datetime.timedelta | None = None) -> dict[str, NDArray[Any]]:
         """Get the statistical tendencies of the dataset.
 
         Parameters
@@ -176,7 +173,7 @@ class Select(Forwards):
             delta = self.frequency
         return {k: v[self.indices] for k, v in self.dataset.statistics_tendencies(delta).items()}
 
-    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+    def metadata_specific(self, **kwargs: Any) -> dict[str, Any]:
         """Get the specific metadata of the dataset.
 
         Parameters
@@ -216,7 +213,7 @@ class Select(Forwards):
         """
         return Node(self, [self.dataset.tree()], **self.reason)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the subclass.
 
         Returns
@@ -231,7 +228,7 @@ class Select(Forwards):
 class Rename(Forwards):
     """Class to rename variables in a dataset."""
 
-    def __init__(self, dataset: Dataset, rename: Dict[str, str]) -> None:
+    def __init__(self, dataset: Dataset, rename: dict[str, str]) -> None:
         """Initialize the Rename class.
 
         Parameters
@@ -251,17 +248,17 @@ class Rename(Forwards):
         self.rename = rename
 
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Get the renamed variables."""
         return self._variables
 
     @property
-    def variables_metadata(self) -> Dict[str, Any]:
+    def variables_metadata(self) -> dict[str, Any]:
         """Get the renamed variables metadata."""
         return self._variables_metadata
 
     @cached_property
-    def name_to_index(self) -> Dict[str, int]:
+    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)}
 
@@ -273,7 +270,7 @@ class Rename(Forwards):
         """
         return Node(self, [self.forward.tree()], rename=self.rename)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the subclass.
 
         Returns:

anemoi/datasets/data/statistics.py

@@ -12,9 +12,6 @@ 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
 
@@ -56,11 +53,11 @@ class Statistics(Forwards):
             )
 
     @cached_property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Get the statistics."""
         return self._statistic.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]]:
         """Get the statistics tendencies.
 
         Parameters
@@ -77,7 +74,7 @@ class Statistics(Forwards):
             delta = self.frequency
         return self._statistic.statistics_tendencies(delta)
 
-    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+    def forwards_subclass_metadata_specific(self) -> dict[str, Any]:
         """Get the metadata specific to the forwards subclass.
 
         Returns
@@ -97,7 +94,7 @@ class Statistics(Forwards):
         """
         return Node(self, [self.forward.tree()])
 
-    def get_dataset_names(self, names: Set[str]) -> None:
+    def get_dataset_names(self, names: set[str]) -> None:
         """Get the dataset names.
 
         Parameters

anemoi/datasets/data/stores.py

@@ -15,11 +15,6 @@ import tempfile
 import warnings
 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 Union
 from urllib.parse import urlparse
 
 import numpy as np
@@ -90,22 +85,20 @@ class S3Store(ReadOnlyStore):
     options using the anemoi configs.
     """
 
-    def __init__(self, url: str, region: Optional[str] = None) -> None:
-        """Initialize the S3Store with a URL and optional region."""
-        from anemoi.utils.remote.s3 import s3_client
+    def __init__(self, url: str) -> None:
+        """Initialize the S3Store with a URL."""
 
-        _, _, self.bucket, self.key = url.split("/", 3)
-        self.s3 = s3_client(self.bucket, region=region)
+        self.url = url
 
     def __getitem__(self, key: str) -> bytes:
         """Retrieve an item from the store."""
+        from anemoi.utils.remote.s3 import get_object
+
         try:
-            response = self.s3.get_object(Bucket=self.bucket, Key=self.key + "/" + key)
-        except self.s3.exceptions.NoSuchKey:
+            return get_object(os.path.join(self.url, key))
+        except FileNotFoundError:
             raise KeyError(key)
 
-        return response["Body"].read()
-
 
 class DebugStore(ReadOnlyStore):
     """A store to debug the zarr loading."""
@@ -199,7 +192,7 @@ def open_zarr(path: str, dont_fail: bool = False, cache: int = None) -> zarr.hie
 class Zarr(Dataset):
     """A zarr dataset."""
 
-    def __init__(self, path: Union[str, zarr.hierarchy.Group]) -> None:
+    def __init__(self, path: str | zarr.hierarchy.Group) -> None:
         """Initialize the Zarr dataset with a path or zarr group."""
         if isinstance(path, zarr.hierarchy.Group):
             self.was_zarr = True
@@ -215,7 +208,7 @@ class Zarr(Dataset):
         self._missing = set()
 
     @property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Return the missing dates of the dataset."""
         return self._missing
 
@@ -236,7 +229,7 @@ class Zarr(Dataset):
         """Retrieve an item from the dataset."""
         return self.data[n]
 
-    def _unwind(self, index: Union[int, slice, list, tuple], rest: list, shape: tuple, axis: int, axes: list) -> iter:
+    def _unwind(self, index: int | slice | list | tuple, rest: list, shape: tuple, axis: int, axes: list) -> iter:
         """Unwind the index for multi-dimensional indexing."""
         if not isinstance(index, (int, slice, list, tuple)):
             try:
@@ -298,7 +291,7 @@ class Zarr(Dataset):
             return self.z.longitude[:]
 
     @property
-    def statistics(self) -> Dict[str, NDArray[Any]]:
+    def statistics(self) -> dict[str, NDArray[Any]]:
         """Return the statistics of the dataset."""
         return dict(
             mean=self.z.mean[:],
@@ -307,7 +300,7 @@ class Zarr(Dataset):
             minimum=self.z.minimum[:],
         )
 
-    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 statistical tendencies of the dataset."""
         if delta is None:
             delta = self.frequency
@@ -354,14 +347,14 @@ class Zarr(Dataset):
         return dates[1].astype(object) - dates[0].astype(object)
 
     @property
-    def name_to_index(self) -> Dict[str, int]:
+    def name_to_index(self) -> dict[str, int]:
         """Return the name to index mapping of the dataset."""
         if "variables" in self.z.attrs:
             return {n: i for i, n in enumerate(self.z.attrs["variables"])}
         return self.z.attrs["name_to_index"]
 
     @property
-    def variables(self) -> List[str]:
+    def variables(self) -> list[str]:
         """Return the variables of the dataset."""
         return [
             k
@@ -372,7 +365,7 @@ class Zarr(Dataset):
         ]
 
     @cached_property
-    def constant_fields(self) -> List[str]:
+    def constant_fields(self) -> list[str]:
         """Return the constant fields of the dataset."""
         result = self.z.attrs.get("constant_fields")
         if result is None:
@@ -380,7 +373,7 @@ class Zarr(Dataset):
         return self.computed_constant_fields()
 
     @property
-    def variables_metadata(self) -> Dict[str, Any]:
+    def variables_metadata(self) -> dict[str, Any]:
         """Return the metadata of the variables."""
         return self.z.attrs.get("variables_metadata", {})
 
@@ -392,7 +385,7 @@ class Zarr(Dataset):
         """Return the end date of the statistics."""
         return self.dates[-1]
 
-    def metadata_specific(self, **kwargs: Any) -> Dict[str, Any]:
+    def metadata_specific(self, **kwargs: Any) -> dict[str, Any]:
         """Return the specific metadata of the dataset."""
         return super().metadata_specific(
             attrs=dict(self.z.attrs),
@@ -416,7 +409,7 @@ class Zarr(Dataset):
         """Return the tree representation of the dataset."""
         return Node(self, [], path=self.path)
 
-    def get_dataset_names(self, names: Set[str]) -> None:
+    def get_dataset_names(self, names: set[str]) -> None:
         """Get the names of the datasets."""
         name, _ = os.path.splitext(os.path.basename(self.path))
         names.add(name)
@@ -433,17 +426,17 @@ class Zarr(Dataset):
 class ZarrWithMissingDates(Zarr):
     """A zarr dataset with missing dates."""
 
-    def __init__(self, path: Union[str, zarr.hierarchy.Group]) -> None:
+    def __init__(self, path: str | zarr.hierarchy.Group) -> None:
         """Initialize the ZarrWithMissingDates dataset with a path or zarr group."""
         super().__init__(path)
 
         missing_dates = self.z.attrs.get("missing_dates", [])
-        missing_dates = set([np.datetime64(x, "s") for x in missing_dates])
+        missing_dates = {np.datetime64(x, "s") for x in missing_dates}
         self.missing_to_dates = {i: d for i, d in enumerate(self.dates) if d in missing_dates}
         self._missing = set(self.missing_to_dates)
 
     @property
-    def missing(self) -> Set[int]:
+    def missing(self) -> set[int]:
         """Return the missing dates of the dataset."""
         return self._missing
 
@@ -506,7 +499,7 @@ class ZarrWithMissingDates(Zarr):
 QUIET = set()
 
 
-def zarr_lookup(name: str, fail: bool = True) -> Optional[str]:
+def zarr_lookup(name: str, fail: bool = True) -> str | None:
     """Look up a zarr dataset by name."""
 
     config = load_config()["datasets"]