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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (81) hide show
  1. anemoi_utils-0.4.5/.github/ci-hpc-config.yml +8 -0
  2. anemoi_utils-0.4.5/.github/workflows/changelog-pr-update.yml +18 -0
  3. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.github/workflows/changelog-release-update.yml +1 -0
  4. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.pre-commit-config.yaml +3 -3
  5. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/CHANGELOG.md +17 -2
  6. anemoi_utils-0.4.5/CONTRIBUTORS.md +13 -0
  7. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/PKG-INFO +2 -1
  8. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/conf.py +2 -2
  9. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/pyproject.toml +8 -4
  10. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/__init__.py +3 -1
  11. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/__main__.py +2 -3
  12. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/_version.py +2 -2
  13. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/checkpoints.py +74 -9
  14. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/cli.py +14 -2
  15. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/commands/__init__.py +2 -3
  16. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/commands/config.py +0 -1
  17. anemoi_utils-0.4.5/src/anemoi/utils/compatibility.py +76 -0
  18. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/config.py +3 -2
  19. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/dates.py +7 -2
  20. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/mars/__init__.py +3 -1
  21. anemoi_utils-0.4.5/src/anemoi/utils/registry.py +129 -0
  22. anemoi_utils-0.4.5/src/anemoi/utils/remote/__init__.py +328 -0
  23. {anemoi_utils-0.4.3/src/anemoi/utils → anemoi_utils-0.4.5/src/anemoi/utils/remote}/s3.py +42 -216
  24. anemoi_utils-0.4.5/src/anemoi/utils/remote/ssh.py +133 -0
  25. anemoi_utils-0.4.5/src/anemoi/utils/s3.py +63 -0
  26. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi_utils.egg-info/PKG-INFO +2 -1
  27. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi_utils.egg-info/SOURCES.txt +14 -1
  28. anemoi_utils-0.4.5/tests/test-transfer-data/directory/b/c/x +1 -0
  29. anemoi_utils-0.4.5/tests/test-transfer-data/directory/b/y +1 -0
  30. anemoi_utils-0.4.5/tests/test-transfer-data/directory/exotic filename ;^/"'[=.,#]()/303/252/303/274/303/247/303/262/342/234/205.txt" +1 -0
  31. anemoi_utils-0.4.5/tests/test-transfer-data/directory/z +1 -0
  32. anemoi_utils-0.4.5/tests/test-transfer-data/file +1 -0
  33. anemoi_utils-0.4.5/tests/test_compatibility.py +32 -0
  34. anemoi_utils-0.4.5/tests/test_remote.py +175 -0
  35. anemoi_utils-0.4.3/.github/ci-hpc-config.yml +0 -8
  36. anemoi_utils-0.4.3/.github/workflows/changelog-pr-update.yml +0 -18
  37. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.gitattributes +0 -0
  38. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.github/CODEOWNERS +0 -0
  39. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.github/workflows/ci.yml +0 -0
  40. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.github/workflows/label-public-pr.yml +0 -0
  41. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.github/workflows/python-publish.yml +0 -0
  42. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.github/workflows/python-pull-request.yml +0 -0
  43. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.github/workflows/readthedocs-pr-update.yml +0 -0
  44. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.gitignore +0 -0
  45. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/.readthedocs.yaml +0 -0
  46. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/LICENSE +0 -0
  47. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/README.md +0 -0
  48. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/Makefile +0 -0
  49. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/_static/logo.png +0 -0
  50. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/_static/style.css +0 -0
  51. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/_templates/.gitkeep +0 -0
  52. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/index.rst +0 -0
  53. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/installing.rst +0 -0
  54. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/modules/checkpoints.rst +0 -0
  55. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/modules/config.rst +0 -0
  56. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/modules/dates.rst +0 -0
  57. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/modules/grib.rst +0 -0
  58. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/modules/humanize.rst +0 -0
  59. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/modules/provenance.rst +0 -0
  60. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/modules/s3.rst +0 -0
  61. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/docs/modules/text.rst +0 -0
  62. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/setup.cfg +0 -0
  63. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/caching.py +0 -0
  64. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/grib.py +0 -0
  65. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/hindcasts.py +0 -0
  66. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/humanize.py +0 -0
  67. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/mars/mars.yaml +0 -0
  68. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/provenance.py +0 -0
  69. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/sanitise.py +0 -0
  70. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/sanitize.py +0 -0
  71. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/text.py +0 -0
  72. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi/utils/timer.py +0 -0
  73. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi_utils.egg-info/dependency_links.txt +0 -0
  74. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi_utils.egg-info/entry_points.txt +0 -0
  75. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi_utils.egg-info/requires.txt +0 -0
  76. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/src/anemoi_utils.egg-info/top_level.txt +0 -0
  77. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/tests/test_dates.py +0 -0
  78. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/tests/test_frequency.py +0 -0
  79. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/tests/test_provenance.py +0 -0
  80. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/tests/test_sanetise.py +0 -0
  81. {anemoi_utils-0.4.3 → anemoi_utils-0.4.5}/tests/test_utils.py +0 -0
