anemoi-datasets 0.4.4__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/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

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

@@ -14,34 +14,32 @@ from functools import cached_property
 import numpy as np
 from earthkit.data.utils.array import ensure_backend
 
-from anemoi.datasets.create.functions.sources.xarray.metadata import MDMapping
-
 from .field import XArrayField
 
 LOG = logging.getLogger(__name__)
 
 
 class Variable:
-    def __init__(self, *, ds, var, coordinates, grid, time, metadata, mapping=None, array_backend=None):
+    def __init__(
+        self,
+        *,
+        ds,
+        var,
+        coordinates,
+        grid,
+        time,
+        metadata,
+        array_backend=None,
+    ):
         self.ds = ds
         self.var = var
 
         self.grid = grid
         self.coordinates = coordinates
 
-        # print("Variable", var.name)
-        # for c in coordinates:
-        #     print(" ", c)
-
         self._metadata = metadata.copy()
-        # self._metadata.update(var.attrs)
         self._metadata.update({"variable": var.name})
 
-        # self._metadata.setdefault("level", None)
-        # self._metadata.setdefault("number", 0)
-        # self._metadata.setdefault("levtype", "sfc")
-        self._mapping = mapping
-
         self.time = time
 
         self.shape = tuple(len(c.variable) for c in coordinates if c.is_dim and not c.scalar and not c.is_grid)
@@ -51,23 +49,6 @@ class Variable:
         self.length = math.prod(self.shape)
         self.array_backend = ensure_backend(array_backend)
 
-    def update_metadata_mapping(self, kwargs):
-
-        result = {}
-
-        for k, v in kwargs.items():
-            if k == "param":
-                result[k] = "variable"
-                continue
-
-            for c in self.coordinates:
-                if k in c.mars_names:
-                    for v in c.mars_names:
-                        result[v] = c.variable.name
-                    break
-
-        self._mapping = MDMapping(result)
-
     @property
     def name(self):
         return self.var.name
@@ -111,17 +92,11 @@ class Variable:
         kwargs = {k: v for k, v in zip(self.names, coords)}
         return XArrayField(self, self.var.isel(kwargs))
 
-    @property
-    def mapping(self):
-        return self._mapping
-
     def sel(self, missing, **kwargs):
 
         if not kwargs:
             return self
 
-        kwargs = self._mapping.from_user(kwargs)
-
         k, v = kwargs.popitem()
 
         c = self.by_name.get(k)
@@ -147,13 +122,15 @@ class Variable:
             grid=self.grid,
             time=self.time,
             metadata=metadata,
-            mapping=self.mapping,
         )
 
         return variable.sel(missing, **kwargs)
 
     def match(self, **kwargs):
-        kwargs = self._mapping.from_user(kwargs)
+
+        if "param" in kwargs:
+            assert "variable" not in kwargs
+            kwargs["variable"] = kwargs.pop("param")
 
         if "variable" in kwargs:
             name = kwargs.pop("variable")

anemoi/datasets/create/input.py

@@ -106,30 +106,32 @@ def _data_request(data):
     area = grid = None
 
     for field in data:
-        if not hasattr(field, "as_mars"):
-            continue
-
-        if date is None:
-            date = field.datetime()["valid_time"]
-
-        if field.datetime()["valid_time"] != date:
-            continue
+        try:
+            if date is None:
+                date = field.datetime()["valid_time"]
 
-        as_mars = field.metadata(namespace="mars")
-        step = as_mars.get("step")
-        levtype = as_mars.get("levtype", "sfc")
-        param = as_mars["param"]
-        levelist = as_mars.get("levelist", None)
-        area = field.mars_area
-        grid = field.mars_grid
+            if field.datetime()["valid_time"] != date:
+                continue
 
-        if levelist is None:
-            params_levels[levtype].add(param)
-        else:
-            params_levels[levtype].add((param, levelist))
+            as_mars = field.metadata(namespace="mars")
+            if not as_mars:
+                continue
+            step = as_mars.get("step")
+            levtype = as_mars.get("levtype", "sfc")
+            param = as_mars["param"]
+            levelist = as_mars.get("levelist", None)
+            area = field.mars_area
+            grid = field.mars_grid
+
+            if levelist is None:
+                params_levels[levtype].add(param)
+            else:
+                params_levels[levtype].add((param, levelist))
 
-        if step:
-            params_steps[levtype].add((param, step))
+            if step:
+                params_steps[levtype].add((param, step))
+        except Exception:
+            LOG.error(f"Error in retrieving metadata (cannot build data request info) for {field}", exc_info=True)
 
     def sort(old_dic):
         new_dic = {}
@@ -288,7 +290,6 @@ class Result:
                 names += list(a.keys())
 
         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, progress_bar=False)