ensemble perturbations and clean up

Author: floriankrb

ecml_tools/create/functions/__init__.py

@@ -0,0 +1,193 @@
+# (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.
+#
+
+import warnings
+
+import numpy as np
+import tqdm
+from climetlab import load_source
+from climetlab.core.temporary import temp_file
+from climetlab.readers.grib.output import new_grib_output
+
+from ecml_tools.create.check import check_data_values
+
+
+def get_unique_field(ds, selection):
+    ds = ds.sel(**selection)
+    assert len(ds) == 1, (ds, selection)
+    return ds[0]
+
+
+def normalise_number(number):
+    if isinstance(number, (tuple, list, int)):
+        return number
+
+    assert isinstance(number, str), (type(number), number)
+
+    number = number.split("/")
+    if len(number) > 4 and (number[1] == "to" and number[3] == "by"):
+        return list(range(int(number[0]), int(number[2]) + 1, int(number[4])))
+
+    if len(number) > 2 and number[1] == "to":
+        return list(range(int(number[0]), int(number[2]) + 1))
+
+    assert isinstance(number, list), (type(number), number)
+    return number
+
+
+def ensembles_perturbations(ensembles, center, mean, remapping={}, patches={}):
+    n_ensembles = len(normalise_number(ensembles["number"]))
+
+    print(f"Retrieving ensemble data with {ensembles}")
+    ensembles = load_source(**ensembles)
+    print(f"Retrieving center data with {center}")
+    center = load_source(**center)
+    print(f"Retrieving mean data with {mean}")
+    mean = load_source(**mean)
+
+    assert len(mean) * n_ensembles == len(ensembles), (
+        len(mean),
+        n_ensembles,
+        len(ensembles),
+    )
+    assert len(center) * n_ensembles == len(ensembles), (
+        len(center),
+        n_ensembles,
+        len(ensembles),
+    )
+
+    tmp = temp_file()
+    path = tmp.path
+    out = new_grib_output(path)
+
+    keys = ["param", "level", "valid_datetime", "number", "date", "time", "step"]
+
+    ensembles_coords = ensembles.unique_values(*keys)
+    center_coords = center.unique_values(*keys)
+    mean_coords = mean.unique_values(*keys)
+
+    for k in keys:
+        if k == "number":
+            assert len(mean_coords[k]) == 1
+            assert len(center_coords[k]) == 1
+            assert len(ensembles_coords[k]) == n_ensembles
+            continue
+        assert set(center_coords[k]) == set(ensembles_coords[k]), (
+            k,
+            center_coords[k],
+            ensembles_coords[k],
+        )
+        assert set(center_coords[k]) == set(mean_coords[k]), (
+            k,
+            center_coords[k],
+            mean_coords[k],
+        )
+
+    for field in tqdm.tqdm(center):
+        param = field.metadata("param")
+        grid = field.metadata("grid")
+
+        selection = dict(
+            valid_datetime=field.metadata("valid_datetime"),
+            param=field.metadata("param"),
+            level=field.metadata("level"),
+            date=field.metadata("date"),
+            time=field.metadata("time"),
+            step=field.metadata("step"),
+        )
+        mean_field = get_unique_field(mean, selection)
+        assert mean_field.metadata("grid") == grid, (mean_field.metadata("grid"), grid)
+
+        m = mean_field.to_numpy()
+        c = field.to_numpy()
+        assert m.shape == c.shape, (m.shape, c.shape)
+
+        for number in ensembles_coords["number"]:
+            ensembles_field = get_unique_field(ensembles.sel(number=number), selection)
+            assert ensembles_field.metadata("grid") == grid, (
+                ensembles_field.metadata("grid"),
+                grid,
+            )
+
+            e = ensembles_field.to_numpy()
+            assert c.shape == e.shape, (c.shape, e.shape)
+
+            x = c + m - e
+            if param == "q":
+                warnings.warn("Clipping q")
+                x = np.maximum(x, 0)
+
+            assert x.shape == c.shape, (x.shape, c.shape)
+
+            check_data_values(x, name=param)
+            out.write(x, template=ensembles_field)
+
+    out.close()
+
+    ds = load_source("file", path)
+    assert len(ds) == len(ensembles), (len(ds), len(ensembles))
+    ds._tmp = tmp
+
+    assert len(mean) * n_ensembles == len(ensembles)
+    assert len(center) * n_ensembles == len(ensembles)
+
+    final_coords = ds.unique_values(*keys)
+    assert len(final_coords["number"]) == n_ensembles, final_coords
+    return ds
+
+
+execute = ensembles_perturbations
+
+if __name__ == "__main__":
+    import yaml
+
+    config = yaml.safe_load(
+        """
+
+    common: &common
+        name: mars
+        # marser is the MARS containing ERA5 reanalysis dataset, avoid hitting the FDB server for nothing
+        database: marser
+        class: ea
+        # date: $datetime_format($dates,%Y%m%d)
+        # time: $datetime_format($dates,%H%M)
+        date: 20221230/to/20230103
+        time: '0000/1200'
+        expver: '0001'
+        grid: 20.0/20.0
+        levtype: sfc
+        param: [2t]
+        # levtype: pl
+        # param: [10u, 10v, 2d, 2t, lsm, msl, sdor, skt, slor, sp, tcw, z]
+
+    config:
+        ensembles: # the ensemble data has one additional dimension
+          <<: *common
+          stream: enda
+          type: an
+          number: [0, 1]
+          # number: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+
+        center: # the new center of the data
+          <<: *common
+          stream: oper
+          type: an
+
+        mean: # the previous center of the data
+          <<: *common
+          stream: enda
+          type: em
+
+    """
+    )["config"]
+    for k, v in config.items():
+        print(k, v)
+
+    for f in ensembles_perturbations(**config):
+        print(f, f.to_numpy().mean())

