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

anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3'
-__version_tuple__ = version_tuple = (0, 4, 3)
+__version__ = version = '0.4.4'
+__version_tuple__ = version_tuple = (0, 4, 4)

anemoi/datasets/compute/recentre.py

@@ -64,7 +64,7 @@ def recentre(
 
     keys = ["param", "level", "valid_datetime", "date", "time", "step", "number"]
 
-    number_list = members.unique_values("number")["number"]
+    number_list = members.unique_values("number", progress_bar=False)["number"]
     n_numbers = len(number_list)
 
     assert None not in number_list

anemoi/datasets/create/functions/sources/accumulations.py

@@ -15,11 +15,10 @@ import earthkit.data as ekd
 import numpy as np
 from earthkit.data.core.temporary import temp_file
 from earthkit.data.readers.grib.output import new_grib_output
-from earthkit.data.utils.availability import Availability
 
 from anemoi.datasets.create.utils import to_datetime_list
 
-from .mars import use_grib_paramid
+from .mars import mars
 
 LOG = logging.getLogger(__name__)
 
@@ -307,26 +306,11 @@ def _compute_accumulations(
     for date, time, steps in mars_date_time_steps:
         for p in param:
             for n in number:
-                requests.append(
-                    patch(
-                        {
-                            "param": p,
-                            "date": date,
-                            "time": time,
-                            "step": sorted(steps),
-                            "number": n,
-                        }
-                    )
-                )
-
-    compressed = Availability(requests)
-    ds = ekd.from_source("empty")
-    for r in compressed.iterate():
-        request.update(r)
-        if context.use_grib_paramid and "param" in request:
-            request = use_grib_paramid(request)
-        print("🌧️", request)
-        ds = ds + ekd.from_source("mars", **request)
+                r = dict(request, param=p, date=date, time=time, step=sorted(steps), number=n)
+
+                requests.append(patch(r))
+
+    ds = mars(context, dates, *requests, request_already_using_valid_datetime=True)
 
     accumulations = {}
     for a in [AccumulationClass(out, frequency=frequency, **r) for r in requests]:

anemoi/datasets/create/functions/sources/hindcasts.py

@@ -7,20 +7,16 @@
 # nor does it submit to any jurisdiction.
 #
 import datetime
+import logging
+
+from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.create.functions.sources.mars import mars
 
+LOGGER = logging.getLogger(__name__)
 DEBUG = True
 
 
-def _member(field):
-    # Bug in eccodes has number=0 randomly
-    number = field.metadata("number")
-    if number is None:
-        number = 0
-    return number
-
-
 def _to_list(x):
     if isinstance(x, (list, tuple)):
         return x
@@ -63,9 +59,19 @@ class HindcastCompute:
 def use_reference_year(reference_year, request):
     request = request.copy()
     hdate = request.pop("date")
-    date = datetime.datetime(reference_year, hdate.month, hdate.day)
+
+    if hdate.year >= reference_year:
+        return None, False
+
+    try:
+        date = datetime.datetime(reference_year, hdate.month, hdate.day)
+    except ValueError:
+        if hdate.month == 2 and hdate.day == 29:
+            return None, False
+        raise
+
     request.update(date=date.strftime("%Y-%m-%d"), hdate=hdate.strftime("%Y-%m-%d"))
-    return request
+    return request, True
 
 
 def hindcasts(context, dates, **request):
@@ -89,9 +95,18 @@ def hindcasts(context, dates, **request):
     requests = []
     for d in dates:
         req = c.compute_hindcast(d)
-        req = use_reference_year(reference_year, req)
+        req, ok = use_reference_year(reference_year, req)
+        if ok:
+            requests.append(req)
+
+    # print("HINDCASTS requests", reference_year, base_times, available_steps)
+    # print("HINDCASTS dates", compress_dates(dates))
+
+    if len(requests) == 0:
+        # print("HINDCASTS no requests")
+        return MultiFieldList([])
 
-        requests.append(req)
+    # print("HINDCASTS requests", requests)
 
     return mars(
         context,

anemoi/datasets/create/functions/sources/tendencies.py

@@ -56,7 +56,7 @@ def tendencies(dates, time_increment, **kwargs):
 
     ds = mars(dates=all_dates, **kwargs)
 
-    dates_in_data = ds.unique_values("valid_datetime")["valid_datetime"]
+    dates_in_data = ds.unique_values("valid_datetime", progress_bar=False)["valid_datetime"]
     for d in all_dates:
         assert d.isoformat() in dates_in_data, d
 

anemoi/datasets/create/input.py

@@ -277,6 +277,9 @@ class Result:
         if len(args) == 1 and isinstance(args[0], (list, tuple)):
             args = args[0]
 
+        # print("Executing", self.action_path)
+        # print("Dates:", compress_dates(self.dates))
+
         names = []
         for a in args:
             if isinstance(a, str):
@@ -287,12 +290,12 @@ class Result:
         print(f"Building a {len(names)}D hypercube using", names)
 
         ds = ds.order_by(*args, remapping=remapping, patches=patches)
-        user_coords = ds.unique_values(*names, remapping=remapping, patches=patches)
+        user_coords = ds.unique_values(*names, remapping=remapping, patches=patches, progress_bar=False)
 
         print()
         print("Number of unique values found for each coordinate:")
         for k, v in user_coords.items():
-            print(f"  {k:20}:", len(v))
+            print(f"  {k:20}:", len(v), shorten_list(v, max_length=10))
         print()
         user_shape = tuple(len(v) for k, v in user_coords.items())
         print("Shape of the hypercube           :", user_shape)
@@ -305,13 +308,18 @@ class Result:
 
         remapping = build_remapping(remapping, patches)
         expected = set(itertools.product(*user_coords.values()))
+        extra = set()
 
         if math.prod(user_shape) > len(ds):
             print(f"This means that all the fields in the datasets do not exists for all combinations of {names}.")
 
             for f in ds:
                 metadata = remapping(f.metadata)
-                expected.remove(tuple(metadata(n) for n in names))
+                key = tuple(metadata(n, default=None) for n in names)
+                if key in expected:
+                    expected.remove(key)
+                else:
+                    extra.add(key)
 
             print("Missing fields:")
             print()
@@ -321,7 +329,35 @@ class Result:
                     print("...", len(expected) - i - 1, "more")
                     break
 
+            print("Extra fields:")
+            print()
+            for i, f in enumerate(sorted(extra)):
+                print(" ", f)
+                if i >= 9 and len(extra) > 10:
+                    print("...", len(extra) - i - 1, "more")
+                    break
+
             print()
+            print("Missing values:")
+            per_name = defaultdict(set)
+            for e in expected:
+                for n, v in zip(names, e):
+                    per_name[n].add(v)
+
+            for n, v in per_name.items():
+                print(" ", n, len(v), shorten_list(sorted(v), max_length=10))
+            print()
+
+            print("Extra values:")
+            per_name = defaultdict(set)
+            for e in extra:
+                for n, v in zip(names, e):
+                    per_name[n].add(v)
+
+            for n, v in per_name.items():
+                print(" ", n, len(v), shorten_list(sorted(v), max_length=10))
+            print()
+
             print("To solve this issue, you can:")
             print(
                 "  - Provide a better selection, like 'step: 0' or 'level: 1000' to "

anemoi/datasets/create/loaders.py

@@ -17,6 +17,7 @@ import numpy as np
 import tqdm
 import zarr
 from anemoi.utils.config import DotDict
+from anemoi.utils.dates import as_datetime
 from anemoi.utils.humanize import seconds_to_human
 
 from anemoi.datasets import MissingDateError
@@ -24,6 +25,7 @@ from anemoi.datasets import open_dataset
 from anemoi.datasets.create.persistent import build_storage
 from anemoi.datasets.data.misc import as_first_date
 from anemoi.datasets.data.misc import as_last_date
+from anemoi.datasets.dates import compress_dates
 from anemoi.datasets.dates.groups import Groups
 
 from .check import DatasetName
@@ -442,14 +444,24 @@ class ContentLoader(Loader):
         dates = result.dates
 
         cube = result.get_cube()
-        assert cube.extended_user_shape[0] == len(dates), (
-            cube.extended_user_shape[0],
-            len(dates),
-        )
-
         shape = cube.extended_user_shape
         dates_in_data = cube.user_coords["valid_datetime"]
 
+        if cube.extended_user_shape[0] != len(dates):
+            print(f"Cube shape does not match the number of dates {cube.extended_user_shape[0]}, {len(dates)}")
+            print("Requested dates", compress_dates(dates))
+            print("Cube dates", compress_dates(dates_in_data))
+
+            a = set(as_datetime(_) for _ in dates)
+            b = set(as_datetime(_) for _ in dates_in_data)
+
+            print("Missing dates", compress_dates(a - b))
+            print("Extra dates", compress_dates(b - a))
+
+            raise ValueError(
+                f"Cube shape does not match the number of dates {cube.extended_user_shape[0]}, {len(dates)}"
+            )
+
         LOG.debug(f"Loading {shape=} in {self.data_array.shape=}")
 
         def check_dates_in_data(lst, lst2):

anemoi/datasets/dates/__init__.py

@@ -9,6 +9,45 @@
 import datetime
 import warnings
 
+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)
@@ -30,6 +69,27 @@ def normalize_date(x):
     return x
 
 
+def extend(x):
+
+    if isinstance(x, (list, tuple)):
+        for y in x:
+            yield from extend(y)
+        return
+
+    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)
+            while start <= end:
+                yield start
+                start += datetime.timedelta(hours=step)
+            return
+
+    yield normalize_date(x)
+
+
 class Dates:
     """Base class for date generation.
 
@@ -59,7 +119,7 @@ class Dates:
     def __init__(self, missing=None):
         if not missing:
             missing = []
-        self.missing = [normalize_date(x) for x in missing]
+        self.missing = list(extend(missing))
         if set(self.missing) - set(self.values):
             warnings.warn(f"Missing dates {self.missing} not in list.")
 
@@ -145,3 +205,9 @@ class StartEndDates(Dates):
             "end": self.end.isoformat(),
             "frequency": f"{self.frequency}h",
         }
+
+
+if __name__ == "__main__":
+    print_dates([datetime.datetime(2023, 1, 1, 0, 0)])
+    s = StartEndDates(start="2023-01-01 00:00", end="2023-01-02 00:00", frequency=1)
+    print_dates(list(s))

anemoi/datasets/dates/groups.py

@@ -9,7 +9,7 @@
 import itertools
 
 from anemoi.datasets.dates import Dates
-from anemoi.datasets.dates import no_time_zone
+from anemoi.datasets.dates import normalize_date
 
 
 class Groups:
@@ -67,7 +67,7 @@ class Groups:
 
 class Filter:
     def __init__(self, missing):
-        self.missing = [no_time_zone(m) for m in missing]
+        self.missing = [normalize_date(m) for m in missing]
 
     def __call__(self, dates):
         return [d for d in dates if d not in self.missing]

{anemoi_datasets-0.4.3.dist-info → anemoi_datasets-0.4.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-datasets
-Version: 0.4.3
+Version: 0.4.4
 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

{anemoi_datasets-0.4.3.dist-info → anemoi_datasets-0.4.4.dist-info}/RECORD

@@ -1,6 +1,6 @@
 anemoi/datasets/__init__.py,sha256=Z1gqZWhecLcT0RZQqYBLlz01MUlUZd0kWEj_RavbITM,782
 anemoi/datasets/__main__.py,sha256=cLA2PidDTOUHaDGzd0_E5iioKYNe-PSTv567Y2fuwQk,723
-anemoi/datasets/_version.py,sha256=mDTy6z4iq5XKGigbX2Q9ub9_BzxGHksMlMjpkDAl_hE,411
+anemoi/datasets/_version.py,sha256=t0Mfy7vCENQWSYE4xRfluKBHgAYWK48fjSubsSGEQEI,411
 anemoi/datasets/grids.py,sha256=3YBMMJodgYhavarXPAlMZHaMtDT9v2IbTmAXZTqf8Qo,8481
 anemoi/datasets/commands/__init__.py,sha256=qAybFZPBBQs0dyx7dZ3X5JsLpE90pwrqt1vSV7cqEIw,706
 anemoi/datasets/commands/compare.py,sha256=svEhyR7pOS1847_RJr1I6vF7ZDPB9AVlcrhy_gxQVms,3263
@@ -9,13 +9,13 @@ anemoi/datasets/commands/create.py,sha256=iAXm7b1zeFEbVyTSqgSdCXwYtgIi8XUN96muv-
 anemoi/datasets/commands/inspect.py,sha256=rEBzR5LmrPzkNrlGS471Q7Ma9O8Djiml2mxAlxvSl0U,18676
 anemoi/datasets/commands/scan.py,sha256=MaTdne4JrtlqO3LhOUr43DZhZ6O-RZwC7uQ7C6PG7Os,2910
 anemoi/datasets/compute/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-anemoi/datasets/compute/recentre.py,sha256=oIqhENljB9ad-wnMJCJ3P0Xf1v76zQjqYCu1TuySgSI,4919
+anemoi/datasets/compute/recentre.py,sha256=CMvRxTMv2H1hF1KqaZHT3Maa9P01GtsgAxr1YHbvH_o,4939
 anemoi/datasets/create/__init__.py,sha256=_XMODLz0xpVf06Tu7x2c1HhFcuBYJ0xeux6dn9dgTXg,6830
 anemoi/datasets/create/check.py,sha256=TAZOrq8QbS9F5kuR779FAuZPJXA7MXsS_DU1kVotq5Q,5838
 anemoi/datasets/create/chunks.py,sha256=1Inh3sIBn-2sNguErb-BsLns6W_HtDiOJAjIb29lp-U,2442
 anemoi/datasets/create/config.py,sha256=vZ-t50mHgrEnxJoaooWrZKzVGeEQrg__rCOrhi5Q4nU,7122
-anemoi/datasets/create/input.py,sha256=BJiGFR-imZcChdij3fodyz8XSHMUs-__ZfoCeMHuhUo,32472
-anemoi/datasets/create/loaders.py,sha256=tosx8Zms7OAEScw9JxHw6BLq6xZSDCNPYKeqFWP5O3Y,32504
+anemoi/datasets/create/input.py,sha256=CNOhqHqW6C2owewJYXkLZC2W_hwGZj794m2_wXXKbyA,33714
+anemoi/datasets/create/loaders.py,sha256=dEGj7FxKHgxXTVBSHB-7kzOK7sjbiWOuzE2y5eDXla0,33130
 anemoi/datasets/create/patch.py,sha256=c8cWoqxFzcY9mKREosLjuQCUeJMJL6sbNNkoVvhCZDA,3800
 anemoi/datasets/create/persistent.py,sha256=wfNnNSY8VLO2EyEOQL9umI3sbQqETjJP7bPCiVF6Cko,4291
 anemoi/datasets/create/size.py,sha256=k4COEjs3wud0oKHH5P3n8Fap35xddXs002ucPjpBC88,1040
@@ -32,18 +32,18 @@ anemoi/datasets/create/functions/filters/rename.py,sha256=02p6zj2g0Qp866RrXeZG9D
 anemoi/datasets/create/functions/filters/rotate_winds.py,sha256=E0P5scdX0lwTMdcFDYfBzQ_X_4A6EvnrtFvF55-56Hk,2414
 anemoi/datasets/create/functions/filters/unrotate_winds.py,sha256=hiIwgWi_2lk_ntxsPFMyZ6Ku8_5p91ht36VN_2kHYDA,2414
 anemoi/datasets/create/functions/sources/__init__.py,sha256=83G1-DD2IV4VJP3MVN9512_CN4D3IVDZHUKa2ghVgKo,1246
-anemoi/datasets/create/functions/sources/accumulations.py,sha256=Rzn9ocBR4CHdws4s6GuKiV4IqVQo0rCQLyPQNBsTGXY,12493
+anemoi/datasets/create/functions/sources/accumulations.py,sha256=nz09rezfNvXdTBtSQAKnsOeoiWwB8v1ecAayrJZq9rE,11947
 anemoi/datasets/create/functions/sources/constants.py,sha256=9MNxjkXAjtIq7X-T7GgKGVzH-V-FBcTxj0gLLJYoXTI,903
 anemoi/datasets/create/functions/sources/empty.py,sha256=ZrXGs8Y3VrLSV8C8YlJTJcHV7Bmi7xPiUlrq8R0JZQY,485
 anemoi/datasets/create/functions/sources/forcings.py,sha256=tF3EyIs5AGF1Ppvp6dIExONM-kGF-wcnMO1sZc_wDuo,646
 anemoi/datasets/create/functions/sources/grib.py,sha256=MbHSAJAkvUsQT5uo_-mEqzN-0-aOygMwFIds42kXcqg,1750
-anemoi/datasets/create/functions/sources/hindcasts.py,sha256=TRe25E6Xb_fBUsyOnzB_Bm1cTZ4H99viplf1ccWzZv8,3154
+anemoi/datasets/create/functions/sources/hindcasts.py,sha256=xgE4a-diiy8rP8cmgAT12cBsQZiGk0cRzUswtZM7i_0,3607
 anemoi/datasets/create/functions/sources/mars.py,sha256=Wrxm15-7QUQ_xCe7eqmxHLLkYo5RDM0GKUoeEP8BZS4,6603
 anemoi/datasets/create/functions/sources/netcdf.py,sha256=8uug0oAdGBJIKKws-EflA4ZgjQye_sWCBymUVv1TEW4,532
 anemoi/datasets/create/functions/sources/opendap.py,sha256=0Bs0PytUvI1WZGn2OdmnJuFDGAQypxN8v44nl106QjY,531
 anemoi/datasets/create/functions/sources/recentre.py,sha256=t07LIXG3Hp9gmPkPriILVt86TxubsHyS1EL1lzwgtXY,1810
 anemoi/datasets/create/functions/sources/source.py,sha256=J3O4M0nB1a-67IJuY_aWqDDqyNGXB_uzxVbicFldO4U,1422
-anemoi/datasets/create/functions/sources/tendencies.py,sha256=OIYKdRTabDvxnBXCcffWWBcorVlhZVNj3VkcVuQjIYA,4088
+anemoi/datasets/create/functions/sources/tendencies.py,sha256=4hFPGU51KnxXp-KteSp5WLb8v586sDq2ePPd8CfUFjA,4108
 anemoi/datasets/create/functions/sources/xarray_kerchunk.py,sha256=JsuGjTs5BQdJIPRI8TuUNde680UQkT4GbUhOt6wYy38,1432
 anemoi/datasets/create/functions/sources/xarray_zarr.py,sha256=qI_fEKo28hu_B_qMPx67gigysw8qw5ePMU9P0wbDcIc,531
 anemoi/datasets/create/functions/sources/zenodo.py,sha256=n7_sfZHJVDjVnSwc07muO9kS7TwOM4S8DIlJf5pnDiU,1237
@@ -75,12 +75,12 @@ anemoi/datasets/data/statistics.py,sha256=lZCcKw9s7ttMBEp6ANyxtbXoZZvchhE7SClq-D
 anemoi/datasets/data/stores.py,sha256=LNCJBBQ--3ng4LGvdZ24ZXyOz7i15MvhtArt3_CxM_Q,11934
 anemoi/datasets/data/subset.py,sha256=9urVTXdnwCgqn0_BRYquMi8oiXn4ubAf0n4586hWfKw,3814
 anemoi/datasets/data/unchecked.py,sha256=xhdMg-ToI1UfBWHNsWyn1y2meZWngZtHx-33L0KqKp8,4037
-anemoi/datasets/dates/__init__.py,sha256=BRrwOapeP20tu7hN0m8hG04INp94tigwEfMClLl0ww4,4635
-anemoi/datasets/dates/groups.py,sha256=bYgw02eCLkjkZjJdyLpcFVMqtH8XO22cBdHOfG7tCBc,3471
+anemoi/datasets/dates/__init__.py,sha256=ZRsUJCA7zHkgKzB7gOq3WndTGcCju63WGgB-jzSJDgc,6155
+anemoi/datasets/dates/groups.py,sha256=xwDkUye3PuMg3c6ao0QvZ5bcMTxAhTcmeIZsLUKYJJY,3475
 anemoi/datasets/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-anemoi_datasets-0.4.3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-anemoi_datasets-0.4.3.dist-info/METADATA,sha256=2tNctMmd44eD_Lc7EDuP-ZWP6XJueMQ76MT53wQ8sbc,16716
-anemoi_datasets-0.4.3.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
-anemoi_datasets-0.4.3.dist-info/entry_points.txt,sha256=yR-o-4uiPEA_GLBL81SkMYnUoxq3CAV3hHulQiRtGG0,66
-anemoi_datasets-0.4.3.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_datasets-0.4.3.dist-info/RECORD,,
+anemoi_datasets-0.4.4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+anemoi_datasets-0.4.4.dist-info/METADATA,sha256=gZz_PkwuY4NrJvcsBWFn5xwm7jOjnbRm-k6cSIUMdJ8,16716
+anemoi_datasets-0.4.4.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+anemoi_datasets-0.4.4.dist-info/entry_points.txt,sha256=yR-o-4uiPEA_GLBL81SkMYnUoxq3CAV3hHulQiRtGG0,66
+anemoi_datasets-0.4.4.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_datasets-0.4.4.dist-info/RECORD,,

{anemoi_datasets-0.4.3.dist-info → anemoi_datasets-0.4.4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (71.1.0)
+Generator: setuptools (72.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any