anemoi-datasets 0.5.7__py3-none-any.whl → 0.5.10__py3-none-any.whl

anemoi/datasets/__init__.py

@@ -1,19 +1,27 @@
-# (C) Copyright 2023 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.
 
-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
 
+try:
+    # NOTE: the `_version.py` file must not be present in the git repository
+    #   as it is generated by setuptools at install time
+    from ._version import __version__  # type: ignore
+except ImportError:  # pragma: no cover
+    # Local copy or not installed with setuptools
+    __version__ = "999"
+
 __all__ = [
-    "__version__",
     "add_dataset_path",
     "add_named_dataset",
     "list_dataset_names",

anemoi/datasets/__main__.py

@@ -1,12 +1,11 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (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.
-#
 
 from anemoi.utils.cli import cli_main
 from anemoi.utils.cli import make_parser

anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.7'
-__version_tuple__ = version_tuple = (0, 5, 7)
+__version__ = version = '0.5.10'
+__version_tuple__ = version_tuple = (0, 5, 10)

anemoi/datasets/commands/__init__.py

@@ -1,12 +1,11 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (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 os
 

anemoi/datasets/commands/cleanup.py

@@ -1,3 +1,12 @@
+# (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 logging
 import time
 

anemoi/datasets/commands/compare.py

@@ -1,12 +1,12 @@
-#!/usr/bin/env python
-# (C) Copyright 2024 ECMWF.
+# (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 numpy as np
 import tqdm

anemoi/datasets/commands/copy.py

@@ -1,20 +1,22 @@
-# (C) Copyright 2023 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.
 
+
 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 anemoi.utils.remote import Transfer
+from anemoi.utils.remote import TransferMethodNotImplementedError
 
 from . import Command
 
@@ -26,54 +28,7 @@ except AttributeError:
     isatty = False
 
 
-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
-
-    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,
-        )
-
-
-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
-
-    def run(self):
-        upload(
-            self.source,
-            self.target,
-            overwrite=self.overwrite,
-            resume=self.resume,
-            verbosity=self.verbosity,
-            threads=self.transfers,
-        )
-
-
-class DefaultCopier:
+class ZarrCopier:
     def __init__(self, source, target, transfers, block_size, overwrite, resume, verbosity, nested, rechunk, **kwargs):
         self.source = source
         self.target = target
@@ -87,6 +42,14 @@ class DefaultCopier:
 
         self.rechunking = rechunk.split(",") if rechunk else []
 
+        source_is_ssh = self.source.startswith("ssh://")
+        target_is_ssh = self.target.startswith("ssh://")
+
+        if source_is_ssh or target_is_ssh:
+            if self.rechunk:
+                raise NotImplementedError("Rechunking with SSH not implemented.")
+            assert NotImplementedError("SSH not implemented.")
+
     def _store(self, path, nested=False):
         if nested:
             import zarr
@@ -334,26 +297,33 @@ class CopyMixin:
         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)
-
+        if not args.rechunk:
+            # rechunking is only supported for ZARR datasets, it is implemented in this package
+            try:
+                if args.source.startswith("s3://") and not args.source.endswith("/"):
+                    args.source = args.source + "/"
+                copier = Transfer(
+                    args.source,
+                    args.target,
+                    overwrite=args.overwrite,
+                    resume=args.resume,
+                    verbosity=args.verbosity,
+                    threads=args.transfers,
+                )
+                copier.run()
+                return
+            except TransferMethodNotImplementedError:
+                # DataTransfer relies on anemoi-utils which is agnostic to the source and target format
+                # it transfers file and folders, ignoring that it is zarr data
+                # if it is not implemented, we fallback to the ZarrCopier
+                pass
+
+        copier = ZarrCopier(**vars(args))
         copier.run()
+        return
 
 
 class Copy(CopyMixin, Command):

anemoi/datasets/commands/create.py

@@ -1,3 +1,12 @@
+# (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 datetime
 import logging
 import time
@@ -14,9 +23,7 @@ LOG = logging.getLogger(__name__)
 
 
 def task(what, options, *args, **kwargs):
-    """
-    Make sure `import Creator` is done in the sub-processes, and not in the main one.
-    """
+    """Make sure `import Creator` is done in the sub-processes, and not in the main one."""
 
     now = datetime.datetime.now()
     LOG.info(f"🎬 Task {what}({args},{kwargs}) starting")
@@ -117,7 +124,9 @@ class Create(Command):
                 opt["parts"] = f"{n+1}/{total}"
                 futures.append(executor.submit(task, "load", opt))
 
-            for future in tqdm.tqdm(as_completed(futures), desc="Loading", total=len(futures), colour="green", position=parallel + 1):  # fmt: skip
+            for future in tqdm.tqdm(
+                as_completed(futures), desc="Loading", total=len(futures), colour="green", position=parallel + 1
+            ):
                 future.result()
 
         with ExecutorClass(max_workers=1) as executor:
@@ -133,7 +142,13 @@ class Create(Command):
             for n in range(total):
                 futures.append(executor.submit(task, "load-additions", opt))
 
-            for future in tqdm.tqdm(as_completed(futures), desc="Computing additions", total=len(futures), colour="green", position=parallel + 1):  # fmt: skip
+            for future in tqdm.tqdm(
+                as_completed(futures),
+                desc="Computing additions",
+                total=len(futures),
+                colour="green",
+                position=parallel + 1,
+            ):
                 future.result()
 
         with ExecutorClass(max_workers=1) as executor:

anemoi/datasets/commands/finalise-additions.py

@@ -1,3 +1,12 @@
+# (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 logging
 import time
 

anemoi/datasets/commands/finalise.py

@@ -1,3 +1,12 @@
+# (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 logging
 import time
 

anemoi/datasets/commands/init-additions.py

@@ -1,3 +1,12 @@
+# (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 logging
 import time
 

anemoi/datasets/commands/init.py

@@ -1,3 +1,12 @@
+# (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 logging
 import time
 

anemoi/datasets/commands/inspect.py

@@ -1,6 +1,8 @@
-# (C) Copyright 2023 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/datasets/commands/load-additions.py

@@ -1,3 +1,12 @@
+# (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 logging
 import time
 

anemoi/datasets/commands/load.py

@@ -1,3 +1,12 @@
+# (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 logging
 import time
 

anemoi/datasets/commands/patch.py

@@ -1,3 +1,12 @@
+# (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 logging
 import time
 

anemoi/datasets/commands/publish.py

@@ -1,3 +1,12 @@
+# (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 logging
 
 from . import Command

anemoi/datasets/commands/scan.py

@@ -1,3 +1,12 @@
+# (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 fnmatch
 import os
 import sys

anemoi/datasets/compute/__init__.py

@@ -0,0 +1,8 @@
+# (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/datasets/compute/recentre.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (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 logging
 

anemoi/datasets/create/__init__.py

@@ -1,11 +1,11 @@
-# (C) Copyright 2023 ECMWF.
+# (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 datetime
 import json
@@ -14,9 +14,9 @@ import os
 import time
 import uuid
 import warnings
-from copy import deepcopy
 from functools import cached_property
 
+import cftime
 import numpy as np
 import tqdm
 from anemoi.utils.config import DotDict as DotDict
@@ -67,6 +67,19 @@ def json_tidy(o):
     if isinstance(o, datetime.timedelta):
         return frequency_to_string(o)
 
+    if isinstance(o, cftime.DatetimeJulian):
+        import pandas as pd
+
+        o = pd.Timestamp(
+            o.year,
+            o.month,
+            o.day,
+            o.hour,
+            o.minute,
+            o.second,
+        )
+        return o.isoformat()
+
     raise TypeError(repr(o) + " is not JSON serializable")
 
 
@@ -94,10 +107,6 @@ def build_statistics_dates(dates, start, end):
     return (start.isoformat(), end.isoformat())
 
 
-def _ignore(*args, **kwargs):
-    pass
-
-
 def _path_readable(path):
     import zarr
 
@@ -278,6 +287,16 @@ class Size(Actor):
         metadata = compute_directory_sizes(self.path)
         self.update_metadata(**metadata)
 
+        # Look for constant fields
+        ds = open_dataset(self.path)
+        constants = ds.computed_constant_fields()
+
+        variables_metadata = self.dataset.zarr_metadata.get("variables_metadata", {}).copy()
+        for k in constants:
+            variables_metadata[k]["constant_in_time"] = True
+
+        self.update_metadata(constant_fields=constants, variables_metadata=variables_metadata)
+
 
 class HasRegistryMixin:
     @cached_property
@@ -328,7 +347,20 @@ def build_input_(main_config, output_config):
 
 class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixin):
     dataset_class = NewDataset