ecml_tools/create/functions/ensemble_perturbations.py

@@ -7,7 +7,9 @@
 # nor does it submit to any jurisdiction.
 #
 import datetime
+import importlib
 import logging
+import os
 import time
 from collections import defaultdict
 from copy import deepcopy
@@ -334,9 +336,10 @@ def __getitem__(self, key):
         raise KeyError(key)
 
 
-class SourceResult(Result):
+class FunctionResult(Result):
     def __init__(self, context, dates, action, previous_sibling=None):
         super().__init__(context, dates)
+        assert isinstance(action, Action), type(action)
         self.action = action
 
         _args = self.action.args
@@ -349,17 +352,19 @@ def __init__(self, context, dates, action, previous_sibling=None):
 
     @cached_property
     def datasource(self):
-        from climetlab import load_source
-
         print(f"loading source with {self.args} {self.kwargs}")
-        return load_source(*self.args, **self.kwargs)
+        return self.action.function(*self.args, **self.kwargs)
 
     def __repr__(self):
         content = " ".join([f"{v}" for v in self.args])
         content += " ".join([f"{k}={v}" for k, v in self.kwargs.items()])
 
         return super().__repr__(content)
 
+    @property
+    def function(self):
+        raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
+
 
 class JoinResult(Result):
     def __init__(self, context, dates, results, **kwargs):
@@ -396,7 +401,7 @@ def __repr__(self):
         return super().__repr__(_inline_=self.name, _indent_=" ")
 
 
-class SourceAction(Action):
+class BaseFunctionAction(Action):
     def __repr__(self):
         content = ""
         content += ",".join([self._short_str(a) for a in self.args])
@@ -407,7 +412,32 @@ def __repr__(self):
         return super().__repr__(_inline_=content, _indent_=" ")
 
     def select(self, dates):
-        return SourceResult(self.context, dates, action=self)
+        return FunctionResult(self.context, dates, action=self)
+
+
+class SourceAction(BaseFunctionAction):
+    @property
+    def function(self):
+        from climetlab import load_source
+
+        return load_source
+
+
+class FunctionAction(BaseFunctionAction):
+    def __init__(self, context, name, **kwargs):
+        super().__init__(context, **kwargs)
+        self.name = name
+
+    @property
+    def function(self):
+        here = os.path.dirname(__file__)
+        path = os.path.join(here, "functions", f"{self.name}.py")
+        spec = importlib.util.spec_from_file_location(self.name, path)
+        module = spec.loader.load_module()
+        # TODO: this fails here, fix this.
+        #   getattr(module, self.name)
+        #   self.action.kwargs
+        return module.execute
 
 
 class ConcatResult(Result):
@@ -610,6 +640,7 @@ def action_factory(config, context):
         label=LabelAction,
         pipe=PipeAction,
         source=SourceAction,
+        function=FunctionAction,
         dates=DateAction,
     )[key]
 
@@ -688,5 +719,8 @@ def select(self, dates):
         """This changes the context."""
         return self._action.select(dates)
 
+    def __repr__(self):
+        return repr(self._action)
+
 
 build_input = InputBuilder

ecml_tools/create/input.py

@@ -16,7 +16,6 @@
 from collections import defaultdict
 
 import numpy as np
-from prepml.utils.text import table
 
 from ecml_tools.provenance import gather_provenance_info
 
@@ -274,19 +273,12 @@ def check(self):
                 raise
 
     def __str__(self):
-        stats = [self[name] for name in self.STATS_NAMES]
-
-        rows = []
+        header = ["Variables"] + [self[name] for name in self.STATS_NAMES]
+        out = " ".join(header)
 
         for i, v in enumerate(self["variables_names"]):
-            rows.append([i, v] + [x[i] for x in stats])
-
-        return table(
-            rows,
-            header=["Index", "Variable", "Min", "Max", "Mean", "Stdev"],
-            align=[">", "<", ">", ">", ">", ">"],
-            margin=3,
-        )
+            out += " ".join([v] + [f"{x[i]:.2f}" for x in self.values()])
+        return out
 
     def save(self, filename, provenance=None):
         assert filename.endswith(".json"), filename