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

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/.gitignore

@@ -185,3 +185,4 @@ _build/
 ~*
 *.sync
 _version.py
+*.code-workspace

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

@@ -61,6 +61,7 @@ repos:
   rev: v0.0.14
   hooks:
     - id: rstfmt
+      exclude: 'cli/.*' # Because we use argparse
 
 - repo: https://github.com/b8raoult/pre-commit-docconvert
   rev: "0.1.4"

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.8
+Version: 0.2.0
 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
@@ -223,6 +223,8 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: tomli
+Requires-Dist: pyyaml
+Requires-Dist: tqdm
 Provides-Extra: provenance
 Requires-Dist: GitPython; extra == "provenance"
 Requires-Dist: nvsmi; extra == "provenance"
@@ -234,6 +236,11 @@ Provides-Extra: docs
 Requires-Dist: tomli; extra == "docs"
 Requires-Dist: termcolor; extra == "docs"
 Requires-Dist: requests; extra == "docs"
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx_rtd_theme; extra == "docs"
+Requires-Dist: nbsphinx; extra == "docs"
+Requires-Dist: pandoc; extra == "docs"
+Requires-Dist: sphinx_argparse; extra == "docs"
 Provides-Extra: all
 Requires-Dist: tomli; extra == "all"
 Requires-Dist: GitPython; extra == "all"
@@ -246,7 +253,3 @@ Requires-Dist: GitPython; extra == "dev"
 Requires-Dist: nvsmi; extra == "dev"
 Requires-Dist: termcolor; extra == "dev"
 Requires-Dist: requests; extra == "dev"
-Requires-Dist: sphinx; extra == "dev"
-Requires-Dist: sphinx_rtd_theme; extra == "dev"
-Requires-Dist: nbsphinx; extra == "dev"
-Requires-Dist: pandoc; extra == "dev"

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/docs/conf.py

@@ -14,14 +14,9 @@ import datetime
 import os
 import sys
 
-sys.path.insert(0, os.path.join(os.path.abspath(".."), "src"))
-
-
 read_the_docs_build = os.environ.get("READTHEDOCS", None) == "True"
 
-# top = os.path.realpath(os.path.dirname(os.path.dirname(__file__)))
-# sys.path.insert(0, top)
-
+sys.path.insert(0, os.path.join(os.path.abspath(".."), "src"))
 
 source_suffix = ".rst"
 master_doc = "index"
@@ -32,7 +27,7 @@ html_logo = "_static/logo.png"
 
 # -- Project information -----------------------------------------------------
 
-project = "Anemoi"
+project = "Anemoi Utils"
 
 author = "ECMWF"
 
@@ -45,9 +40,9 @@ else:
 copyright = "%s, ECMWF" % (years,)
 
 try:
-    import anemoi.utils
+    from anemoi.utils._version import __version__
 
-    release = anemoi.utils.__version__
+    release = __version__
 except ImportError:
     release = "0.0.0"
 
@@ -65,6 +60,7 @@ extensions = [
     "sphinx.ext.intersphinx",
     "sphinx.ext.autodoc",
     "sphinx.ext.napoleon",
+    "sphinxarg.ext",
 ]
 
 # Add any paths that contain templates here, relative to this directory.

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/docs/index.rst

@@ -2,9 +2,9 @@
 
 .. _index-page:
 
-####################################
- Welcome to Anemoi's documentation!
-####################################
+##########################################
+ Welcome to `anemoi-utils` documentation!
+##########################################
 
 .. warning::
 

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/docs/requirements.txt

@@ -2,6 +2,7 @@
 sphinx
 sphinx_rtd_theme
 nbsphinx
+sphinx_argparse
 
 # Also requires `brew install pandoc` on Mac
 pandoc

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/pyproject.toml

@@ -40,7 +40,9 @@ classifiers = [
 ]
 
 dependencies = [
-    "tomli", # Only needed before 3.11
+    "tomli",  # Only needed before 3.11
+    "pyyaml",
+    "tqdm",
 ]
 
 [project.optional-dependencies]
