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

anemoi/datasets/create/check.py

@@ -12,6 +12,7 @@ import re
 import warnings
 
 import numpy as np
+from anemoi.utils.dates import frequency_to_string
 
 LOG = logging.getLogger(__name__)
 
@@ -56,10 +57,11 @@ class DatasetName:
             raise ValueError(self.error_message)
 
     def _parse(self, name):
-        pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h)-v(\d+)-?([a-zA-Z0-9-]+)$"
+        pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h)-v(\d+)-?([a-zA-Z0-9-]+)?$"
         match = re.match(pattern, name)
 
-        assert match, (name, pattern)
+        if not match:
+            raise ValueError(f"the dataset name '{name}' does not follow naming convention. Does not match {pattern}")
 
         parsed = {}
         if match:
@@ -105,7 +107,7 @@ class DatasetName:
     def check_frequency(self, frequency):
         if frequency is None:
             return
-        frequency_str = f"{frequency}h"
+        frequency_str = frequency_to_string(frequency)
         self._check_missing("frequency", frequency_str)
         self._check_mismatch("frequency", frequency_str)
 

anemoi/datasets/create/config.py

@@ -16,6 +16,8 @@ from anemoi.utils.config import DotDict
 from anemoi.utils.config import load_any_dict_format
 from earthkit.data.core.order import normalize_order_by
 
+from anemoi.datasets.dates.groups import Groups
+
 LOG = logging.getLogger(__name__)
 
 
@@ -153,6 +155,8 @@ class LoadersConfig(Config):
             raise ValueError("statistics_end is not supported anymore. Use 'statistics:end:' instead")
 
         self.setdefault("statistics", Config())
+        if "allow_nans" not in self.statistics:
+            self.statistics.allow_nans = []
 
         check_dict_value_and_set(self.output, "flatten_grid", True)
         check_dict_value_and_set(self.output, "ensemble_dimension", 2)
@@ -207,8 +211,49 @@ def _prepare_serialisation(o):
     return str(o)
 
 
-def loader_config(config):
+def set_to_test_mode(cfg):
+    NUMBER_OF_DATES = 4
+
+    dates = cfg["dates"]
+    LOG.warn(f"Running in test mode. Changing the list of dates to use only {NUMBER_OF_DATES}.")
+    groups = Groups(**LoadersConfig(cfg).dates)
+    dates = groups.dates
+    cfg["dates"] = dict(
+        start=dates[0],
+        end=dates[NUMBER_OF_DATES - 1],
+        frequency=dates.frequency,
+        group_by=NUMBER_OF_DATES,
+    )
+
+    def set_element_to_test(obj):
+        if isinstance(obj, (list, tuple)):
+            for v in obj:
+                set_element_to_test(v)
+            return
+        if isinstance(obj, (dict, DotDict)):
+            if "grid" in obj:
+                previous = obj["grid"]
+                obj["grid"] = "20./20."
+                LOG.warn(f"Running in test mode. Setting grid to {obj['grid']} instead of {previous}")
+            if "number" in obj:
+                if isinstance(obj["number"], (list, tuple)):
+                    previous = obj["number"]
+                    obj["number"] = previous[0:3]
+                    LOG.warn(f"Running in test mode. Setting number to {obj['number']} instead of {previous}")
+            for k, v in obj.items():
+                set_element_to_test(v)
+            if "constants" in obj:
+                constants = obj["constants"]
+                if "param" in constants and isinstance(constants["param"], list):
+                    constants["param"] = ["cos_latitude"]
+
+    set_element_to_test(cfg)
+
+
+def loader_config(config, is_test=False):
     config = Config(config)
+    if is_test:
+        set_to_test_mode(config)
     obj = LoadersConfig(config)
 
     # yaml round trip to check that serialisation works as expected
@@ -216,7 +261,13 @@ def loader_config(config):
     copy = yaml.load(yaml.dump(copy), Loader=yaml.SafeLoader)
     copy = Config(copy)
     copy = LoadersConfig(config)
-    assert yaml.dump(obj) == yaml.dump(copy), (obj, copy)
+
+    a = yaml.dump(obj)
+    b = yaml.dump(copy)
+    if a != b:
+        print(a)
+        print(b)
+        raise ValueError("Serialisation failed")
 
     return copy
 

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/functions/sources/xarray/__init__.py

@@ -52,9 +52,19 @@ def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs)
     result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
 
     if len(result) == 0:
