anemoi-utils 0.4.10__py3-none-any.whl → 0.4.12__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.4.10'
-__version_tuple__ = version_tuple = (0, 4, 10)
+__version__ = version = '0.4.12'
+__version_tuple__ = version_tuple = (0, 4, 12)

anemoi/utils/caching.py

@@ -14,47 +14,31 @@ import os
 import time
 from threading import Lock
 
+import numpy as np
+
 LOCK = Lock()
 CACHE = {}
 
 
-def cache(key, proc, collection="default", expires=None):
-
-    key = json.dumps(key, sort_keys=True)
-    m = hashlib.md5()
-    m.update(key.encode("utf-8"))
-    m = m.hexdigest()
-
-    if m in CACHE:
-        return CACHE[m]
-
-    path = os.path.join(os.path.expanduser("~"), ".cache", "anemoi", collection)
-    os.makedirs(path, exist_ok=True)
-
-    filename = os.path.join(path, m)
-    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
+def _get_cache_path(collection):
+    return os.path.join(os.path.expanduser("~"), ".cache", "anemoi", collection)
 
-    with open(filename, "w") as f:
-        json.dump(data, f)
 
-    CACHE[m] = value
-    return value
+def clean_cache(collection="default"):
+    global CACHE
+    CACHE = {}
+    path = _get_cache_path(collection)
+    if not os.path.exists(path):
+        return
+    for filename in os.listdir(path):
+        os.remove(os.path.join(path, filename))
 
 
-class cached:
-    """Decorator to cache the result of a function."""
+class Cacher:
+    """This class implements a simple caching mechanism.
+    Private class, do not use directly"""
 
-    def __init__(self, collection="default", expires=None):
+    def __init__(self, collection, expires):
         self.collection = collection
         self.expires = expires
 
@@ -64,11 +48,77 @@ class cached:
 
         def wrapped(*args, **kwargs):
             with LOCK:
