anemoi-datasets 0.5.7__py3-none-any.whl → 0.5.10__py3-none-any.whl

anemoi/datasets/create/input/result.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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 itertools
 import logging
 import math
@@ -33,9 +34,61 @@ LOG = logging.getLogger(__name__)
 def _fields_metatata(variables, cube):
     assert isinstance(variables, tuple), variables
 
-    result = {}
-    for i, c in enumerate(cube.iterate_cubelets()):
-        assert c._coords_names[1] == variables[i], (c._coords_names[1], variables[i])
+    KNOWN = {
+        "cos_julian_day": dict(computed_forcing=True, constant_in_time=False),
+        "cos_latitude": dict(computed_forcing=True, constant_in_time=True),
+        "cos_local_time": dict(computed_forcing=True, constant_in_time=False),
+        "cos_longitude": dict(computed_forcing=True, constant_in_time=True),
+        "cos_solar_zenith_angle": dict(computed_forcing=True, constant_in_time=False),
+        "insolation": dict(computed_forcing=True, constant_in_time=False),
+        "latitude": dict(computed_forcing=True, constant_in_time=True),
+        "longitude": dict(computed_forcing=True, constant_in_time=True),
+        "sin_julian_day": dict(computed_forcing=True, constant_in_time=False),
+        "sin_latitude": dict(computed_forcing=True, constant_in_time=True),
+        "sin_local_time": dict(computed_forcing=True, constant_in_time=False),
+        "sin_longitude": dict(computed_forcing=True, constant_in_time=True),
+    }
+
+    def _merge(md1, md2):
+        assert set(md1.keys()) == set(md2.keys()), (set(md1.keys()), set(md2.keys()))
+        result = {}
+        for k in md1.keys():
+            v1 = md1[k]
+            v2 = md2[k]
+
+            if v1 == v2:
+                result[k] = v1
+                continue
+
+            if isinstance(v1, list):
+                assert v2 not in v1, (v1, v2)
+                result[k] = sorted(v1 + [v2])
+                continue
+
+            if isinstance(v2, list):
+                assert v1 not in v2, (v1, v2)
+                result[k] = sorted(v2 + [v1])
+                continue
+
+            result[k] = sorted([v1, v2])
+
+        return result
+
+    mars = {}
+    other = defaultdict(dict)
+    i = -1
+    date = None
+    for c in cube.iterate_cubelets():
+
+        if date is None:
+            date = c._coords_names[0]
+
+        if date != c._coords_names[0]:
+            continue
+
+        if i == -1 or c._coords_names[1] != variables[i]:
+            i += 1
+
         f = cube[c.coords]
         md = f.metadata(namespace="mars")
         if not md:
@@ -49,7 +102,102 @@ def _fields_metatata(variables, cube):
             md["param"] = str(f.metadata("paramId", default="unknown"))
             # assert md['param'] != 'unknown', (md, f.metadata('param'))
 
-        result[variables[i]] = md
+        startStep = f.metadata("startStep", default=None)
+        assert startStep is None or isinstance(startStep, int), (startStep, type(f))
+
+        endStep = f.metadata("endStep", default=None)
+        assert endStep is None or isinstance(endStep, int), endStep
+
+        stepTypeForConversion = f.metadata("stepTypeForConversion", default=None)
+        typeOfStatisticalProcessing = f.metadata("typeOfStatisticalProcessing", default=None)
+        timeRangeIndicator = f.metadata("timeRangeIndicator", default=None)
+
+        # GRIB1 precipitation accumulations are not correctly encoded
+        if startStep == endStep and stepTypeForConversion == "accum":
+            startStep = 0
+
+        if startStep != endStep:
+            # https://codes.ecmwf.int/grib/format/grib2/ctables/4/10/
+            TYPE_OF_STATISTICAL_PROCESSING = {
+                None: None,
+                0: "average",
+                1: "accumulation",
+                2: "maximum",
+                3: "minimum",
+                4: "difference(end-start)",
+                5: "root_mean_square",
+                6: "standard_deviation",
+                7: "covariance",
+                8: "difference(start-end)",
+                9: "ratio",
+                10: "standardized_anomaly",
+                11: "summation",
+                100: "severity",
+                101: "mode",
+            }
+
+            # https://codes.ecmwf.int/grib/format/grib1/ctable/5/
+
+            TIME_RANGE_INDICATOR = {
+                4: "accumulation",
+                3: "average",
+            }
+
+            STEP_TYPE_FOR_CONVERSION = {
+                "min": "minimum",
+                "max": "maximum",
+                "accum": "accumulation",
+            }
+
+            #
+            # A few patches
+            #
+
+            PATCHES = {
+                "10fg6": "maximum",
+                "mntpr3": "minimum",  # Not in param db
+                "mntpr6": "minimum",  # Not in param db
+                "mxtpr3": "maximum",  # Not in param db
+                "mxtpr6": "maximum",  # Not in param db
+            }
+
+            process = TYPE_OF_STATISTICAL_PROCESSING.get(typeOfStatisticalProcessing)
+            if process is None:
+                process = TIME_RANGE_INDICATOR.get(timeRangeIndicator)
+            if process is None:
+                process = STEP_TYPE_FOR_CONVERSION.get(stepTypeForConversion)
+            if process is None:
+                process = PATCHES.get(md["param"])
+                if process is not None:
+                    LOG.error(f"Unknown process {stepTypeForConversion} for {md['param']}, using {process} instead")
+
+            if process is None:
+                raise ValueError(
+                    f"Unknown for {md['param']}:"
+                    f" {stepTypeForConversion=} ({STEP_TYPE_FOR_CONVERSION.get('stepTypeForConversion')}),"
+                    f" {typeOfStatisticalProcessing=} ({TYPE_OF_STATISTICAL_PROCESSING.get(typeOfStatisticalProcessing)}),"
+                    f" {timeRangeIndicator=} ({TIME_RANGE_INDICATOR.get(timeRangeIndicator)})"
+                )
+
+            # print(md["param"], "startStep", startStep, "endStep", endStep, "process", process, "typeOfStatisticalProcessing", typeOfStatisticalProcessing)
+            other[variables[i]]["process"] = process
+            other[variables[i]]["period"] = (startStep, endStep)
+
+        for k in md.copy().keys():
+            if k.startswith("_"):
+                md.pop(k)
+
+        if variables[i] in mars:
+            mars[variables[i]] = _merge(md, mars[variables[i]])
+        else:
+            mars[variables[i]] = md
+
+    result = {}
+    for k, v in mars.items():
+        result[k] = dict(mars=v) if v else {}
+        result[k].update(other[k])
+        result[k].update(KNOWN.get(k, {}))
+        assert result[k], k
 
     assert i + 1 == len(variables), (i + 1, len(variables))
     return result