@@ -0,0 +1,8 @@
1
+ build:
2
+ python: '3.10'
3
+ modules:
4
+ - ninja
5
+ parallel: 64
6
+
7
+ pytest_cmd: |
8
+ python -m pytest -vv -m 'not notebook and not no_cache_init and not skip_on_hpc' --cov=. --cov-report=xml
@@ -0,0 +1,18 @@
1
+ # name: Check Changelog Update on PR
2
+ # on:
3
+ # pull_request:
4
+ # types: [assigned, opened, synchronize, reopened, labeled, unlabeled]
5
+ # branches:
6
+ # - main
7
+ # - develop
8
+ # paths-ignore:
9
+ # - .pre-commit-config.yaml
10
+ # - .readthedocs.yaml
11
+ # jobs:
12
+ # Check-Changelog:
13
+ # name: Check Changelog Action
14
+ # runs-on: ubuntu-20.04
15
+ # steps:
16
+ # - uses: tarides/changelog-check-action@v2
17
+ # with:
18
+ # changelog: CHANGELOG.md
@@ -25,6 +25,7 @@ jobs:
25
25
  with:
26
26
  latest-version: ${{ github.event.release.tag_name }}
27
27
  heading-text: ${{ github.event.release.name }}
28
+ release-notes: ${{ github.event.release.body }}
28
29
 
29
30
  - name: Create Pull Request
30
31
  uses: peter-evans/create-pull-request@v6
@@ -27,7 +27,7 @@ repos:
27
27
  - id: python-check-blanket-noqa # Check for # noqa: all
28
28
  - id: python-no-log-warn # Check for log.warn
29
29
  - repo: https://github.com/psf/black-pre-commit-mirror
30
- rev: 24.8.0
30
+ rev: 24.10.0
31
31
  hooks:
32
32
  - id: black
33
33
  args: [--line-length=120]
@@ -40,7 +40,7 @@ repos:
40
40
  - --force-single-line-imports
41
41
  - --profile black
42
42
  - repo: https://github.com/astral-sh/ruff-pre-commit
43
- rev: v0.6.9
43
+ rev: v0.7.2
44
44
  hooks:
45
45
  - id: ruff
46
46
  args:
@@ -65,7 +65,7 @@ repos:
65
65
  - id: docconvert
66
66
  args: ["numpy"]
67
67
  - repo: https://github.com/tox-dev/pyproject-fmt
68
- rev: "2.2.4"
68
+ rev: "v2.5.0"
69
69
  hooks:
70
70
  - id: pyproject-fmt
71
71
  - repo: https://github.com/jshwi/docsig # Check docstrings against function sig
@@ -8,14 +8,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
8
8
  Please add your functional changes to the appropriate section in the PR.
9
9
  Keep it human-readable, your future self will thank you!
10
10
 
