anemoi-datasets 0.4.5__py3-none-any.whl → 0.5.5__py3-none-any.whl

anemoi/datasets/dates/__init__.py

@@ -10,8 +10,11 @@ import datetime
 import warnings
 
 # from anemoi.utils.dates import as_datetime
+from anemoi.utils.dates import DateTimes
 from anemoi.utils.dates import as_datetime
+from anemoi.utils.dates import frequency_to_string
 from anemoi.utils.dates import frequency_to_timedelta
+from anemoi.utils.hindcasts import HindcastDatesTimes
 from anemoi.utils.humanize import print_dates
 
 
@@ -30,32 +33,32 @@ def extend(x):
             step = frequency_to_timedelta(step)
             while start <= end:
                 yield start
-                start += datetime.timedelta(hours=step)
+                start += step
             return
 
     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"]))
@@ -67,12 +70,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):
@@ -89,7 +98,7 @@ class Dates:
         return f"📅 {self.values[0]} ... {self.values[-1]}"
 
 
-class ValuesDates(Dates):
+class ValuesDates(DatesProvider):
     def __init__(self, values, **kwargs):
         self.values = sorted([as_datetime(_) for _ in values])
         super().__init__(**kwargs)
@@ -101,8 +110,9 @@ class ValuesDates(Dates):
         return {"values": self.values[0]}
 
 
-class StartEndDates(Dates):
-    def __init__(self, start, end, frequency=1, months=None, **kwargs):
+class StartEndDates(DatesProvider):
+    def __init__(self, start, end, frequency=1, **kwargs):
+
         frequency = frequency_to_timedelta(frequency)
         assert isinstance(frequency, datetime.timedelta), frequency
 
@@ -123,35 +133,108 @@ class StartEndDates(Dates):
         start = as_datetime(start)
         end = as_datetime(end)
 
-        # if end <= start:
-        #     raise ValueError(f"End date {end} must be after start date {start}")
-
-        increment = frequency
-
         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__":

anemoi/datasets/dates/groups.py

@@ -7,11 +7,42 @@
 
 
 import itertools
+from functools import cached_property
 
-from anemoi.datasets.dates import Dates
+from anemoi.datasets.dates import DatesProvider
 from anemoi.datasets.dates import as_datetime
 
 
