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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.7
+Version: 0.1.8
 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

{anemoi_utils-0.1.7 → anemoi_utils-0.1.8}/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.8'
+__version_tuple__ = version_tuple = (0, 1, 8)

anemoi_utils-0.1.8/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.8}/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.7 → anemoi_utils-0.1.8}/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.8
 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

{anemoi_utils-0.1.7 → anemoi_utils-0.1.8}/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

{anemoi_utils-0.1.7 → anemoi_utils-0.1.8}/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()