anemoi/datasets/create/input/step.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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
 
@@ -59,6 +60,7 @@ class StepAction(Action):
         )
 
     def __repr__(self):
+        # raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
         return super().__repr__(self.previous_step, _inline_=str(self.kwargs))
 
 

anemoi/datasets/create/input/template.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2023 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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
 import re

anemoi/datasets/create/input/trace.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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
 import textwrap

anemoi/datasets/create/patch.py

@@ -1,4 +1,12 @@
-#!/usr/bin/env python3
+# (C) Copyright 2024 Anemoi contributors.
+#
+# 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 json
 import logging
 import os

anemoi/datasets/create/persistent.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2023 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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 glob
 import hashlib

anemoi/datasets/create/size.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2023 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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
 import os

anemoi/datasets/create/statistics/__init__.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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 datetime
 import glob
 import hashlib

anemoi/datasets/create/statistics/summary.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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 json
 from collections import defaultdict
 

anemoi/datasets/create/utils.py

@@ -1,14 +1,16 @@
-# (C) Copyright 2023 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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 datetime
 import os
+import warnings
 from contextlib import contextmanager
 
 import numpy as np
@@ -31,12 +33,23 @@ def cache_context(dirname):
 def to_datetime_list(*args, **kwargs):
     from earthkit.data.utils.dates import to_datetime_list as to_datetime_list_
 
+    warnings.warn(
+        "to_datetime_list() is deprecated. Call earthkit.data.utils.dates.to_datetime_list() instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     return to_datetime_list_(*args, **kwargs)
 
 
 def to_datetime(*args, **kwargs):
     from earthkit.data.utils.dates import to_datetime as to_datetime_
 
+    warnings.warn(
+        "to_datetime() is deprecated. Call earthkit.data.utils.dates.to_datetime() instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
     return to_datetime_(*args, **kwargs)
 
 

anemoi/datasets/create/writer.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2023 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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
 

anemoi/datasets/create/zarr.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2023 ECMWF.
+# (C) Copyright 2024 Anemoi contributors.
 #
 # 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 datetime
 import logging
 import shutil

anemoi/datasets/data/__init__.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2023 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # 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.
@@ -25,7 +27,31 @@ class MissingDateError(Exception):
     pass
 
 
+def _convert(x):
+
+    if isinstance(x, list):
+        return [_convert(a) for a in x]
+
+    if isinstance(x, tuple):
+        return tuple(_convert(a) for a in x)
+
+    if isinstance(x, dict):
+        return {k: _convert(v) for k, v in x.items()}
+
+    if x.__class__.__name__ in ("DictConfig", "ListConfig"):
+        from omegaconf import OmegaConf
+
+        return OmegaConf.to_container(x, resolve=True)
+
+    return x
+
+
 def open_dataset(*args, **kwargs):
+
+    # That will get rid of OmegaConf objects
+
+    args, kwargs = _convert(args), _convert(kwargs)
+
     ds = _open_dataset(*args, **kwargs)
     ds = ds.mutate()
     ds.arguments = {"args": args, "kwargs": kwargs}

anemoi/datasets/data/concat.py

@@ -1,10 +1,13 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (C) Copyright 2024 Anemoi contributors.
+#
 # 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
 
@@ -148,6 +151,7 @@ def concat_factory(args, kwargs):
 
     datasets = kwargs.pop("concat")
     fill_missing_gaps = kwargs.pop("fill_missing_gaps", False)
+
     assert isinstance(datasets, (list, tuple))
     assert len(args) == 0