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

anemoi/datasets/create/functions/__init__.py

@@ -1,11 +1,11 @@
-# (C) Copyright 2020 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 importlib
@@ -21,6 +21,8 @@ def assert_is_fieldlist(obj):
 
 def import_function(name, kind):
 
+    from anemoi.transform.filters import filter_registry
+
     name = name.replace("-", "_")
 
     plugins = {}
@@ -30,8 +32,21 @@ def import_function(name, kind):
     if name in plugins:
         return plugins[name].load()
 
-    module = importlib.import_module(
-        f".{kind}.{name}",
-        package=__name__,
-    )
-    return module.execute
+    try:
+        module = importlib.import_module(
+            f".{kind}.{name}",
+            package=__name__,
+        )
+        return module.execute
+    except ModuleNotFoundError:
+        pass
+
+    if kind == "filters":
+        if filter_registry.lookup(name, return_none=True):
+
+            def proc(context, data, *args, **kwargs):
+                return filter_registry.create(name, *args, **kwargs)(data)
+
+            return proc
+
+    raise ValueError(f"Unknown {kind} '{name}'")

anemoi/datasets/create/functions/filters/__init__.py

@@ -1,7 +1,8 @@
-# (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.

anemoi/datasets/create/functions/filters/empty.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 earthkit.data as ekd
 

anemoi/datasets/create/functions/filters/noop.py

@@ -1,11 +1,11 @@
-# (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.
-#
 
 
 def execute(context, input, *args, **kwargs):

anemoi/datasets/create/functions/filters/pressure_level_relative_humidity_to_specific_humidity.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.
-#
+
 
 from collections import defaultdict
 

anemoi/datasets/create/functions/filters/pressure_level_specific_humidity_to_relative_humidity.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.
-#
+
 
 from collections import defaultdict
 

anemoi/datasets/create/functions/filters/rename.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 re
 
@@ -24,7 +25,9 @@ class RenamedFieldMapping:
     def __init__(self, field, what, renaming):
         self.field = field
         self.what = what
-        self.renaming = renaming
+        self.renaming = {}
+        for k, v in renaming.items():
+            self.renaming[k] = {str(a): str(b) for a, b in v.items()}
 
     def metadata(self, key=None, **kwargs):
         if key is None:
