anemoi-utils 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

anemoi/utils/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.6'
-__version_tuple__ = version_tuple = (0, 1, 6)
+__version__ = version = '0.1.8'
+__version_tuple__ = version_tuple = (0, 1, 8)

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/checkpoints.py

@@ -6,8 +6,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-"""
-Read and write extra metadata in PyTorch checkpoints files. These files
+"""Read and write extra metadata in PyTorch checkpoints files. These files
 are zip archives containing the model weights.
 """
 

anemoi/utils/grib.py

@@ -11,12 +11,26 @@ See https://codes.ecmwf.int/grib/param-db/ for more information.
 
 """
 
+import logging
 import re
 
 import requests
 
+from .caching import cached
 
-def _search(name):
+LOG = logging.getLogger(__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()
@@ -26,7 +40,12 @@ def _search(name):
 
     if len(results) > 1:
         names = [f'{r.get("id")} ({r.get("name")})' for r in results]
-        raise ValueError(f"{name} is ambiguous: {', '.join(names)}")
+        dissemination = [r for r in results if "dissemination" in r.get("access_ids", [])]
+        if len(dissemination) == 1:
+            return dissemination[0]
+
+        results = sorted(results, key=lambda x: x["id"])
+        LOG.warning(f"{name} is ambiguous: {', '.join(names)}. Using param_id={results[0]['id']}")
 
     return results[0]
 
@@ -44,12 +63,11 @@ def shortname_to_paramid(shortname: str) -> int:
     int
         Parameter id.
 
-
     >>> shortname_to_paramid("2t")
     167
 
     """
-    return _search(shortname)["id"]
+    return _search_param(shortname)["id"]
 
 
 def paramid_to_shortname(paramid: int) -> str:
@@ -65,9 +83,50 @@ def paramid_to_shortname(paramid: int) -> str:
     str
         Parameter shortname.
 
-
     >>> paramid_to_shortname(167)
     '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/provenance.py

@@ -5,8 +5,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-"""
-Collect information about the current environment, like:
+"""Collect information about the current environment, like:
 
  - The Python version
  - The versions of the modules which are currently loaded
@@ -215,7 +214,6 @@ def git_check(*args):
     dict
         An object with the git information for the given arguments.
 
-
     >>> {
             "anemoi.utils": {
                 "sha1": "c999d83ae283bcbb99f68d92c42d24315922129f",

anemoi/utils/text.py

@@ -5,9 +5,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-"""
-Text utilities
-"""
+"""Text utilities"""
 
 import sys
 from collections import defaultdict
@@ -58,9 +56,6 @@ def boxed(text, min_width=80, max_width=80) -> str:
     str
         A boxed version of the input text
 
-
-
-
     """
 
     lines = text.split("\n")
@@ -266,7 +261,6 @@ def table(rows, header, align, margin=0):
     margin : int, optional
         Extra spaces on the left side of the table, by default 0
 
-
     Returns
     -------
     str
@@ -339,7 +333,6 @@ def progress(done, todo, width=80) -> str:
     str
         _description_
 
-
     """
     done = min(int(done / todo * width + 0.5), width)
     return green("█" * done) + red("█" * (width - done))

{anemoi_utils-0.1.6.dist-info → anemoi_utils-0.1.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.6
+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.8.dist-info/RECORD

@@ -0,0 +1,15 @@
+anemoi/utils/__init__.py,sha256=zZZpbKIoGWwdCOuo6YSruLR7C0GzvzI1Wzhyqaa0K7M,456
+anemoi/utils/_version.py,sha256=PdJ7dZoz_SyEgX0MdrMfQYBFlGcwpemv6ibF8NKALBY,411
+anemoi/utils/caching.py,sha256=HrC9aFHlcCTaM2Z5u0ivGIXz7eFu35UQQhUuwwuG2pk,1743
+anemoi/utils/checkpoints.py,sha256=IR86FFNh5JR_uQVlgybnZG74PyU0CNLhyocqARwZIrs,2069
+anemoi/utils/config.py,sha256=XEesqODvkuE3ZA7dnEnZ-ooBRtU6ecPmkfP65FtialA,2147
+anemoi/utils/dates.py,sha256=ZabXLecHMio9ecIJbUxNcfOFoGfkcotSwy_xw1kNQ9w,6474
+anemoi/utils/grib.py,sha256=gVfo4KYQv31iRyoqRDwk5tiqZDUgOIvhag_kO0qjYD0,3067
+anemoi/utils/humanize.py,sha256=LD6dGnqChxA5j3tMhSybsAGRQzi33d_qS9pUoUHubkc,10330
+anemoi/utils/provenance.py,sha256=v54L9jF1JgYcclOhg3iojRl1v3ajbiWz_oc289xTgO4,9574
+anemoi/utils/text.py,sha256=pGWtDvRFoDxAnSuZJiA-GOGJOJLHsw2dAm0tfVvPKno,8599
+anemoi_utils-0.1.8.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+anemoi_utils-0.1.8.dist-info/METADATA,sha256=bEA_bMbWJubncfBmc6tMi6cJLtFgZrgth899Z60i5Yg,15079
+anemoi_utils-0.1.8.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+anemoi_utils-0.1.8.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_utils-0.1.8.dist-info/RECORD,,

anemoi_utils-0.1.6.dist-info/RECORD

@@ -1,14 +0,0 @@
-anemoi/utils/__init__.py,sha256=zZZpbKIoGWwdCOuo6YSruLR7C0GzvzI1Wzhyqaa0K7M,456
-anemoi/utils/_version.py,sha256=L9DFp_i-1xztMRdin6ZHqLFMlDKUn5bsVgheqXFwTTc,411
-anemoi/utils/checkpoints.py,sha256=j1QMmffpo6Blo6kmAJDdUzK3dqq30UWIsWghFqsY4QU,2070
-anemoi/utils/config.py,sha256=XEesqODvkuE3ZA7dnEnZ-ooBRtU6ecPmkfP65FtialA,2147
-anemoi/utils/dates.py,sha256=ZabXLecHMio9ecIJbUxNcfOFoGfkcotSwy_xw1kNQ9w,6474
-anemoi/utils/grib.py,sha256=8fxp2O7tXrFvvp0h68IZ0OhDkyH7GQFTX_6qwJ-hzBs,1692
-anemoi/utils/humanize.py,sha256=LD6dGnqChxA5j3tMhSybsAGRQzi33d_qS9pUoUHubkc,10330
-anemoi/utils/provenance.py,sha256=RiDpevwTbCzTAkU0oL7pNwST37xt4W9GOkeNNDqbAgk,9576
-anemoi/utils/text.py,sha256=mBCPH9NwLyWCU9GkXzB0K4xiDN6UQ1XGlD4T5d8x8Vk,8606
-anemoi_utils-0.1.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-anemoi_utils-0.1.6.dist-info/METADATA,sha256=gGK2zftLCNffyYnA0FhMXja0k5MOKzoixlL24O29dhk,15079
-anemoi_utils-0.1.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-anemoi_utils-0.1.6.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_utils-0.1.6.dist-info/RECORD,,