@@ -53,9 +55,14 @@ grib = ["requests"]
 # Loaded by read-the-docs
 # `pip install .[docs]`
 docs = [
-    "tomli",     # Only needed before 3.11
+    "tomli",            # Only needed before 3.11
     "termcolor",
     "requests",
+    "sphinx",
+    "sphinx_rtd_theme",
+    "nbsphinx",
+    "pandoc",
+    "sphinx_argparse",
 ]
 
 all = [
@@ -72,10 +79,6 @@ dev = [
     "nvsmi",
     "termcolor",
     "requests",
-    "sphinx",
-    "sphinx_rtd_theme",
-    "nbsphinx",
-    "pandoc",
 ]
 
 [project.urls]
@@ -85,6 +88,11 @@ Repository = "https://github.com/ecmwf/anemoi-utils/"
 Issues = "https://github.com/ecmwf/anemoi-utils/issues"
 # Changelog = "https://github.com/ecmwf/anemoi-utils/CHANGELOG.md"
 
+[project.scripts]
+anemoi-utils = "anemoi.utils.__main__:main"
 
 [tool.setuptools_scm]
 version_file = "src/anemoi/utils/_version.py"
+
+[tool.setuptools.package-data]
+"anemoi.utils.mars" = ["*.yaml"]

anemoi_utils-0.2.0/src/anemoi/utils/__main__.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# 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 argparse
+import logging
+import sys
+import traceback
+
+from . import __version__
+from .commands import COMMANDS
+
+LOG = logging.getLogger(__name__)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+
+    parser.add_argument(
+        "--version",
+        "-V",
+        action="store_true",
+        help="show the version and exit",
+    )
+    parser.add_argument(
+        "--debug",
+        "-d",
+        action="store_true",
+        help="Debug mode",
+    )
+
+    subparsers = parser.add_subparsers(help="commands:", dest="command")
+    for name, command in COMMANDS.items():
+        command_parser = subparsers.add_parser(name, help=command.__doc__)
+        command.add_arguments(command_parser)
+
+    args = parser.parse_args()
+
+    if args.version:
+        print(__version__)
+        return
+
+    if args.command is None:
+        parser.print_help()
+        return
+
+    cmd = COMMANDS[args.command]
+
+    logging.basicConfig(
+        format="%(asctime)s %(levelname)s %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        level=logging.DEBUG if args.debug else logging.INFO,
+    )
+
+    try:
+        cmd.run(args)
+    except ValueError as e:
+        traceback.print_exc()
+        LOG.error("\n💣 %s", str(e).lstrip())
+        LOG.error("💣 Exiting")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/src/anemoi/utils/_version.py

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

anemoi_utils-0.2.0/src/anemoi/utils/checkpoints.py

@@ -0,0 +1,179 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+
+"""Read and write extra metadata in PyTorch checkpoints files. These files
+are zip archives containing the model weights.
+"""
+
+import json
+import logging
+import os
+import time
+import zipfile
+from tempfile import TemporaryDirectory
+
+import tqdm
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_NAME = "ai-models.json"
+DEFAULT_FOLDER = "anemoi-metadata"
+
+
+def has_metadata(path: str, name: str = DEFAULT_NAME) -> bool:
+    """Check if a checkpoint file has a metadata file
+
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    name : str, optional
+        The name of the metadata file in the zip archive
+
+    Returns
+    -------
+    bool
+        True if the metadata file is found
+    """
+    with zipfile.ZipFile(path, "r") as f:
+        for b in f.namelist():
+            if os.path.basename(b) == name:
+                return True
+    return False
+
+
+def load_metadata(path: str, name: str = DEFAULT_NAME):
+    """Load metadata from a checkpoint file
+
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    name : str, optional
+        The name of the metadata file in the zip archive
+
+    Returns
+    -------
+    JSON
+        The content of the metadata file
+
+    Raises
+    ------
+    ValueError
+        If the metadata file is not found
+    """
+    with zipfile.ZipFile(path, "r") as f:
+        metadata = None
+        for b in f.namelist():
+            if os.path.basename(b) == name:
+                if metadata is not None:
+                    raise ValueError(f"Found two or more '{name}' in {path}.")
+                metadata = b
+
+    if metadata is not None:
+        with zipfile.ZipFile(path, "r") as f:
+            return json.load(f.open(metadata, "r"))
+    else:
+        raise ValueError(f"Could not find '{name}' in {path}.")
+
+
+def save_metadata(path, metadata, name=DEFAULT_NAME, folder=DEFAULT_FOLDER):
+    """Save metadata to a checkpoint file
+
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    metadata : JSON
+        A JSON serializable object
+    name : str, optional
+        The name of the metadata file in the zip archive
+    """
+    with zipfile.ZipFile(path, "a") as zipf:
+
+        directories = set()
+
+        for b in zipf.namelist():
+            directory = os.path.dirname(b)
+            while os.path.dirname(directory) not in (".", ""):
+                directory = os.path.dirname(directory)
+            directories.add(directory)
+
+            if os.path.basename(b) == name:
+                raise ValueError(f"'{name}' already in {path}")
+
+        if len(directories) != 1:
+            # PyTorch checkpoints should have a single directory
+            # otherwise PyTorch will complain
+            raise ValueError(f"No or multiple directories in the checkpoint {path}, directories={directories}")
+
+        directory = list(directories)[0]
+
+        LOG.info("Saving metadata to %s/%s/%s", directory, folder, name)
+
+        zipf.writestr(
+            f"{directory}/{folder}/{name}",
+            json.dumps(metadata),
+        )
+
+
+def _edit_metadata(path, name, callback):
+    new_path = f"{path}.anemoi-edit-{time.time()}-{os.getpid()}.tmp"
+
+    found = False
+
+    with TemporaryDirectory() as temp_dir:
+        zipfile.ZipFile(path, "r").extractall(temp_dir)
+        total = 0
+        for root, dirs, files in os.walk(temp_dir):
+            for f in files:
+                total += 1
+                full = os.path.join(root, f)
+                if f == name:
+                    found = True
+                    callback(full)
+
+        if not found:
+            raise ValueError(f"Could not find '{name}' in {path}")
+
+        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):
+                    for f in files:
+                        full = os.path.join(root, f)
+                        rel = os.path.relpath(full, temp_dir)
+                        zipf.write(full, rel)
+                        pbar.update(1)
+
+    os.rename(new_path, path)
+    LOG.info("Updated metadata in %s", path)
+
+
+def replace_metadata(path, metadata, name=DEFAULT_NAME):
+
+    if not isinstance(metadata, dict):
+        raise ValueError(f"metadata must be a dict, got {type(metadata)}")
+
+    if "version" not in metadata:
+        raise ValueError("metadata must have a 'version' key")
+
+    def callback(full):
+        with open(full, "w") as f:
+            json.dump(metadata, f)
+
+    _edit_metadata(path, name, callback)
+
+
+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)