@@ -51,16 +54,19 @@ class RenamedFieldFormat:
         format (str): A string that defines the new name of the field.
     """
 
-    def __init__(self, field, format):
+    def __init__(self, field, what, format):
         self.field = field
+        self.what = what
         self.format = format
         self.bits = re.findall(r"{(\w+)}", format)
 
-    def metadata(self, key, **kwargs):
-        value = self.field.metadata(key, **kwargs)
-        if "{" + key + "}" in self.format:
-            bits = {b: self.field.metadata(b, **kwargs) for b in self.bits}
-            return self.format.format(**bits)
+    def metadata(self, *args, **kwargs):
+        value = self.field.metadata(*args, **kwargs)
+        if args:
+            assert len(args) == 1
+            if args[0] == self.what:
+                bits = {b: self.field.metadata(b, **kwargs) for b in self.bits}
+                return self.format.format(**bits)
         return value
 
     def __getattr__(self, name):
@@ -69,6 +75,6 @@ class RenamedFieldFormat:
 
 def execute(context, input, what="param", **kwargs):
     if what in kwargs and isinstance(kwargs[what], str):
-        return FieldArray([RenamedFieldFormat(fs, kwargs[what]) for fs in input])
+        return FieldArray([RenamedFieldFormat(fs, what, kwargs[what]) for fs in input])
 
     return FieldArray([RenamedFieldMapping(fs, what, kwargs) for fs in input])

anemoi/datasets/create/functions/filters/rotate_winds.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.
-#
+
 
 from collections import defaultdict
 

anemoi/datasets/create/functions/filters/single_level_dewpoint_to_relative_humidity.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.
-#
+
 
 from collections import defaultdict
 

anemoi/datasets/create/functions/filters/single_level_relative_humidity_to_dewpoint.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.
-#
+
 
 from collections import defaultdict
 

anemoi/datasets/create/functions/filters/single_level_relative_humidity_to_specific_humidity.py

@@ -1,11 +1,11 @@
-# (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 numpy as np

anemoi/datasets/create/functions/filters/single_level_specific_humidity_to_relative_humidity.py

@@ -1,11 +1,11 @@
-# (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 numpy as np

anemoi/datasets/create/functions/filters/speeddir_to_uv.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.
-#
+
 
 from collections import defaultdict
 

anemoi/datasets/create/functions/filters/unrotate_winds.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.
-#
+
 
 from collections import defaultdict
 

anemoi/datasets/create/functions/filters/uv_to_speeddir.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.
-#
+
 
 from collections import defaultdict
 

anemoi/datasets/create/functions/sources/__init__.py

@@ -1,11 +1,11 @@
-# (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 glob
 import logging

anemoi/datasets/create/functions/sources/accumulations.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 logging
 import warnings
@@ -327,7 +328,9 @@ def _compute_accumulations(
             _member(field),
         )
         values = field.values  # optimisation
-        assert accumulations[key], key
+        if key not in accumulations:
+            raise ValueError(f"Key not found: {key}. Is it an accumulation field?")
+
         for a in accumulations[key]:
             a.add(field, values)
 
@@ -370,12 +373,15 @@ def accumulations(context, dates, **request):
 
     user_accumulation_period = request.pop("accumulation_period", 6)
 
+    # If `data_accumulation_period` is not set, this means that the accumulations are from the start
+    # of the forecast.
+
     KWARGS = {
         ("od", "oper"): dict(patch=_scda),
         ("od", "elda"): dict(base_times=(6, 18)),
         ("ea", "oper"): dict(data_accumulation_period=1, base_times=(6, 18)),
         ("ea", "enda"): dict(data_accumulation_period=3, base_times=(6, 18)),
-        ("rr", "oper"): dict(data_accumulation_period=3, base_times=(0, 3, 6, 9, 12, 15, 18, 21)),
+        ("rr", "oper"): dict(base_times=(0, 3, 6, 9, 12, 15, 18, 21)),
         ("l5", "oper"): dict(data_accumulation_period=1, base_times=(0,)),
     }
 

anemoi/datasets/create/functions/sources/constants.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.
-#
+
 from earthkit.data import from_source
 
 

anemoi/datasets/create/functions/sources/empty.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 earthkit.data as ekd
 

anemoi/datasets/create/functions/sources/forcings.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.
-#
+
 from earthkit.data import from_source
 
 

anemoi/datasets/create/functions/sources/grib.py

@@ -1,11 +1,11 @@
-# (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 glob

anemoi/datasets/create/functions/sources/hindcasts.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 earthkit.data.core.fieldlist import MultiFieldList

anemoi/datasets/create/functions/sources/mars.py

@@ -1,12 +1,14 @@
-# (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 re
 
 from anemoi.utils.humanize import did_you_mean
 from earthkit.data import from_source
@@ -31,6 +33,25 @@ def _date_to_datetime(d):
     return datetime.datetime.fromisoformat(d)
 
 
+def expand_to_by(x):
+
+    if isinstance(x, (str, int)):
+        return expand_to_by(str(x).split("/"))
+
+    if len(x) == 3 and x[1] == "to":
+        start = int(x[0])
+        end = int(x[2])
+        return list(range(start, end + 1))
+
+    if len(x) == 5 and x[1] == "to" and x[3] == "by":
+        start = int(x[0])
+        end = int(x[2])
+        by = int(x[4])
+        return list(range(start, end + 1, by))
+
+    return x
+
+
 def normalise_time_delta(t):
     if isinstance(t, datetime.timedelta):
         assert t == datetime.timedelta(hours=t.hours), t
@@ -42,25 +63,48 @@ def normalise_time_delta(t):
     return t
 
 
+def _normalise_time(t):
+    t = int(t)
+    if t < 100:
+        t * 100
+    return "{:04d}".format(t)
+
+
 def _expand_mars_request(request, date, request_already_using_valid_datetime=False, date_key="date"):
     requests = []
-    step = to_list(request.get("step", [0]))
-    for s in step:
+
+    user_step = to_list(expand_to_by(request.get("step", [0])))
+    user_time = None
+    user_date = None
+
+    if not request_already_using_valid_datetime:
+        user_time = request.get("time")
+        if user_time is not None:
+            user_time = to_list(user_time)
+            user_time = [_normalise_time(t) for t in user_time]
+
+        user_date = request.get(date_key)
+        if user_date is not None:
+            assert isinstance(user_date, str), user_date
+            user_date = re.compile("^{}$".format(user_date.replace("-", "").replace("?", ".")))
+
+    for step in user_step:
         r = request.copy()
 
         if not request_already_using_valid_datetime:
 
-            if isinstance(s, str) and "-" in s:
-                assert s.count("-") == 1, s
+            if isinstance(step, str) and "-" in step:
+                assert step.count("-") == 1, step
+
             # this takes care of the cases where the step is a period such as 0-24 or 12-24
-            hours = int(str(s).split("-")[-1])
+            hours = int(str(step).split("-")[-1])
 
             base = date - datetime.timedelta(hours=hours)
             r.update(
                 {
                     date_key: base.strftime("%Y%m%d"),
                     "time": base.strftime("%H%M"),
-                    "step": s,
+                    "step": step,
                 }
             )
 
@@ -69,12 +113,28 @@ def _expand_mars_request(request, date, request_already_using_valid_datetime=Fal
                 if isinstance(r[pproc], (list, tuple)):
                     r[pproc] = "/".join(str(x) for x in r[pproc])
 
+        if user_date is not None:
+            if not user_date.match(r[date_key]):
+                continue
+
+        if user_time is not None:
+            # It time is provided by the user, we only keep the requests that match the time
+            if r["time"] not in user_time:
+                continue
+
         requests.append(r)
 
+    # assert requests, requests
+
     return requests
 
 
-def factorise_requests(dates, *requests, request_already_using_valid_datetime=False, date_key="date"):
+def factorise_requests(
+    dates,
+    *requests,
+    request_already_using_valid_datetime=False,
+    date_key="date",
+):
     updates = []
     for req in requests:
         # req = normalise_request(req)
@@ -87,6 +147,9 @@ def factorise_requests(dates, *requests, request_already_using_valid_datetime=Fa
                 date_key=date_key,
             )
 
+    if not updates:
+        return
+
     compressed = Availability(updates)
     for r in compressed.iterate():
         for k, v in r.items():
@@ -177,13 +240,24 @@ MARS_KEYS = [
 ]
 
 
-def mars(context, dates, *requests, request_already_using_valid_datetime=False, date_key="date", **kwargs):
+def mars(
+    context,
+    dates,
+    *requests,
+    request_already_using_valid_datetime=False,
+    date_key="date",
+    **kwargs,
+):
+
     if not requests:
         requests = [kwargs]
 
     for r in requests:
+        param = r.get("param", [])
+        if not isinstance(param, (list, tuple)):
+            param = [param]
         # check for "Norway bug" where yaml transforms 'no' into False, etc.
-        for p in r.get("param", []):
+        for p in param:
             if p is False:
                 raise ValueError(
                     "'param' cannot be 'False'. If you wrote 'param: no' or 'param: off' in yaml, you may want to use quotes?"
@@ -197,12 +271,18 @@ def mars(context, dates, *requests, request_already_using_valid_datetime=False,
                     "'param' cannot be 'True'. If you wrote 'param: on' in yaml, you may want to use quotes?"
                 )
 
-    requests = factorise_requests(
-        dates,
-        *requests,
-        request_already_using_valid_datetime=request_already_using_valid_datetime,
-        date_key=date_key,
-    )
+    if len(dates) == 0:  # When using `repeated_dates`
+        assert len(requests) == 1, requests
+        assert "date" in requests[0], requests[0]
+        if isinstance(requests[0]["date"], datetime.date):
+            requests[0]["date"] = requests[0]["date"].strftime("%Y%m%d")
+    else:
+        requests = factorise_requests(
+            dates,
+            *requests,
+            request_already_using_valid_datetime=request_already_using_valid_datetime,
+            date_key=date_key,
+        )
 
     requests = list(requests)
 

anemoi/datasets/create/functions/sources/netcdf.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.
-#
+
 
 from .xarray import load_many
 

anemoi/datasets/create/functions/sources/opendap.py

@@ -1,11 +1,11 @@
-# (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.
-#
 
 
 from .xarray import load_many