-                return cache(
+                return self.cache(
                     (full, args, kwargs),
                     lambda: func(*args, **kwargs),
-                    self.collection,
-                    self.expires,
                 )
 
         return wrapped
+
+    def cache(self, key, proc):
+
+        key = json.dumps(key, sort_keys=True)
+        m = hashlib.md5()
+        m.update(key.encode("utf-8"))
+        m = m.hexdigest()
+
+        if m in CACHE:
+            return CACHE[m]
+
+        path = _get_cache_path(self.collection)
+
+        filename = os.path.join(path, m) + self.ext
+        if os.path.exists(filename):
+            data = self.load(filename)
+            if self.expires is None or data["expires"] > time.time():
+                if data["key"] == key:
+                    return data["value"]
+
+        value = proc()
+        data = {"key": key, "value": value}
+        if self.expires is not None:
+            data["expires"] = time.time() + self.expires
+
+        os.makedirs(path, exist_ok=True)
+        temp_filename = self.save(filename, data)
+        os.rename(temp_filename, filename)
+
+        CACHE[m] = value
+        return value
+
+
+class JsonCacher(Cacher):
+    ext = ""
+
+    def save(self, path, data):
+        temp_path = path + ".tmp"
+        with open(temp_path, "w") as f:
+            json.dump(data, f)
+        return temp_path
+
+    def load(self, path):
+        with open(path, "r") as f:
+            return json.load(f)
+
+
+class NpzCacher(Cacher):
+    ext = ".npz"
+
+    def save(self, path, data):
+        temp_path = path + ".tmp.npz"
+        np.savez(temp_path, **data)
+        return temp_path
+
+    def load(self, path):
+        return np.load(path, allow_pickle=True)
+
+
+# PUBLIC API
+def cached(collection="default", expires=None, encoding="json"):
+    """Decorator to cache the result of a function.
+
+    Default is to use a json file to store the cache, but you can also use npz files
+    to cache dict of numpy arrays.
+
+    """
+    return dict(json=JsonCacher, npz=NpzCacher)[encoding](collection, expires)

anemoi/utils/commands/requests.py

@@ -6,6 +6,7 @@
 # nor does it submit to any jurisdiction.
 
 import json
+import sys
 
 from anemoi.utils.mars.requests import print_request
 
@@ -22,8 +23,11 @@ class Requests(Command):
         command_parser.add_argument("--only-one-field", action="store_true")
 
     def run(self, args):
-        with open(args.input) as f:
-            requests = json.load(f)
+        if args.input == "-":
+            requests = json.load(sys.stdin)
+        else:
+            with open(args.input) as f:
+                requests = json.load(f)
 
         if args.only_one_field:
             for r in requests:

anemoi/utils/config.py

@@ -223,6 +223,23 @@ def load_any_dict_format(path) -> dict:
         if path.endswith(".toml"):
             with open(path, "rb") as f:
                 return tomllib.load(f)
+
+        if path == "-":
+            import sys
+
+            config = sys.stdin.read()
+
+            parsers = [(yaml.safe_load, "yaml"), (json.loads, "json"), (tomllib.loads, "toml")]
+
+            for parser, parser_type in parsers:
+                try:
+                    LOG.debug(f"Trying {parser_type} parser for stdin")
+                    return parser(config)
+                except Exception:
+                    pass
+
+            raise ValueError("Failed to parse configuration from stdin")
+
     except (json.JSONDecodeError, yaml.YAMLError, tomllib.TOMLDecodeError) as e:
         LOG.warning(f"Failed to parse config file {path}", exc_info=e)
         raise ValueError(f"Failed to parse config file {path} [{e}]")

anemoi/utils/dates.py

@@ -155,12 +155,19 @@ def as_timedelta(frequency) -> datetime.timedelta:
         unit = {"h": "hours", "d": "days", "s": "seconds", "m": "minutes"}[unit]
         return datetime.timedelta(**{unit: v})
 
-    m = frequency.split(":")
-    if len(m) == 2:
-        return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]))
-
-    if len(m) == 3:
-        return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]), seconds=int(m[2]))
+    if re.match(r"^\d+:\d+(:\d+)?$", frequency):
+        m = frequency.split(":")
+        if len(m) == 2:
+            return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]))
+
+        if len(m) == 3:
+            return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]), seconds=int(m[2]))
+
+    if re.match(r"^\d+ days?, \d+:\d+:\d+$", frequency):
+        m = frequency.split(", ")
+        days = int(m[0].split()[0])
+        hms = m[1].split(":")
+        return datetime.timedelta(days=days, hours=int(hms[0]), minutes=int(hms[1]), seconds=int(hms[2]))
 
     # ISO8601
     try:

anemoi/utils/devtools.py

@@ -0,0 +1,83 @@
+# (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 cartopy.crs as ccrs
+import cartopy.feature as cfeature
+import matplotlib.pyplot as plt
+import matplotlib.tri as tri
+import numpy as np
+
+"""FOR DEVELOPMENT PURPOSES ONLY
+
+This module contains
+
+"""
+
+# TODO: use earthkit-plots
+
+
+def fix(lons):
+    return np.where(lons > 180, lons - 360, lons)
+
+
+def plot_values(
+    values, latitudes, longitudes, title=None, missing_value=None, min_value=None, max_value=None, **kwargs
+):
+
+    _, ax = plt.subplots(subplot_kw={"projection": ccrs.PlateCarree()})
+    ax.coastlines()
+    ax.add_feature(cfeature.BORDERS, linestyle=":")
+
+    missing_values = np.isnan(values)
+
+    if missing_value is None:
+        values = values[~missing_values]
+        longitudes = longitudes[~missing_values]
+        latitudes = latitudes[~missing_values]
+    else:
+        values = np.where(missing_values, missing_value, values)
+
+    if max_value is not None:
+        values = np.where(values > max_value, max_value, values)
+
+    if min_value is not None:
+        values = np.where(values < min_value, min_value, values)
+
+    triangulation = tri.Triangulation(fix(longitudes), latitudes)
+
+    levels = kwargs.pop("levels", 10)
+
+    _ = ax.tricontourf(triangulation, values, levels=levels, transform=ccrs.PlateCarree())
+
+    options = dict(
+        levels=levels,
+        colors="black",
+        linewidths=0.5,
+        transform=ccrs.PlateCarree(),
+    )
+
+    options.update(kwargs)
+
+    ax.tricontour(
+        triangulation,
+        values,
+        **options,
+    )
+
+    if title is not None:
+        ax.set_title(title)
+
+    return ax
+
+
+def plot_field(field, title=None, **kwargs):
+    values = field.to_numpy(flatten=True)
+    latitudes, longitudes = field.grid_points()
+    return plot_values(values, latitudes, longitudes, title=title, **kwargs)

