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

anemoi/datasets/data/stores.py

@@ -5,6 +5,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import logging
 import os
 import warnings
@@ -13,6 +14,7 @@ from urllib.parse import urlparse
 
 import numpy as np
 import zarr
+from anemoi.utils.dates import frequency_to_timedelta
 
 from . import MissingDateError
 from .dataset import Dataset
@@ -82,6 +84,8 @@ class S3Store(ReadOnlyStore):
 
 
 class DebugStore(ReadOnlyStore):
+    """A store to debug the zarr loading."""
+
     def __init__(self, store):
         assert not isinstance(store, DebugStore)
         self.store = store
@@ -147,6 +151,8 @@ def open_zarr(path, dont_fail=False, cache=None):
 
 
 class Zarr(Dataset):
+    """A zarr dataset."""
+
     def __init__(self, path):
         if isinstance(path, zarr.hierarchy.Group):
             self.was_zarr = True
@@ -243,14 +249,20 @@ class Zarr(Dataset):
             delta = self.frequency
         if isinstance(delta, int):
             delta = f"{delta}h"
-        from anemoi.datasets.create.loaders import TendenciesStatisticsAddition
+        from anemoi.utils.dates import frequency_to_string
+        from anemoi.utils.dates import frequency_to_timedelta
+
+        delta = frequency_to_timedelta(delta)
+        delta = frequency_to_string(delta)
+
+        def func(k):
+            return f"statistics_tendencies_{delta}_{k}"
 
-        func = TendenciesStatisticsAddition.final_storage_name_from_delta
         return dict(
-            mean=self.z[func("mean", delta)][:],
-            stdev=self.z[func("stdev", delta)][:],
-            maximum=self.z[func("maximum", delta)][:],
-            minimum=self.z[func("minimum", delta)][:],
+            mean=self.z[func("mean")][:],
+            stdev=self.z[func("stdev")][:],
+            maximum=self.z[func("maximum")][:],
+            minimum=self.z[func("minimum")][:],
         )
 
     @property
@@ -268,12 +280,11 @@ class Zarr(Dataset):
     @property
     def frequency(self):
         try:
-            return self.z.attrs["frequency"]
+            return frequency_to_timedelta(self.z.attrs["frequency"])
         except KeyError:
             LOG.warning("No 'frequency' in %r, computing from 'dates'", self)
         dates = self.dates
-        delta = dates[1].astype(object) - dates[0].astype(object)
-        return int(delta.total_seconds() / 3600)
+        return dates[1].astype(object) - dates[0].astype(object)
 
     @property
     def name_to_index(self):
@@ -322,11 +333,13 @@ class Zarr(Dataset):
 
 
 class ZarrWithMissingDates(Zarr):
+    """A zarr dataset with missing dates."""
+
     def __init__(self, path):
         super().__init__(path)
 
         missing_dates = self.z.attrs.get("missing_dates", [])
