anemoi-datasets 0.3.6__tar.gz → 0.3.8__tar.gz

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/.pre-commit-config.yaml

@@ -38,7 +38,7 @@ repos:
   hooks:
   - id: ruff
     # Next line if for documenation cod snippets
-    exclude: '^(dev/.*|[A-Za-z].*_)\.py$'
+    exclude: '^[^_].*_\.py$'
     args:
     - --line-length=120
     - --fix

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-datasets
-Version: 0.3.6
+Version: 0.3.8
 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,7 +223,7 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.9
 License-File: LICENSE
-Requires-Dist: anemoi-utils[provenance]>=0.3
+Requires-Dist: anemoi-utils[provenance]>=0.3.4
 Requires-Dist: numpy
 Requires-Dist: pyyaml
 Requires-Dist: semantic-version

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/docs/conf.py

@@ -99,6 +99,14 @@ intersphinx_mapping = {
         "https://anemoi-inference.readthedocs.io/en/latest/",
         ("../../anemoi-inference/docs/_build/html/objects.inv", None),
     ),
+    "anemoi-graphs": (
+        "https://anemoi-graphs.readthedocs.io/en/latest/",
+        ("../../anemoi-graphs/docs/_build/html/objects.inv", None),
+    ),
+    "anemoi-registry": (
+        "https://anemoi-registry.readthedocs.io/en/latest/",
+        ("../../anemoi-registry/docs/_build/html/objects.inv", None),
+    ),
 }
 
 

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/docs/index.rst

