anemoi-datasets 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

anemoi/datasets/data/dataset.py

@@ -5,24 +5,41 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+import datetime
+import json
 import logging
 import os
+import pprint
 import warnings
 from functools import cached_property
 
+from anemoi.utils.dates import frequency_to_seconds
+from anemoi.utils.dates import frequency_to_string
+from anemoi.utils.dates import frequency_to_timedelta
+
 LOG = logging.getLogger(__name__)
 
 
 class Dataset:
     arguments = {}
 
+    def mutate(self) -> "Dataset":
+        """
+        Give an opportunity to a subclass to return a new Dataset
+        object of a different class, if needed.
+        """
+        return self
+
+    def swap_with_parent(self, parent):
+        return parent
+
     @cached_property
     def _len(self):
         return len(self)
 
     def _subset(self, **kwargs):
         if not kwargs:
-            return self
+            return self.mutate()
 
         if "start" in kwargs or "end" in kwargs:
             start = kwargs.pop("start", None)
@@ -30,37 +47,58 @@ class Dataset:
 
             from .subset import Subset
 
-            return Subset(self, self._dates_to_indices(start, end), dict(start=start, end=end))._subset(**kwargs)
+            return (
+                Subset(self, self._dates_to_indices(start, end), dict(start=start, end=end))._subset(**kwargs).mutate()
+            )
 
         if "frequency" in kwargs:
             from .subset import Subset
 
+            if "interpolate_frequency" in kwargs:
+                raise ValueError("Cannot use both `frequency` and `interpolate_frequency`")
+
             frequency = kwargs.pop("frequency")
-            return Subset(self, self._frequency_to_indices(frequency), dict(frequency=frequency))._subset(**kwargs)
+            return (
+                Subset(self, self._frequency_to_indices(frequency), dict(frequency=frequency))
+                ._subset(**kwargs)
+                .mutate()
+            )
+
+        if "interpolate_frequency" in kwargs:
+            from .interpolate import InterpolateFrequency
+
+            interpolate_frequency = kwargs.pop("interpolate_frequency")
+            return InterpolateFrequency(self, interpolate_frequency)._subset(**kwargs).mutate()
 
         if "select" in kwargs:
             from .select import Select
 
             select = kwargs.pop("select")
-            return Select(self, self._select_to_columns(select), {"select": select})._subset(**kwargs)
+            return Select(self, self._select_to_columns(select), {"select": select})._subset(**kwargs).mutate()
 
         if "drop" in kwargs:
             from .select import Select
 
             drop = kwargs.pop("drop")
-            return Select(self, self._drop_to_columns(drop), {"drop": drop})._subset(**kwargs)
+            return Select(self, self._drop_to_columns(drop), {"drop": drop})._subset(**kwargs).mutate()
 
         if "reorder" in kwargs:
             from .select import Select
 
             reorder = kwargs.pop("reorder")
-            return Select(self, self._reorder_to_columns(reorder), {"reoder": reorder})._subset(**kwargs)
+            return Select(self, self._reorder_to_columns(reorder), {"reoder": reorder})._subset(**kwargs).mutate()
 
         if "rename" in kwargs:
             from .select import Rename
 
             rename = kwargs.pop("rename")
-            return Rename(self, rename)._subset(**kwargs)
+            return Rename(self, rename)._subset(**kwargs).mutate()
+
+        if "rescale" in kwargs:
+            from .rescale import Rescale
+
+            rescale = kwargs.pop("rescale")
+            return Rescale(self, rescale)._subset(**kwargs).mutate()
 
         if "statistics" in kwargs:
             from ..data import open_dataset
@@ -68,20 +106,38 @@ class Dataset:
 
             statistics = kwargs.pop("statistics")
 
-            return Statistics(self, open_dataset(statistics))._subset(**kwargs)
+            return Statistics(self, open_dataset(statistics))._subset(**kwargs).mutate()
 
         if "thinning" in kwargs:
             from .masked import Thinning
 
             thinning = kwargs.pop("thinning")
             method = kwargs.pop("method", "every-nth")
-            return Thinning(self, thinning, method)._subset(**kwargs)
+            return Thinning(self, thinning, method)._subset(**kwargs).mutate()
 
         if "area" in kwargs:
             from .masked import Cropping
 
             bbox = kwargs.pop("area")
