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

anemoi/datasets/create/input/step.py

@@ -0,0 +1,99 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 copy import deepcopy
+
+from anemoi.utils.dates import as_datetime as as_datetime
+from anemoi.utils.dates import frequency_to_timedelta as frequency_to_timedelta
+
+from anemoi.datasets.dates import DatesProvider as DatesProvider
+from anemoi.datasets.fields import FieldArray as FieldArray
+from anemoi.datasets.fields import NewValidDateTimeField as NewValidDateTimeField
+
+from .action import Action
+from .context import Context
+from .misc import is_function
+from .result import Result
+from .template import notify_result
+from .trace import trace_datasource
+from .trace import trace_select
+
+LOG = logging.getLogger(__name__)
+
+
+class StepResult(Result):
+    def __init__(self, context, action_path, group_of_dates, action, upstream_result):
+        super().__init__(context, action_path, group_of_dates)
+        assert isinstance(upstream_result, Result), type(upstream_result)
+        self.upstream_result = upstream_result
+        self.action = action
+
+    @property
+    @notify_result
+    @trace_datasource
+    def datasource(self):
+        raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
+
+
+class StepAction(Action):
+    result_class = None
+
+    def __init__(self, context, action_path, previous_step, *args, **kwargs):
+        super().__init__(context, action_path, *args, **kwargs)
+        self.previous_step = previous_step
+
+    @trace_select
+    def select(self, group_of_dates):
+        return self.result_class(
+            self.context,
+            self.action_path,
+            group_of_dates,
+            self,
+            self.previous_step.select(group_of_dates),
+        )
+
+    def __repr__(self):
+        return super().__repr__(self.previous_step, _inline_=str(self.kwargs))
+
+
+def step_factory(config, context, action_path, previous_step):
+
+    from .filter import FilterStepAction
+    from .filter import FunctionStepAction
+
+    assert isinstance(context, Context), (type, context)
+    if not isinstance(config, dict):
+        raise ValueError(f"Invalid input config {config}")
+
+    config = deepcopy(config)
+    assert len(config) == 1, config
+
+    key = list(config.keys())[0]
+    cls = dict(
+        filter=FilterStepAction,
+        # rename=RenameAction,
+        # remapping=RemappingAction,
+    ).get(key)
+
+    if isinstance(config[key], list):
+        args, kwargs = config[key], {}
+
+    if isinstance(config[key], dict):
+        args, kwargs = [], config[key]
+
+    if isinstance(config[key], str):
+        args, kwargs = [config[key]], {}
+
+    if cls is None:
+        if not is_function(key, "filters"):
+            raise ValueError(f"Unknown step {key}")
+        cls = FunctionStepAction
+        args = [key] + args
+
+    return cls(context, action_path, previous_step, *args, **kwargs)

anemoi/datasets/create/{template.py → input/template.py}

@@ -9,14 +9,8 @@
 
 import logging
 import re
-import textwrap
 from functools import wraps
 
-from anemoi.utils.humanize import plural
-
-from .trace import step
-from .trace import trace
-
 LOG = logging.getLogger(__name__)
 
 
@@ -30,42 +24,6 @@ def notify_result(method):
     return wrapper
 
 
-class Context:
-    def __init__(self):
-        # used_references is a set of reference paths that will be needed
-        self.used_references = set()
-        # results is a dictionary of reference path -> obj
-        self.results = {}
-
-    def will_need_reference(self, key):
-        assert isinstance(key, (list, tuple)), key
-        key = tuple(key)
-        self.used_references.add(key)
-
-    def notify_result(self, key, result):
-        trace(
-            "🎯",
-            step(key),
-            "notify result",
-            textwrap.shorten(repr(result).replace(",", ", "), width=40),
-            plural(len(result), "field"),
-        )
-        assert isinstance(key, (list, tuple)), key
-        key = tuple(key)
-        if key in self.used_references:
-            if key in self.results:
-                raise ValueError(f"Duplicate result {key}")
-            self.results[key] = result
-
-    def get_result(self, key):
-        assert isinstance(key, (list, tuple)), key
-        key = tuple(key)
-        if key in self.results:
-            return self.results[key]
-        all_keys = sorted(list(self.results.keys()))
-        raise ValueError(f"Cannot find result {key} in {all_keys}")
-
-
 class Substitution:
     pass
 

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/statistics/__init__.py

@@ -155,7 +155,7 @@ def compute_statistics(array, check_variables_names=None, allow_nans=False):
             check_data_values(values[j, :], name=name, allow_nans=allow_nans)
             if np.isnan(values[j, :]).all():
                 # LOG.warning(f"All NaN values for {name} ({j}) for date {i}")
-                raise ValueError(f"All NaN values for {name} ({j}) for date {i}")
+                LOG.warning(f"All NaN values for {name} ({j}) for date {i}")
 
         # Ignore NaN values
         minimum[i] = np.nanmin(values, axis=1)

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/create/zarr.py

@@ -128,7 +128,7 @@ class ZarrBuiltRegistry:
     def add_to_history(self, action, **kwargs):
         new = dict(
             action=action,
-            timestamp=datetime.datetime.utcnow().isoformat(),
+            timestamp=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None).isoformat(),
         )
         new.update(kwargs)
 
@@ -151,7 +151,9 @@ class ZarrBuiltRegistry:
 
     def set_flag(self, i, value=True):
         z = self._open_write()
-        z.attrs["latest_write_timestamp"] = datetime.datetime.utcnow().isoformat()
+        z.attrs["latest_write_timestamp"] = (
+            datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None).isoformat()
+        )
         z["_build"][self.name_flags][i] = value
 
     def ready(self):

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)