-        missing_dates = [np.datetime64(x) for x in missing_dates]
+        missing_dates = set([np.datetime64(x) 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)
 

anemoi/datasets/data/subset.py

@@ -9,6 +9,7 @@ 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 Source
@@ -23,13 +24,51 @@ from .indexing import update_tuple
 LOG = logging.getLogger(__name__)
 
 
+def _default(a, b, dates):
+    return [a, b]
+
+
+def _start(a, b, dates):
+    from .misc import as_first_date
+
+    c = as_first_date(a, dates)
+    d = as_first_date(b, dates)
+    if c < d:
+        return b
+    else:
+        return a
+
+
+def _end(a, b, dates):
+    from .misc import as_last_date
+
+    c = as_last_date(a, dates)
+    d = as_last_date(b, dates)
+    if c < d:
+        return a
+    else:
+        return b
+
+
+def _combine_reasons(reason1, reason2, dates):
+
+    reason = reason1.copy()
+    for k, v in reason2.items():
+        if k not in reason:
+            reason[k] = v
+        else:
+            func = globals().get(f"_{k}", _default)
+            reason[k] = func(reason[k], v, dates)
+    return reason
+
+
 class Subset(Forwards):
     """Select a subset of the dates."""
 
     def __init__(self, dataset, indices, reason):
         while isinstance(dataset, Subset):
             indices = [dataset.indices[i] for i in indices]
-            reason = {**reason, **dataset.reason}
+            reason = _combine_reasons(reason, dataset.reason, dataset.dates)
             dataset = dataset.dataset
 
         self.dataset = dataset
@@ -39,6 +78,12 @@ class Subset(Forwards):
         # Forward other properties to the super dataset
         super().__init__(dataset)
 
+    def clone(self, dataset):
+        return self.__class__(dataset, self.indices, self.reason).mutate()
+
+    def mutate(self):
+        return self.forward.swap_with_parent(parent=self)
+
     @debug_indexing
     def __getitem__(self, n):
         if isinstance(n, tuple):
@@ -66,10 +111,8 @@ class Subset(Forwards):
     @expand_list_indexing
     def _get_tuple(self, n):
         index, changes = index_to_slices(n, self.shape)
-        # print('INDEX', index, changes)
         indices = [self.indices[i] for i in range(*index[0].indices(self._len))]
         indices = make_slice_or_index_from_list_or_tuple(indices)
-        # print('INDICES', indices)
         index, _ = update_tuple(index, 0, indices)
         result = self.dataset[index]
         result = apply_index_to_slices_changes(result, changes)
@@ -89,8 +132,7 @@ class Subset(Forwards):
     @cached_property
     def frequency(self):
         dates = self.dates
-        delta = dates[1].astype(object) - dates[0].astype(object)
-        return int(delta.total_seconds() / 3600)
+        return frequency_to_timedelta(dates[1].astype(object) - dates[0].astype(object))
 
     def source(self, index):
         return Source(self, index, self.forward.source(index))

anemoi/datasets/data/unchecked.py

@@ -104,22 +104,29 @@ class Unchecked(Combined):
     def shape(self):
         raise NotImplementedError()
 
-    @property
-    def dtype(self):
-        raise NotImplementedError()
+    # @property
+    # def field_shape(self):
+    #     return tuple(d.shape for d in self.datasets)
 
-    @property
-    def grids(self):
-        raise NotImplementedError()
+    # @property
+    # def latitudes(self):
+    #     return tuple(d.latitudes for d in self.datasets)
 
+    # @property
+    # def longitudes(self):
+    #     return tuple(d.longitudes for d in self.datasets)
 
-class Zip(Unchecked):
+    # @property
+    # def statistics(self):
+    #     return tuple(d.statistics for d in self.datasets)
 
-    def __len__(self):
-        return min(len(d) for d in self.datasets)
+    # @property
+    # def resolution(self):
+    #     return tuple(d.resolution for d in self.datasets)
 
-    def __getitem__(self, n):
-        return tuple(d[n] for d in self.datasets)
+    # @property
+    # def name_to_index(self):
+    #     return tuple(d.name_to_index for d in self.datasets)
 
     @cached_property
     def missing(self):
@@ -142,17 +149,8 @@ class Chain(ConcatMixin, Unchecked):
     def dates(self):
         raise NotImplementedError()
 
-
-def zip_factory(args, kwargs):
-
-    zip = kwargs.pop("zip")
-    assert len(args) == 0
-    assert isinstance(zip, (list, tuple))
-
-    datasets = [_open(e) for e in zip]
-    datasets, kwargs = _auto_adjust(datasets, kwargs)
-
-    return Zip(datasets)._subset(**kwargs)
+    def dataset_metadata(self):
+        return {"multiple": [d.dataset_metadata() for d in self.datasets]}
 
 
 def chain_factory(args, kwargs):

anemoi/datasets/data/xy.py

@@ -0,0 +1,125 @@
+# (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
+
+from .debug import Node
+from .forwards import Combined
+from .misc import _auto_adjust
+from .misc import _open
+
+LOG = logging.getLogger(__name__)
+
+
+class ZipBase(Combined):
+
+    def swap_with_parent(self, parent):
+        new_parents = [parent.clone(ds) for ds in self.datasets]
+        return self.clone(new_parents)
+
+    def clone(self, datasets):
+        return self.__class__(datasets)
+
+    def tree(self):
+        return Node(self, [d.tree() for d in self.datasets])
+
+    def __len__(self):
+        return min(len(d) for d in self.datasets)
+
+    def __getitem__(self, n):
+        return tuple(d[n] for d in self.datasets)
+
+    def check_same_resolution(self, d1, d2):
+        pass
+
+    def check_same_grid(self, d1, d2):
+        pass
+
+    def check_same_variables(self, d1, d2):
+        pass
+
+    @cached_property
+    def missing(self):
+        result = set()
+        for d in self.datasets:
+            result = result | d.missing
+        return result
+
+    @property
+    def shape(self):
+        return tuple(d.shape for d in self.datasets)
+
+    @property
+    def field_shape(self):
+        return tuple(d.shape for d in self.datasets)
+
+    @property
+    def latitudes(self):
+        return tuple(d.latitudes for d in self.datasets)
+
+    @property
+    def longitudes(self):
+        return tuple(d.longitudes for d in self.datasets)
+
+    @property
+    def dtype(self):
+        return tuple(d.dtype for d in self.datasets)
+
+    @property
+    def grids(self):
+        return tuple(d.grids for d in self.datasets)
+
+    @property
+    def statistics(self):
+        return tuple(d.statistics for d in self.datasets)
+
+    @property
+    def resolution(self):
+        return tuple(d.resolution for d in self.datasets)
+
+    @property
+    def name_to_index(self):
+        return tuple(d.name_to_index for d in self.datasets)
+
+
+class Zip(ZipBase):
+    pass
+
+
+class XY(ZipBase):
+    pass
+
+
+def xy_factory(args, kwargs):
+
+    if "xy" in kwargs:
+        xy = kwargs.pop("xy")
+    else:
+        xy = [kwargs.pop("x"), kwargs.pop("y")]
+
+    assert len(args) == 0
+    assert isinstance(xy, (list, tuple))
+
+    datasets = [_open(e) for e in xy]
+    datasets, kwargs = _auto_adjust(datasets, kwargs)
+
+    assert len(datasets) == 2
+
+    return XY(datasets)._subset(**kwargs)
+
+
+def zip_factory(args, kwargs):
+
+    zip = kwargs.pop("zip")
+    assert len(args) == 0
+    assert isinstance(zip, (list, tuple))
+
+    datasets = [_open(e) for e in zip]
+    datasets, kwargs = _auto_adjust(datasets, kwargs)
+
+    return Zip(datasets)._subset(**kwargs)

anemoi/datasets/dates/__init__.py

@@ -9,64 +9,12 @@
 import datetime
 import warnings
 
+# from anemoi.utils.dates import as_datetime
+from anemoi.utils.dates import DateTimes
 from anemoi.utils.dates import as_datetime
-
-
-def _compress_dates(dates):
-    dates = sorted(dates)
-    if len(dates) < 3:
-        yield dates
-        return
-
-    prev = first = dates.pop(0)
-    curr = dates.pop(0)
-    delta = curr - prev
-    while curr - prev == delta:
-        prev = curr
-        if not dates:
-            break
-        curr = dates.pop(0)
-
-    yield (first, prev, delta)
-    if dates:
-        yield from _compress_dates([curr] + dates)
-
-
-def compress_dates(dates):
-    dates = [as_datetime(_) for _ in dates]
-    result = []
-
-    for n in _compress_dates(dates):
-        if isinstance(n, list):
-            result.extend([str(_) for _ in n])
-        else:
-            result.append(" ".join([str(n[0]), "to", str(n[1]), "by", str(n[2])]))
-
-    return result
-
-
-def print_dates(dates):
-    print(compress_dates(dates))
-
-
-def no_time_zone(date):
-    return date.replace(tzinfo=None)
-
-
-def frequency_to_hours(frequency):
-    if isinstance(frequency, int):
-        return frequency
-    assert isinstance(frequency, str), (type(frequency), frequency)
-
-    unit = frequency[-1].lower()
-    v = int(frequency[:-1])
-    return {"h": v, "d": v * 24}[unit]
-
-
-def normalize_date(x):
-    if isinstance(x, str):
-        return no_time_zone(datetime.datetime.fromisoformat(x))
-    return x
+from anemoi.utils.dates import frequency_to_timedelta
+from anemoi.utils.hindcasts import HindcastDatesTimes
+from anemoi.utils.humanize import print_dates
 
 
 def extend(x):
@@ -79,37 +27,37 @@ def extend(x):
     if isinstance(x, str):
         if "/" in x:
             start, end, step = x.split("/")
-            start = normalize_date(start)
-            end = normalize_date(end)
-            step = frequency_to_hours(step)
+            start = as_datetime(start)
+            end = as_datetime(end)
+            step = frequency_to_timedelta(step)
             while start <= end:
                 yield start
-                start += datetime.timedelta(hours=step)
+                start += step
             return
 
-    yield normalize_date(x)
+    yield as_datetime(x)
 
 
-class Dates:
+class DatesProvider:
     """Base class for date generation.
 
-    >>> Dates.from_config(**{"start": "2023-01-01 00:00", "end": "2023-01-02 00:00", "frequency": "1d"}).values
+    >>> 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)]
 
-    >>> Dates.from_config(**{"start": "2023-01-01 00:00", "end": "2023-01-03 00:00", "frequency": "18h"}).values
+    >>> DatesProvider.from_config(**{"start": "2023-01-01 00:00", "end": "2023-01-03 00:00", "frequency": "18h"}).values
     [datetime.datetime(2023, 1, 1, 0, 0), datetime.datetime(2023, 1, 1, 18, 0), datetime.datetime(2023, 1, 2, 12, 0)]
 
-    >>> Dates.from_config(start="2023-01-01 00:00", end="2023-01-02 00:00", frequency=6).as_dict()
+    >>> DatesProvider.from_config(start="2023-01-01 00:00", end="2023-01-02 00:00", frequency=6).as_dict()
     {'start': '2023-01-01T00:00:00', 'end': '2023-01-02T00:00:00', 'frequency': '6h'}
 
-    >>> len(Dates.from_config(start="2023-01-01 00:00", end="2023-01-02 00:00", frequency=12))
+    >>> len(DatesProvider.from_config(start="2023-01-01 00:00", end="2023-01-02 00:00", frequency=12))
     3
-    >>> len(Dates.from_config(start="2023-01-01 00:00",
+    >>> len(DatesProvider.from_config(start="2023-01-01 00:00",
     ...                   end="2023-01-02 00:00",
     ...                   frequency=12,
     ...                   missing=["2023-01-01 12:00"]))
     3
-    >>> len(Dates.from_config(start="2023-01-01 00:00",
+    >>> len(DatesProvider.from_config(start="2023-01-01 00:00",
     ...                   end="2023-01-02 00:00",
     ...                   frequency=12,
     ...                   missing=["2099-01-01 12:00"]))
@@ -121,12 +69,18 @@ class Dates:
             missing = []
         self.missing = list(extend(missing))
         if set(self.missing) - set(self.values):
-            warnings.warn(f"Missing dates {self.missing} not in list.")
+            diff = set(self.missing) - set(self.values)
+            warnings.warn(f"Missing dates {len(diff)=} not in list.")
 
     @classmethod
     def from_config(cls, **kwargs):
+
+        if kwargs.pop("hindcasts", False):
+            return HindcastsDates(**kwargs)
+
         if "values" in kwargs:
             return ValuesDates(**kwargs)
+
         return StartEndDates(**kwargs)
 
     def __iter__(self):
@@ -143,9 +97,9 @@ class Dates:
         return f"📅 {self.values[0]} ... {self.values[-1]}"
 
 
-class ValuesDates(Dates):
+class ValuesDates(DatesProvider):
     def __init__(self, values, **kwargs):
-        self.values = sorted([no_time_zone(_) for _ in values])
+        self.values = sorted([as_datetime(_) for _ in values])
         super().__init__(**kwargs)
 
     def __repr__(self):
@@ -155,9 +109,11 @@ class ValuesDates(Dates):
         return {"values": self.values[0]}
 
 
-class StartEndDates(Dates):
-    def __init__(self, start, end, frequency=1, months=None, **kwargs):
-        frequency = frequency_to_hours(frequency)
+class StartEndDates(DatesProvider):
+    def __init__(self, start, end, frequency=1, **kwargs):
+
+        frequency = frequency_to_timedelta(frequency)
+        assert isinstance(frequency, datetime.timedelta), frequency
 
         def _(x):
             if isinstance(x, str):
@@ -173,38 +129,111 @@ class StartEndDates(Dates):
         if isinstance(end, datetime.date) and not isinstance(end, datetime.datetime):
             end = datetime.datetime(end.year, end.month, end.day)
 
-        start = no_time_zone(start)
-        end = no_time_zone(end)
-
-        # if end <= start:
-        #     raise ValueError(f"End date {end} must be after start date {start}")
-
-        increment = datetime.timedelta(hours=frequency)
+        start = as_datetime(start)
+        end = as_datetime(end)
 
         self.start = start
         self.end = end
         self.frequency = frequency
 
-        date = start
-        self.values = []
-        while date <= end:
+        missing = kwargs.pop("missing", [])
 
-            if months is not None:
-                if date.month not in months:
-                    date += increment
-                    continue
+        self.values = list(DateTimes(start, end, increment=frequency, **kwargs))
+        self.kwargs = kwargs
 
-            self.values.append(date)
-            date += increment
-
-        super().__init__(**kwargs)
+        super().__init__(missing=missing)
 
     def as_dict(self):
         return {
             "start": self.start.isoformat(),
             "end": self.end.isoformat(),
-            "frequency": f"{self.frequency}h",
-        }
+            "frequency": frequency_to_string(self.frequency),
+        }.update(self.kwargs)
+
+
+class Hindcast:
+
+    def __init__(self, date, refdate, hdate, step):
+        self.date = date
+        self.refdate = refdate
+        self.hdate = hdate
+        self.step = step
+
+
+class HindcastsDates(DatesProvider):
+    def __init__(self, start, end, steps=[0], years=20, **kwargs):
+
+        if not isinstance(start, list):
+            start = [start]
+            end = [end]
+
+        reference_dates = []
+        for s, e in zip(start, end):
+            reference_dates.extend(list(DateTimes(s, e, increment=24, **kwargs)))
+        # reference_dates = list(DateTimes(start, end, increment=24, **kwargs))
+        dates = []
+
+        seen = {}
+
+        for hdate, refdate in HindcastDatesTimes(reference_dates=reference_dates, years=years):
+            assert refdate - hdate >= datetime.timedelta(days=365), (refdate - hdate, refdate, hdate)
+            for step in steps:
+
+                date = hdate + datetime.timedelta(hours=step)
+
+                if date in seen:
+                    raise ValueError(f"Duplicate date {date}={hdate}+{step} for {refdate} and {seen[date]}")
+
+                seen[date] = Hindcast(date, refdate, hdate, step)
+
+                assert refdate - date > datetime.timedelta(days=360), (refdate - date, refdate, date, hdate, step)
+
+                dates.append(date)
+
+        dates = sorted(dates)
+
+        mindelta = None
+        for a, b in zip(dates, dates[1:]):
+            delta = b - a
+            assert isinstance(delta, datetime.timedelta), delta
+            if mindelta is None:
+                mindelta = delta
+            else:
+                mindelta = min(mindelta, delta)
+
+        self.frequency = mindelta
+        assert mindelta.total_seconds() > 0, mindelta
+
+        print("🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥", dates[0], dates[-1], mindelta)
+
+        # Use all values between start and end by frequency, and set the ones that are missing
+        self.values = []
+        missing = []
+        date = dates[0]
+        last = date
+        print("------", date, dates[-1])
+        dateset = set(dates)
+        while date <= dates[-1]:
+            self.values.append(date)
+            if date not in dateset:
+                missing.append(date)
+                seen[date] = seen[last]
+            else:
+                last = date
+            date = date + mindelta
+
+        self.mapping = seen
+
+        print("🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥", self.values[0], self.values[-1], mindelta)
+        print("🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥🔥", f"{len(self.values)=} - {len(missing)=}")
+
+        super().__init__(missing=missing)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.values[0]}..{self.values[-1]})"
+
+    def as_dict(self):
+        return {"hindcasts": self.hindcasts}
 
 
 if __name__ == "__main__":