anemoi-utils 0.4.3__py3-none-any.whl → 0.4.4__py3-none-any.whl

anemoi/utils/__init__.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# (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.

anemoi/utils/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3'
-__version_tuple__ = version_tuple = (0, 4, 3)
+__version__ = version = '0.4.4'
+__version_tuple__ = version_tuple = (0, 4, 4)

anemoi/utils/checkpoints.py

@@ -27,7 +27,7 @@ DEFAULT_NAME = "ai-models.json"
 DEFAULT_FOLDER = "anemoi-metadata"
 
 
-def has_metadata(path: str, name: str = DEFAULT_NAME) -> bool:
+def has_metadata(path: str, *, name: str = DEFAULT_NAME) -> bool:
     """Check if a checkpoint file has a metadata file
 
     Parameters
@@ -49,13 +49,26 @@ def has_metadata(path: str, name: str = DEFAULT_NAME) -> bool:
     return False
 
 
-def load_metadata(path: str, name: str = DEFAULT_NAME) -> dict:
+def metadata_root(path: str, *, name: str = DEFAULT_NAME) -> bool:
+
+    with zipfile.ZipFile(path, "r") as f:
+        for b in f.namelist():
+            if os.path.basename(b) == name:
+                return os.path.dirname(b)
+    raise ValueError(f"Could not find '{name}' in {path}.")
+
+
+def load_metadata(path: str, *, supporting_arrays=False, name: str = DEFAULT_NAME) -> dict:
     """Load metadata from a checkpoint file
 
     Parameters
     ----------
     path : str
         The path to the checkpoint file
+
+    supporting_arrays: bool, optional
+        If True, the function will return a dictionary with the supporting arrays
+
     name : str, optional
         The name of the metadata file in the zip archive
 
@@ -79,12 +92,29 @@ def load_metadata(path: str, name: str = DEFAULT_NAME) -> dict:
 
     if metadata is not None:
         with zipfile.ZipFile(path, "r") as f:
-            return json.load(f.open(metadata, "r"))
+            metadata = json.load(f.open(metadata, "r"))
+            if supporting_arrays:
+                metadata["supporting_arrays"] = load_supporting_arrays(f, metadata.get("supporting_arrays", {}))
+                return metadata, supporting_arrays
+
+            return metadata
     else:
         raise ValueError(f"Could not find '{name}' in {path}.")
 
 
