Merge branch 'refactor-tendencies' of github.com:ecmwf-lab/ecml-tools…

… into refactor-tendencies

ecml_tools/create/config.py

@@ -9,7 +9,6 @@
 import datetime
 import logging
 import os
-import warnings
 from copy import deepcopy
 
 import yaml
@@ -143,11 +142,21 @@ def __init__(self, config, *args, **kwargs):
 
         # deprecated/obsolete
         if "order" in self.output:
-            raise ValueError(f"Do not use 'order'. Use order_by in {self}")
+            raise ValueError(
+                f"Do not use 'order'. Use order_by instead. {list(self.keys())}"
+            )
         if "loops" in self:
-            assert "loop" not in self
-            warnings.warn("Should use loop instead of loops in config")
-            self.loop = self.pop("loops")
+            raise ValueError(
+                f"Do not use 'loops'. Use dates instead. {list(self.keys())}"
+            )
+        if "loop" in self:
+            print(f"Do not use 'loop'. Use 'dates' instead. {list(self.keys())}")
+            self.dates = self.pop("loop")
+
+        if isinstance(self.dates, dict):
+            self.dates = [self.dates]
+
+        self.loop = self.pop("dates")
 
         self.normalise()
 

ecml_tools/create/functions/accumulation.py

@@ -0,0 +1,86 @@
+# (C) Copyright 2020 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.
+#
+from copy import deepcopy
+
+from climetlab import load_source
+
+from ecml_tools.create.functions.mars import factorise_requests
+from ecml_tools.create.utils import to_datetime_list
+
+DEBUG = True
+
+
+def to_list(x):
+    if isinstance(x, (list, tuple)):
+        return x
+    return [x]
+
+
+def normalise_time_to_hours(r):
+    r = deepcopy(r)
+    if "time" not in r:
+        return r
+
+    times = []
+    for t in to_list(r["time"]):
+        assert len(t) == 4, r
+        assert t.endswith("00"), r
+        times.append(int(t) // 100)
+    r["time"] = tuple(times)
+    return r
+
+
+def accumulations(context, dates, request, **kwargs):
+    to_list(request["param"])
+    class_ = request["class"]
+
+    source_name = dict(
+        ea="era5-accumulations",
+        oper="oper-accumulations",
+        ei="oper-accumulations",
+    )[class_]
+
+    requests = factorise_requests(dates, request)
+
+    ds = load_source("empty")
+    for r in requests:
+        r = {k: v for k, v in r.items() if v != ("-",)}
+        r = normalise_time_to_hours(r)
+
+        if DEBUG:
+            print(f"✅ load_source({source_name},  {r}")
+        ds = ds + load_source(source_name, **r)
+    return ds
+
+
+execute = accumulations
+
+if __name__ == "__main__":
+    import yaml
+
+    config = yaml.safe_load(
+        """
+    - class: ea
+      expver: '0001'
+      grid: 20.0/20.0
+      levtype: sfc
+#      number: [0, 1]
+#      stream: enda
+      param: [cp, tp]
+#      accumulation_period: 6h
+    """
+    )
+    dates = yaml.safe_load(
+        "[2022-12-30 18:00, 2022-12-31 00:00, 2022-12-31 06:00, 2022-12-31 12:00]"
+    )
+    dates = to_datetime_list(dates)
+
+    DEBUG = True
+    for f in accumulations(None, dates, *config):
+        print(f, f.to_numpy().mean())

ecml_tools/create/functions/constants.py

@@ -0,0 +1,85 @@
+# (C) Copyright 2020 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.
+#
+from copy import deepcopy
+
+from climetlab import load_source
+
+from ecml_tools.create.utils import to_datetime_list
+
+DEBUG = True
+
+
+def to_list(x):
+    if isinstance(x, (list, tuple)):
+        return x
+    return [x]
+
+
+def get_template_field(request):
+    template_request = {
+        "date": "20200101",
+        "time": "0000",
+        "levtype": "sfc",
+        "param": "2t",
+    }
+    for k in ["area", "grid", "class"]:  # is class needed?
+        if k in request:
+            template_request[k] = request[k]
+    template = load_source("mars", template_request)
+    assert len(template) == 1, (len(template), template_request)
+    return template
+
+
+def normalise_time_to_hours(r):
+    r = deepcopy(r)
+    if "time" not in r:
+        return r
+
+    times = []
+    for t in to_list(r["time"]):
+        assert len(t) == 4, r
+        assert t.endswith("00"), r
+        times.append(int(t) // 100)
+    r["time"] = tuple(times)
+    return r
+
+
+def constants(context, dates, request, **kwargs):
+    to_list(request["param"])
+
+    template = get_template_field(request)
+
+    print(f"✅ load_source(constants, {template}, {request}")
+    return load_source("constants", source_or_dataset=template, **request)
+
+
+execute = constants
+
+if __name__ == "__main__":
+    import yaml
+
+    config = yaml.safe_load(
+        """
+    - class: ea
+      expver: '0001'
+      grid: 20.0/20.0
+      levtype: sfc
+      # param: [10u, 10v, 2d, 2t, lsm, msl, sdor, skt, slor, sp, tcw, z]
+      number: [0, 1]
+      param: [cos_latitude]
+    """
+    )
+    dates = yaml.safe_load(
+        "[2022-12-30 18:00, 2022-12-31 00:00, 2022-12-31 06:00, 2022-12-31 12:00]"
+    )
+    dates = to_datetime_list(dates)
+
+    DEBUG = True
+    for f in constants(None, dates, *config):
+        print(f, f.to_numpy().mean())

ecml_tools/create/functions/mars.py

@@ -71,7 +71,7 @@ def factorise_requests(dates, *requests):
     return compressed.iterate()
 
 
-def mars(dates, *requests, **kwargs):
+def mars(dates, *requests, _source_name="mars", **kwargs):
     if not requests:
         requests = [kwargs]
 
@@ -80,8 +80,8 @@ def mars(dates, *requests, **kwargs):
     for r in requests:
         r = {k: v for k, v in r.items() if v != ("-",)}
         if DEBUG:
-            print(f"✅ load_source(mars, {r}")
-        ds = ds + load_source("mars", **r)
+            print(f"✅ load_source({_source_name}, {r}")
+        ds = ds + load_source(_source_name, **r)
     return ds
 
 

ecml_tools/create/functions/tendencies.py

@@ -8,7 +8,6 @@
 #
 import datetime
 from collections import defaultdict
-from copy import deepcopy
 
 from climetlab.core.temporary import temp_file
 from climetlab.readers.grib.output import new_grib_output

ecml_tools/create/input.py

@@ -81,7 +81,9 @@ def sort(old_dic):
     params_steps = sort(params_steps)
     params_levels = sort(params_levels)
 
-    return dict(param_level=params_levels, param_step=params_steps, area=area, grid=grid)
+    return dict(
+        param_level=params_levels, param_step=params_steps, area=area, grid=grid
+    )
 
 
 class Cache:
@@ -114,7 +116,14 @@ def _build_coords(self):
         ensembles_key = list(from_config.keys())[2]
 
         if isinstance(from_config[variables_key], (list, tuple)):
-            assert all([v == w for v, w in zip(from_data[variables_key], from_config[variables_key])]), (
+            assert all(
+                [
+                    v == w
+                    for v, w in zip(
+                        from_data[variables_key], from_config[variables_key]
+                    )
+                ]
+            ), (
                 from_data[variables_key],
                 from_config[variables_key],
             )
@@ -347,7 +356,7 @@ def __init__(self, context, dates, action, previous_sibling=None):
     @cached_property
     def datasource(self):
         print(f"loading source with {self.args} {self.kwargs}")
-        return self.action.function(*self.args, **self.kwargs)
+        return self.action.function(self.dates, *self.args, **self.kwargs)
 
     def __repr__(self):
         content = " ".join([f"{v}" for v in self.args])
@@ -415,7 +424,9 @@ class BaseFunctionAction(Action):
     def __repr__(self):
         content = ""
         content += ",".join([self._short_str(a) for a in self.args])
-        content += " ".join([self._short_str(f"{k}={v}") for k, v in self.kwargs.items()])
+        content += " ".join(
+            [self._short_str(f"{k}={v}") for k, v in self.kwargs.items()]
+        )
         content = self._short_str(content)
         return super().__repr__(_inline_=content, _indent_=" ")
 
@@ -565,7 +576,15 @@ def __repr__(self):
         return super().__repr__(self.content, _inline_=str(self.kwargs))
 
 
-class FilterResult(StepResult):
+class StepFunctionResult(StepAction):
+    @property
+    def datasource(self):
+        return self.function(
+            self.context, self.content.datasource, self.dates, **self.kwargs
+        )
+
+
+class FilterStepResult(StepResult):
     @property
     def datasource(self):
         ds = self.content.datasource
@@ -575,8 +594,8 @@ def datasource(self):
         return ds
 
 
-class FilterAction(StepAction):
-    result_class = FilterResult
+class FilterStepAction(StepAction):
+    result_class = FilterStepResult
 
 
 class ConcatAction(ActionWithList):
@@ -647,7 +666,9 @@ def action_factory(config, context):
         )
 
     if len(config) != 1:
-        raise ValueError(f"Invalid input config. Expecting dict with only one key, got {list(config.keys())}")
+        raise ValueError(
+            f"Invalid input config. Expecting dict with only one key, got {list(config.keys())}"
+        )
 
     config = deepcopy(config)
     key = list(config.keys())[0]
@@ -685,7 +706,7 @@ def step_factory(config, context, _upstream_action):
 
     key = list(config.keys())[0]
     cls = dict(
-        filter=FilterAction,
+        filter=FilterStepAction,
         # rename=RenameAction,
         # remapping=RemappingAction,
     )[key]
@@ -704,7 +725,8 @@ def step_factory(config, context, _upstream_action):
 
 
 class Context:
-    def __init__(self, /, order_by, flatten_grid, remapping):
+    def __init__(self, /, dates, order_by, flatten_grid, remapping):
+        self.dates = dates
         self.order_by = order_by
         self.flatten_grid = flatten_grid
         self.remapping = build_remapping(remapping)
@@ -732,18 +754,17 @@ def __init__(self, config, **kwargs):
         self.kwargs = kwargs
         self.config = config
 
-    @property
-    def _action(self):
-        context = Context(**self.kwargs)
-        return action_factory(self.config, context)
-
     def select(self, dates):
         """This changes the context."""
         dates = build_groups(dates)
-        return self._action.select(dates)
+        context = Context(dates=dates, **self.kwargs)
+        action = action_factory(self.config, context)
+        return action.select(dates)
 
     def __repr__(self):
-        return repr(self._action)
+        context = Context(dates=None, **self.kwargs)
+        a = action_factory(self.config, context)
+        return repr(a)
 
 
 build_input = InputBuilder

ecml_tools/create/loaders.py

@@ -330,8 +330,8 @@ def load(self, parts):
             self.print(f" -> Processing {igroup} total={n_groups}")
             assert isinstance(group[0], datetime.datetime), group
 
-            inputs = self.input.select(dates=group)
-            data_writer.write(inputs, igroup)
+            result = self.input.select(dates=group)
+            data_writer.write(result, igroup, group)
 
         self.registry.add_to_history("loading_data_end", parts=parts)
         self.registry.add_provenance(name="provenance_load")