anemoi_utils-0.2.0/src/anemoi/utils/commands/__init__.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# 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 argparse
+import importlib
+import logging
+import os
+import sys
+
+LOG = logging.getLogger(__name__)
+
+
+def register(here, package, select, fail=None):
+    result = {}
+    not_available = {}
+
+    for p in os.listdir(here):
+        full = os.path.join(here, p)
+        if p.startswith("_"):
+            continue
+        if not (p.endswith(".py") or (os.path.isdir(full) and os.path.exists(os.path.join(full, "__init__.py")))):
+            continue
+
+        name, _ = os.path.splitext(p)
+
+        try:
+            imported = importlib.import_module(
+                f".{name}",
+                package=package,
+            )
+        except ImportError as e:
+            not_available[name] = e
+            continue
+
+        obj = select(imported)
+        if obj is not None:
+            result[name] = obj
+
+    for name, e in not_available.items():
+        if fail is None:
+            pass
+        if callable(fail):
+            result[name] = fail(name, e)
+
+    return result
+
+
+class Command:
+    def run(self, args):
+        raise NotImplementedError(f"Command not implemented: {args.command}")
+
+
+class Failed(Command):
+    def __init__(self, name, error):
+        self.name = name
+        self.error = error
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument("x", nargs=argparse.REMAINDER)
+
+    def run(self, args):
+        print(f"Command '{self.name}' not available: {self.error}")
+        sys.exit(1)
+
+
+COMMANDS = register(
+    os.path.dirname(__file__),
+    __name__,
+    lambda x: x.command(),
+    lambda name, error: Failed(name, error),
+)