-            return Cropping(self, bbox)._subset(**kwargs)
+            return Cropping(self, bbox)._subset(**kwargs).mutate()
+
+        if "missing_dates" in kwargs:
+            from .missing import MissingDates
+
+            missing_dates = kwargs.pop("missing_dates")
+            return MissingDates(self, missing_dates)._subset(**kwargs).mutate()
+
+        if "skip_missing_dates" in kwargs:
+            from .missing import SkipMissingDates
+
+            if "expected_access" not in kwargs:
+                raise ValueError("`expected_access` is required with `skip_missing_dates`")
+
+            skip_missing_dates = kwargs.pop("skip_missing_dates")
+            expected_access = kwargs.pop("expected_access")
+
+            if skip_missing_dates:
+                return SkipMissingDates(self, expected_access)._subset(**kwargs).mutate()
 
         # Keep last
         if "shuffle" in kwargs:
@@ -90,15 +146,14 @@ class Dataset:
             shuffle = kwargs.pop("shuffle")
 
             if shuffle:
-                return Subset(self, self._shuffle_indices(), dict(shuffle=True))._subset(**kwargs)
+                return Subset(self, self._shuffle_indices(), dict(shuffle=True))._subset(**kwargs).mutate()
 
         raise NotImplementedError("Unsupported arguments: " + ", ".join(kwargs))
 
     def _frequency_to_indices(self, frequency):
-        from .misc import _frequency_to_hours
 
-        requested_frequency = _frequency_to_hours(frequency)
-        dataset_frequency = _frequency_to_hours(self.frequency)
+        requested_frequency = frequency_to_seconds(frequency)
+        dataset_frequency = frequency_to_seconds(self.frequency)
         assert requested_frequency % dataset_frequency == 0
         # Question: where do we start? first date, or first date that is a multiple of the frequency?
         step = requested_frequency // dataset_frequency
@@ -171,37 +226,71 @@ class Dataset:
         import anemoi
 
         def tidy(v):
-            if isinstance(v, (list, tuple)):
+            if isinstance(v, (list, tuple, set)):
                 return [tidy(i) for i in v]
             if isinstance(v, dict):
                 return {k: tidy(v) for k, v in v.items()}
             if isinstance(v, str) and v.startswith("/"):
                 return os.path.basename(v)
+            if isinstance(v, datetime.datetime):
+                return v.isoformat()
+            if isinstance(v, datetime.date):
+                return v.isoformat()
+            if isinstance(v, datetime.timedelta):
+                return frequency_to_string(v)
+
+            if isinstance(v, Dataset):
+                # That can happen in the `arguments`
+                # if a dataset is passed as an argument
+                return repr(v)
+
+            if isinstance(v, slice):
+                return (v.start, v.stop, v.step)
+
             return v
 
