anemoi-utils 0.4.3__tar.gz → 0.4.4__tar.gz

anemoi_utils-0.4.4/.github/ci-hpc-config.yml

@@ -0,0 +1,8 @@
+build:
+  python: '3.10'
+  modules:
+    - ninja
+  parallel: 64
+
+  pytest_cmd: |
+    python -m pytest -vv -m 'not notebook and not no_cache_init and not skip_on_hpc' --cov=. --cov-report=xml

{anemoi_utils-0.4.3 → anemoi_utils-0.4.4}/CHANGELOG.md

@@ -8,14 +8,27 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Please add your functional changes to the appropriate section in the PR.
 Keep it human-readable, your future self will thank you!
 
-## [Unreleased](https://github.com/ecmwf/anemoi-utils/compare/0.4.1...HEAD)
+## [Unreleased](https://github.com/ecmwf/anemoi-utils/compare/0.4.3...HEAD)
+
+## [0.4.3](https://github.com/ecmwf/anemoi-utils/compare/0.4.1...0.4.3) - 2024-10-26
+
+## [0.4.2](https://github.com/ecmwf/anemoi-utils/compare/0.4.1...0.4.2) - 2024-10-25
+
+### Added
+- Add supporting_arrays to checkpoints
+- Add factories registry
+- Optional renaming of subcommands via `command` attribute [#34](https://github.com/ecmwf/anemoi-utils/pull/34)
+- `skip_on_hpc` pytest marker for tests that should not be run on HPC [36](https://github.com/ecmwf/anemoi-utils/pull/36)
+
 
 ## [0.4.1](https://github.com/ecmwf/anemoi-utils/compare/0.4.0...0.4.1) - 2024-10-23
 
 ## Fixed
+
 - Fix `__version__` import in init
 
 ### Changed
+
 - Fix: resolve mounted filesystems in provenance
 - Fix pre-commit regex
 - ci: extend python versions [#23] (https://github.com/ecmwf/anemoi-utils/pull/23)
@@ -26,6 +39,7 @@ Keep it human-readable, your future self will thank you!
 ### Added
 
 - Add anemoi-transform link to documentation
+- Add CONTRIBUTORS.md (#33)
 
 ## [0.3.17](https://github.com/ecmwf/anemoi-utils/compare/0.3.13...0.3.17) - 2024-10-01
 

anemoi_utils-0.4.4/CONTRIBUTORS.md

@@ -0,0 +1,13 @@
+## How to Contribute
+
+Please see the [read the docs](https://anemoi-training.readthedocs.io/en/latest/dev/contributing.html).
+
+
+## Contributors
+
+Thank you to all the wonderful people who have contributed to Anemoi. Contributions can come in many forms, including code, documentation, bug reports, feature suggestions, design, and more. A list of code-based contributors can be found [here](https://github.com/ecmwf/anemoi-utils/graphs/contributors).
+
+
+## Contributing Organisations
+
+Significant contributions have been made by the following organisations: [DWD](https://www.dwd.de/), [MET Norway](https://www.met.no/), [MeteoSwiss](https://www.meteoswiss.admin.ch/), [RMI](https://www.meteo.be/) & [ECMWF](https://www.ecmwf.int/)

{anemoi_utils-0.4.3 → anemoi_utils-0.4.4}/PKG-INFO

@@ -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 → anemoi_utils-0.4.4}/docs/conf.py

@@ -29,7 +29,7 @@ html_logo = "_static/logo.png"
 
 project = "Anemoi Utils"
 
-author = "ECMWF"
+author = "Anemoi contributors"
 
 year = datetime.datetime.now().year
 if year == 2024:
@@ -37,7 +37,7 @@ if year == 2024:
 else:
     years = "2024-%s" % (year,)
 
-copyright = "%s, ECMWF" % (years,)
+copyright = "%s, Anemoi contributors" % (years,)
 
 try:
     from anemoi.utils._version import __version__

{anemoi_utils-0.4.3 → anemoi_utils-0.4.4}/pyproject.toml

@@ -82,3 +82,8 @@ scripts.anemoi-utils = "anemoi.utils.__main__:main"
 
 [tool.setuptools_scm]
 version_file = "src/anemoi/utils/_version.py"
+
+[tool.pytest.ini_options]
+markers = [
+  "skip_on_hpc: mark a test that should not be run on HPC",
+]

{anemoi_utils-0.4.3 → anemoi_utils-0.4.4}/src/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-0.4.3 → anemoi_utils-0.4.4}/src/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-0.4.3 → anemoi_utils-0.4.4}/src/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-0.4.3 → anemoi_utils-0.4.4}/src/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-0.4.3 → anemoi_utils-0.4.4}/src/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-0.4.3 → anemoi_utils-0.4.4}/src/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-0.4.4/src/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 → anemoi_utils-0.4.4}/src/anemoi_utils.egg-info/PKG-INFO

@@ -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 → anemoi_utils-0.4.4}/src/anemoi_utils.egg-info/SOURCES.txt

@@ -3,6 +3,7 @@
 .pre-commit-config.yaml
 .readthedocs.yaml
 CHANGELOG.md
+CONTRIBUTORS.md
 LICENSE
 README.md
 pyproject.toml
@@ -42,6 +43,7 @@ src/anemoi/utils/grib.py
 src/anemoi/utils/hindcasts.py
 src/anemoi/utils/humanize.py
 src/anemoi/utils/provenance.py
+src/anemoi/utils/registry.py
 src/anemoi/utils/s3.py
 src/anemoi/utils/sanitise.py
 src/anemoi/utils/sanitize.py

anemoi_utils-0.4.3/.github/ci-hpc-config.yml

@@ -1,8 +0,0 @@
-build:
-  python: '3.10'
-  modules:
-    - ninja
-  parallel: 64
-
-  pytest_cmd: |
-    python -m pytest -vv -m 'not notebook and not no_cache_init' --cov=. --cov-report=xml