anemoi_utils-0.2.0/src/anemoi/utils/commands/checkpoint.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# 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 json
+
+from . import Command
+
+
+def visit(x, path, name, value):
+    if isinstance(x, dict):
+        for k, v in x.items():
+            if k == name:
+                print(".".join(path), k, v)
+
+            if v == value:
+                print(".".join(path), k, v)
+
+            path.append(k)
+            visit(v, path, name, value)
+            path.pop()
+
+    if isinstance(x, list):
+        for i, v in enumerate(x):
+            path.append(str(i))
+            visit(v, path, name, value)
+            path.pop()
+
+
+class Checkpoint(Command):
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument("path", help="Path to the checkpoint.")
+        command_parser.add_argument("--name", help="Search for a specific name.")
+        command_parser.add_argument("--value", help="Search for a specific value.")
+
+    def run(self, args):
+        from anemoi.utils.checkpoints import load_metadata
+
+        checkpoint = load_metadata(args.path, "*.json")
+
+        if args.name or args.value:
+            visit(
+                checkpoint,
+                [],
+                args.name if args.name is not None else object(),
+                args.value if args.value is not None else object(),
+            )
+            return
+
+        print(json.dumps(checkpoint, sort_keys=True, indent=4))
+
+
+command = Checkpoint

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/src/anemoi/utils/dates.py

@@ -10,6 +10,16 @@ import calendar
 import datetime
 
 
+def normalise_frequency(frequency):
+    if isinstance(frequency, int):
+        return frequency
+    assert isinstance(frequency, str), (type(frequency), frequency)
+
+    unit = frequency[-1].lower()
+    v = int(frequency[:-1])
+    return {"h": v, "d": v * 24}[unit]
+
+
 def no_time_zone(date):
     """Remove time zone information from a date.
 
@@ -27,6 +37,7 @@ def no_time_zone(date):
     return date.replace(tzinfo=None)
 
 
+# this function is use in anemoi-datasets
 def as_datetime(date):
     """Convert a date to a datetime object, removing any time zone information.
 
@@ -162,11 +173,15 @@ class HindcastDatesTimes:
         """
 
         self.reference_dates = reference_dates
-        self.years = (1, years + 1)
+
+        if isinstance(years, list):
+            self.years = years
+        else:
+            self.years = range(1, years + 1)
 
     def __iter__(self):
         for reference_date in self.reference_dates:
-            for year in range(*self.years):
+            for year in self.years:
                 if reference_date.month == 2 and reference_date.day == 29:
                     date = datetime.datetime(reference_date.year - year, 2, 28)
                 else:
@@ -246,3 +261,61 @@ class Autumn(DateTimes):
             _description_
         """
         super().__init__(datetime.datetime(year, 9, 1), datetime.datetime(year, 11, 30), **kwargs)