anemoi/utils/grids.py

@@ -0,0 +1,97 @@
+# (C) Copyright 2025 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.
+
+
+"""Utilities for working with grids.
+
+"""
+
+import logging
+import os
+from io import BytesIO
+
+import numpy as np
+import requests
+
+from .caching import cached
+
+LOG = logging.getLogger(__name__)
+
+
+GRIDS_URL_PATTERN = "https://get.ecmwf.int/repository/anemoi/grids/grid-{name}.npz"
+
+
+def xyz_to_latlon(x, y, z):
+    return (
+        np.rad2deg(np.arcsin(np.minimum(1.0, np.maximum(-1.0, z)))),
+        np.rad2deg(np.arctan2(y, x)),
+    )
+
+
+def latlon_to_xyz(lat, lon, radius=1.0):
+    # https://en.wikipedia.org/wiki/Geographic_coordinate_conversion#From_geodetic_to_ECEF_coordinates
+    # We assume that the Earth is a sphere of radius 1 so N(phi) = 1
+    # We assume h = 0
+    #
+    phi = np.deg2rad(lat)
+    lda = np.deg2rad(lon)
+
+    cos_phi = np.cos(phi)
+    cos_lda = np.cos(lda)
+    sin_phi = np.sin(phi)
+    sin_lda = np.sin(lda)
+
+    x = cos_phi * cos_lda * radius
+    y = cos_phi * sin_lda * radius
+    z = sin_phi * radius
+
+    return x, y, z
+
+
+def nearest_grid_points(source_latitudes, source_longitudes, target_latitudes, target_longitudes):
+    from scipy.spatial import cKDTree
+
+    source_xyz = latlon_to_xyz(source_latitudes, source_longitudes)
+    source_points = np.array(source_xyz).transpose()
+
+    target_xyz = latlon_to_xyz(target_latitudes, target_longitudes)
+    target_points = np.array(target_xyz).transpose()
+
+    _, indices = cKDTree(source_points).query(target_points, k=1)
+    return indices
+
+
+@cached(collection="grids", encoding="npz")
+def _grids(name):
+    from anemoi.utils.config import load_config
+
+    user_path = load_config().get("utils", {}).get("grids_path")
+    if user_path:
+        path = os.path.expanduser(os.path.join(user_path, f"grid-{name}.npz"))
+        if os.path.exists(path):
+            LOG.warning("Loading grids from custom user path %s", path)
+            with open(path, "rb") as f:
+                return f.read()
+        else:
+            LOG.warning("Custom user path %s does not exist", path)
+
+    url = GRIDS_URL_PATTERN.format(name=name.lower())
+    LOG.warning("Downloading grids from %s", url)
+    response = requests.get(url)
+    response.raise_for_status()
+    return response.content
+
+
+def grids(name):
+    if name.endswith(".npz"):
+        return dict(np.load(name))
+
+    data = _grids(name)
+    npz = np.load(BytesIO(data))
+    return dict(npz)

anemoi/utils/humanize.py