-        LOG.warning(f"No data found for {dataset} and dates {dates}")
+        LOG.warning(f"No data found for {dataset} and dates {dates} and {kwargs}")
         LOG.warning(f"Options: {options}")
-        LOG.warning(data)
+
+        for i, k in enumerate(fs):
+            a = ["valid_datetime", k.metadata("valid_datetime", default=None)]
+            for n in kwargs.keys():
+                a.extend([n, k.metadata(n, default=None)])
+            print([str(x) for x in a])
+
+            if i > 16:
+                break
+
+        # LOG.warning(data)
 
     return result
 

anemoi/datasets/create/functions/sources/xarray/coordinates.py

@@ -55,6 +55,7 @@ class Coordinate:
     is_time = False
     is_step = False
     is_date = False
+    is_member = False
 
     def __init__(self, variable):
         self.variable = variable
@@ -201,8 +202,14 @@ class LevelCoordinate(Coordinate):
 
 
 class EnsembleCoordinate(Coordinate):
+    is_member = True
     mars_names = ("number",)
 
+    def normalise(self, value):
+        if int(value) == value:
+            return int(value)
+        return value
+
 
 class LongitudeCoordinate(Coordinate):
     is_grid = True

anemoi/datasets/create/functions/sources/xarray/field.py

@@ -80,7 +80,7 @@ class XArrayField(Field):
         return values.reshape(self.shape)
 
     def _make_metadata(self):
-        return XArrayMetadata(self, self.owner.mapping)
+        return XArrayMetadata(self)
 
     def grid_points(self):
         return self.owner.grid_points()

anemoi/datasets/create/functions/sources/xarray/fieldlist.py

@@ -134,8 +134,6 @@ class XarrayFieldList(FieldList):
 
         for v in self.variables:
 
-            v.update_metadata_mapping(kwargs)
-
             # First, select matching variables
             # This will consume 'param' or 'variable' from kwargs
             # and return the rest

anemoi/datasets/create/functions/sources/xarray/flavour.py

@@ -9,6 +9,7 @@
 
 
 from .coordinates import DateCoordinate
+from .coordinates import EnsembleCoordinate
 from .coordinates import LatitudeCoordinate
 from .coordinates import LevelCoordinate
 from .coordinates import LongitudeCoordinate
@@ -135,6 +136,17 @@ class CoordinateGuesser:
         if d is not None:
             return d
 