-    def __init__(self, path, config, check_name=False, overwrite=False, use_threads=False, statistics_temp_dir=None, progress=None, test=False, cache=None, **kwargs):  # fmt: skip
+
+    def __init__(
+        self,
+        path,
+        config,
+        check_name=False,
+        overwrite=False,
+        use_threads=False,
+        statistics_temp_dir=None,
+        progress=None,
+        test=False,
+        cache=None,
+        **kwargs,
+    ):
         if _path_readable(path) and not overwrite:
             raise Exception(f"{path} already exists. Use overwrite=True to overwrite.")
 
@@ -412,7 +444,24 @@ class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         metadata.update(self.main_config.get("add_metadata", {}))
 
         metadata["_create_yaml_config"] = self.main_config.get_serialisable_dict()
-        metadata["recipe"] = sanitise(self.main_config.get_serialisable_dict())
+
+        recipe = sanitise(self.main_config.get_serialisable_dict())
+
+        # Remove stuff added by prepml
+        for k in [
+            "build_dataset",
+            "config_format_version",
+            "config_path",
+            "dataset_status",
+            "ecflow",
+            "metadata",
+            "platform",
+            "reading_chunks",
+            "upload",
+        ]:
+            recipe.pop(k, None)
+
+        metadata["recipe"] = recipe
 
         metadata["description"] = self.main_config.description
         metadata["licence"] = self.main_config["licence"]