@@ -689,3 +689,36 @@ def print_dates(dates) -> None:
         A list of dates, as datetime objects or strings.
     """
     print(compress_dates(dates))
+
+
+def make_list_int(value) -> list:
+    """Convert a string like "1/2/3" or "1/to/3" or "1/to/10/by/2" to a list of integers.
+
+    Parameters
+    ----------
+    value : str, list, tuple, int
+        The value to convert to a list of integers.
+
+    Returns
+    -------
+    list
+        A list of integers.
+    """
+    if isinstance(value, str):
+        if "/" not in value:
+            return [int(value)]
+        bits = value.split("/")
+        if len(bits) == 3 and bits[1].lower() == "to":
+            value = list(range(int(bits[0]), int(bits[2]) + 1, 1))
+
+        elif len(bits) == 5 and bits[1].lower() == "to" and bits[3].lower() == "by":
+            value = list(range(int(bits[0]), int(bits[2]) + int(bits[4]), int(bits[4])))
+
+    if isinstance(value, list):
+        return value
+    if isinstance(value, tuple):
+        return value
+    if isinstance(value, int):
+        return [value]
+
+    raise ValueError(f"Cannot make list from {value}")

{anemoi_utils-0.4.10.dist-info → anemoi_utils-0.4.12.dist-info}/LICENSE

@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright [yyyy] [name of copyright owner]
+   Copyright 2024-2025 Anemoi Contributors
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

{anemoi_utils-0.4.10.dist-info → anemoi_utils-0.4.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: anemoi-utils
-Version: 0.4.10
+Version: 0.4.12
 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
@@ -191,7 +191,7 @@ License:                                  Apache License
               same "printed page" as the copyright notice for easier
               identification within third-party archives.
         
-           Copyright [yyyy] [name of copyright owner]
+           Copyright 2024-2025 Anemoi Contributors
         
            Licensed under the Apache License, Version 2.0 (the "License");
            you may not use this file except in compliance with the License.
@@ -226,6 +226,7 @@ Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: aniso8601
 Requires-Dist: importlib-metadata; python_version < "3.10"
+Requires-Dist: numpy
 Requires-Dist: python-dateutil
 Requires-Dist: pyyaml
 Requires-Dist: tomli; python_version < "3.11"

{anemoi_utils-0.4.10.dist-info → anemoi_utils-0.4.12.dist-info}/RECORD

@@ -1,15 +1,17 @@
 anemoi/utils/__init__.py,sha256=0u0eIdu5-H1frf6V4KHpNmlh_SS-bJnxjzIejlsLqdw,702
 anemoi/utils/__main__.py,sha256=5NW2A3OgTimB4ptwYThivIRSeCrvabMuvnr8mmnVx0E,715
-anemoi/utils/_version.py,sha256=OK_y1wkdggDGF27DpKS31C6p5FSDPep8dDgBhyKlfYE,413
-anemoi/utils/caching.py,sha256=0cznpvaaox14NSVi-Q3PqumfuGtXo0YNcEFwDPxvMZw,1948
+anemoi/utils/_version.py,sha256=sg-lGJqAzZSE-Cr-6X2XM20x4P3h_gsFXHtA5vFf4qs,413
+anemoi/utils/caching.py,sha256=UkHQOKcBoB6xni83qgTEGfkyam7iqu1YiuBLZQIn9RM,3208
 anemoi/utils/checkpoints.py,sha256=q8QqKlZ6qChjzEfq7KM1gVXuyqgsVRGIb4dJFtkGk58,7774
 anemoi/utils/cli.py,sha256=rmMP60VY3em99rQP6TCrKibMngWwVe5h_0GDcf16c5U,4117
 anemoi/utils/compatibility.py,sha256=0_nIcbdQbNMrS6AkqrBgJGJlSJXW8R23ncaZaDwdJ4c,2190
-anemoi/utils/config.py,sha256=c7dh9jRKHQ4DmwNTuniOZcrqZQgxU5VLQqH6rHSE7bI,10460
-anemoi/utils/dates.py,sha256=wwYD5_QI7EWY_jhpENNYtL5O7fjwYkzmqHkNoayvmrY,12452
+anemoi/utils/config.py,sha256=ADCksBl6qrJpU9eebVebwkthNanzgCOxBgzlqCUTwbg,10962
+anemoi/utils/dates.py,sha256=brsBEje1w8sJ1RbVJl-_nxbSUYsn1h0zaF6EY54zpDs,12785
+anemoi/utils/devtools.py,sha256=Mns5vU9o2HrO4zS1e0-W4gBIhk8xHrhcB7wLR_q6OiA,2172
 anemoi/utils/grib.py,sha256=zBICyOsYtR_9px1C5UDT6wL_D6kiIhUi_00kjFmas5c,3047
+anemoi/utils/grids.py,sha256=tqhH8ZiRS9Re7xQZHVKtl8bBtqH_kDVOGHfDTxu3RuI,2708
 anemoi/utils/hindcasts.py,sha256=TEYDmrZUajuhp_dfWeg6z5c6XfntE-mwugUQJyAgUco,1419
-anemoi/utils/humanize.py,sha256=tSQkiUHiDj3VYk-DeruHp9P79sJO1b0whsPBphqy9qA,16627
+anemoi/utils/humanize.py,sha256=ZD5UMD7m79I1h_IoIQFnd4FZR5K9VARw9cVaNrD1QdM,17579
 anemoi/utils/logs.py,sha256=o0xXiO2BdG_bZkljxxI2TKlCiA5QbWHgAUlYM53lirE,1058
 anemoi/utils/provenance.py,sha256=SqOiNoY1y36Zec83Pjt7OhihbwxMyknscfmogHCuriA,10894
 anemoi/utils/registry.py,sha256=Iit_CfTGuoVffXkZA2A5mUXb4AdGIUX9TpnUqWT4HJ0,4291
@@ -20,16 +22,16 @@ anemoi/utils/text.py,sha256=Xfr_3wvsjg7m-BwvdJVz1bV6f5KNMnGIIFRtXaiMfbs,10496
 anemoi/utils/timer.py,sha256=Twnr3GZu-n0WzgboELRKJWs87qyDYqy6Dwr9cQ_JG18,1803
 anemoi/utils/commands/__init__.py,sha256=O5W3yHZywRoAqmRUioAr3zMCh0hGVV18wZYGvc00ioM,698
 anemoi/utils/commands/config.py,sha256=zt4PFATYJ-zs0C5mpUlrQ4Fj5m1kM3CcsszUP1VBbzA,816
-anemoi/utils/commands/requests.py,sha256=7joRYnJUzJh5O8Pqkqa-s9M9woHy-Z86czp00uCZXGc,1448
+anemoi/utils/commands/requests.py,sha256=NZPtQ2-PEtj7w5nNzTq5Tvj5axtsVAeuTUOKIQt-Faw,1555
 anemoi/utils/mars/__init__.py,sha256=kvbu-gSaYI9jSNEzfQltrtHPVIameYGoLjOJKwI7x_U,1723
 anemoi/utils/mars/mars.yaml,sha256=R0dujp75lLA4wCWhPeOQnzJ45WZAYLT8gpx509cBFlc,66
 anemoi/utils/mars/requests.py,sha256=0khe_mbq4GNueR_B8fGPTBoWHtCfjQvtoKXOSVm6La4,759
 anemoi/utils/remote/__init__.py,sha256=-_AA1xm9GpagW5zP0PGpz-3SRKEUjw_AGSNd_bhuh7g,11639
 anemoi/utils/remote/s3.py,sha256=hykbVlh1_aFI00FWjgm_FWIMfVCTFiQf_cq8_gAo31s,11976
 anemoi/utils/remote/ssh.py,sha256=3lqFpY9CEor_DvIK9ZxSmj3rND-366Sm9R3Vw61sWSs,4695
-anemoi_utils-0.4.10.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-anemoi_utils-0.4.10.dist-info/METADATA,sha256=CHhddg_iojS8zfV67wBSG0DIlGeZq7ShPf7UiSAdu3Y,15237
-anemoi_utils-0.4.10.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-anemoi_utils-0.4.10.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
-anemoi_utils-0.4.10.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_utils-0.4.10.dist-info/RECORD,,
+anemoi_utils-0.4.12.dist-info/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
+anemoi_utils-0.4.12.dist-info/METADATA,sha256=-9E31lmUv9bhpu_e4ZlLRgWtZdrRdeAo5KFzfJg42HU,15255
+anemoi_utils-0.4.12.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+anemoi_utils-0.4.12.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
+anemoi_utils-0.4.12.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_utils-0.4.12.dist-info/RECORD,,

{anemoi_utils-0.4.10.dist-info → anemoi_utils-0.4.12.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any