anemoi-datasets 0.4.5__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/field.py

@@ -7,6 +7,7 @@
 # nor does it submit to any jurisdiction.
 #
 
+import datetime
 import logging
 
 from earthkit.data.core.fieldlist import Field
@@ -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/metadata.py

@@ -70,15 +70,17 @@ 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
@@ -135,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/input.py

@@ -23,7 +23,7 @@ from earthkit.data.core.fieldlist import FieldList
 from earthkit.data.core.fieldlist import MultiFieldList
 from earthkit.data.core.order import build_remapping
 
-from anemoi.datasets.dates import Dates
+from anemoi.datasets.dates import DatesProvider
 
 from .functions import import_function
 from .template import Context
@@ -75,7 +75,7 @@ def time_delta_to_string(delta):
 
 
 def is_function(name, kind):
-    name, delta = parse_function_name(name)  # noqa
+    name, _ = parse_function_name(name)
     try:
         import_function(name, kind)
         return True
@@ -204,11 +204,15 @@ class Result:
     _coords_already_built = False
 
     def __init__(self, context, action_path, dates):
+        from anemoi.datasets.dates.groups import GroupOfDates
+
+        assert isinstance(dates, GroupOfDates), dates
+
         assert isinstance(context, ActionContext), type(context)
         assert isinstance(action_path, list), action_path
 
         self.context = context
-        self.dates = dates
+        self.group_of_dates = dates
         self.action_path = action_path
 
     @property
@@ -405,10 +409,10 @@ class Result:
         more += ",".join([f"{k}={v}"[:5000] for k, v in kwargs.items()])
 
         dates = " no-dates"
-        if self.dates is not None:
-            dates = f" {len(self.dates)} dates"
+        if self.group_of_dates is not None:
+            dates = f" {len(self.group_of_dates)} dates"
             dates += " ("
-            dates += "/".join(d.strftime("%Y-%m-%d:%H") for d in self.dates)
+            dates += "/".join(d.strftime("%Y-%m-%d:%H") for d in self.group_of_dates)
             if len(dates) > 100:
                 dates = dates[:100] + "..."
             dates += ")"
@@ -423,7 +427,7 @@ class Result:
         raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
 
     def _trace_datasource(self, *args, **kwargs):
-        return f"{self.__class__.__name__}({shorten(self.dates)})"
+        return f"{self.__class__.__name__}({self.group_of_dates})"
 
     def build_coords(self):
         if self._coords_already_built:
@@ -513,7 +517,7 @@ class Result:
     @cached_property
     def shape(self):
         return [
-            len(self.dates),
+            len(self.group_of_dates),
             len(self.variables),
             len(self.ensembles),
             len(self.grid_values),
@@ -522,7 +526,7 @@ class Result:
     @cached_property
     def coords(self):
         return {
-            "dates": self.dates,
+            "dates": list(self.group_of_dates),
             "variables": self.variables,
             "ensembles": self.ensembles,
             "values": self.grid_values,
@@ -573,7 +577,7 @@ class FunctionResult(Result):
         self.args, self.kwargs = substitute(context, (self.action.args, self.action.kwargs))
 
     def _trace_datasource(self, *args, **kwargs):
-        return f"{self.action.name}({shorten(self.dates)})"
+        return f"{self.action.name}({self.group_of_dates})"
 
     @cached_property
     @assert_fieldlist
@@ -583,14 +587,21 @@ class FunctionResult(Result):
         args, kwargs = resolve(self.context, (self.args, self.kwargs))
 
         try:
-            return _tidy(self.action.function(FunctionContext(self), self.dates, *args, **kwargs))
+            return _tidy(
+                self.action.function(
+                    FunctionContext(self),
+                    list(self.group_of_dates),  # Will provide a list of datetime objects
+                    *args,
+                    **kwargs,
+                )
+            )
         except Exception:
             LOG.error(f"Error in {self.action.function.__name__}", exc_info=True)
             raise
 
     def __repr__(self):
         try:
-            return f"{self.action.name}({shorten(self.dates)})"
+            return f"{self.action.name}({self.group_of_dates})"
         except Exception:
             return f"{self.__class__.__name__}(unitialised)"
 
@@ -609,7 +620,7 @@ class JoinResult(Result):
     @notify_result
     @trace_datasource
     def datasource(self):
-        ds = EmptyResult(self.context, self.action_path, self.dates).datasource
+        ds = EmptyResult(self.context, self.action_path, self.group_of_dates).datasource
         for i in self.results:
             ds += i.datasource
         return _tidy(ds)
@@ -824,7 +835,7 @@ class ConcatResult(Result):
     @notify_result
     @trace_datasource
     def datasource(self):
-        ds = EmptyResult(self.context, self.action_path, self.dates).datasource
+        ds = EmptyResult(self.context, self.action_path, self.group_of_dates).datasource
         for i in self.results:
             ds += i.datasource
         return _tidy(ds)
@@ -904,7 +915,7 @@ class ConcatAction(Action):
             cfg = deepcopy(cfg)
             dates_cfg = cfg.pop("dates")
             assert isinstance(dates_cfg, dict), dates_cfg
-            filtering_dates = Dates.from_config(**dates_cfg)
+            filtering_dates = DatesProvider.from_config(**dates_cfg)
             action = action_factory(cfg, context, action_path + [str(i)])
             parts.append((filtering_dates, action))
         self.parts = parts
@@ -915,9 +926,11 @@ class ConcatAction(Action):
 
     @trace_select
     def select(self, dates):
+        from anemoi.datasets.dates.groups import GroupOfDates
+
         results = []
         for filtering_dates, action in self.parts:
-            newdates = sorted(set(dates) & set(filtering_dates))
+            newdates = GroupOfDates(sorted(set(dates) & set(filtering_dates)), dates.provider)
             if newdates:
                 results.append(action.select(newdates))
         if not results:
@@ -953,8 +966,10 @@ def action_factory(config, context, action_path):
 
     if isinstance(config[key], list):
         args, kwargs = config[key], {}
-    if isinstance(config[key], dict):
+    elif isinstance(config[key], dict):
         args, kwargs = [], config[key]
+    else:
+        raise ValueError(f"Invalid input config {config[key]} ({type(config[key])}")
 
     cls = {
         # "date_shift": DateShiftAction,
@@ -1021,6 +1036,13 @@ class FunctionContext:
     def trace(self, emoji, *args):
         trace(emoji, *args)
 
+    def info(self, *args, **kwargs):
+        LOG.info(*args, **kwargs)
+
+    @property
+    def dates_provider(self):
+        return self.owner.group_of_dates.provider
+
 
 class ActionContext(Context):
     def __init__(self, /, order_by, flatten_grid, remapping, use_grib_paramid):

anemoi/datasets/create/persistent.py

@@ -68,7 +68,7 @@ class PersistentDict:
         path = os.path.join(self.dirname, f"{h}.pickle")
 
         if os.path.exists(path):
-            LOG.warn(f"{path} already exists")
+            LOG.warning(f"{path} already exists")
 
         tmp_path = path + f".tmp-{os.getpid()}-on-{socket.gethostname()}"
         with open(tmp_path, "wb") as f:

anemoi/datasets/create/utils.py

@@ -62,6 +62,9 @@ def make_list_int(value):
 
 
 def normalize_and_check_dates(dates, start, end, frequency, dtype="datetime64[s]"):
+
+    dates = [d.hdate if hasattr(d, "hdate") else d for d in dates]
+
     assert isinstance(frequency, datetime.timedelta), frequency
     start = np.datetime64(start)
     end = np.datetime64(end)

anemoi/datasets/data/dataset.py

@@ -23,7 +23,11 @@ LOG = logging.getLogger(__name__)
 class Dataset:
     arguments = {}
 
-    def mutate(self):
+    def mutate(self) -> "Dataset":
+        """
+        Give an opportunity to a subclass to return a new Dataset
+        object of a different class, if needed.
+        """
         return self
 
     def swap_with_parent(self, parent):
@@ -90,6 +94,12 @@ class Dataset:
             rename = kwargs.pop("rename")
             return Rename(self, rename)._subset(**kwargs).mutate()
 
+        if "rescale" in kwargs:
+            from .rescale import Rescale
+
+            rescale = kwargs.pop("rescale")
+            return Rescale(self, rescale)._subset(**kwargs).mutate()
+
         if "statistics" in kwargs:
             from ..data import open_dataset
             from .statistics import Statistics

anemoi/datasets/data/debug.py

@@ -209,10 +209,14 @@ def _debug_indexing(method):
     return wrapper
 
 
+def _identity(x):
+    return x
+
+
 if DEBUG_ZARR_INDEXING:
     debug_indexing = _debug_indexing
 else:
-    debug_indexing = lambda x: x  # noqa
+    debug_indexing = _identity
 
 
 def debug_zarr_loading(on_off):

anemoi/datasets/data/masked.py

@@ -112,5 +112,5 @@ class Cropping(Masked):
     def tree(self):
         return Node(self, [self.forward.tree()], area=self.area)
 
-    def metadata_specific(self, **kwargs):
-        return super().metadata_specific(area=self.area, **kwargs)
+    def subclass_metadata_specific(self):
+        return dict(area=self.area)

anemoi/datasets/data/rescale.py

@@ -0,0 +1,147 @@
+# (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.
+
+import logging
+from functools import cached_property
+
+import numpy as np
+
+from .debug import Node
+from .debug import debug_indexing
+from .forwards import Forwards
+from .indexing import apply_index_to_slices_changes
+from .indexing import expand_list_indexing
+from .indexing import index_to_slices
+from .indexing import update_tuple
+
+LOG = logging.getLogger(__name__)
+
+
+def make_rescale(variable, rescale):
+
+    if isinstance(rescale, (tuple, list)):
+
+        assert len(rescale) == 2, rescale
+
+        if isinstance(rescale[0], (int, float)):
+            return rescale
+
+        from cfunits import Units
+
+        u0 = Units(rescale[0])
+        u1 = Units(rescale[1])
+
+        x1, x2 = 0.0, 1.0
+        y1, y2 = Units.conform([x1, x2], u0, u1)
+
+        a = (y2 - y1) / (x2 - x1)
+        b = y1 - a * x1
+
+        return a, b
+
+        return rescale
+
+    if isinstance(rescale, dict):
+        assert "scale" in rescale, rescale
+        assert "offset" in rescale, rescale
+        return rescale["scale"], rescale["offset"]
+
+    assert False
+
+
+class Rescale(Forwards):
+    def __init__(self, dataset, rescale):
+        super().__init__(dataset)
+        for n in rescale:
+            assert n in dataset.variables, n
+
+        variables = dataset.variables
+
+        self._a = np.ones(len(variables))
+        self._b = np.zeros(len(variables))
+
+        self.rescale = {}
+        for i, v in enumerate(variables):
+            if v in rescale:
+                a, b = make_rescale(v, rescale[v])
+                self.rescale[v] = a, b
+                self._a[i], self._b[i] = a, b
+
+        self._a = self._a[np.newaxis, :, np.newaxis, np.newaxis]
+        self._b = self._b[np.newaxis, :, np.newaxis, np.newaxis]
+
+        self._a = self._a.astype(self.forward.dtype)
+        self._b = self._b.astype(self.forward.dtype)
+
+    def tree(self):
+        return Node(self, [self.forward.tree()], rescale=self.rescale)
+
+    def subclass_metadata_specific(self):
+        return dict(rescale=self.rescale)
+
+    @debug_indexing
+    @expand_list_indexing
+    def _get_tuple(self, index):
+        index, changes = index_to_slices(index, self.shape)
+        index, previous = update_tuple(index, 1, slice(None))
+        result = self.forward[index]
+        result = result * self._a + self._b
+        result = result[:, previous]
+        result = apply_index_to_slices_changes(result, changes)
+        return result
+
+    @debug_indexing
+    def __get_slice_(self, n):
+        data = self.forward[n]
+        return data * self._a + self._b
+
+    @debug_indexing
+    def __getitem__(self, n):
+
+        if isinstance(n, tuple):
+            return self._get_tuple(n)
+
+        if isinstance(n, slice):
+            return self.__get_slice_(n)
+
+        data = self.forward[n]
+
+        return data * self._a[0] + self._b[0]
+
+    @cached_property
+    def statistics(self):
+        result = {}
+        a = self._a.squeeze()
+        assert np.all(a >= 0)
+
+        b = self._b.squeeze()
+        for k, v in self.forward.statistics.items():
+            if k in ("maximum", "minimum", "mean"):
+                result[k] = v * a + b
+                continue
+
+            if k in ("stdev",):
+                result[k] = v * a
+                continue
+
+            raise NotImplementedError("rescale statistics", k)
+
+        return result
+
+    def statistics_tendencies(self, delta=None):
+        result = {}
+        a = self._a.squeeze()
+        assert np.all(a >= 0)
+
+        for k, v in self.forward.statistics_tendencies(delta).items():
+            if k in ("maximum", "minimum", "mean", "stdev"):
+                result[k] = v * a
+                continue
+
+            raise NotImplementedError("rescale tendencies statistics", k)
+
+        return result

anemoi/datasets/data/stores.py

@@ -5,6 +5,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+
 import logging
 import os
 import warnings
@@ -83,6 +84,8 @@ class S3Store(ReadOnlyStore):
 
 
 class DebugStore(ReadOnlyStore):
+    """A store to debug the zarr loading."""
+
     def __init__(self, store):
         assert not isinstance(store, DebugStore)
         self.store = store
@@ -148,6 +151,8 @@ def open_zarr(path, dont_fail=False, cache=None):
 
 
 class Zarr(Dataset):
+    """A zarr dataset."""
+
     def __init__(self, path):
         if isinstance(path, zarr.hierarchy.Group):
             self.was_zarr = True
@@ -244,14 +249,20 @@ class Zarr(Dataset):
             delta = self.frequency
         if isinstance(delta, int):
             delta = f"{delta}h"
-        from anemoi.datasets.create.loaders import TendenciesStatisticsAddition
+        from anemoi.utils.dates import frequency_to_string
+        from anemoi.utils.dates import frequency_to_timedelta
+
+        delta = frequency_to_timedelta(delta)
+        delta = frequency_to_string(delta)
+
+        def func(k):
+            return f"statistics_tendencies_{delta}_{k}"
 
-        func = TendenciesStatisticsAddition.final_storage_name_from_delta
         return dict(
-            mean=self.z[func("mean", delta)][:],
-            stdev=self.z[func("stdev", delta)][:],
-            maximum=self.z[func("maximum", delta)][:],
-            minimum=self.z[func("minimum", delta)][:],
+            mean=self.z[func("mean")][:],
+            stdev=self.z[func("stdev")][:],
+            maximum=self.z[func("maximum")][:],
+            minimum=self.z[func("minimum")][:],
         )
 
     @property
@@ -322,11 +333,13 @@ class Zarr(Dataset):
 
 
 class ZarrWithMissingDates(Zarr):
+    """A zarr dataset with missing dates."""
+
     def __init__(self, path):
         super().__init__(path)
 
         missing_dates = self.z.attrs.get("missing_dates", [])
-        missing_dates = [np.datetime64(x) for x in missing_dates]
+        missing_dates = set([np.datetime64(x) for x in missing_dates])
         self.missing_to_dates = {i: d for i, d in enumerate(self.dates) if d in missing_dates}
         self.missing = set(self.missing_to_dates)