+
+
+class ConcatDateTimes:
+    def __init__(self, *dates):
+        if len(dates) == 1 and isinstance(dates[0], list):
+            dates = dates[0]
+
+        self.dates = dates
+
+    def __iter__(self):
+        for date in self.dates:
+            yield from date
+
+
+class EnumDateTimes:
+    def __init__(self, dates):
+        self.dates = dates
+
+    def __iter__(self):
+        for date in self.dates:
+            yield as_datetime(date)
+
+
+def datetimes_factory(*args, **kwargs):
+    if args and kwargs:
+        raise ValueError("Cannot provide both args and kwargs for a list of dates")
+
+    if not args and not kwargs:
+        raise ValueError("No dates provided")
+
+    if kwargs:
+        name = kwargs.get("name")
+
+        if name == "hindcast":
+            reference_dates = kwargs["reference_dates"]
+            reference_dates = datetimes_factory(reference_dates)
+            years = kwargs["years"]
+            return HindcastDatesTimes(reference_dates=reference_dates, years=years)
+
+        kwargs = kwargs.copy()
+        if "frequency" in kwargs:
+            freq = kwargs.pop("frequency")
+            kwargs["increment"] = normalise_frequency(freq)
+        return DateTimes(**kwargs)
+
+    if not any((isinstance(x, dict) or isinstance(x, list)) for x in args):
+        return EnumDateTimes(args)
+
+    if len(args) == 1:
+        a = args[0]
+
+        if isinstance(a, dict):
+            return datetimes_factory(**a)
+
+        if isinstance(a, list):
+            return datetimes_factory(*a)
+
+    return ConcatDateTimes(*[datetimes_factory(a) for a in args])

anemoi_utils-0.2.0/src/anemoi/utils/mars/__init__.py

