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

anemoi/datasets/create/utils.py

@@ -96,48 +96,6 @@ def to_datetime(*args: Any, **kwargs: Any) -> datetime.datetime:
     return to_datetime_(*args, **kwargs)
 
 
-def make_list_int(value: 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/data/dataset.py

@@ -293,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
 

anemoi/datasets/data/grids.py

@@ -21,167 +21,15 @@ from .dataset import FullIndex
 from .dataset import Shape
 from .dataset import TupleIndex
 from .debug import Node
-from .debug import debug_indexing
-from .forwards import Combined
 from .forwards import GivenAxis
 from .indexing import apply_index_to_slices_changes
-from .indexing import expand_list_indexing
 from .indexing import index_to_slices
-from .indexing import length_to_slices
-from .indexing import update_tuple
 from .misc import _auto_adjust
 from .misc import _open
 
 LOG = logging.getLogger(__name__)
 
 
-class Concat(Combined):
-    """A class to represent concatenated datasets."""
-
-    def __len__(self) -> int:
-        """Returns the total length of the concatenated datasets.
-
-        Returns
-        -------
-        int
-            Total length of the concatenated datasets.
-        """
-        return sum(len(i) for i in self.datasets)
-
-    @debug_indexing
-    @expand_list_indexing
-    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
-        """Retrieves a tuple of data from the concatenated datasets based on the given index.
-
-        Parameters
-        ----------
-        index : TupleIndex
-            Index specifying the data to retrieve.
-
-        Returns
-        -------
-        NDArray[Any]
-            Concatenated data array from the specified index.
-        """
-        index, changes = index_to_slices(index, self.shape)
-        # print(index, changes)
-        lengths = [d.shape[0] for d in self.datasets]
-        slices = length_to_slices(index[0], lengths)
-        # print("slies", slices)
-        result = [d[update_tuple(index, 0, i)[0]] for (d, i) in zip(self.datasets, slices) if i is not None]
-        result = np.concatenate(result, axis=0)
-        return apply_index_to_slices_changes(result, changes)
-
-    @debug_indexing
-    def __getitem__(self, n: FullIndex) -> NDArray[Any]:
-        """Retrieves data from the concatenated datasets based on the given index.
-
-        Parameters
-        ----------
-        n : FullIndex
-            Index specifying the data to retrieve.
-
-        Returns
-        -------
-        NDArray[Any]
-            Data array from the concatenated datasets based on the index.
-        """
-        if isinstance(n, tuple):
-            return self._get_tuple(n)
-
-        if isinstance(n, slice):
-            return self._get_slice(n)
-
-        # TODO: optimize
-        k = 0
-        while n >= self.datasets[k]._len:
-            n -= self.datasets[k]._len
-            k += 1
-        return self.datasets[k][n]
-
-    @debug_indexing
-    def _get_slice(self, s: slice) -> NDArray[Any]:
-        """Retrieves a slice of data from the concatenated datasets.
-
-        Parameters
-        ----------
-        s : slice
-            Slice object specifying the range of data to retrieve.
-
-        Returns
-        -------
-        NDArray[Any]
-            Concatenated data array from the specified slice.
-        """
-        result = []
-
-        lengths = [d.shape[0] for d in self.datasets]
-        slices = length_to_slices(s, lengths)
-
-        result = [d[i] for (d, i) in zip(self.datasets, slices) if i is not None]
-
-        return np.concatenate(result)
-
-    def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
-        """Check the compatibility of two datasets for concatenation.
-
-        Parameters
-        ----------
-        d1 : Dataset
-            The first dataset.
-        d2 : Dataset
-            The second dataset.
-        """
-        super().check_compatibility(d1, d2)
-        self.check_same_sub_shapes(d1, d2, drop_axis=0)
-
-    def check_same_lengths(self, d1: Dataset, d2: Dataset) -> None:
-        """Check if the lengths of two datasets are the same.
-
-        Parameters
-        ----------
-        d1 : Dataset
-            The first dataset.
-        d2 : Dataset
-            The second dataset.
-        """
-        # Turned off because we are concatenating along the first axis
-        pass
-
-    def check_same_dates(self, d1: Dataset, d2: Dataset) -> None:
-        """Check if the dates of two datasets are the same.
-
-        Parameters
-        ----------
-        d1 : Dataset
-            The first dataset.
-        d2 : Dataset
-            The second dataset.
-        """
-        # Turned off because we are concatenating along the dates axis
-        pass
-
-    @property
-    def dates(self) -> NDArray[np.datetime64]:
-        """Returns the concatenated dates of all datasets."""
-        return np.concatenate([d.dates for d in self.datasets])
-
-    @property
-    def shape(self) -> Shape:
-        """Returns the shape of the concatenated datasets."""
-        return (len(self),) + self.datasets[0].shape[1:]
-
-    def tree(self) -> Node:
-        """Generates a hierarchical tree structure for the concatenated datasets.
-
-        Returns
-        -------
-        Node
-            A Node object representing the concatenated datasets.
-        """
-        return Node(self, [d.tree() for d in self.datasets])
-
-
 class GridsBase(GivenAxis):
     """A base class for handling grids in datasets."""
 

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/stores.py

@@ -85,22 +85,20 @@ class S3Store(ReadOnlyStore):
     options using the anemoi configs.
     """
 
-    def __init__(self, url: str, region: str | None = 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."""

anemoi/datasets/dates/__init__.py

@@ -58,6 +58,8 @@ def extend(x: str | list[Any] | tuple[Any, ...]) -> Iterator[datetime.datetime]:
 class DatesProvider:
     """Base class for date generation.
 
+    Examples
+    --------
     >>> DatesProvider.from_config(**{"start": "2023-01-01 00:00", "end": "2023-01-02 00:00", "frequency": "1d"}).values
     [datetime.datetime(2023, 1, 1, 0, 0), datetime.datetime(2023, 1, 2, 0, 0)]
 

anemoi/datasets/dumper.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 datetime
+import io
+import logging
+
+import ruamel.yaml
+
+LOG = logging.getLogger(__name__)
+
+
+def represent_date(dumper, data):
+
+    if isinstance(data, datetime.datetime):
+        if data.tzinfo is None:
+            data = data.replace(tzinfo=datetime.timezone.utc)
+        data = data.astimezone(datetime.timezone.utc)
+        iso_str = data.replace(tzinfo=None).isoformat(timespec="seconds") + "Z"
+    else:
+        iso_str = data.isoformat()
+
+    return dumper.represent_scalar("tag:yaml.org,2002:timestamp", iso_str)
+
+
+# --- Represent multiline strings with | style ---
+def represent_multiline_str(dumper, data):
+    if "\n" in data:
+        return dumper.represent_scalar("tag:yaml.org,2002:str", data.strip(), style="|")
+    return dumper.represent_scalar("tag:yaml.org,2002:str", data)
+
+
+# --- Represent short lists inline (flow style) ---
+def represent_inline_list(dumper, data):
+
+    if not all(isinstance(i, (str, int, float, bool, type(None))) for i in data):
+        return dumper.represent_sequence("tag:yaml.org,2002:seq", data)
+
+    return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True)
+
+
+def yaml_dump(obj, order=None, stream=None, **kwargs):
+
+    if order:
+
+        def _ordering(k):
+            return order.index(k) if k in order else len(order)
+
+        obj = {k: v for k, v in sorted(obj.items(), key=lambda item: _ordering(item[0]))}
+
+    yaml = ruamel.yaml.YAML()
+    yaml.width = 120  # wrap long flow sequences
+
+    yaml.Representer.add_representer(datetime.date, represent_date)
+    yaml.Representer.add_representer(datetime.datetime, represent_date)
+    yaml.Representer.add_representer(str, represent_multiline_str)
+    yaml.Representer.add_representer(list, represent_inline_list)
+
+    data = ruamel.yaml.comments.CommentedMap()
+    for i, (k, v) in enumerate(obj.items()):
+        data[k] = v
+        if i > 0:
+            data.yaml_set_comment_before_after_key(key=k, before="\n")
+
+    if stream:
+        yaml.dump(data, stream=stream, **kwargs)
+
+    stream = io.StringIO()
+    yaml.dump(data, stream=stream, **kwargs)
+    return stream.getvalue()

anemoi/datasets/grids.py

@@ -8,11 +8,11 @@
 # nor does it submit to any jurisdiction.
 
 
-import base64
 import logging
 from typing import Any
 
 import numpy as np
+from anemoi.utils.grids import latlon_to_xyz
 from numpy.typing import NDArray
 
 LOG = logging.getLogger(__name__)
@@ -88,71 +88,6 @@ def plot_mask(
         plt.savefig(path + "-global-zoomed.png")
 
 
-# TODO: Use the one from anemoi.utils.grids instead
-# from anemoi.utils.grids import ...
-def xyz_to_latlon(x: NDArray[Any], y: NDArray[Any], z: NDArray[Any]) -> tuple[NDArray[Any], NDArray[Any]]:
-    """Convert Cartesian coordinates to latitude and longitude.
-
-    Parameters
-    ----------
-    x : NDArray[Any]
-        X coordinates.
-    y : NDArray[Any]
-        Y coordinates.
-    z : NDArray[Any]
-        Z coordinates.
-
-    Returns
-    -------
-    Tuple[NDArray[Any], NDArray[Any]]
-        Latitude and longitude coordinates.
-    """
-    return (
-        np.rad2deg(np.arcsin(np.minimum(1.0, np.maximum(-1.0, z)))),
-        np.rad2deg(np.arctan2(y, x)),
-    )
-
-
-# TODO: Use the one from anemoi.utils.grids instead
-# from anemoi.utils.grids import ...
-def latlon_to_xyz(
-    lat: NDArray[Any], lon: NDArray[Any], radius: float = 1.0
-) -> tuple[NDArray[Any], NDArray[Any], NDArray[Any]]:
-    """Convert latitude and longitude to Cartesian coordinates.
-
-    Parameters
-    ----------
-    lat : NDArray[Any]
-        Latitude coordinates.
-    lon : NDArray[Any]
-        Longitude coordinates.
-    radius : float, optional
-        Radius of the sphere. Defaults to 1.0.
-
-    Returns
-    -------
-    Tuple[NDArray[Any], NDArray[Any], NDArray[Any]]
-        X, Y, and Z coordinates.
-    """
-    # https://en.wikipedia.org/wiki/Geographic_coordinate_conversion#From_geodetic_to_ECEF_coordinates
-    # We assume that the Earth is a sphere of radius 1 so N(phi) = 1
-    # We assume h = 0
-    #
-    phi = np.deg2rad(lat)
-    lda = np.deg2rad(lon)
-
-    cos_phi = np.cos(phi)
-    cos_lda = np.cos(lda)
-    sin_phi = np.sin(phi)
-    sin_lda = np.sin(lda)
-
-    x = cos_phi * cos_lda * radius
-    y = cos_phi * sin_lda * radius
-    z = sin_phi * radius
-
-    return x, y, z
-
-
 class Triangle3D:
     """A class to represent a 3D triangle and perform intersection tests with rays."""
 
@@ -509,92 +444,6 @@ def outline(lats: NDArray[Any], lons: NDArray[Any], neighbours: int = 5) -> list
     return outside
 
 
-def deserialise_mask(encoded: str) -> NDArray[Any]:
-    """Deserialise a mask from a base64 encoded string.
-
-    Parameters
-    ----------
-    encoded : str
-        Base64 encoded string.
-
-    Returns
-    -------
-    NDArray[Any]
-        Deserialised mask array.
-    """
-    import pickle
-    import zlib
-
-    packed = pickle.loads(zlib.decompress(base64.b64decode(encoded)))
-
-    mask = []
-    value = False
-    for count in packed:
-        mask.extend([value] * count)
-        value = not value
-    return np.array(mask, dtype=bool)
-
-
-def _serialise_mask(mask: NDArray[Any]) -> str:
-    """Serialise a mask to a base64 encoded string.
-
-    Parameters
-    ----------
-    mask : NDArray[Any]
-        Mask array.
-
-    Returns
-    -------
-    str
-        Base64 encoded string.
-    """
-    import pickle
-    import zlib
-
-    assert len(mask.shape) == 1
-    assert len(mask)
-
-    packed = []
-    last = mask[0]
-    count = 1
-
-    for value in mask[1:]:
-        if value == last:
-            count += 1
-        else:
-            packed.append(count)
-            last = value
-            count = 1
-
-    packed.append(count)
-
-    # We always start with an 'off' value
-    # So if the first value is 'on', we need to add a zero
-    if mask[0]:
-        packed.insert(0, 0)
-
-    return base64.b64encode(zlib.compress(pickle.dumps(packed))).decode("utf-8")
-
-
-def serialise_mask(mask: NDArray[Any]) -> str:
-    """Serialise a mask and ensure it can be deserialised.
-
-    Parameters
-    ----------
-    mask : NDArray[Any]
-        Mask array.
-
-    Returns
-    -------
-    str
-        Base64 encoded string.
-    """
-    result = _serialise_mask(mask)
-    # Make sure we can deserialise it
-    assert np.all(mask == deserialise_mask(result))
-    return result
-
-
 def nearest_grid_points(
     source_latitudes: NDArray[Any],
     source_longitudes: NDArray[Any],
@@ -640,29 +489,3 @@ def nearest_grid_points(
     else:
         distances, indices = cKDTree(source_points).query(target_points, k=k, distance_upper_bound=max_distance)
     return distances, indices
-
-
-if __name__ == "__main__":
-    global_lats, global_lons = np.meshgrid(
-        np.linspace(90, -90, 90),
-        np.linspace(-180, 180, 180),
-    )
-    global_lats = global_lats.flatten()
-    global_lons = global_lons.flatten()
-
-    lats, lons = np.meshgrid(
-        np.linspace(50, 40, 100),
-        np.linspace(-10, 15, 100),
-    )
-    lats = lats.flatten()
-    lons = lons.flatten()
-
-    mask = cutout_mask(lats, lons, global_lats, global_lons, cropping_distance=5.0)
-
-    import matplotlib.pyplot as plt
-
-    fig = plt.figure(figsize=(10, 5))
-    plt.scatter(global_lons, global_lats, s=0.01, marker="o", c="r")
-    plt.scatter(global_lons[mask], global_lats[mask], s=0.1, c="k")
-    # plt.scatter(lons, lats, s=0.01)
-    plt.savefig("cutout.png")