+        d = self._is_number(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
         if c.shape in ((1,), tuple()):
             return ScalarCoordinate(c)
 
@@ -249,9 +261,13 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
         if standard_name == "depth":
             return LevelCoordinate(c, "depth")
 
-        if name == "pressure":
+        if name == "vertical" and units == "hPa":
             return LevelCoordinate(c, "pl")
 
+    def _is_number(self, c, *, axis, name, long_name, standard_name, units):
+        if name in ("realization", "number"):
+            return EnsembleCoordinate(c)
+
 
 class FlavourCoordinateGuesser(CoordinateGuesser):
     def __init__(self, ds, flavour):
@@ -328,3 +344,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
             return self.flavour["levtype"]
 
         raise NotImplementedError(f"levtype for {c=}")
+
+    def _is_number(self, c, *, axis, name, long_name, standard_name, units):
+        if self._match(c, "number", locals()):
+            return DateCoordinate(c)

anemoi/datasets/create/functions/sources/xarray/metadata.py

@@ -10,29 +10,37 @@
 import logging
 from functools import cached_property
 
+from anemoi.utils.dates import as_datetime
 from earthkit.data.core.geography import Geography
 from earthkit.data.core.metadata import RawMetadata
-from earthkit.data.utils.dates import to_datetime
 from earthkit.data.utils.projections import Projection
 
 LOG = logging.getLogger(__name__)
 
 
-class MDMapping:
+class _MDMapping:
 
-    def __init__(self, mapping):
-        self.user_to_internal = mapping
+    def __init__(self, variable):
+        self.variable = variable
+        self.time = variable.time
+        self.mapping = dict(param="variable")
+        for c in variable.coordinates:
+            for v in c.mars_names:
+                assert v not in self.mapping, f"Duplicate key '{v}' in {c}"
+                self.mapping[v] = c.variable.name
 
-    def from_user(self, kwargs):
-        if isinstance(kwargs, str):
-            return self.user_to_internal.get(kwargs, kwargs)
-        return {self.user_to_internal.get(k, k): v for k, v in kwargs.items()}
+    def _from_user(self, key):
+        return self.mapping.get(key, key)
 
-    def __len__(self):
-        return len(self.user_to_internal)
+    def from_user(self, kwargs):
+        print("from_user", kwargs, self)
+        return {self._from_user(k): v for k, v in kwargs.items()}
 
     def __repr__(self):
-        return f"MDMapping({self.user_to_internal})"
+        return f"MDMapping({self.mapping})"
+
+    def fill_time_metadata(self, field, md):
+        md["valid_datetime"] = as_datetime(self.variable.time.fill_time_metadata(field._md, md)).isoformat()
 
 
 class XArrayMetadata(RawMetadata):
@@ -40,23 +48,11 @@ class XArrayMetadata(RawMetadata):
     NAMESPACES = ["default", "mars"]
     MARS_KEYS = ["param", "step", "levelist", "levtype", "number", "date", "time"]
 
-    def __init__(self, field, mapping):
+    def __init__(self, field):
         self._field = field
         md = field._md.copy()
-
-        self._mapping = mapping
-        if mapping is None:
-            time_coord = [c for c in field.owner.coordinates if c.is_time]
-            if len(time_coord) == 1:
-                time_key = time_coord[0].name
-            else:
-                time_key = "time"
-        else:
-            time_key = mapping.from_user("valid_datetime")
-        self._time = to_datetime(md.pop(time_key))
-        self._field.owner.time.fill_time_metadata(self._time, md)
-        md["valid_datetime"] = self._time.isoformat()
-
+        self._mapping = _MDMapping(field.owner)
+        self._mapping.fill_time_metadata(field, md)
         super().__init__(md)
 
     @cached_property
@@ -88,10 +84,13 @@ class XArrayMetadata(RawMetadata):
         return self._field.forecast_reference_time
 
     def _valid_datetime(self):
-        return self._time
+        return self._get("valid_datetime")
 
     def _get(self, key, **kwargs):
 
+        if key in self._d:
+            return self._d[key]
+
         if key.startswith("mars."):
             key = key[5:]
             if key not in self.MARS_KEYS:
@@ -100,8 +99,7 @@ class XArrayMetadata(RawMetadata):
                 else:
                     return kwargs.get("default", None)
 
-        if self._mapping is not None:
-            key = self._mapping.from_user(key)
+        key = self._mapping._from_user(key)
 
         return super()._get(key, **kwargs)
 

anemoi/datasets/create/functions/sources/xarray/time.py

@@ -10,8 +10,11 @@
 
 import datetime
 
+from anemoi.utils.dates import as_datetime
+
 
 class Time:
+
     @classmethod
     def from_coordinates(cls, coordinates):
         time_coordinate = [c for c in coordinates if c.is_time]
@@ -19,16 +22,16 @@ class Time:
         date_coordinate = [c for c in coordinates if c.is_date]
 
         if len(date_coordinate) == 0 and len(time_coordinate) == 1 and len(step_coordinate) == 1:
-            return ForecasstFromValidTimeAndStep(step_coordinate[0])
+            return ForecastFromValidTimeAndStep(time_coordinate[0], step_coordinate[0])
 
         if len(date_coordinate) == 0 and len(time_coordinate) == 1 and len(step_coordinate) == 0:
-            return Analysis()
+            return Analysis(time_coordinate[0])
 
         if len(date_coordinate) == 0 and len(time_coordinate) == 0 and len(step_coordinate) == 0:
             return Constant()
 
         if len(date_coordinate) == 1 and len(time_coordinate) == 1 and len(step_coordinate) == 0:
-            return ForecastFromValidTimeAndBaseTime(date_coordinate[0])
+            return ForecastFromValidTimeAndBaseTime(date_coordinate[0], time_coordinate[0])
 
         if len(date_coordinate) == 1 and len(time_coordinate) == 0 and len(step_coordinate) == 1:
             return ForecastFromBaseTimeAndDate(date_coordinate[0], step_coordinate[0])
@@ -38,61 +41,91 @@ class Time:
 
 class Constant(Time):
 
-    def fill_time_metadata(self, time, metadata):
-        metadata["date"] = time.strftime("%Y%m%d")
-        metadata["time"] = time.strftime("%H%M")
-        metadata["step"] = 0
+    def fill_time_metadata(self, coords_values, metadata):
+        raise NotImplementedError("Constant time not implemented")
+        # print("Constant", coords_values, metadata)
+        # metadata["date"] = time.strftime("%Y%m%d")
+        # metadata["time"] = time.strftime("%H%M")
+        # metadata["step"] = 0
 
 
 class Analysis(Time):
 
-    def fill_time_metadata(self, time, metadata):
-        metadata["date"] = time.strftime("%Y%m%d")
-        metadata["time"] = time.strftime("%H%M")
+    def __init__(self, time_coordinate):
+        self.time_coordinate_name = time_coordinate.variable.name
+
+    def fill_time_metadata(self, coords_values, metadata):
+        valid_datetime = coords_values[self.time_coordinate_name]
+
+        metadata["date"] = as_datetime(valid_datetime).strftime("%Y%m%d")
+        metadata["time"] = as_datetime(valid_datetime).strftime("%H%M")
         metadata["step"] = 0
 
+        return valid_datetime
 
-class ForecasstFromValidTimeAndStep(Time):
-    def __init__(self, step_coordinate):
-        self.step_name = step_coordinate.variable.name
 
-    def fill_time_metadata(self, time, metadata):
-        step = metadata.pop(self.step_name)
+class ForecastFromValidTimeAndStep(Time):
+
+    def __init__(self, time_coordinate, step_coordinate):
+        self.time_coordinate_name = time_coordinate.variable.name
+        self.step_coordinate_name = step_coordinate.variable.name
+
+    def fill_time_metadata(self, coords_values, metadata):
+        valid_datetime = coords_values[self.time_coordinate_name]
+        step = coords_values[self.step_coordinate_name]
+
         assert isinstance(step, datetime.timedelta)
-        base = time - step
+        base_datetime = valid_datetime - step
 
         hours = step.total_seconds() / 3600
         assert int(hours) == hours
 
-        metadata["date"] = base.strftime("%Y%m%d")
-        metadata["time"] = base.strftime("%H%M")
+        metadata["date"] = as_datetime(base_datetime).strftime("%Y%m%d")
+        metadata["time"] = as_datetime(base_datetime).strftime("%H%M")
         metadata["step"] = int(hours)
+        return valid_datetime
 
 
 class ForecastFromValidTimeAndBaseTime(Time):
-    def __init__(self, date_coordinate):
-        self.date_coordinate = date_coordinate
 
-    def fill_time_metadata(self, time, metadata):
+    def __init__(self, date_coordinate, time_coordinate):
+        self.date_coordinate.name = date_coordinate.name
+        self.time_coordinate.name = time_coordinate.name
+
+    def fill_time_metadata(self, coords_values, metadata):
+        valid_datetime = coords_values[self.time_coordinate_name]
+        base_datetime = coords_values[self.date_coordinate_name]
 
-        step = time - self.date_coordinate
+        step = valid_datetime - base_datetime
 
         hours = step.total_seconds() / 3600
         assert int(hours) == hours
 
-        metadata["date"] = self.date_coordinate.single_value.strftime("%Y%m%d")
-        metadata["time"] = self.date_coordinate.single_value.strftime("%H%M")
+        metadata["date"] = as_datetime(base_datetime).strftime("%Y%m%d")
+        metadata["time"] = as_datetime(base_datetime).strftime("%H%M")
         metadata["step"] = int(hours)
 
+        return valid_datetime
+
 
 class ForecastFromBaseTimeAndDate(Time):
+
     def __init__(self, date_coordinate, step_coordinate):
-        self.date_coordinate = date_coordinate
-        self.step_coordinate = step_coordinate
+        self.date_coordinate_name = date_coordinate.name
+        self.step_coordinate_name = step_coordinate.name
+
+    def fill_time_metadata(self, coords_values, metadata):
+
+        date = coords_values[self.date_coordinate_name]
+        step = coords_values[self.step_coordinate_name]
+        assert isinstance(step, datetime.timedelta)
+
+        metadata["date"] = as_datetime(date).strftime("%Y%m%d")
+        metadata["time"] = as_datetime(date).strftime("%H%M")
+
+        hours = step.total_seconds() / 3600
 
-    def fill_time_metadata(self, time, metadata):
-        metadata["date"] = time.strftime("%Y%m%d")
-        metadata["time"] = time.strftime("%H%M")
-        hours = metadata[self.step_coordinate.name].total_seconds() / 3600
         assert int(hours) == hours
         metadata["step"] = int(hours)
+
+        return date + step