+def _shorten(dates):
+    if isinstance(dates, (list, tuple)):
+        dates = [d.isoformat() for d in dates]
+        if len(dates) > 5:
+            return f"{dates[0]}...{dates[-1]}"
+    return dates
+
+
+class GroupOfDates:
+    def __init__(self, dates, provider, partial_ok=False):
+        assert isinstance(provider, DatesProvider), type(provider)
+        assert isinstance(dates, list)
+
+        self.dates = dates
+        self.provider = provider
+        self.partial_ok = partial_ok
+
+    def __len__(self):
+        return len(self.dates)
+
+    def __iter__(self):
+        return iter(self.dates)
+
+    def __repr__(self) -> str:
+        return f"GroupOfDates(dates={_shorten(self.dates)})"
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, GroupOfDates) and self.dates == other.dates
+
+
 class Groups:
     """>>> list(Groups(group_by="daily", start="2023-01-01 00:00", end="2023-01-05 00:00", frequency=12))[0]
     [datetime.datetime(2023, 1, 1, 0, 0), datetime.datetime(2023, 1, 1, 12, 0)]
@@ -41,33 +72,48 @@ class Groups:
 
     def __init__(self, **kwargs):
         group_by = kwargs.pop("group_by")
-        self.dates = Dates.from_config(**kwargs)
-        self.grouper = Grouper.from_config(group_by)
-        self.filter = Filter(self.dates.missing)
+        self._dates = DatesProvider.from_config(**kwargs)
+        self._grouper = Grouper.from_config(group_by)
+        self._filter = Filter(self._dates.missing)
+
+    @property
+    def provider(self):
+        return self._dates
 
     def __iter__(self):
-        for dates in self.grouper(self.dates):
-            dates = self.filter(dates)
+        for go in self._grouper(self._dates):
+            dates = self._filter(go.dates)
             if not dates:
                 continue
-            yield dates
+            yield GroupOfDates(dates, go.provider)
 
     def __len__(self):
-        count = 0
-        for dates in self.grouper(self.dates):
-            dates = self.filter(dates)
+        return self._len
+
+    @cached_property
+    def _len(self):
+        n = 0
+        for go in self._grouper(self._dates):
+            dates = self._filter(go.dates)
             if not dates:
                 continue
-            count += 1
-        return count
+            n += 1
+        return n
 
     def __repr__(self):
-        return f"{self.__class__.__name__}(dates={len(self)})"
+        return f"{self.__class__.__name__}(dates={len(self)},{_shorten(self._dates)})"
+
+    def describe(self):
+        return self.dates.summary
+
+    def one_date(self):
+        go = next(iter(self))
+        return GroupOfDates([go.dates[0]], go.provider)
 
 
 class Filter:
     def __init__(self, missing):
-        self.missing = [as_datetime(m) for m in missing]
+        self.missing = set(as_datetime(m) for m in missing)
 
     def __call__(self, dates):
         return [d for d in dates if d not in self.missing]
@@ -76,10 +122,16 @@ class Filter:
 class Grouper:
     @classmethod
     def from_config(cls, group_by):
+
         if isinstance(group_by, int) and group_by > 0:
             return GrouperByFixedSize(group_by)
+
         if group_by is None:
             return GrouperOneGroup()
+
+        if group_by == "reference_date":
+            return ReferenceDateGroup()
+
         key = {
             "monthly": lambda dt: (dt.year, dt.month),
             "daily": lambda dt: (dt.year, dt.month, dt.day),
@@ -89,30 +141,51 @@ class Grouper:
         return GrouperByKey(key)
 
 
+class ReferenceDateGroup(Grouper):
+    def __call__(self, dates):
+        assert isinstance(dates, DatesProvider), type(dates)
+
+        mapping = dates.mapping
+
+        def same_refdate(dt):
+            return mapping[dt].refdate
+
+        for _, g in itertools.groupby(sorted(dates, key=same_refdate), key=same_refdate):
+            yield GroupOfDates(list(g), dates)
+
+
 class GrouperOneGroup(Grouper):
     def __call__(self, dates):
-        yield dates.values
+        assert isinstance(dates, DatesProvider), type(dates)
+
+        yield GroupOfDates(dates.values, dates)
 
 
 class GrouperByKey(Grouper):
+    """Group dates by a key."""
+
     def __init__(self, key):
         self.key = key
 
     def __call__(self, dates):
-        for _, g in itertools.groupby(dates, key=self.key):
-            yield list(g)
+        for _, g in itertools.groupby(sorted(dates, key=self.key), key=self.key):
+            yield GroupOfDates(list(g), dates)
 
 
 class GrouperByFixedSize(Grouper):
+    """Group dates by a fixed size."""
+
     def __init__(self, size):
         self.size = size
 
     def __call__(self, dates):
         batch = []
+
         for d in dates:
             batch.append(d)
             if len(batch) == self.size:
-                yield batch
+                yield GroupOfDates(batch, dates)
                 batch = []
+
         if batch:
-            yield batch
+            yield GroupOfDates(batch, dates)

anemoi/datasets/fields.py

@@ -0,0 +1,66 @@
+# (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.
+
+
+from earthkit.data.indexing.fieldlist import FieldArray
+
+
+def list_to_fieldlist(fields):
+    return FieldArray(fields)
+
+
+def empty_fieldlist():
+    return FieldArray([])
+
+
+class WrappedField:
+    def __init__(self, field):
+        self._field = field
+
+    def __getattr__(self, name):
+        return getattr(self._field, name)
+
+    def __repr__(self) -> str:
+        return repr(self._field)
+
+
+class NewDataField(WrappedField):
+    def __init__(self, field, data):
+        super().__init__(field)
+        self._data = data
+        self.shape = data.shape
+
+    def to_numpy(self, flatten=False, dtype=None, index=None):
+        data = self._data
+        if dtype is not None:
+            data = data.astype(dtype)
+        if flatten:
+            data = data.flatten()
+        if index is not None:
+            data = data[index]
+        return data
+
+
+class NewMetadataField(WrappedField):
+    def __init__(self, field, **kwargs):
+        super().__init__(field)
+        self._metadata = kwargs
+
+    def metadata(self, *args, **kwargs):
+        if len(args) == 1 and args[0] in self._metadata:
+            return self._metadata[args[0]]
+        return self._field.metadata(*args, **kwargs)
+
+
+class NewValidDateTimeField(NewMetadataField):
+    def __init__(self, field, valid_datetime):
+        date = valid_datetime.date().strftime("%Y%m%d")
+        time = valid_datetime.time().strftime("%H%M")
+
+        self.valid_datetime = valid_datetime
+
+        super().__init__(field, date=date, time=time, step=0, valid_datetime=valid_datetime.isoformat())

anemoi/datasets/utils/fields.py

@@ -0,0 +1,47 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+#
+
+
+class WrappedField:
+    def __init__(self, field):
+        self._field = field
+
+    def __getattr__(self, name):
+        return getattr(self._field, name)
+
+    def __repr__(self) -> str:
+        return repr(self._field)
+
+
+class NewDataField(WrappedField):
+    def __init__(self, field, data):
+        super().__init__(field)
+        self._data = data
+        self.shape = data.shape
+
+    def to_numpy(self, flatten=False, dtype=None, index=None):
+        data = self._data
+        if dtype is not None:
+            data = data.astype(dtype)
+        if flatten:
+            data = data.flatten()
+        if index is not None:
+            data = data[index]
+        return data
+
+
+class NewMetadataField(WrappedField):
+    def __init__(self, field, **kwargs):
+        super().__init__(field)
+        self._metadata = kwargs
+
+    def metadata(self, *args, **kwargs):
+        if len(args) == 1 and args[0] in self._metadata:
+            return self._metadata[args[0]]
+        return self._field.metadata(*args, **kwargs)

{anemoi_datasets-0.4.5.dist-info → anemoi_datasets-0.5.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-datasets
-Version: 0.4.5
+Version: 0.5.5
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License: Apache License
@@ -224,38 +224,36 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: anemoi-utils[provenance] >=0.3.15
+Requires-Dist: cfunits
 Requires-Dist: numpy
 Requires-Dist: pyyaml
 Requires-Dist: semantic-version
 Requires-Dist: tqdm
-Requires-Dist: zarr
+Requires-Dist: zarr <=2.17
 Provides-Extra: all
-Requires-Dist: aiohttp ; extra == 'all'
 Requires-Dist: boto3 ; extra == 'all'
 Requires-Dist: earthkit-data[mars] >=0.9 ; extra == 'all'
 Requires-Dist: earthkit-geo >=0.2 ; extra == 'all'
 Requires-Dist: earthkit-meteo ; extra == 'all'
-Requires-Dist: eccodes >=2.37 ; extra == 'all'
+Requires-Dist: ecmwflibs >=0.6.3 ; extra == 'all'
 Requires-Dist: entrypoints ; extra == 'all'
 Requires-Dist: gcsfs ; extra == 'all'
 Requires-Dist: kerchunk ; extra == 'all'
 Requires-Dist: pyproj ; extra == 'all'
 Requires-Dist: requests ; extra == 'all'
-Requires-Dist: s3fs ; extra == 'all'
 Provides-Extra: create
 Requires-Dist: earthkit-data[mars] >=0.9 ; extra == 'create'
 Requires-Dist: earthkit-geo >=0.2 ; extra == 'create'
 Requires-Dist: earthkit-meteo ; extra == 'create'
-Requires-Dist: eccodes >=2.37 ; extra == 'create'
+Requires-Dist: ecmwflibs >=0.6.3 ; extra == 'create'
 Requires-Dist: entrypoints ; extra == 'create'
 Requires-Dist: pyproj ; extra == 'create'
 Provides-Extra: dev
-Requires-Dist: aiohttp ; extra == 'dev'
 Requires-Dist: boto3 ; extra == 'dev'
 Requires-Dist: earthkit-data[mars] >=0.9 ; extra == 'dev'
 Requires-Dist: earthkit-geo >=0.2 ; extra == 'dev'
 Requires-Dist: earthkit-meteo ; extra == 'dev'
-Requires-Dist: eccodes >=2.37 ; extra == 'dev'
+Requires-Dist: ecmwflibs >=0.6.3 ; extra == 'dev'
 Requires-Dist: entrypoints ; extra == 'dev'
 Requires-Dist: gcsfs ; extra == 'dev'
 Requires-Dist: kerchunk ; extra == 'dev'
@@ -264,27 +262,20 @@ Requires-Dist: pandoc ; extra == 'dev'
 Requires-Dist: pyproj ; extra == 'dev'
 Requires-Dist: pytest ; extra == 'dev'
 Requires-Dist: requests ; extra == 'dev'
-Requires-Dist: rstfmt ; extra == 'dev'
-Requires-Dist: s3fs ; extra == 'dev'
 Requires-Dist: sphinx ; extra == 'dev'
-Requires-Dist: sphinx-argparse <0.5 ; extra == 'dev'
 Requires-Dist: sphinx-rtd-theme ; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: nbsphinx ; extra == 'docs'
 Requires-Dist: pandoc ; extra == 'docs'
-Requires-Dist: rstfmt ; extra == 'docs'
 Requires-Dist: sphinx ; extra == 'docs'
-Requires-Dist: sphinx-argparse <0.5 ; extra == 'docs'
+Requires-Dist: sphinx-argparse ; extra == 'docs'
 Requires-Dist: sphinx-rtd-theme ; extra == 'docs'
-Provides-Extra: kerchunk
-Requires-Dist: gcsfs ; extra == 'kerchunk'
-Requires-Dist: kerchunk ; extra == 'kerchunk'
-Requires-Dist: s3fs ; extra == 'kerchunk'
 Provides-Extra: remote
-Requires-Dist: aiohttp ; extra == 'remote'
 Requires-Dist: boto3 ; extra == 'remote'
 Requires-Dist: requests ; extra == 'remote'
-Requires-Dist: s3fs ; extra == 'remote'
 Provides-Extra: tests
 Requires-Dist: pytest ; extra == 'tests'
+Provides-Extra: xarray
+Requires-Dist: gcsfs ; extra == 'xarray'
+Requires-Dist: kerchunk ; extra == 'xarray'