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

{anemoi_utils-0.1.6 → anemoi_utils-0.1.8}/.pre-commit-config.yaml

@@ -62,8 +62,8 @@ repos:
   hooks:
     - id: rstfmt
 
-# - repo: ../pre-commit-docconvert
-#   rev: 07f5748
-#   hooks:
-#   - id: docconvert
-#     args: ["--in-place", "--output", "numpy"]
+- repo: https://github.com/b8raoult/pre-commit-docconvert
+  rev: "0.1.4"
+  hooks:
+  - id: docconvert
+    args: ["numpy"]

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

@@ -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.6 → 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.6'
-__version_tuple__ = version_tuple = (0, 1, 6)
+__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.6 → anemoi_utils-0.1.8}/src/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-0.1.8/src/anemoi/utils/grib.py

@@ -0,0 +1,132 @@
+# (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 GRIB parameters.
+
+See https://codes.ecmwf.int/grib/param-db/ for more information.
+
+"""
+
+import logging
+import re
+
+import requests
+
+from .caching import cached
+
+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()
+    results = r.json()
+    if len(results) == 0:
+        raise KeyError(name)
+
+    if len(results) > 1:
+        names = [f'{r.get("id")} ({r.get("name")})' for r in results]
+        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]
+
+
+def shortname_to_paramid(shortname: str) -> int:
+    """Return the GRIB parameter id given its shortname.
+
+    Parameters
+    ----------
+    shortname : str
+        Parameter shortname.
+
+    Returns
+    -------
+    int
+        Parameter id.
+
+    >>> shortname_to_paramid("2t")
+    167
+
+    """
+    return _search_param(shortname)["id"]
+
+
+def paramid_to_shortname(paramid: int) -> str:
+    """Return the shortname of a GRIB parameter given its id.
+
+    Parameters
+    ----------
+    paramid : int
+        Parameter id.
+
+    Returns
+    -------
+    str
+        Parameter shortname.
+
+    >>> paramid_to_shortname(167)
+    '2t'
+
+    """
+    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.6 → anemoi_utils-0.1.8}/src/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-0.1.6 → anemoi_utils-0.1.8}/src/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 → 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.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.6 → 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.6 → 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()

anemoi_utils-0.1.6/src/anemoi/utils/grib.py

@@ -1,73 +0,0 @@
-# (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 GRIB parameters.
-
-See https://codes.ecmwf.int/grib/param-db/ for more information.
-
-"""
-
-import re
-
-import requests
-
-
-def _search(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()
-    results = r.json()
-    if len(results) == 0:
-        raise KeyError(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)}")
-
-    return results[0]
-
-
-def shortname_to_paramid(shortname: str) -> int:
-    """Return the GRIB parameter id given its shortname.
-
-    Parameters
-    ----------
-    shortname : str
-        Parameter shortname.
-
-    Returns
-    -------
-    int
-        Parameter id.
-
-
-    >>> shortname_to_paramid("2t")
-    167
-
-    """
-    return _search(shortname)["id"]
-
-
-def paramid_to_shortname(paramid: int) -> str:
-    """Return the shortname of a GRIB parameter given its id.
-
-    Parameters
-    ----------
-    paramid : int
-        Parameter id.
-
-    Returns
-    -------
-    str
-        Parameter shortname.
-
-
-    >>> paramid_to_shortname(167)
-    '2t'
-
-    """
-    return _search(str(paramid))["shortname"]