-def save_metadata(path, metadata, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> None:
+def load_supporting_arrays(zipf, entries) -> dict:
+    import numpy as np
+
+    supporting_arrays = {}
+    for key, entry in entries.items():
+        supporting_arrays[key] = np.frombuffer(
+            zipf.read(entry["path"]),
+            dtype=entry["dtype"],
+        ).reshape(entry["shape"])
+    return supporting_arrays
+
+
+def save_metadata(path, metadata, *, supporting_arrays=None, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> None:
     """Save metadata to a checkpoint file
 
     Parameters
@@ -93,6 +123,8 @@ def save_metadata(path, metadata, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> N
         The path to the checkpoint file
     metadata : JSON
         A JSON serializable object
+    supporting_arrays: dict, optional
+        A dictionary of supporting NumPy arrays
     name : str, optional
         The name of the metadata file in the zip archive
     folder : str, optional
@@ -118,19 +150,41 @@ def save_metadata(path, metadata, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> N
 
         directory = list(directories)[0]
 
+        LOG.info("Adding extra information to checkpoint %s", path)
         LOG.info("Saving metadata to %s/%s/%s", directory, folder, name)
 
+        metadata = metadata.copy()
+        if supporting_arrays is not None:
+            metadata["supporting_arrays_paths"] = {
+                key: dict(path=f"{directory}/{folder}/{key}.numpy", shape=value.shape, dtype=str(value.dtype))
+                for key, value in supporting_arrays.items()
+            }
+        else:
+            metadata["supporting_arrays_paths"] = {}
+
         zipf.writestr(
             f"{directory}/{folder}/{name}",
             json.dumps(metadata),
         )
 
+        for name, entry in metadata["supporting_arrays_paths"].items():
+            value = supporting_arrays[name]
+            LOG.info(
+                "Saving supporting array `%s` to %s (shape=%s, dtype=%s)",
+                name,
+                entry["path"],
+                entry["shape"],
+                entry["dtype"],
+            )
+            zipf.writestr(entry["path"], value.tobytes())
+
 
-def _edit_metadata(path, name, callback):
+def _edit_metadata(path, name, callback, supporting_arrays=None):
     new_path = f"{path}.anemoi-edit-{time.time()}-{os.getpid()}.tmp"
 
     found = False
 
+    directory = None
     with TemporaryDirectory() as temp_dir:
         zipfile.ZipFile(path, "r").extractall(temp_dir)
         total = 0
@@ -141,10 +195,21 @@ def _edit_metadata(path, name, callback):
                 if f == name:
                     found = True
                     callback(full)
+                    directory = os.path.dirname(full)
 
         if not found:
             raise ValueError(f"Could not find '{name}' in {path}")
 
+        if supporting_arrays is not None:
+
+            for key, entry in supporting_arrays.items():
+                value = entry.tobytes()
+                fname = os.path.join(directory, f"{key}.numpy")
+                os.makedirs(os.path.dirname(fname), exist_ok=True)
+                with open(fname, "wb") as f:
+                    f.write(value)
+                    total += 1
+
         with zipfile.ZipFile(new_path, "w", zipfile.ZIP_DEFLATED) as zipf:
             with tqdm.tqdm(total=total, desc="Rebuilding checkpoint") as pbar:
                 for root, dirs, files in os.walk(temp_dir):
@@ -158,7 +223,7 @@ def _edit_metadata(path, name, callback):
     LOG.info("Updated metadata in %s", path)
 
 
-def replace_metadata(path, metadata, name=DEFAULT_NAME):
+def replace_metadata(path, metadata, supporting_arrays=None, *, name=DEFAULT_NAME):
 
     if not isinstance(metadata, dict):
         raise ValueError(f"metadata must be a dict, got {type(metadata)}")
@@ -170,14 +235,14 @@ def replace_metadata(path, metadata, name=DEFAULT_NAME):
         with open(full, "w") as f:
             json.dump(metadata, f)
 
-    _edit_metadata(path, name, callback)
+    return _edit_metadata(path, name, callback, supporting_arrays)
 
 
-def remove_metadata(path, name=DEFAULT_NAME):
+def remove_metadata(path, *, name=DEFAULT_NAME):
 
     LOG.info("Removing metadata '%s' from %s", name, path)
 
     def callback(full):
         os.remove(full)
 
-    _edit_metadata(path, name, callback)
+    return _edit_metadata(path, name, callback)

anemoi/utils/cli.py

@@ -96,8 +96,20 @@ def register_commands(here, package, select, fail=None):
             continue
 
         obj = select(imported)
-        if obj is not None:
-            result[name] = obj
+        if obj is None:
+            continue
+
+        if hasattr(obj, "command"):
+            name = obj.command
+
+        if name in result:
+            msg = f"Duplicate command '{name}', please choose a different command name for {type(obj)}"
+            raise ValueError(msg)
+        if " " in name:
+            msg = f"Commands cannot contain spaces: '{name}' in {type(obj)}"
+            raise ValueError(msg)
+
+        result[name] = obj
 
     for name, e in not_available.items():
         if fail is None:

anemoi/utils/config.py

@@ -358,7 +358,7 @@ def check_config_mode(name="settings.toml", secrets_name=None, secrets=None) ->
         CHECKED[name] = True
 
 
-def find(metadata, what, result=None):
+def find(metadata, what, result=None, *, select: callable = None):
     if result is None:
         result = []
 
@@ -369,7 +369,8 @@ def find(metadata, what, result=None):
 
     if isinstance(metadata, dict):
         if what in metadata:
-            result.append(metadata[what])
+            if select is None or select(metadata[what]):
+                result.append(metadata[what])
 
         for k, v in metadata.items():
             find(v, what, result)

anemoi/utils/dates.py

@@ -107,8 +107,8 @@ def as_datetime_list(date, default_increment=1):
     return list(_as_datetime_list(date, default_increment))
 
 
-def frequency_to_timedelta(frequency) -> datetime.timedelta:
-    """Convert a frequency to a timedelta object.
+def as_timedelta(frequency) -> datetime.timedelta:
+    """Convert anything to a timedelta object.
 
     Parameters
     ----------
@@ -171,6 +171,11 @@ def frequency_to_timedelta(frequency) -> datetime.timedelta:
     raise ValueError(f"Cannot convert frequency {frequency} to timedelta")
 
 
+def frequency_to_timedelta(frequency) -> datetime.timedelta:
+    """Convert a frequency to a timedelta object."""
+    return as_timedelta(frequency)
+
+
 def frequency_to_string(frequency) -> str:
     """Convert a frequency (i.e. a datetime.timedelta) to a string.
 

anemoi/utils/registry.py

@@ -0,0 +1,98 @@
+# (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 importlib
+import logging
+import os
+import sys
+
+import entrypoints
+
+LOG = logging.getLogger(__name__)
+
+
+class Wrapper:
+    """A wrapper for the registry"""
+
+    def __init__(self, name, registry):
+        self.name = name
+        self.registry = registry
+
+    def __call__(self, factory):
+        self.registry.register(self.name, factory)
+        return factory
+
+
+class Registry:
+    """A registry of factories"""
+
+    def __init__(self, package):
+
+        self.package = package
+        self.registered = {}
+        self.kind = package.split(".")[-1]
+
+    def register(self, name: str, factory: callable = None):
+
+        if factory is None:
+            return Wrapper(name, self)
+
+        self.registered[name] = factory
+
+    def _load(self, file):
+        name, _ = os.path.splitext(file)
+        try:
+            importlib.import_module(f".{name}", package=self.package)
+        except Exception:
+            LOG.warning(f"Error loading filter '{self.package}.{name}'", exc_info=True)
+
+    def lookup(self, name: str) -> callable:
+        if name in self.registered:
+            return self.registered[name]
+
+        directory = sys.modules[self.package].__path__[0]
+
+        for file in os.listdir(directory):
+
+            if file[0] == ".":
+                continue
+
+            if file == "__init__.py":
+                continue
+
+            full = os.path.join(directory, file)
+            if os.path.isdir(full):
+                if os.path.exists(os.path.join(full, "__init__.py")):
+                    self._load(file)
+                continue
+
+            if file.endswith(".py"):
+                self._load(file)
+
+        entrypoint_group = f"anemoi.{self.kind}"
+        for entry_point in entrypoints.get_group_all(entrypoint_group):
+            if entry_point.name == name:
+                if name in self.registered:
+                    LOG.warning(
+                        f"Overwriting builtin '{name}' from {self.package} with plugin '{entry_point.module_name}'"
+                    )
+                self.registered[name] = entry_point.load()
+
+        if name not in self.registered:
+            raise ValueError(f"Cannot load '{name}' from {self.package}")
+
+        return self.registered[name]
+
+    def create(self, name: str, *args, **kwargs):
+        factory = self.lookup(name)
+        return factory(*args, **kwargs)
+
+    def __call__(self, name: str, *args, **kwargs):
+        return self.create(name, *args, **kwargs)

{anemoi_utils-0.4.3.dist-info → anemoi_utils-0.4.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.4.3
+Version: 0.4.4
 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.4.3.dist-info → anemoi_utils-0.4.4.dist-info}/RECORD

@@ -1,15 +1,16 @@
-anemoi/utils/__init__.py,sha256=se9GXELQQbXBN_5O8lolt4vAPuIavGDKzNWUMmsERWI,729
+anemoi/utils/__init__.py,sha256=0u0eIdu5-H1frf6V4KHpNmlh_SS-bJnxjzIejlsLqdw,702
 anemoi/utils/__main__.py,sha256=cLA2PidDTOUHaDGzd0_E5iioKYNe-PSTv567Y2fuwQk,723
-anemoi/utils/_version.py,sha256=mDTy6z4iq5XKGigbX2Q9ub9_BzxGHksMlMjpkDAl_hE,411
+anemoi/utils/_version.py,sha256=t0Mfy7vCENQWSYE4xRfluKBHgAYWK48fjSubsSGEQEI,411
 anemoi/utils/caching.py,sha256=0cznpvaaox14NSVi-Q3PqumfuGtXo0YNcEFwDPxvMZw,1948
-anemoi/utils/checkpoints.py,sha256=irFLb0R2M7XkVNduYGrcs7lR1rZ8r6J6csQQyvyYqYs,5292
-anemoi/utils/cli.py,sha256=teCHTgo77hcF75Tn0hHdfsPZ5_YEYJNN2VqKl30wWV8,3734
-anemoi/utils/config.py,sha256=VtHzYO_fcRpXH2zrPJQDv3C3Bi2XhPZD7SvNEsvxelA,9784
-anemoi/utils/dates.py,sha256=hlO32nkCVu8PCGyVfV2RRREqSujfpwEn75xhN1IZaTs,12314
+anemoi/utils/checkpoints.py,sha256=719HjvY8zyseQxwk-08rMB3X3vI_o26Sq_AiBFkZ8Fk,7802
+anemoi/utils/cli.py,sha256=rmMP60VY3em99rQP6TCrKibMngWwVe5h_0GDcf16c5U,4117
+anemoi/utils/config.py,sha256=Fpy4wzj3dahTwwO75Iet6zmQQFGmTvhXml6-EsTEvgk,9873
+anemoi/utils/dates.py,sha256=wwYD5_QI7EWY_jhpENNYtL5O7fjwYkzmqHkNoayvmrY,12452
 anemoi/utils/grib.py,sha256=zBICyOsYtR_9px1C5UDT6wL_D6kiIhUi_00kjFmas5c,3047
 anemoi/utils/hindcasts.py,sha256=OUOY2nDa3LBnzJ3ncgANzJDapouh82KgVyofDAu7K_Q,1426
 anemoi/utils/humanize.py,sha256=tSQkiUHiDj3VYk-DeruHp9P79sJO1b0whsPBphqy9qA,16627
 anemoi/utils/provenance.py,sha256=SqOiNoY1y36Zec83Pjt7OhihbwxMyknscfmogHCuriA,10894
+anemoi/utils/registry.py,sha256=m7jNJKTkMPOv_muUCn1RPeVW8D8DtggRIhY36RtcQfU,2847
 anemoi/utils/s3.py,sha256=LMljA5OoaVcgZcg_rmH-_LOX4uicMZl1FY64Bx4uOO8,18694
 anemoi/utils/sanitise.py,sha256=MqEMLwVZ1jSemLDBoQXuJyXKIfyR0gzYi7DoITBcir8,2866
 anemoi/utils/sanitize.py,sha256=43ZKDcfVpeXSsJ9TFEc9aZnD6oe2cUh151XnDspM98M,462
@@ -19,9 +20,9 @@ anemoi/utils/commands/__init__.py,sha256=qAybFZPBBQs0dyx7dZ3X5JsLpE90pwrqt1vSV7c
 anemoi/utils/commands/config.py,sha256=cAt6yYF3rN1shr57c8wXsgMEvQMRN08l_fCdA1Ux9gE,839
 anemoi/utils/mars/__init__.py,sha256=RAeY8gJ7ZvsPlcIvrQ4fy9xVHs3SphTAPw_XJDtNIKo,1750
 anemoi/utils/mars/mars.yaml,sha256=R0dujp75lLA4wCWhPeOQnzJ45WZAYLT8gpx509cBFlc,66
-anemoi_utils-0.4.3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-anemoi_utils-0.4.3.dist-info/METADATA,sha256=U7EY8vK4PeGqk4-p7FR3nd7_6m6ytmKXqwKIXyOeJ9o,15171
-anemoi_utils-0.4.3.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-anemoi_utils-0.4.3.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
-anemoi_utils-0.4.3.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_utils-0.4.3.dist-info/RECORD,,
+anemoi_utils-0.4.4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+anemoi_utils-0.4.4.dist-info/METADATA,sha256=9Tnb0OBqtTnOcJPBg3yRRYKLEgPW0FByS_GkAke-mGg,15171
+anemoi_utils-0.4.4.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+anemoi_utils-0.4.4.dist-info/entry_points.txt,sha256=LENOkn88xzFQo-V59AKoA_F_cfYQTJYtrNTtf37YgHY,60
+anemoi_utils-0.4.4.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_utils-0.4.4.dist-info/RECORD,,

{anemoi_utils-0.4.3.dist-info → anemoi_utils-0.4.4.dist-info}/WHEEL

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