anemoi-utils 0.1.7__tar.gz → 0.1.9__tar.gz

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/.gitignore

@@ -185,3 +185,4 @@ _build/
 ~*
 *.sync
 _version.py
+*.code-workspace

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.7
+Version: 0.1.9
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License
@@ -223,6 +223,7 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: tomli
+Requires-Dist: pyyaml
 Provides-Extra: provenance
 Requires-Dist: GitPython; extra == "provenance"
 Requires-Dist: nvsmi; extra == "provenance"

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/pyproject.toml

@@ -41,6 +41,7 @@ classifiers = [
 
 dependencies = [
     "tomli", # Only needed before 3.11
+    "pyyaml",
 ]
 
 [project.optional-dependencies]
@@ -88,3 +89,6 @@ Issues = "https://github.com/ecmwf/anemoi-utils/issues"
 
 [tool.setuptools_scm]
 version_file = "src/anemoi/utils/_version.py"
+
+[tool.setuptools.package-data]
+"anemoi.utils.mars" = ["*.yaml"]

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/src/anemoi/utils/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.7'
-__version_tuple__ = version_tuple = (0, 1, 7)
+__version__ = version = '0.1.9'
+__version_tuple__ = version_tuple = (0, 1, 9)

anemoi_utils-0.1.9/src/anemoi/utils/caching.py

@@ -0,0 +1,59 @@
+# (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 hashlib
+import json
+import os
+import time
+
+
+def cache(key, proc, collection="default", expires=None):
+    path = os.path.join(os.path.expanduser("~"), ".cache", "anemoi", collection)
+    os.makedirs(path, exist_ok=True)
+
+    key = json.dumps(key, sort_keys=True)
+    m = hashlib.md5()
+    m.update(key.encode("utf-8"))
+
+    filename = os.path.join(path, m.hexdigest())
+    if os.path.exists(filename):
+        with open(filename, "r") as f:
+            data = json.load(f)
+            if expires is None or data["expires"] > time.time():
+                if data["key"] == key:
+                    return data["value"]
+
+    value = proc()
+    data = {"key": key, "value": value}
+    if expires is not None:
+        data["expires"] = time.time() + expires
+
+    with open(filename, "w") as f:
+        json.dump(data, f)
+
+    return value
+
+
+class cached:
+
+    def __init__(self, collection="default", expires=None):
+        self.collection = collection
+        self.expires = expires
+
+    def __call__(self, func):
+
+        full = f"{func.__module__}.{func.__name__}"
+
+        def wrapped(*args, **kwargs):
+            return cache(
+                (full, args, kwargs),
+                lambda: func(*args, **kwargs),
+                self.collection,
+                self.expires,
+            )
+
+        return wrapped

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/src/anemoi/utils/dates.py

@@ -10,6 +10,16 @@ import calendar
 import datetime
 
 
+def normalise_frequency(frequency):
+    if isinstance(frequency, int):
+        return frequency
+    assert isinstance(frequency, str), (type(frequency), frequency)
+
+    unit = frequency[-1].lower()
+    v = int(frequency[:-1])
+    return {"h": v, "d": v * 24}[unit]
+
+
 def no_time_zone(date):
     """Remove time zone information from a date.
 
@@ -27,6 +37,7 @@ def no_time_zone(date):
     return date.replace(tzinfo=None)
 
 
+# this function is use in anemoi-datasets
 def as_datetime(date):
     """Convert a date to a datetime object, removing any time zone information.
 
@@ -162,11 +173,15 @@ class HindcastDatesTimes:
         """
 
         self.reference_dates = reference_dates
-        self.years = (1, years + 1)
+
+        if isinstance(years, list):
+            self.years = years
+        else:
+            self.years = range(1, years + 1)
 
     def __iter__(self):
         for reference_date in self.reference_dates:
-            for year in range(*self.years):
+            for year in self.years:
                 if reference_date.month == 2 and reference_date.day == 29:
                     date = datetime.datetime(reference_date.year - year, 2, 28)
                 else:
@@ -246,3 +261,61 @@ class Autumn(DateTimes):
             _description_
         """
         super().__init__(datetime.datetime(year, 9, 1), datetime.datetime(year, 11, 30), **kwargs)
+
+
+class ConcatDateTimes:
+    def __init__(self, *dates):
+        if len(dates) == 1 and isinstance(dates[0], list):
+            dates = dates[0]
+
+        self.dates = dates
+
+    def __iter__(self):
+        for date in self.dates:
+            yield from date
+
+
+class EnumDateTimes:
+    def __init__(self, dates):
+        self.dates = dates
+
+    def __iter__(self):
+        for date in self.dates:
+            yield as_datetime(date)
+
+
+def datetimes_factory(*args, **kwargs):
+    if args and kwargs:
+        raise ValueError("Cannot provide both args and kwargs for a list of dates")
+
+    if not args and not kwargs:
+        raise ValueError("No dates provided")
+
+    if kwargs:
+        name = kwargs.get("name")
+
+        if name == "hindcast":
+            reference_dates = kwargs["reference_dates"]
+            reference_dates = datetimes_factory(reference_dates)
+            years = kwargs["years"]
+            return HindcastDatesTimes(reference_dates=reference_dates, years=years)
+
+        kwargs = kwargs.copy()
+        if "frequency" in kwargs:
+            freq = kwargs.pop("frequency")
+            kwargs["increment"] = normalise_frequency(freq)
+        return DateTimes(**kwargs)
+
+    if not any((isinstance(x, dict) or isinstance(x, list)) for x in args):
+        return EnumDateTimes(args)
+
+    if len(args) == 1:
+        a = args[0]
+
+        if isinstance(a, dict):
+            return datetimes_factory(**a)
+
+        if isinstance(a, list):
+            return datetimes_factory(*a)
+
+    return ConcatDateTimes(*[datetimes_factory(a) for a in args])

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/src/anemoi/utils/grib.py

@@ -16,10 +16,21 @@ import re
 
 import requests
 
+from .caching import cached
+
 LOG = logging.getLogger(__name__)
 
 
-def _search(name):
+@cached(collection="grib", expires=30 * 24 * 60 * 60)
+def _units():
+    r = requests.get("https://codes.ecmwf.int/parameter-database/api/v1/unit/")
+    r.raise_for_status()
+    units = r.json()
+    return {str(u["id"]): u["name"] for u in units}
+
+
+@cached(collection="grib", expires=30 * 24 * 60 * 60)
+def _search_param(name):
     name = re.escape(name)
     r = requests.get(f"https://codes.ecmwf.int/parameter-database/api/v1/param/?search=^{name}$&regex=true")
     r.raise_for_status()
@@ -56,7 +67,7 @@ def shortname_to_paramid(shortname: str) -> int:
     167
 
     """
-    return _search(shortname)["id"]
+    return _search_param(shortname)["id"]
 
 
 def paramid_to_shortname(paramid: int) -> str:
@@ -76,4 +87,46 @@ def paramid_to_shortname(paramid: int) -> str:
     '2t'
 
     """
-    return _search(str(paramid))["shortname"]
+    return _search_param(str(paramid))["shortname"]
+
+
+def units(param) -> str:
+    """Return the units of a GRIB parameter given its name or id.
+
+    Parameters
+    ----------
+    paramid : int or str
+        Parameter id ir name.
+
+    Returns
+    -------
+    str
+        Parameter unit.
+
+    >>> unit(167)
+    'K'
+
+    """
+
+    unit_id = str(_search_param(str(param))["unit_id"])
+    return _units()[unit_id]
+
+
+def must_be_positive(param):
+    """Check if a parameter must be positive.
+
+    Parameters
+    ----------
+    param : int or str
+        Parameter id or shortname.
+
+    Returns
+    -------
+    bool
+        True if the parameter must be positive.
+
+    >>> must_be_positive("tp")
+    True
+
+    """
+    return units(param) in ["m", "kg kg**-1", "m of water equivalent"]

anemoi_utils-0.1.9/src/anemoi/utils/mars/__init__.py

@@ -0,0 +1,76 @@
+# (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.
+
+
+"""Utilities for working with Mars requests.
+
+Has some konwledge of how certain streams are organised in Mars.
+
+"""
+
+import datetime
+import logging
+import os
+
+import yaml
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_MARS_LABELLING = {
+    "class": "od",
+    "type": "an",
+    "stream": "oper",
+    "expver": "0001",
+}
+
+
+def _expand_mars_labelling(request):
+    """Expand the request with the default Mars labelling.
+
+    The default Mars labelling is:
+
+    {'class': 'od',
+     'type': 'an',
+     'stream': 'oper',
+     'expver': '0001'}
+
+    """
+    result = DEFAULT_MARS_LABELLING.copy()
+    result.update(request)
+    return result
+
+
+STREAMS = None
+
+
+def _lookup_mars_stream(request):
+    global STREAMS
+
+    if STREAMS is None:
+
+        with open(os.path.join(os.path.dirname(__file__), "mars.yaml")) as f:
+            STREAMS = yaml.safe_load(f)
+
+    request = _expand_mars_labelling(request)
+    for s in STREAMS:
+        match = s["match"]
+        if all(request.get(k) == v for k, v in match.items()):
+            return s["info"]
+
+
+def recenter(date, center, members):
+
+    center = _lookup_mars_stream(center)
+    members = _lookup_mars_stream(members)
+
+    return (center, members)
+
+
+if __name__ == "__main__":
+    date = datetime.datetime(2024, 5, 9, 0)
+
+    print(recenter(date, {"type": "an"}, {"stream": "elda"}))

anemoi_utils-0.1.9/src/anemoi/utils/mars/mars.yaml

@@ -0,0 +1,5 @@
+- match:
+    class: od
+    stream: elda
+  info:
+    runs: [6, 18]

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/src/anemoi_utils.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.7
+Version: 0.1.9
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License
@@ -223,6 +223,7 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: tomli
+Requires-Dist: pyyaml
 Provides-Extra: provenance
 Requires-Dist: GitPython; extra == "provenance"
 Requires-Dist: nvsmi; extra == "provenance"

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/src/anemoi_utils.egg-info/SOURCES.txt

@@ -22,6 +22,7 @@ docs/modules/provenance.rst
 docs/modules/text.rst
 src/anemoi/utils/__init__.py
 src/anemoi/utils/_version.py
+src/anemoi/utils/caching.py
 src/anemoi/utils/checkpoints.py
 src/anemoi/utils/config.py
 src/anemoi/utils/dates.py
@@ -29,10 +30,13 @@ src/anemoi/utils/grib.py
 src/anemoi/utils/humanize.py
 src/anemoi/utils/provenance.py
 src/anemoi/utils/text.py
+src/anemoi/utils/mars/__init__.py
+src/anemoi/utils/mars/mars.yaml
 src/anemoi_utils.egg-info/PKG-INFO
 src/anemoi_utils.egg-info/SOURCES.txt
 src/anemoi_utils.egg-info/dependency_links.txt
 src/anemoi_utils.egg-info/requires.txt
 src/anemoi_utils.egg-info/top_level.txt
 tests/requirements.txt
+tests/test_dates.py
 tests/test_utils.py

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/src/anemoi_utils.egg-info/requires.txt

@@ -1,4 +1,5 @@
 tomli
+pyyaml
 
 [all]
 tomli

anemoi_utils-0.1.9/tests/test_dates.py

@@ -0,0 +1,113 @@
+import datetime
+from textwrap import dedent
+
+import yaml
+
+from anemoi.utils.dates import datetimes_factory
+
+
+def _(txt):
+    txt = dedent(txt)
+    config = yaml.safe_load(txt)
+    return datetimes_factory(config)
+
+
+def test_date_1():
+    d = _(
+        """
+          - 2023-01-01
+          - 2023-01-02
+          - 2023-01-03
+    """
+    )
+    assert len(list(d)) == 3
+
+
+def test_date_2():
+    d = _(
+        """
+        start: 2023-01-01
+        end: 2023-01-07
+        frequency: 12
+        day_of_week: [monday, friday]
+    """
+    )
+    assert len(list(d)) == 4
+
+
+def test_date_3():
+    d = _(
+        """
+        - start: 2023-01-01
+          end: 2023-01-03
+          frequency: 24
+        - start: 2024-01-01T06:00:00
+          end: 2024-01-02T18:00:00
+          frequency: 6h
+    """
+    )
+    assert datetime.datetime(2023, 1, 1, 0) in d
+    assert datetime.datetime(2023, 1, 2, 0) in d
+    assert datetime.datetime(2023, 1, 3, 0) in d
+    assert datetime.datetime(2024, 1, 1, 6) in d
+    assert datetime.datetime(2024, 1, 1, 12) in d
+    assert datetime.datetime(2024, 1, 1, 18) in d
+    assert datetime.datetime(2024, 1, 2, 0) in d
+    assert datetime.datetime(2024, 1, 2, 6) in d
+    assert datetime.datetime(2024, 1, 2, 12) in d
+    assert datetime.datetime(2024, 1, 2, 18) in d
+    assert len(list(d)) == 10
+
+
+def test_date_hindcast_1():
+    d = _(
+        """
+        - name: hindcast
+          reference_dates:
+            start: 2023-01-01
+            end: 2023-01-03
+            frequency: 24
+          years: 20
+    """
+    )
+    assert len(list(d)) == 60
+
+
+def test_date_hindcast_2():
+    d = _(
+        """
+        - name: hindcast
+          reference_dates:
+            start: 2023-01-01
+            end: 2023-01-03
+            frequency: 24
+          years: [2018, 2019, 2020, 2021]
+    """
+    )
+    assert len(list(d)) == 12
+
+
+def test_date_hindcast_3():
+    d = _(
+        """
+        - name: hindcast
+          reference_dates:
+            start: 2022-12-25 00:00:00
+            end: 2022-12-31 12:00:00
+            frequency: 12h
+            day_of_week: tuesday
+          years: [2018, 2019, 2020, 2021]
+    """
+    )
+    print(list(d))
+    assert len(list(d)) == 8
+
+
+if __name__ == "__main__":
+    test_functions = [
+        obj for name, obj in globals().items() if name.startswith("test_") and isinstance(obj, type(lambda: 0))
+    ]
+    for test_func in test_functions:
+        print(f"Running test: {test_func.__name__}")
+        test_func()
+    print("All tests passed!")

{anemoi_utils-0.1.7 → anemoi_utils-0.1.9}/tests/test_utils.py

@@ -7,6 +7,8 @@
 
 
 from anemoi.utils.config import DotDict
+from anemoi.utils.grib import paramid_to_shortname
+from anemoi.utils.grib import shortname_to_paramid
 
 
 def test_dotdict():
@@ -26,5 +28,10 @@ def test_dotdict():
     assert d.d.x == 6
 
 
+def test_grib():
+    assert shortname_to_paramid("2t") == 167
+    assert paramid_to_shortname(167) == "2t"
+
+
 if __name__ == "__main__":
-    test_dotdict()
+    test_grib()