@@ -116,8 +116,10 @@ datasets <building-introduction>`.
 -  :ref:`anemoi-utils <anemoi-utils:index-page>`
 -  :ref:`anemoi-datasets <anemoi-datasets:index-page>`
 -  :ref:`anemoi-models <anemoi-models:index-page>`
+-  :ref:`anemoi-graphs <anemoi-graphs:index-page>`
 -  :ref:`anemoi-training <anemoi-training:index-page>`
 -  :ref:`anemoi-inference <anemoi-inference:index-page>`
+-  :ref:`anemoi-registry <anemoi-registry:index-page>`
 
 *********
  License

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/docs/using/configuration.rst

@@ -9,7 +9,8 @@ end with ``.zarr`` or ``.zip``, it is considered a dataset name and not
 a path or a URL.
 
 In that case, the *Anemoi* configuration is read from
-``~/.anemoi.toml``. Below is an example of such a configuration:
+``~/.config/anemoi/settings.toml``. Below is an example of such a
+configuration:
 
 .. literalinclude:: configuration.toml
    :language: toml

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/pyproject.toml

@@ -50,7 +50,7 @@ dynamic = [
   "version",
 ]
 dependencies = [
-  "anemoi-utils[provenance]>=0.3",
+  "anemoi-utils[provenance]>=0.3.4",
   "numpy",
   "pyyaml",
   "semantic-version",

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/src/anemoi/datasets/__init__.py

@@ -9,11 +9,14 @@ from ._version import __version__
 from .data import MissingDateError
 from .data import add_dataset_path
 from .data import add_named_dataset
+from .data import list_dataset_names
 from .data import open_dataset
 
 __all__ = [
-    "open_dataset",
-    "MissingDateError",
+    "__version__",
     "add_dataset_path",
     "add_named_dataset",
+    "list_dataset_names",
+    "MissingDateError",
+    "open_dataset",
 ]

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/src/anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.3.6'
-__version_tuple__ = version_tuple = (0, 3, 6)
+__version__ = version = '0.3.8'
+__version_tuple__ = version_tuple = (0, 3, 8)

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/src/anemoi/datasets/commands/copy.py

@@ -7,11 +7,14 @@
 
 import logging
 import os
+import shutil
 import sys
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import as_completed
 
 import tqdm
+from anemoi.utils.s3 import download
+from anemoi.utils.s3 import upload
 
 from . import Command
 
@@ -22,34 +25,63 @@ try:
 except AttributeError:
     isatty = False
 
-"""
 
-~/.aws/credentials
+class S3Downloader:
+    def __init__(self, source, target, transfers, overwrite, resume, verbosity, **kwargs):
+        self.source = source
+        self.target = target
+        self.transfers = transfers
+        self.overwrite = overwrite
+        self.resume = resume
+        self.verbosity = verbosity
 
-[default]
-endpoint_url = https://object-store.os-api.cci1.ecmwf.int
-aws_access_key_id=xxx
-aws_secret_access_key=xxxx
+    def run(self):
+        if self.target == ".":
+            self.target = os.path.basename(self.source)
+
+        if self.overwrite and os.path.exists(self.target):
+            LOG.info(f"Deleting {self.target}")
+            shutil.rmtree(self.target)
+
+        download(
+            self.source + "/" if not self.source.endswith("/") else self.source,
+            self.target,
+            overwrite=self.overwrite,
+            resume=self.resume,
+            verbosity=self.verbosity,
+            threads=self.transfers,
+        )
 
-Then:
 
-anemoi-datasets copy aifs-ea-an-oper-0001-mars-o96-1979-2022-1h-v3.zarr/
-    s3://ml-datasets/stable/aifs-ea-an-oper-0001-mars-o96-1979-2022-1h-v3.zarr
+class S3Uploader:
+    def __init__(self, source, target, transfers, overwrite, resume, verbosity, **kwargs):
+        self.source = source
+        self.target = target
+        self.transfers = transfers
+        self.overwrite = overwrite
+        self.resume = resume
+        self.verbosity = verbosity
 
-zinfo https://object-store.os-api.cci1.ecmwf.int/
-    ml-datasets/stable/aifs-ea-an-oper-0001-mars-o96-1979-2022-1h-v3.zarr
-"""
+    def run(self):
+        upload(
+            self.source,
+            self.target,
+            overwrite=self.overwrite,
+            resume=self.resume,
+            verbosity=self.verbosity,
+            threads=self.transfers,
+        )
 
 
-class Copier:
-    def __init__(self, source, target, transfers, block_size, overwrite, resume, progress, nested, rechunk, **kwargs):
+class DefaultCopier:
+    def __init__(self, source, target, transfers, block_size, overwrite, resume, verbosity, nested, rechunk, **kwargs):
         self.source = source
         self.target = target
         self.transfers = transfers
         self.block_size = block_size
         self.overwrite = overwrite
         self.resume = resume
-        self.progress = progress
+        self.verbosity = verbosity
         self.nested = nested
         self.rechunk = rechunk
 
@@ -62,7 +94,7 @@ class Copier:
             return zarr.storage.NestedDirectoryStore(path)
         return path
 
-    def copy_chunk(self, n, m, source, target, _copy, progress):
+    def copy_chunk(self, n, m, source, target, _copy, verbosity):
         if _copy[n:m].all():
             LOG.info(f"Skipping {n} to {m}")
             return None
@@ -82,7 +114,7 @@ class Copier:
                 range(n, m),
                 desc=f"Copying {n} to {m}",
                 leave=False,
-                disable=not isatty and not progress,
+                disable=not isatty and not verbosity,
             ):
                 target[i] = source[i]
 
@@ -107,7 +139,7 @@ class Copier:
             #    raise NotImplementedError("Rechunking with multiple transfers is not implemented")
         return chunks
 
-    def copy_data(self, source, target, _copy, progress):
+    def copy_data(self, source, target, _copy, verbosity):
         LOG.info("Copying data")
         source_data = source["data"]
 
@@ -121,6 +153,7 @@ class Copier:
                 shape=source_data.shape,
                 chunks=self.data_chunks,
                 dtype=source_data.dtype,
+                fill_value=source_data.fill_value,
             )
         )
 
@@ -136,7 +169,7 @@ class Copier:
                     source_data,
                     target_data,
                     _copy,
-                    progress,
+                    verbosity,
                 )
             )
             n += self.block_size
@@ -151,7 +184,7 @@ class Copier:
 
         LOG.info("Copied data")
 
-    def copy_array(self, name, source, target, _copy, progress):
+    def copy_array(self, name, source, target, _copy, verbosity):
         for k, v in source.attrs.items():
             target.attrs[k] = v
 
@@ -159,14 +192,14 @@ class Copier:
             return
 
         if name == "data":
-            self.copy_data(source, target, _copy, progress)
+            self.copy_data(source, target, _copy, verbosity)
             return
 
         LOG.info(f"Copying {name}")
         target[name] = source[name]
         LOG.info(f"Copied {name}")
 
-    def copy_group(self, source, target, _copy, progress):
+    def copy_group(self, source, target, _copy, verbosity):
         import zarr
 
         for k, v in source.attrs.items():
@@ -179,7 +212,7 @@ class Copier:
                     source[name],
                     group,
                     _copy,
-                    progress,
+                    verbosity,
                 )
             else:
                 self.copy_array(
@@ -187,10 +220,10 @@ class Copier:
                     source,
                     target,
                     _copy,
-                    progress,
+                    verbosity,
                 )
 
-    def copy(self, source, target, progress):
+    def copy(self, source, target, verbosity):
         import zarr
 
         if "_copy" not in target:
@@ -201,7 +234,7 @@ class Copier:
         _copy = target["_copy"]
         _copy_np = _copy[:]
 
-        self.copy_group(source, target, _copy_np, progress)
+        self.copy_group(source, target, _copy_np, verbosity)
         del target["_copy"]
 
     def run(self):
@@ -260,7 +293,7 @@ class Copier:
         assert target is not None, target
 
         source = zarr.open(self._store(self.source), mode="r")
-        self.copy(source, target, self.progress)
+        self.copy(source, target, self.verbosity)
 
 
 class CopyMixin:
@@ -279,7 +312,10 @@ class CopyMixin:
         )
         command_parser.add_argument("--transfers", type=int, default=8, help="Number of parallel transfers.")
         command_parser.add_argument(
-            "--progress", action="store_true", help="Force show progress bar, even if not in an interactive shell."
+            "--verbosity",
+            type=int,
+            help="Verbosity level. 0 is silent, 1 is normal, 2 is verbose.",
+            default=1,
         )
         command_parser.add_argument("--nested", action="store_true", help="Use ZARR's nested directpry backend.")
         command_parser.add_argument(
@@ -295,7 +331,29 @@ class CopyMixin:
         command_parser.add_argument("target", help="Target location.")
 
     def run(self, args):
-        Copier(**vars(args)).run()
+        if args.source == args.target:
+            raise ValueError("Source and target are the same.")
+
+        kwargs = vars(args)
+
+        if args.overwrite and args.resume:
+            raise ValueError("Cannot use --overwrite and --resume together.")
+
+        source_in_s3 = args.source.startswith("s3://")
+        target_in_s3 = args.target.startswith("s3://")
+
+        copier = None
+
+        if args.rechunk or (source_in_s3 and target_in_s3):
+            copier = DefaultCopier(**kwargs)
+        else:
+            if source_in_s3:
+                copier = S3Downloader(**kwargs)
+
+            if target_in_s3:
+                copier = S3Uploader(**kwargs)
+
+        copier.run()
 
 
 class Copy(CopyMixin, Command):

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/src/anemoi/datasets/create/__init__.py

@@ -97,13 +97,14 @@ class Creator:
 
         apply_patch(self.path, **kwargs)
 
-    def init_additions(self, delta=[1, 3, 6, 12, 24]):
+    def init_additions(self, delta=[1, 3, 6, 12, 24], statistics=True):
         from .loaders import StatisticsAddition
         from .loaders import TendenciesStatisticsAddition
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
-        a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
-        a.initialise()
+        if statistics:
+            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a.initialise()
 
         for d in delta:
             try:
@@ -112,13 +113,14 @@ class Creator:
             except TendenciesStatisticsDeltaNotMultipleOfFrequency:
                 self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
 
-    def run_additions(self, parts=None, delta=[1, 3, 6, 12, 24]):
+    def run_additions(self, parts=None, delta=[1, 3, 6, 12, 24], statistics=True):
         from .loaders import StatisticsAddition
         from .loaders import TendenciesStatisticsAddition
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
-        a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
-        a.run(parts)
+        if statistics:
+            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a.run(parts)
 
         for d in delta:
             try:
@@ -127,13 +129,14 @@ class Creator:
             except TendenciesStatisticsDeltaNotMultipleOfFrequency:
                 self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
 
-    def finalise_additions(self, delta=[1, 3, 6, 12, 24]):
+    def finalise_additions(self, delta=[1, 3, 6, 12, 24], statistics=True):
         from .loaders import StatisticsAddition
         from .loaders import TendenciesStatisticsAddition
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
-        a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
-        a.finalise()
+        if statistics:
+            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a.finalise()
 
         for d in delta:
             try:

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/src/anemoi/datasets/create/functions/sources/accumulations.py

@@ -19,6 +19,8 @@ from climetlab.utils.availability import Availability
 
 from anemoi.datasets.create.utils import to_datetime_list
 
+from .mars import use_grib_paramid
+
 LOG = logging.getLogger(__name__)
 
 
@@ -85,6 +87,7 @@ class Accumulation:
             stepType="accum",
             startStep=self.startStep,
             endStep=self.endStep,
+            check_nans=True,
         )
         self.values = None
         self.done = True
@@ -230,6 +233,7 @@ def identity(x):
 
 
 def compute_accumulations(
+    context,
     dates,
     request,
     user_accumulation_period=6,
@@ -306,7 +310,10 @@ def compute_accumulations(
     ds = cml.load_source("empty")
     for r in compressed.iterate():
         request.update(r)
+        if context.use_grib_paramid and "param" in request:
+            request = use_grib_paramid(request)
         print("🌧️", request)
+
         ds = ds + cml.load_source("mars", **request)
 
     accumulations = {}
@@ -395,7 +402,7 @@ def accumulations(context, dates, **request):
     class_ = request.get("class", "od")
     stream = request.get("stream", "oper")
 
-    user_accumulation_period = request.get("accumulation_period", 6)
+    user_accumulation_period = request.pop("accumulation_period", 6)
 
     KWARGS = {
         ("od", "oper"): dict(patch=scda),
@@ -409,6 +416,7 @@ def accumulations(context, dates, **request):
     context.trace("🌧️", f"accumulations {request} {user_accumulation_period} {kwargs}")
 
     return compute_accumulations(
+        context,
         dates,
         request,
         user_accumulation_period=user_accumulation_period,

{anemoi_datasets-0.3.6 → anemoi_datasets-0.3.8}/src/anemoi/datasets/create/functions/sources/mars.py

@@ -9,6 +9,7 @@
 import datetime
 from copy import deepcopy
 
+from anemoi.utils.humanize import did_you_mean
 from climetlab import load_source
 from climetlab.utils.availability import Availability
 
@@ -102,6 +103,74 @@ def use_grib_paramid(r):
     return r
 
 
+MARS_KEYS = [
+    "accuracy",
+    "activity",
+    "anoffset",
+    "area",
+    "bitmap",
+    "channel",
+    "class",
+    "database",
+    "dataset",
+    "date",
+    "diagnostic",
+    "direction",
+    "domain",
+    "expect",
+    "experiment",
+    "expver",
+    "fcmonth",
+    "fcperiod",
+    "fieldset",
+    "filter",
+    "format",
+    "frame",
+    "frequency",
+    "gaussian",
+    "generation",
+    "grid",
+    "hdate",
+    "ident",
+    "instrument",
+    "interpolation",
+    "intgrid",
+    "iteration",
+    "level",
+    "levelist",
+    "levtype",
+    "method",
+    "model",
+    "month",
+    "number",
+    "obsgroup",
+    "obstype",
+    "offsetdate",
+    "offsettime",
+    "optimise",
+    "origin",
+    "packing",
+    "padding",
+    "param",
+    "quantile",
+    "realization",
+    "reference",
+    "reportype",
+    "repres",
+    "resol",
+    "resolution",
+    "rotation",
+    "step",
+    "stream",
+    "system",
+    "target",
+    "time",
+    "truncation",
+    "type",
+    "year",
+]
+
+
 def mars(context, dates, *requests, date_key="date", **kwargs):
     if not requests:
         requests = [kwargs]
@@ -117,6 +186,11 @@ def mars(context, dates, *requests, date_key="date", **kwargs):
         if DEBUG:
             context.trace("✅", f"load_source(mars, {r}")
 
+        for k, v in r.items():
+            if k not in MARS_KEYS:
+                raise ValueError(
+                    f"⚠️ Unknown key {k}={v} in MARS request. Did you mean '{did_you_mean(k, MARS_KEYS)}' ?"
+                )
         ds = ds + load_source("mars", **r)
     return ds