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

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

@@ -6,7 +6,6 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 #
-import datetime
 import logging
 
 from earthkit.data.core.fieldlist import MultiFieldList
@@ -14,7 +13,6 @@ from earthkit.data.core.fieldlist import MultiFieldList
 from anemoi.datasets.create.functions.sources.mars import mars
 
 LOGGER = logging.getLogger(__name__)
-DEBUG = True
 
 
 def _to_list(x):
@@ -23,91 +21,34 @@ def _to_list(x):
     return [x]
 
 
-class HindcastCompute:
-    def __init__(self, base_times, available_steps, request):
-        self.base_times = base_times
-        self.available_steps = available_steps
-        self.request = request
-
-    def compute_hindcast(self, date):
-        result = []
-        for step in sorted(self.available_steps):  # Use the shortest step
-            start_date = date - datetime.timedelta(hours=step)
-            hours = start_date.hour
-            if hours in self.base_times:
-                r = self.request.copy()
-                r["date"] = start_date
-                r["time"] = f"{start_date.hour:02d}00"
-                r["step"] = step
-                result.append(r)
-
-        if not result:
-            raise ValueError(
-                f"Cannot find data for {self.request} for {date} (base_times={self.base_times}, "
-                f"available_steps={self.available_steps})"
-            )
-
-        if len(result) > 1:
-            raise ValueError(
-                f"Multiple requests for {self.request} for {date} (base_times={self.base_times}, "
-                f"available_steps={self.available_steps})"
-            )
-
-        return result[0]
-
-
-def use_reference_year(reference_year, request):
-    request = request.copy()
-    hdate = request.pop("date")
-
-    if hdate.year >= reference_year:
-        return None, False
+def hindcasts(context, dates, **request):
 
-    try:
-        date = datetime.datetime(reference_year, hdate.month, hdate.day)
-    except ValueError:
-        if hdate.month == 2 and hdate.day == 29:
-            return None, False
-        raise
+    from anemoi.datasets.dates import HindcastsDates
 
-    request.update(date=date.strftime("%Y-%m-%d"), hdate=hdate.strftime("%Y-%m-%d"))
-    return request, True
+    provider = context.dates_provider
+    assert isinstance(provider, HindcastsDates)
 
+    context.trace("H️", f"hindcasts {len(dates)=}")
 
-def hindcasts(context, dates, **request):
     request["param"] = _to_list(request["param"])
-    request["step"] = _to_list(request["step"])
+    request["step"] = _to_list(request.get("step", 0))
     request["step"] = [int(_) for _ in request["step"]]
 
-    if request.get("stream") == "enfh" and "base_times" not in request:
-        request["base_times"] = [0]
-
-    available_steps = request.pop("step")
-    available_steps = _to_list(available_steps)
-
-    base_times = request.pop("base_times")
-
-    reference_year = request.pop("reference_year")
+    context.trace("H️", f"hindcast {request}")
 
-    context.trace("H️", f"hindcast {request} {base_times} {available_steps} {reference_year}")
-
-    c = HindcastCompute(base_times, available_steps, request)
     requests = []
     for d in dates:
-        req = c.compute_hindcast(d)
-        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))
+        r = request.copy()
+        hindcast = provider.mapping[d]
+        r["hdate"] = hindcast.hdate.strftime("%Y-%m-%d")
+        r["date"] = hindcast.refdate.strftime("%Y-%m-%d")
+        r["time"] = hindcast.refdate.strftime("%H")
+        r["step"] = hindcast.step
+        requests.append(r)
 
     if len(requests) == 0:
-        # print("HINDCASTS no requests")
         return MultiFieldList([])
 
-    # print("HINDCASTS requests", requests)
-
     return mars(
         context,
         dates,

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

@@ -203,16 +203,22 @@ def mars(context, dates, *requests, request_already_using_valid_datetime=False,
         request_already_using_valid_datetime=request_already_using_valid_datetime,
         date_key=date_key,
     )
+
+    requests = list(requests)
+
     ds = from_source("empty")
+    context.trace("✅", f"{[str(d) for d in dates]}")
+    context.trace("✅", f"Will run {len(requests)} requests")
+    for r in requests:
+        r = {k: v for k, v in r.items() if v != ("-",)}
+        context.trace("✅", f"mars {r}")
+
     for r in requests:
         r = {k: v for k, v in r.items() if v != ("-",)}
 
         if context.use_grib_paramid and "param" in r:
             r = use_grib_paramid(r)
 
-        if DEBUG:
-            context.trace("✅", f"from_source(mars, {r}")
-
         for k, v in r.items():
             if k not in MARS_KEYS:
                 raise ValueError(

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

@@ -7,6 +7,7 @@
 # nor does it submit to any jurisdiction.
 #
 
+import datetime
 import logging
 
 from earthkit.data.core.fieldlist import Field
@@ -80,7 +81,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()
@@ -103,7 +104,12 @@ class XArrayField(Field):
 
     @property
     def forecast_reference_time(self):
-        return self.owner.forecast_reference_time
+        date, time = self.metadata("date", "time")
+        assert len(time) == 4, time
+        assert len(date) == 8, date
+        yyyymmdd = int(date)
+        time = int(time) // 100
+        return datetime.datetime(yyyymmdd // 10000, yyyymmdd // 100 % 100, yyyymmdd % 100, time)
 
     def __repr__(self):
         return repr(self._metadata)

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
@@ -74,24 +70,29 @@ class XArrayMetadata(RawMetadata):
             return self._as_mars()
 
     def _as_mars(self):
-        return dict(
-            param=self["variable"],
-            step=self["step"],
-            levelist=self["level"],
-            levtype=self["levtype"],
-            number=self["number"],
-            date=self["date"],
-            time=self["time"],
-        )
+        return {}
+        # p =  dict(
+        #     param=self.get("variable", self.get("param")),
+        #     step=self.get("step"),
+        #     levelist=self.get("levelist", self.get("level")),
+        #     levtype=self.get("levtype"),
+        #     number=self.get("number"),
+        #     date=self.get("date"),
+        #     time=self.get("time"),
+        # )
+        # return {k: v for k, v in p.items() if v is not None}
 
     def _base_datetime(self):
         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 +101,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)
 
@@ -137,12 +137,12 @@ class XArrayFieldGeography(Geography):
         # TODO: implement resolution
         return None
 
-    @property
+    # @property
     def mars_grid(self):
         # TODO: implement mars_grid
         return None
 
-    @property
+    # @property
     def mars_area(self):
         # TODO: code me
         # return [self.north, self.west, self.south, self.east]

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")