@@ -496,7 +545,9 @@ class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
 
 
 class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixin):
-    def __init__(self, path, parts=None,  use_threads=False, statistics_temp_dir=None, progress=None, cache=None, **kwargs):  # fmt: skip
+    def __init__(
+        self, path, parts=None, use_threads=False, statistics_temp_dir=None, progress=None, cache=None, **kwargs
+    ):
         super().__init__(path, cache=cache)
         self.use_threads = use_threads
         self.statistics_temp_dir = statistics_temp_dir

anemoi/datasets/create/check.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2023 ECMWF.
+# (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 logging
 import re
@@ -88,7 +89,7 @@ class DatasetName:
             self.messages.append(
                 f"the dataset name {self} does not follow naming convention. "
                 "See here for details: "
-                "https://confluence.ecmwf.int/display/DWF/Datasets+available+as+zarr"
+                "https://anemoi-registry.readthedocs.io/en/latest/naming-conventions.html"
             )
 
     def check_resolution(self, resolution):

anemoi/datasets/create/chunks.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2024 ECMWF.
+# (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 logging
 import warnings
 

anemoi/datasets/create/config.py

@@ -1,11 +1,12 @@
-# (C) Copyright 2023 ECMWF.
+# (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 datetime
 import logging
 import os
@@ -214,15 +215,14 @@ def _prepare_serialisation(o):
 def set_to_test_mode(cfg):
     NUMBER_OF_DATES = 4
 
-    dates = cfg["dates"]
     LOG.warning(f"Running in test mode. Changing the list of dates to use only {NUMBER_OF_DATES}.")
     groups = Groups(**LoadersConfig(cfg).dates)
 
-    dates = groups.dates
+    dates = groups.provider.values
     cfg["dates"] = dict(
         start=dates[0],
         end=dates[NUMBER_OF_DATES - 1],
-        frequency=dates.frequency,
+        frequency=groups.provider.frequency,
         group_by=NUMBER_OF_DATES,
     )