11
- ## [Unreleased](https://github.com/ecmwf/anemoi-utils/compare/0.4.1...HEAD)
11
+ ## [0.4.4](https://github.com/ecmwf/anemoi-utils/compare/0.4.3...0.4.4) - 2024-11-01
12
+
13
+ ## [0.4.3](https://github.com/ecmwf/anemoi-utils/compare/0.4.1...0.4.3) - 2024-10-26
14
+
15
+ ## [0.4.2](https://github.com/ecmwf/anemoi-utils/compare/0.4.1...0.4.2) - 2024-10-25
16
+
17
+ ### Added
18
+ - Add supporting_arrays to checkpoints
19
+ - Add factories registry
20
+ - Optional renaming of subcommands via `command` attribute [#34](https://github.com/ecmwf/anemoi-utils/pull/34)
21
+ - `skip_on_hpc` pytest marker for tests that should not be run on HPC [36](https://github.com/ecmwf/anemoi-utils/pull/36)
12
22
 
13
23
  ## [0.4.1](https://github.com/ecmwf/anemoi-utils/compare/0.4.0...0.4.1) - 2024-10-23
14
24
 
15
25
  ## Fixed
26
+
16
27
  - Fix `__version__` import in init
17
28
 
18
29
  ### Changed
30
+
19
31
  - Fix: resolve mounted filesystems in provenance
20
32
  - Fix pre-commit regex
21
33
  - ci: extend python versions [#23] (https://github.com/ecmwf/anemoi-utils/pull/23)
@@ -26,6 +38,7 @@ Keep it human-readable, your future self will thank you!
26
38
  ### Added
27
39
 
28
40
  - Add anemoi-transform link to documentation
41
+ - Add CONTRIBUTORS.md (#33)
29
42
 
30
43
  ## [0.3.17](https://github.com/ecmwf/anemoi-utils/compare/0.3.13...0.3.17) - 2024-10-01
31
44
 
@@ -37,7 +50,9 @@ Keep it human-readable, your future self will thank you!
37
50
  - Changelog merge strategy- Codeowners file
38
51
  - Create dependency on wcwidth. MIT licence.
39
52
  - Add distribution name dictionary to provenance [#15](https://github.com/ecmwf/anemoi-utils/pull/15) & [#19](https://github.com/ecmwf/anemoi-utils/pull/19)
40
- - Add anonimize() function.
53
+ - Add anonymize() function.
54
+ - Add transfer to ssh:// target (experimental)
55
+ - Deprecated 'anemoi.utils.s3'
41
56
 
42
57
  ### Changed
43
58
 
@@ -0,0 +1,13 @@
1
+ ## How to Contribute
2
+
3
+ Please see the [read the docs](https://anemoi-training.readthedocs.io/en/latest/dev/contributing.html).
4
+
5
+
6
+ ## Contributors
7
+
8
+ 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).
9
+
10
+
11
+ ## Contributing Organisations
12
+
13
+ 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/)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: anemoi-utils
3
- Version: 0.4.3
3
+ Version: 0.4.5
4
4
  Summary: A package to hold various functions to support training of ML models on ECMWF data.
5
5
  Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
6
6
  License: Apache License
@@ -219,6 +219,7 @@ Classifier: Programming Language :: Python :: 3.9
219
219
  Classifier: Programming Language :: Python :: 3.10
220
220
  Classifier: Programming Language :: Python :: 3.11
221
221
  Classifier: Programming Language :: Python :: 3.12
222
+ Classifier: Programming Language :: Python :: 3.13
222
223
  Classifier: Programming Language :: Python :: Implementation :: CPython
223
224
  Classifier: Programming Language :: Python :: Implementation :: PyPy
224
225
  Requires-Python: >=3.9
@@ -29,7 +29,7 @@ html_logo = "_static/logo.png"
29
29
 
30
30
  project = "Anemoi Utils"
31
31
 
32
- author = "ECMWF"
32
+ author = "Anemoi contributors"
33
33
 
34
34
  year = datetime.datetime.now().year
35
35
  if year == 2024:
@@ -37,7 +37,7 @@ if year == 2024:
37
37
  else:
38
38
  years = "2024-%s" % (year,)
39
39
 
40
- copyright = "%s, ECMWF" % (years,)
40
+ copyright = "%s, Anemoi contributors" % (years,)
41
41
 
42
42
  try:
43
43
  from anemoi.utils._version import __version__
@@ -1,14 +1,12 @@
1
- #!/usr/bin/env python
2
- # (C) Copyright 2024 ECMWF.
1
+ # (C) Copyright 2024 Anemoi contributors.
3
2
  #
4
3
  # This software is licensed under the terms of the Apache Licence Version 2.0
5
4
  # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
5
+ #
6
6
  # In applying this licence, ECMWF does not waive the privileges and immunities
7
7
  # granted to it by virtue of its status as an intergovernmental organisation
8
8
  # nor does it submit to any jurisdiction.
9
9
 
10
- # https://packaging.python.org/en/latest/guides/writing-pyproject-toml/
11
-
12
10
  [build-system]
13
11
  requires = [ "setuptools>=60", "setuptools-scm>=8" ]
14
12
 
@@ -35,6 +33,7 @@ classifiers = [
35
33
  "Programming Language :: Python :: 3.10",
36
34
  "Programming Language :: Python :: 3.11",
37
35
  "Programming Language :: Python :: 3.12",
36
+ "Programming Language :: Python :: 3.13",
38
37
  "Programming Language :: Python :: Implementation :: CPython",
39
38
  "Programming Language :: Python :: Implementation :: PyPy",
40
39
  ]
@@ -82,3 +81,8 @@ scripts.anemoi-utils = "anemoi.utils.__main__:main"
82
81
 
83
82
  [tool.setuptools_scm]
84
83
  version_file = "src/anemoi/utils/_version.py"
84
+
85
+ [tool.pytest.ini_options]
86
+ markers = [
87
+ "skip_on_hpc: mark a test that should not be run on HPC",
88
+ ]
@@ -1,6 +1,8 @@
1
- # (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
1
+ # (C) Copyright 2024 Anemoi contributors.
2
+ #
2
3
  # This software is licensed under the terms of the Apache Licence Version 2.0
3
4
  # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
5
+ #
4
6
  # In applying this licence, ECMWF does not waive the privileges and immunities
5
7
  # granted to it by virtue of its status as an intergovernmental organisation
6
8
  # nor does it submit to any jurisdiction.
@@ -1,12 +1,11 @@
1
- #!/usr/bin/env python
2
- # (C) Copyright 2024 ECMWF.
1
+ # (C) Copyright 2024 Anemoi contributors.
3
2
  #
4
3
  # This software is licensed under the terms of the Apache Licence Version 2.0
5
4
  # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
5
+ #
6
6
  # In applying this licence, ECMWF does not waive the privileges and immunities
7
7
  # granted to it by virtue of its status as an intergovernmental organisation
8
8
  # nor does it submit to any jurisdiction.
9
- #
10
9
 
11
10
  from anemoi.utils.cli import cli_main
12
11
  from anemoi.utils.cli import make_parser
@@ -12,5 +12,5 @@ __version__: str
12
12
  __version_tuple__: VERSION_TUPLE
13
13
  version_tuple: VERSION_TUPLE
14
14
 
15
- __version__ = version = '0.4.3'
16
- __version_tuple__ = version_tuple = (0, 4, 3)
15
+ __version__ = version = '0.4.5'
16
+ __version_tuple__ = version_tuple = (0, 4, 5)
@@ -27,7 +27,7 @@ DEFAULT_NAME = "ai-models.json"
27
27
  DEFAULT_FOLDER = "anemoi-metadata"
28
28
 
29
29
 
30
- def has_metadata(path: str, name: str = DEFAULT_NAME) -> bool:
30
+ def has_metadata(path: str, *, name: str = DEFAULT_NAME) -> bool:
31
31
  """Check if a checkpoint file has a metadata file
32
32
 
33
33
  Parameters
@@ -49,13 +49,26 @@ def has_metadata(path: str, name: str = DEFAULT_NAME) -> bool:
49
49
  return False
50
50
 
51
51
 
52
- def load_metadata(path: str, name: str = DEFAULT_NAME) -> dict:
52
+ def metadata_root(path: str, *, name: str = DEFAULT_NAME) -> bool:
53
+
54
+ with zipfile.ZipFile(path, "r") as f:
55
+ for b in f.namelist():
56
+ if os.path.basename(b) == name:
57
+ return os.path.dirname(b)
58
+ raise ValueError(f"Could not find '{name}' in {path}.")
59
+
60
+
61
+ def load_metadata(path: str, *, supporting_arrays=False, name: str = DEFAULT_NAME) -> dict:
53
62
  """Load metadata from a checkpoint file
54
63
 
55
64
  Parameters
56
65
  ----------
57
66
  path : str
58
67
  The path to the checkpoint file
68
+
69
+ supporting_arrays: bool, optional
70
+ If True, the function will return a dictionary with the supporting arrays
71
+
59
72
  name : str, optional
60
73
  The name of the metadata file in the zip archive
61
74
 
@@ -79,12 +92,29 @@ def load_metadata(path: str, name: str = DEFAULT_NAME) -> dict:
79
92
 
80
93
  if metadata is not None:
81
94
  with zipfile.ZipFile(path, "r") as f:
82
- return json.load(f.open(metadata, "r"))
95
+ metadata = json.load(f.open(metadata, "r"))
96
+ if supporting_arrays:
97
+ arrays = load_supporting_arrays(f, metadata.get("supporting_arrays_paths", {}))
98
+ return metadata, arrays
99
+
100
+ return metadata
83
101
  else:
84
102
  raise ValueError(f"Could not find '{name}' in {path}.")
85
103
 
86
104
 
87
- def save_metadata(path, metadata, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> None:
105
+ def load_supporting_arrays(zipf, entries) -> dict:
106
+ import numpy as np
107
+
108
+ supporting_arrays = {}
109
+ for key, entry in entries.items():
110
+ supporting_arrays[key] = np.frombuffer(
111
+ zipf.read(entry["path"]),
112
+ dtype=entry["dtype"],
113
+ ).reshape(entry["shape"])
114
+ return supporting_arrays
115
+
116
+
117
+ def save_metadata(path, metadata, *, supporting_arrays=None, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> None:
88
118
  """Save metadata to a checkpoint file
89
119
 
90
120
  Parameters
@@ -93,6 +123,8 @@ def save_metadata(path, metadata, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> N
93
123
  The path to the checkpoint file
94
124
  metadata : JSON
95
125
  A JSON serializable object
126
+ supporting_arrays: dict, optional
127
+ A dictionary of supporting NumPy arrays
96
128
  name : str, optional
97
129
  The name of the metadata file in the zip archive
98
130
  folder : str, optional
@@ -118,19 +150,41 @@ def save_metadata(path, metadata, name=DEFAULT_NAME, folder=DEFAULT_FOLDER) -> N
118
150
 
119
151
  directory = list(directories)[0]
120
152
 
153
+ LOG.info("Adding extra information to checkpoint %s", path)
121
154
  LOG.info("Saving metadata to %s/%s/%s", directory, folder, name)
122
155
 
156
+ metadata = metadata.copy()
157
+ if supporting_arrays is not None:
158
+ metadata["supporting_arrays_paths"] = {
159
+ key: dict(path=f"{directory}/{folder}/{key}.numpy", shape=value.shape, dtype=str(value.dtype))
160
+ for key, value in supporting_arrays.items()
161
+ }
162
+ else:
163
+ metadata["supporting_arrays_paths"] = {}
164
+
123
165
  zipf.writestr(
124
166
  f"{directory}/{folder}/{name}",
125
167
  json.dumps(metadata),
126
168
  )
127
169
 
170
+ for name, entry in metadata["supporting_arrays_paths"].items():
171
+ value = supporting_arrays[name]
172
+ LOG.info(
173
+ "Saving supporting array `%s` to %s (shape=%s, dtype=%s)",
174
+ name,
175
+ entry["path"],
176
+ entry["shape"],
177
+ entry["dtype"],
178
+ )
179
+ zipf.writestr(entry["path"], value.tobytes())
180
+
128
181
 
129
- def _edit_metadata(path, name, callback):
182
+ def _edit_metadata(path, name, callback, supporting_arrays=None):
130
183
  new_path = f"{path}.anemoi-edit-{time.time()}-{os.getpid()}.tmp"
131
184
 
132
185
  found = False
133
186
 
187
+ directory = None
134
188
  with TemporaryDirectory() as temp_dir:
135
189
  zipfile.ZipFile(path, "r").extractall(temp_dir)
136
190
  total = 0
@@ -141,10 +195,21 @@ def _edit_metadata(path, name, callback):
141
195
  if f == name:
142
196
  found = True
143
197
  callback(full)
198
+ directory = os.path.dirname(full)
144
199
 
145
200
  if not found:
146
201
  raise ValueError(f"Could not find '{name}' in {path}")
147
202
 
203
+ if supporting_arrays is not None:
204
+
205
+ for key, entry in supporting_arrays.items():
206
+ value = entry.tobytes()
207
+ fname = os.path.join(directory, f"{key}.numpy")
208
+ os.makedirs(os.path.dirname(fname), exist_ok=True)
209
+ with open(fname, "wb") as f:
210
+ f.write(value)
211
+ total += 1
212
+
148
213
  with zipfile.ZipFile(new_path, "w", zipfile.ZIP_DEFLATED) as zipf:
149
214
  with tqdm.tqdm(total=total, desc="Rebuilding checkpoint") as pbar:
150
215
  for root, dirs, files in os.walk(temp_dir):
@@ -158,7 +223,7 @@ def _edit_metadata(path, name, callback):
158
223
  LOG.info("Updated metadata in %s", path)
159
224
 
160
225
 
161
- def replace_metadata(path, metadata, name=DEFAULT_NAME):
226
+ def replace_metadata(path, metadata, supporting_arrays=None, *, name=DEFAULT_NAME):
162
227
 
163
228
  if not isinstance(metadata, dict):
164
229
  raise ValueError(f"metadata must be a dict, got {type(metadata)}")
@@ -170,14 +235,14 @@ def replace_metadata(path, metadata, name=DEFAULT_NAME):
170
235
  with open(full, "w") as f:
171
236
  json.dump(metadata, f)
172
237
 
173
- _edit_metadata(path, name, callback)
238
+ return _edit_metadata(path, name, callback, supporting_arrays)
174
239
 
175
240
 
176
- def remove_metadata(path, name=DEFAULT_NAME):
241
+ def remove_metadata(path, *, name=DEFAULT_NAME):
177
242
 
178
243
  LOG.info("Removing metadata '%s' from %s", name, path)
179
244
 
180
245
  def callback(full):
181
246
  os.remove(full)
182
247
 
183
- _edit_metadata(path, name, callback)
248
+ return _edit_metadata(path, name, callback)
@@ -96,8 +96,20 @@ def register_commands(here, package, select, fail=None):
96
96
  continue
97
97
 
98
98
  obj = select(imported)
99
- if obj is not None:
100
- result[name] = obj
99
+ if obj is None:
100
+ continue
101
+
102
+ if hasattr(obj, "command"):
103
+ name = obj.command
104
+
105
+ if name in result:
106
+ msg = f"Duplicate command '{name}', please choose a different command name for {type(obj)}"
107
+ raise ValueError(msg)
108
+ if " " in name:
109
+ msg = f"Commands cannot contain spaces: '{name}' in {type(obj)}"
110
+ raise ValueError(msg)
111
+
112
+ result[name] = obj
101
113
 
102
114
  for name, e in not_available.items():
103
115
  if fail is None:
@@ -1,12 +1,11 @@
1
- #!/usr/bin/env python
2
- # (C) Copyright 2024 ECMWF.
1
+ # (C) Copyright 2024 Anemoi contributors.
3
2
  #
4
3
  # This software is licensed under the terms of the Apache Licence Version 2.0
5
4
  # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
5
+ #
6
6
  # In applying this licence, ECMWF does not waive the privileges and immunities
7
7
  # granted to it by virtue of its status as an intergovernmental organisation
8
8
  # nor does it submit to any jurisdiction.
9
- #
10
9
 
11
10
  import os
12
11
 
@@ -1,4 +1,3 @@
1
- #!/usr/bin/env python3
2
1
  # (C) Copyright 2024 Anemoi contributors.
3
2
  #
4
3
  # This software is licensed under the terms of the Apache Licence Version 2.0
@@ -0,0 +1,76 @@
1
+ # (C) Copyright 2024 Anemoi contributors.
2
+ #
3
+ # This software is licensed under the terms of the Apache Licence Version 2.0
4
+ # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
5
+ #
6
+ # In applying this licence, ECMWF does not waive the privileges and immunities
7
+ # granted to it by virtue of its status as an intergovernmental organisation
8
+ # nor does it submit to any jurisdiction.
9
+
10
+ from __future__ import annotations
11
+
12
+ import functools
13
+ from typing import Any
14
+ from typing import Callable
15
+
16
+
17
+ def aliases(
18
+ aliases: dict[str, str | list[str]] | None = None, **kwargs: str | list[str]
19
+ ) -> Callable[[Callable], Callable]:
20
+ """Alias keyword arguments in a function call.
21
+
22
+ Allows for dynamically renaming keyword arguments in a function call.
23
+
24
+ Parameters
25
+ ----------
26
+ aliases : dict[str, str | list[str]] | None, optional
27
+ Key, value pair of aliases, with keys being the true name, and value being a str or list of aliases,
28
+ by default None
29
+ **kwargs : str | list[str]
30
+ Kwargs form of aliases
31
+
32
+ Returns
33
+ -------
34
+ Callable
35
+ Decorator function that renames keyword arguments in a function call.
36
+
37
+ Raises
38
+ ------
39
+ ValueError
40
+ If the aliasing would result in duplicate keys.
41
+
42
+ Examples
43
+ --------
44
+ ```python
45
+ @aliases(a="b", c=["d", "e"])
46
+ def func(a, c):
47
+ return a, c
48
+
49
+ func(a=1, c=2) # (1, 2)
50
+ func(b=1, d=2) # (1, 2)
51
+ ```
52
+
53
+ """
54
+
55
+ if aliases is None:
56
+ aliases = {}
57
+ aliases.update(kwargs)
58
+
59
+ aliases = {v: k for k, vs in aliases.items() for v in (vs if isinstance(vs, list) else [vs])}
60
+
61
+ def decorator(func: Callable) -> Callable:
62
+ @functools.wraps(func)
63
+ def wrapper(*args, **kwargs) -> Any:
64
+ keys = kwargs.keys()
65
+ for k in set(keys).intersection(set(aliases.keys())):
66
+ if aliases[k] in keys:
67
+ raise ValueError(
68
+ f"When aliasing {k} with {aliases[k]} duplicate keys were present. Cannot include both."
69
+ )
70
+ kwargs[aliases[k]] = kwargs.pop(k)
71
+
72
+ return func(*args, **kwargs)
73
+
74
+ return wrapper
75
+
76
+ return decorator
@@ -358,7 +358,7 @@ def check_config_mode(name="settings.toml", secrets_name=None, secrets=None) ->
358
358
  CHECKED[name] = True
359
359
 
360
360
 
361
- def find(metadata, what, result=None):
361
+ def find(metadata, what, result=None, *, select: callable = None):
362
362
  if result is None:
363
363
  result = []
364
364
 
@@ -369,7 +369,8 @@ def find(metadata, what, result=None):
369
369
 
370
370
  if isinstance(metadata, dict):
371
371
  if what in metadata:
372
- result.append(metadata[what])
372
+ if select is None or select(metadata[what]):
373
+ result.append(metadata[what])
373
374
 
374
375
  for k, v in metadata.items():
375
376
  find(v, what, result)
@@ -107,8 +107,8 @@ def as_datetime_list(date, default_increment=1):
107
107
  return list(_as_datetime_list(date, default_increment))
108
108
 
109
109
 
110
- def frequency_to_timedelta(frequency) -> datetime.timedelta:
111
- """Convert a frequency to a timedelta object.
110
+ def as_timedelta(frequency) -> datetime.timedelta:
111
+ """Convert anything to a timedelta object.
112
112
 
113
113
  Parameters
114
114
  ----------
@@ -171,6 +171,11 @@ def frequency_to_timedelta(frequency) -> datetime.timedelta:
171
171
  raise ValueError(f"Cannot convert frequency {frequency} to timedelta")
172
172
 
173
173
 
174
+ def frequency_to_timedelta(frequency) -> datetime.timedelta:
175
+ """Convert a frequency to a timedelta object."""
176
+ return as_timedelta(frequency)
177
+
178
+
174
179
  def frequency_to_string(frequency) -> str:
175
180
  """Convert a frequency (i.e. a datetime.timedelta) to a string.
176
181
 
@@ -1,6 +1,8 @@
1
- # (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
1
+ # (C) Copyright 2024 Anemoi contributors.
2
+ #
2
3
  # This software is licensed under the terms of the Apache Licence Version 2.0
3
4
  # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
5
+ #
4
6
  # In applying this licence, ECMWF does not waive the privileges and immunities
5
7
  # granted to it by virtue of its status as an intergovernmental organisation
6
8
  # nor does it submit to any jurisdiction.