anemoi-datasets 0.3.6__py3-none-any.whl → 0.3.7__py3-none-any.whl

anemoi/datasets/__init__.py

@@ -9,6 +9,7 @@ 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__ = [
@@ -16,4 +17,5 @@ __all__ = [
     "MissingDateError",
     "add_dataset_path",
     "add_named_dataset",
+    "list_dataset_names",
 ]

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.7'
+__version_tuple__ = version_tuple = (0, 3, 7)

anemoi/datasets/commands/copy.py

@@ -12,6 +12,8 @@ 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,26 +24,48 @@ try:
 except AttributeError:
     isatty = False
 
-"""
 
-~/.aws/credentials
+class S3Downloader:
+    def __init__(self, source, target, transfers, overwrite, resume, progress, **kwargs):
+        self.source = source
+        self.target = target
+        self.transfers = transfers
+        self.overwrite = overwrite
+        self.resume = resume
+        self.progress = progress
 
-[default]
-endpoint_url = https://object-store.os-api.cci1.ecmwf.int
-aws_access_key_id=xxx
-aws_secret_access_key=xxxx
+    def run(self):
+        download(
+            self.source + "/" if not self.source.endswith("/") else self.source,
+            self.target,
+            overwrite=self.overwrite,
+            ignore_existing=self.resume,
+            threads=self.transfers,
+            show_progress=self.progress,
+        )
 
-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, progress, **kwargs):
+        self.source = source
+        self.target = target
+        self.transfers = transfers
+        self.overwrite = overwrite
+        self.resume = resume
+        self.progress = progress
 
-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,
+            ignore_existing=self.resume,
+            threads=self.transfers,
+            show_progress=self.progress,
+        )
 
 
-class Copier:
+class DefaultCopier:
     def __init__(self, source, target, transfers, block_size, overwrite, resume, progress, nested, rechunk, **kwargs):
         self.source = source
         self.target = target
@@ -295,7 +319,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/data/__init__.py

@@ -30,3 +30,10 @@ def open_dataset(*args, **kwargs):
     ds.arguments = {"args": args, "kwargs": kwargs}
     ds._check()
     return ds
+
+
+def list_dataset_names(*args, **kwargs):
+    ds = _open_dataset(*args, **kwargs)
+    names = set()
+    ds.get_dataset_names(names)
+    return sorted(names)

anemoi/datasets/data/dataset.py

@@ -226,3 +226,6 @@ class Dataset:
     @property
     def label(self):
         return self.__class__.__name__.lower()
+
+    def get_dataset_names(self, names):
+        raise NotImplementedError(self)

anemoi/datasets/data/forwards.py

@@ -103,6 +103,9 @@ class Forwards(Dataset):
             f"subclass_metadata_specific() must be implemented in derived class {self.__class__.__name__}"
         )
 
+    def get_dataset_names(self, names):
+        self.forward.get_dataset_names(names)
+
 
 class Combined(Forwards):
     def __init__(self, datasets):
@@ -193,6 +196,10 @@ class Combined(Forwards):
             offset += len(d)
         return result
 
+    def get_dataset_names(self, names):
+        for d in self.datasets:
+            d.get_dataset_names(names)
+
 
 class GivenAxis(Combined):
     """Given a given axis, combine the datasets along that axis."""

anemoi/datasets/data/masked.py

@@ -70,8 +70,12 @@ class Thinning(Masked):
         self.thinning = thinning
         self.method = method
 
-        latitudes = forward.latitudes.reshape(forward.field_shape)
-        longitudes = forward.longitudes.reshape(forward.field_shape)
+        shape = forward.field_shape
+        if len(shape) != 2:
+            raise ValueError("Thinning only works latitude/longitude fields")
+
+        latitudes = forward.latitudes.reshape(shape)
+        longitudes = forward.longitudes.reshape(shape)
         latitudes = latitudes[::thinning, ::thinning].flatten()
         longitudes = longitudes[::thinning, ::thinning].flatten()
 

anemoi/datasets/data/statistics.py

@@ -39,3 +39,7 @@ class Statistics(Forwards):
 
     def tree(self):
         return Node(self, [self.forward.tree()])
+
+    def get_dataset_names(self, names):
+        super().get_dataset_names(names)
+        self._statistic.get_dataset_names(names)

anemoi/datasets/data/stores.py

@@ -6,6 +6,7 @@
 # nor does it submit to any jurisdiction.
 
 import logging
+import os
 import warnings
 from functools import cached_property
 
@@ -235,7 +236,11 @@ class Zarr(Dataset):
 
     @property
     def field_shape(self):
-        return tuple(self.z.attrs["field_shape"])
+        try:
+            return tuple(self.z.attrs["field_shape"])
+        except KeyError:
+            LOG.warning("No 'field_shape' in %r, assuming 1D fields", self)
+            return (self.shape[-1],)
 
     @property
     def frequency(self):
@@ -288,6 +293,10 @@ class Zarr(Dataset):
     def tree(self):
         return Node(self, [], path=self.path)
 
+    def get_dataset_names(self, names):
+        name, _ = os.path.splitext(os.path.basename(self.path))
+        names.add(name)
+
 
 class ZarrWithMissingDates(Zarr):
     def __init__(self, path):

{anemoi_datasets-0.3.6.dist-info → anemoi_datasets-0.3.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-datasets
-Version: 0.3.6
+Version: 0.3.7
 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.dist-info → anemoi_datasets-0.3.7.dist-info}/RECORD

@@ -1,10 +1,10 @@
-anemoi/datasets/__init__.py,sha256=DC7ttKT--pmhBQALX_Cn7P28dngsJucKi5y-Ydm28QM,700
+anemoi/datasets/__init__.py,sha256=50-v6XPwmNE4LK9PVmyWu3v42F38EhxShd2EaXt1IIA,763
 anemoi/datasets/__main__.py,sha256=cLA2PidDTOUHaDGzd0_E5iioKYNe-PSTv567Y2fuwQk,723
-anemoi/datasets/_version.py,sha256=IKAQ4gPrCQ2FWMXOFRqouULC2EQI1zCb4iXHsnfbmTQ,411
+anemoi/datasets/_version.py,sha256=vVN20516E2VTC9JNgtvqrQNlj5XptaB_a5z2XL8NFxg,411
 anemoi/datasets/grids.py,sha256=3YBMMJodgYhavarXPAlMZHaMtDT9v2IbTmAXZTqf8Qo,8481
 anemoi/datasets/commands/__init__.py,sha256=qAybFZPBBQs0dyx7dZ3X5JsLpE90pwrqt1vSV7cqEIw,706
 anemoi/datasets/commands/compare.py,sha256=p2jQOAC3JhScCLF0GjTCO8goYLWLN8p7vzy_gf5fFcI,1473
-anemoi/datasets/commands/copy.py,sha256=fba-zjD0iTHHXHhPEcm8VhDzsXQXDUxlbtTA1TovyT0,9991
+anemoi/datasets/commands/copy.py,sha256=CHIQuzHUs0uezU9c6jJry6yhilYeGn1wmIubGYatEaY,11450
 anemoi/datasets/commands/create.py,sha256=POdOsVDlvRrHFFkI3SNXNgNIbSxkVUUPMoo660x7Ma0,987
 anemoi/datasets/commands/inspect.py,sha256=G3fzcgiLaU8jln7GKvgamN7Y06-qC_JnFw2SbNn1_E4,18646
 anemoi/datasets/commands/scan.py,sha256=HxsLdCgBMSdEXjlJfPq5M_9LxXHHQIoZ1ZEHO_AoPgA,2881
@@ -45,29 +45,29 @@ anemoi/datasets/create/functions/sources/source.py,sha256=hPQnV_6UIxFw97uRKcTA8T
 anemoi/datasets/create/functions/sources/tendencies.py,sha256=kwS_GZt8R9kpfs5RrvxPb0Gj-5nDP0sgJgfSRCAwwww,4057
 anemoi/datasets/create/statistics/__init__.py,sha256=X50drgE-ltuNe7bSIyvyeC4GeTqGTQGbglh2-2aVWKE,15445
 anemoi/datasets/create/statistics/summary.py,sha256=sgmhA24y3VRyjmDUgTnPIqcHSlWBbFA0qynx6gJ9Xw8,3370
-anemoi/datasets/data/__init__.py,sha256=tacn6K_VZ-pYhLmGePG5sze8kmqGpqscYb-bMyQnWtk,888
+anemoi/datasets/data/__init__.py,sha256=to9L_RZVQ4OgyHUpX6lcvt4GqJdZjBa5HCTaWx1aGKo,1046
 anemoi/datasets/data/concat.py,sha256=AkpyOs16OjW7X0cdyYFQfWSCV6dteXBp-x9WlokO-DI,3550
-anemoi/datasets/data/dataset.py,sha256=UDnidq2amyCT2COH05pGfDCJcmkdMj1ubtHk9cl-qcE,7384
+anemoi/datasets/data/dataset.py,sha256=LBUwWhwcAcovLv0FOLT-rA-yNZhcBFUMc03BfVh7UFc,7465
 anemoi/datasets/data/debug.css,sha256=z2X_ZDSnZ9C3pyZPWnQiEyAxuMxUaxJxET4oaCImTAQ,211
 anemoi/datasets/data/debug.py,sha256=PcyrjgxaLzeb_vf12pvUtPPVvBRHNm1SimythZvqsP4,6303
 anemoi/datasets/data/ensemble.py,sha256=AsP7Xx0ZHLoZs6a4EC0jtyGYIcOvZvvKXhgNsIvqIN8,1137
-anemoi/datasets/data/forwards.py,sha256=t9YQCN7j75VMInt0uP9JUJoh1klF1Z1xnwy5_kDMDQs,7700
+anemoi/datasets/data/forwards.py,sha256=UZOOMUblGS21aaPoFfQa0ONSUaxkqlZQF3KGRhlCr9I,7899
 anemoi/datasets/data/grids.py,sha256=rooOeR6rvjl4U8B4LO3N23fcgxvGE7ZUmhVryk1QS4M,7493
 anemoi/datasets/data/indexing.py,sha256=625m__JG5m_tDMrkz1hB6Vydenwt0oHuyAlc-o3Zwos,4799
 anemoi/datasets/data/join.py,sha256=dtCBbMTicqrRPxfBULi3RwEcQBLhQpIcvCjdN5A3XUU,4892
-anemoi/datasets/data/masked.py,sha256=Fzkehyka70CiS0LYSy_uyVYu2gKLwDSxlbm8GiC_pYs,3742
+anemoi/datasets/data/masked.py,sha256=czAv1ZfZ9q6Wr4RqI2Xj8SEm7yoCgJrwMl-CPDs_wSI,3857
 anemoi/datasets/data/misc.py,sha256=cu2rMNtq8M8yzcR1CGHHQksE0p9C2SLVigOTMH6ilMs,10400
 anemoi/datasets/data/select.py,sha256=U3AEid80mrJKu0SF4lLc-bRWMVcAZwHNUHUHRehvuHU,3680
-anemoi/datasets/data/statistics.py,sha256=PKRgcCiZEb1HjkIveVGhE3TzUy9Qe3AYWGFD72Urah8,1514
-anemoi/datasets/data/stores.py,sha256=oEjUra0zzIysiUvh-RBQRzcbviggejEQiMO5RfpjPyM,10896
+anemoi/datasets/data/statistics.py,sha256=lZCcKw9s7ttMBEp6ANyxtbXoZZvchhE7SClq-D4AUR8,1645
+anemoi/datasets/data/stores.py,sha256=yy914zMHIYKm5q6mHOqGeK0dC_26VFeqKLXyb7x9NXE,11190
 anemoi/datasets/data/subset.py,sha256=9urVTXdnwCgqn0_BRYquMi8oiXn4ubAf0n4586hWfKw,3814
 anemoi/datasets/data/unchecked.py,sha256=xhdMg-ToI1UfBWHNsWyn1y2meZWngZtHx-33L0KqKp8,4037
 anemoi/datasets/dates/__init__.py,sha256=4ItowfLLh90T8L_JOjtv98lE6M7gAaWt7dV3niUrFvk,4473
 anemoi/datasets/dates/groups.py,sha256=iq310Pi7ullglOhcNblv14MmcT8FPgYCD5s45qAfV_s,3383
 anemoi/datasets/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-anemoi_datasets-0.3.6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-anemoi_datasets-0.3.6.dist-info/METADATA,sha256=sEnXtS3eh-ix9Z1aR8RFDqBlnmK-XE6A42ga2uE_Utc,16017
-anemoi_datasets-0.3.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-anemoi_datasets-0.3.6.dist-info/entry_points.txt,sha256=yR-o-4uiPEA_GLBL81SkMYnUoxq3CAV3hHulQiRtGG0,66
-anemoi_datasets-0.3.6.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_datasets-0.3.6.dist-info/RECORD,,
+anemoi_datasets-0.3.7.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+anemoi_datasets-0.3.7.dist-info/METADATA,sha256=MviDtSJZZDW9XL00ALGMtwUjOJzfozG1wrxzq4tb3DA,16019
+anemoi_datasets-0.3.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+anemoi_datasets-0.3.7.dist-info/entry_points.txt,sha256=yR-o-4uiPEA_GLBL81SkMYnUoxq3CAV3hHulQiRtGG0,66
+anemoi_datasets-0.3.7.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_datasets-0.3.7.dist-info/RECORD,,