-        return tidy(
-            dict(
-                version=anemoi.datasets.__version__,
-                shape=self.shape,
-                arguments=self.arguments,
-                specific=self.metadata_specific(),
-                frequency=self.frequency,
-                variables=self.variables,
-                start_date=self.dates[0].astype(str),
-                end_date=self.dates[-1].astype(str),
-            )
+        md = dict(
+            version=anemoi.datasets.__version__,
+            arguments=self.arguments,
+            **self.dataset_metadata(),
+        )
+
+        try:
+            return json.loads(json.dumps(tidy(md)))
+        except Exception:
+            LOG.exception("Failed to serialize metadata")
+            pprint.pprint(md)
+
+            raise
+
+    @property
+    def start_date(self):
+        return self.dates[0]
+
+    @property
+    def end_date(self):
+        return self.dates[-1]
+
+    def dataset_metadata(self):
+        return dict(
+            specific=self.metadata_specific(),
+            frequency=self.frequency,
+            variables=self.variables,
+            shape=self.shape,
+            start_date=self.start_date.astype(str),
+            end_date=self.end_date.astype(str),
         )
 
     def metadata_specific(self, **kwargs):
         action = self.__class__.__name__.lower()
-        assert isinstance(self.frequency, int), (self.frequency, self, action)
+        # assert isinstance(self.frequency, datetime.timedelta), (self.frequency, self, action)
         return dict(
             action=action,
             variables=self.variables,
             shape=self.shape,
-            frequency=self.frequency,
-            start_date=self.dates[0].astype(str),
-            end_date=self.dates[-1].astype(str),
+            frequency=frequency_to_string(frequency_to_timedelta(self.frequency)),
+            start_date=self.start_date.astype(str),
+            end_date=self.end_date.astype(str),
             **kwargs,
         )
 
@@ -220,10 +309,6 @@ class Dataset:
             if n.startswith("_") and not n.startswith("__"):
                 warnings.warn(f"Private method {n} is overriden in {ds.__class__.__name__}")
 
-        # for n in ('metadata_specific', 'tree', 'source'):
-        #     if n not in overriden:
-        #         warnings.warn(f"Method {n} is not overriden in {ds.__class__.__name__}")
-
     def _repr_html_(self):
         return self.tree().html()
 

anemoi/datasets/data/debug.py

@@ -209,10 +209,14 @@ def _debug_indexing(method):
     return wrapper
 
 
+def _identity(x):
+    return x
+
+
 if DEBUG_ZARR_INDEXING:
     debug_indexing = _debug_indexing
 else:
-    debug_indexing = lambda x: x  # noqa
+    debug_indexing = _identity
 
 
 def debug_zarr_loading(on_off):

anemoi/datasets/data/forwards.py

@@ -23,7 +23,7 @@ LOG = logging.getLogger(__name__)
 
 class Forwards(Dataset):
     def __init__(self, forward):
-        self.forward = forward
+        self.forward = forward.mutate()
 
     def __len__(self):
         return len(self.forward)
@@ -118,6 +118,9 @@ class Combined(Forwards):
         # Forward most properties to the first dataset
         super().__init__(datasets[0])
 
+    def mutate(self):
+        return self
+
     def check_same_resolution(self, d1, d2):
         if d1.resolution != d2.resolution:
             raise ValueError(f"Incompatible resolutions: {d1.resolution} and {d2.resolution} ({d1} {d2})")
@@ -187,14 +190,9 @@ class Combined(Forwards):
             **kwargs,
         )
 
-    @cached_property
+    @property
     def missing(self):
-        offset = 0
-        result = set()
-        for d in self.datasets:
-            result.update(offset + m for m in d.missing)
-            offset += len(d)
-        return result
+        raise NotImplementedError("missing() not implemented for Combined")
 
     def get_dataset_names(self, names):
         for d in self.datasets:
@@ -249,3 +247,14 @@ class GivenAxis(Combined):
             return self._get_slice(n)
 
         return np.concatenate([d[n] for d in self.datasets], axis=self.axis - 1)
+
+    @cached_property
+    def missing(self):
+        offset = 0
+        result = set()
+        for d in self.datasets:
+            print("--->", d.missing, d)
+            result.update(offset + m for m in d.missing)
+            if self.axis == 0:  # Advance if axis is time
+                offset += len(d)
+        return result

anemoi/datasets/data/grids.py

@@ -128,7 +128,7 @@ class Grids(GridsBase):
 
 
 class Cutout(GridsBase):
-    def __init__(self, datasets, axis):
+    def __init__(self, datasets, axis, min_distance_km=None, cropping_distance=2.0, neighbours=5, plot=False):
         from anemoi.datasets.grids import cutout_mask
 
         super().__init__(datasets, axis)
@@ -144,7 +144,10 @@ class Cutout(GridsBase):
             self.lam.longitudes,
             self.globe.latitudes,
             self.globe.longitudes,