@@ -0,0 +1,76 @@
+# (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 Mars requests.
+
+Has some konwledge of how certain streams are organised in Mars.
+
+"""
+
+import datetime
+import logging
+import os
+
+import yaml
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_MARS_LABELLING = {
+    "class": "od",
+    "type": "an",
+    "stream": "oper",
+    "expver": "0001",
+}
+
+
+def _expand_mars_labelling(request):
+    """Expand the request with the default Mars labelling.
+
+    The default Mars labelling is:
+
+    {'class': 'od',
+     'type': 'an',
+     'stream': 'oper',
+     'expver': '0001'}
+
+    """
+    result = DEFAULT_MARS_LABELLING.copy()
+    result.update(request)
+    return result
+
+
+STREAMS = None
+
+
+def _lookup_mars_stream(request):
+    global STREAMS
+
+    if STREAMS is None:
+
+        with open(os.path.join(os.path.dirname(__file__), "mars.yaml")) as f:
+            STREAMS = yaml.safe_load(f)
+
+    request = _expand_mars_labelling(request)
+    for s in STREAMS:
+        match = s["match"]
+        if all(request.get(k) == v for k, v in match.items()):
+            return s["info"]
+
+
+def recenter(date, center, members):
+
+    center = _lookup_mars_stream(center)
+    members = _lookup_mars_stream(members)
+
+    return (center, members)
+
+
+if __name__ == "__main__":
+    date = datetime.datetime(2024, 5, 9, 0)
+
+    print(recenter(date, {"type": "an"}, {"stream": "elda"}))

anemoi_utils-0.2.0/src/anemoi/utils/mars/mars.yaml

@@ -0,0 +1,5 @@
+- match:
+    class: od
+    stream: elda
+  info:
+    runs: [6, 18]

anemoi_utils-0.2.0/src/anemoi/utils/timer.py

@@ -0,0 +1,32 @@
+# (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.
+
+"""Logging utilities."""
+
+import logging
+import time
+
+from .humanize import seconds
+
+LOGGER = logging.getLogger(__name__)
+
+
+class Timer:
+    def __init__(self, title, logger=LOGGER):
+        self.title = title
+        self.start = time.time()
+        self.logger = logger
+
+    def __enter__(self):
+        return self
+
+    @property
+    def elapsed(self):
+        return time.time() - self.start
+
+    def __exit__(self, *args):
+        self.logger.info("%s: %s.", self.title, seconds(self.elapsed))

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/src/anemoi_utils.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.8
+Version: 0.2.0
 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
@@ -223,6 +223,8 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: tomli
+Requires-Dist: pyyaml
+Requires-Dist: tqdm
 Provides-Extra: provenance
 Requires-Dist: GitPython; extra == "provenance"
 Requires-Dist: nvsmi; extra == "provenance"
@@ -234,6 +236,11 @@ Provides-Extra: docs
 Requires-Dist: tomli; extra == "docs"
 Requires-Dist: termcolor; extra == "docs"
 Requires-Dist: requests; extra == "docs"
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx_rtd_theme; extra == "docs"
+Requires-Dist: nbsphinx; extra == "docs"
+Requires-Dist: pandoc; extra == "docs"
+Requires-Dist: sphinx_argparse; extra == "docs"
 Provides-Extra: all
 Requires-Dist: tomli; extra == "all"
 Requires-Dist: GitPython; extra == "all"
@@ -246,7 +253,3 @@ Requires-Dist: GitPython; extra == "dev"
 Requires-Dist: nvsmi; extra == "dev"
 Requires-Dist: termcolor; extra == "dev"
 Requires-Dist: requests; extra == "dev"
-Requires-Dist: sphinx; extra == "dev"
-Requires-Dist: sphinx_rtd_theme; extra == "dev"
-Requires-Dist: nbsphinx; extra == "dev"
-Requires-Dist: pandoc; extra == "dev"

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/src/anemoi_utils.egg-info/SOURCES.txt

@@ -21,6 +21,7 @@ docs/modules/humanize.rst
 docs/modules/provenance.rst
 docs/modules/text.rst
 src/anemoi/utils/__init__.py
+src/anemoi/utils/__main__.py
 src/anemoi/utils/_version.py
 src/anemoi/utils/caching.py
 src/anemoi/utils/checkpoints.py
@@ -30,10 +31,17 @@ src/anemoi/utils/grib.py
 src/anemoi/utils/humanize.py
 src/anemoi/utils/provenance.py
 src/anemoi/utils/text.py
+src/anemoi/utils/timer.py
+src/anemoi/utils/commands/__init__.py
+src/anemoi/utils/commands/checkpoint.py
+src/anemoi/utils/mars/__init__.py
+src/anemoi/utils/mars/mars.yaml
 src/anemoi_utils.egg-info/PKG-INFO
 src/anemoi_utils.egg-info/SOURCES.txt
 src/anemoi_utils.egg-info/dependency_links.txt
+src/anemoi_utils.egg-info/entry_points.txt
 src/anemoi_utils.egg-info/requires.txt
 src/anemoi_utils.egg-info/top_level.txt
 tests/requirements.txt
+tests/test_dates.py
 tests/test_utils.py

anemoi_utils-0.2.0/src/anemoi_utils.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+anemoi-utils = anemoi.utils.__main__:main

{anemoi_utils-0.1.8 → anemoi_utils-0.2.0}/src/anemoi_utils.egg-info/requires.txt

@@ -1,4 +1,6 @@
 tomli
+pyyaml
+tqdm
 
 [all]
 tomli
@@ -13,15 +15,16 @@ GitPython
 nvsmi
 termcolor
 requests
-sphinx
-sphinx_rtd_theme
-nbsphinx
-pandoc
 
 [docs]
 tomli
 termcolor
 requests
+sphinx
+sphinx_rtd_theme
+nbsphinx
+pandoc
+sphinx_argparse
 
 [grib]
 requests

anemoi_utils-0.2.0/tests/test_dates.py

@@ -0,0 +1,113 @@
+import datetime
+from textwrap import dedent
+
+import yaml
+
+from anemoi.utils.dates import datetimes_factory
+
+
+def _(txt):
+    txt = dedent(txt)
+    config = yaml.safe_load(txt)
+    return datetimes_factory(config)
+
+
+def test_date_1():
+    d = _(
+        """
+          - 2023-01-01
+          - 2023-01-02
+          - 2023-01-03
+    """
+    )
+    assert len(list(d)) == 3
+
+
+def test_date_2():
+    d = _(
+        """
+        start: 2023-01-01
+        end: 2023-01-07
+        frequency: 12
+        day_of_week: [monday, friday]
+    """
+    )
+    assert len(list(d)) == 4
+
+
+def test_date_3():
+    d = _(
+        """
+        - start: 2023-01-01
+          end: 2023-01-03
+          frequency: 24
+        - start: 2024-01-01T06:00:00
+          end: 2024-01-02T18:00:00
+          frequency: 6h
+    """
+    )
+    assert datetime.datetime(2023, 1, 1, 0) in d
+    assert datetime.datetime(2023, 1, 2, 0) in d
+    assert datetime.datetime(2023, 1, 3, 0) in d
+    assert datetime.datetime(2024, 1, 1, 6) in d
+    assert datetime.datetime(2024, 1, 1, 12) in d
+    assert datetime.datetime(2024, 1, 1, 18) in d
+    assert datetime.datetime(2024, 1, 2, 0) in d
+    assert datetime.datetime(2024, 1, 2, 6) in d
+    assert datetime.datetime(2024, 1, 2, 12) in d
+    assert datetime.datetime(2024, 1, 2, 18) in d
+    assert len(list(d)) == 10
+
+
+def test_date_hindcast_1():
+    d = _(
+        """
+        - name: hindcast
+          reference_dates:
+            start: 2023-01-01
+            end: 2023-01-03
+            frequency: 24
+          years: 20
+    """
+    )
+    assert len(list(d)) == 60
+
+
+def test_date_hindcast_2():
+    d = _(
+        """
+        - name: hindcast
+          reference_dates:
+            start: 2023-01-01
+            end: 2023-01-03
+            frequency: 24
+          years: [2018, 2019, 2020, 2021]
+    """
+    )
+    assert len(list(d)) == 12
+
+
+def test_date_hindcast_3():
+    d = _(
+        """
+        - name: hindcast
+          reference_dates:
+            start: 2022-12-25 00:00:00
+            end: 2022-12-31 12:00:00
+            frequency: 12h
+            day_of_week: tuesday
+          years: [2018, 2019, 2020, 2021]
+    """
+    )
+    print(list(d))
+    assert len(list(d)) == 8
+
+
+if __name__ == "__main__":
+    test_functions = [
+        obj for name, obj in globals().items() if name.startswith("test_") and isinstance(obj, type(lambda: 0))
+    ]
+    for test_func in test_functions:
+        print(f"Running test: {test_func.__name__}")
+        test_func()
+    print("All tests passed!")

anemoi_utils-0.1.8/src/anemoi/utils/checkpoints.py

@@ -1,75 +0,0 @@
-# (C) Copyright 2024 ECMWF.
-#
-# 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.
-
-"""Read and write extra metadata in PyTorch checkpoints files. These files
-are zip archives containing the model weights.
-"""
-
-import json
-import logging
-import os
-import zipfile
-
-LOG = logging.getLogger(__name__)
-
-DEFAULT_NAME = "anemoi-metadata.json"
-
-
-def load_metadata(path: str, name: str = DEFAULT_NAME):
-    """Load metadata from a checkpoint file
-
-    Parameters
-    ----------
-    path : str
-        The path to the checkpoint file
-    name : str, optional
-        The name of the metadata file in the zip archive
-
-    Returns
-    -------
-    JSON
-        The content of the metadata file
-
-    Raises
-    ------
-    ValueError
-        If the metadata file is not found
-    """
-    with zipfile.ZipFile(path, "r") as f:
-        metadata = None
-        for b in f.namelist():
-            if os.path.basename(b) == name:
-                if metadata is not None:
-                    LOG.warning(f"Found two '{name}' if {path}")
-                metadata = b
-
-    if metadata is not None:
-        with zipfile.ZipFile(path, "r") as f:
-            return json.load(f.open(metadata, "r"))
-    else:
-        raise ValueError(f"Could not find {name} in {path}")
-
-
-def save_metadata(path, metadata, name=DEFAULT_NAME):
-    """Save metadata to a checkpoint file
-
-    Parameters
-    ----------
-    path : str
-        The path to the checkpoint file
-    metadata : JSON
-        A JSON serializable object
-    name : str, optional
-        The name of the metadata file in the zip archive
-    """
-    with zipfile.ZipFile(path, "a") as zipf:
-        base, _ = os.path.splitext(os.path.basename(path))
-        zipf.writestr(
-            f"{base}/{name}",
-            json.dumps(metadata),
-        )