-            # plot="cutout",
+            plot=plot,
+            min_distance_km=min_distance_km,
+            cropping_distance=cropping_distance,
+            neighbours=neighbours,
         )
         assert len(self.mask) == self.globe.shape[3], (
             len(self.mask),
@@ -229,6 +232,10 @@ def cutout_factory(args, kwargs):
 
     cutout = kwargs.pop("cutout")
     axis = kwargs.pop("axis", 3)
+    plot = kwargs.pop("plot", None)
+    min_distance_km = kwargs.pop("min_distance_km", None)
+    cropping_distance = kwargs.pop("cropping_distance", 2.0)
+    neighbours = kwargs.pop("neighbours", 5)
 
     assert len(args) == 0
     assert isinstance(cutout, (list, tuple))
@@ -236,4 +243,11 @@ def cutout_factory(args, kwargs):
     datasets = [_open(e) for e in cutout]
     datasets, kwargs = _auto_adjust(datasets, kwargs)
 
-    return Cutout(datasets, axis=axis)._subset(**kwargs)
+    return Cutout(
+        datasets,
+        axis=axis,
+        neighbours=neighbours,
+        min_distance_km=min_distance_km,
+        cropping_distance=cropping_distance,
+        plot=plot,
+    )._subset(**kwargs)

anemoi/datasets/data/interpolate.py

@@ -0,0 +1,133 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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
+
+import numpy as np
+from anemoi.utils.dates import frequency_to_timedelta
+
+from .debug import Node
+from .debug import debug_indexing
+from .forwards import Forwards
+from .indexing import apply_index_to_slices_changes
+from .indexing import expand_list_indexing
+from .indexing import index_to_slices
+from .indexing import update_tuple
+
+LOG = logging.getLogger(__name__)
+
+
+class InterpolateFrequency(Forwards):
+
+    def __init__(self, dataset, frequency):
+        super().__init__(dataset)
+        self._frequency = frequency_to_timedelta(frequency)
+
+        self.seconds = self._frequency.total_seconds()
+        other_seconds = dataset.frequency.total_seconds()
+
+        self.seconds = int(self.seconds)
+        assert self.seconds == self._frequency.total_seconds()
+
+        other_seconds = int(other_seconds)
+        assert other_seconds == dataset.frequency.total_seconds()
+
+        if self.seconds >= other_seconds:
+            raise ValueError(
+                f"Interpolate frequency {self._frequency} must be more frequent than dataset frequency {dataset.frequency}"
+            )
+
+        if other_seconds % self.seconds != 0:
+            raise ValueError(
+                f"Interpolate frequency {self._frequency}  must be a multiple of the dataset frequency {dataset.frequency}"
+            )
+
+        self.ratio = other_seconds // self.seconds
+        self.alphas = np.linspace(0, 1, self.ratio + 1)
+        self.other_len = len(dataset)
+
+    @debug_indexing
+    @expand_list_indexing
+    def _get_tuple(self, index):
+        index, changes = index_to_slices(index, self.shape)
+        index, previous = update_tuple(index, 0, slice(None))
+        result = self._get_slice(previous)
+        return apply_index_to_slices_changes(result[index], changes)
+
+    def _get_slice(self, s):
+        return np.stack([self[i] for i in range(*s.indices(self._len))])
+
+    @debug_indexing
+    def __getitem__(self, n):
+        if isinstance(n, tuple):
+            return self._get_tuple(n)
+
+        if isinstance(n, slice):
+            return self._get_slice(n)
+
+        if n < 0:
+            n += self._len
+
+        if n == self._len - 1:
+            # Special case for the last element
+            return self.forward[-1]
+
+        i = n // self.ratio
+        x = n % self.ratio
+
+        if x == 0:
+            # No interpolation needed
+            return self.forward[i]
+
+        alpha = self.alphas[x]
+
+        assert 0 < alpha < 1, alpha
+        return self.forward[i] * (1 - alpha) + self.forward[i + 1] * alpha
+
+    def __len__(self):
+        return (self.other_len - 1) * self.ratio + 1
+
+    @property
+    def frequency(self):
+        return self._frequency
+
+    @cached_property
+    def dates(self):
+        result = []
+        deltas = [np.timedelta64(self.seconds * i, "s") for i in range(self.ratio)]
+        for d in self.forward.dates[:-1]:
+            for i in deltas:
+                result.append(d + i)
+        result.append(self.forward.dates[-1])
+        return np.array(result)
+
+    @property
+    def shape(self):
+        return (self._len,) + self.forward.shape[1:]
+
+    def tree(self):
+        return Node(self, [self.forward.tree()], frequency=self.frequency)
+
+    @cached_property
+    def missing(self):
+        result = []
+        j = 0
+        for i in range(self.other_len):
+            missing = i in self.forward.missing
+            for _ in range(self.ratio):
+                if missing:
+                    result.append(j)
+                j += 1
+
+        result = set(x for x in result if x < self._len)
+        return result
+
+    def subclass_metadata_specific(self):
+        return {
+            # "frequency": frequency_to_string(self._frequency),
+        }

anemoi/datasets/data/masked.py

@@ -112,5 +112,5 @@ class Cropping(Masked):
     def tree(self):
         return Node(self, [self.forward.tree()], area=self.area)
 
-    def metadata_specific(self, **kwargs):
-        return super().metadata_specific(area=self.area, **kwargs)
+    def subclass_metadata_specific(self):
+        return dict(area=self.area)