anemoi-datasets 0.4.0__tar.gz → 0.4.3__tar.gz

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/.pre-commit-config.yaml

@@ -37,8 +37,8 @@ repos:
   rev: v0.4.6
   hooks:
   - id: ruff
-    # Next line if for documenation cod snippets
-    exclude: '^[^_].*_\.py$'
+    # Next line is to exclude for documentation code snippets
+    exclude: 'docs/(.*/)?[a-z]\w+_.py$'
     args:
     - --line-length=120
     - --fix

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/.readthedocs.yaml

@@ -10,7 +10,6 @@ sphinx:
 
 python:
    install:
-    - requirements: docs/requirements.txt
     - method: pip
       path: .
       extra_requirements:

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-datasets
-Version: 0.4.0
+Version: 0.4.3
 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,19 +223,22 @@ 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.5
+Requires-Dist: anemoi-utils[provenance]>=0.3.13
 Requires-Dist: numpy
 Requires-Dist: pyyaml
 Requires-Dist: semantic-version
 Requires-Dist: tqdm
-Requires-Dist: zarr<=2.17
+Requires-Dist: zarr
 Provides-Extra: all
+Requires-Dist: aiohttp; extra == "all"
 Requires-Dist: boto3; extra == "all"
 Requires-Dist: earthkit-data[mars]>=0.9; extra == "all"
 Requires-Dist: earthkit-geo>=0.2; extra == "all"
 Requires-Dist: earthkit-meteo; extra == "all"
 Requires-Dist: ecmwflibs>=0.6.3; extra == "all"
 Requires-Dist: entrypoints; extra == "all"
+Requires-Dist: gcsfs; extra == "all"
+Requires-Dist: kerchunk; extra == "all"
 Requires-Dist: pyproj; extra == "all"
 Requires-Dist: requests; extra == "all"
 Requires-Dist: s3fs; extra == "all"
@@ -247,26 +250,38 @@ Requires-Dist: ecmwflibs>=0.6.3; extra == "create"
 Requires-Dist: entrypoints; extra == "create"
 Requires-Dist: pyproj; extra == "create"
 Provides-Extra: dev
+Requires-Dist: aiohttp; extra == "dev"
 Requires-Dist: boto3; extra == "dev"
 Requires-Dist: earthkit-data[mars]>=0.9; extra == "dev"
 Requires-Dist: earthkit-geo>=0.2; extra == "dev"
 Requires-Dist: earthkit-meteo; extra == "dev"
 Requires-Dist: ecmwflibs>=0.6.3; extra == "dev"
 Requires-Dist: entrypoints; extra == "dev"
+Requires-Dist: gcsfs; extra == "dev"
+Requires-Dist: kerchunk; extra == "dev"
 Requires-Dist: nbsphinx; extra == "dev"
 Requires-Dist: pandoc; extra == "dev"
 Requires-Dist: pyproj; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: requests; extra == "dev"
+Requires-Dist: rstfmt; extra == "dev"
 Requires-Dist: s3fs; extra == "dev"
 Requires-Dist: sphinx; extra == "dev"
+Requires-Dist: sphinx-argparse<0.5; extra == "dev"
 Requires-Dist: sphinx-rtd-theme; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: nbsphinx; extra == "docs"
 Requires-Dist: pandoc; extra == "docs"
+Requires-Dist: rstfmt; extra == "docs"
 Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx-argparse<0.5; extra == "docs"
 Requires-Dist: sphinx-rtd-theme; extra == "docs"
+Provides-Extra: kerchunk
+Requires-Dist: gcsfs; extra == "kerchunk"
+Requires-Dist: kerchunk; extra == "kerchunk"
+Requires-Dist: s3fs; extra == "kerchunk"
 Provides-Extra: remote
+Requires-Dist: aiohttp; extra == "remote"
 Requires-Dist: boto3; extra == "remote"
 Requires-Dist: requests; extra == "remote"
 Requires-Dist: s3fs; extra == "remote"

anemoi_datasets-0.4.3/docs/building/advanced-options.rst

@@ -0,0 +1,3 @@
+##################
+ Advanced Options
+##################

anemoi_datasets-0.4.3/docs/building/sources/xarray.rst

@@ -0,0 +1,6 @@
+########
+ xarray
+########
+
+.. literalinclude:: xarray.yaml
+   :language: yaml

anemoi_datasets-0.4.3/docs/building/sources/xarray.yaml

@@ -0,0 +1,3 @@
+input:
+  xarray:
+    url: https://...

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/docs/building/sources.rst

@@ -23,6 +23,7 @@ The following `sources` are currently available:
    sources/mars
    sources/grib
    sources/netcdf
+   sources/xarray
    sources/opendap
    sources/forcings
    sources/accumulations

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/docs/index.rst

@@ -89,6 +89,7 @@ datasets <building-introduction>`.
    building/handling-missing-dates
    building/handling-missing-values
    building/statistics
+   building/advanced-options
 
 **Command line tool**
 

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/pyproject.toml

@@ -50,22 +50,24 @@ dynamic = [
   "version",
 ]
 dependencies = [
-  "anemoi-utils[provenance]>=0.3.5",
+  "anemoi-utils[provenance]>=0.3.13",
   "numpy",
   "pyyaml",
   "semantic-version",
   "tqdm",
-  "zarr<=2.17",
-
+  "zarr",
 ]
 
 optional-dependencies.all = [
+  "aiohttp",
   "boto3",
   "earthkit-data[mars]>=0.9",
   "earthkit-geo>=0.2",
   "earthkit-meteo",
   "ecmwflibs>=0.6.3",
   "entrypoints",
+  "gcsfs",
+  "kerchunk",
   "pyproj",
   "requests",
   "s3fs",
@@ -81,37 +83,53 @@ optional-dependencies.create = [
 ]
 
 optional-dependencies.dev = [
+  "aiohttp",
   "boto3",
   "earthkit-data[mars]>=0.9",
   "earthkit-geo>=0.2",
   "earthkit-meteo",
   "ecmwflibs>=0.6.3",
   "entrypoints",
+  "gcsfs",
+  "kerchunk",
   "nbsphinx",
   "pandoc",
   "pyproj",
   "pytest",
   "requests",
+  "rstfmt",
   "s3fs",
   "sphinx",
+  "sphinx-argparse<0.5",
   "sphinx-rtd-theme",
 ]
 
 optional-dependencies.docs = [
   "nbsphinx",
   "pandoc",
+  "rstfmt",
   "sphinx",
+  "sphinx-argparse<0.5",
   "sphinx-rtd-theme",
 ]
 
+optional-dependencies.kerchunk = [
+  "gcsfs",
+  "kerchunk",
+  "s3fs",
+]
+
 optional-dependencies.remote = [
+  "aiohttp",
   "boto3",
   "requests",
   "s3fs",
 ]
+
 optional-dependencies.tests = [
   "pytest",
 ]
+
 urls.Documentation = "https://anemoi-datasets.readthedocs.io/"
 urls.Homepage = "https://github.com/ecmwf/anemoi-datasets/"
 urls.Issues = "https://github.com/ecmwf/anemoi-datasets/issues"

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/src/anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.0'
-__version_tuple__ = version_tuple = (0, 4, 0)
+__version__ = version = '0.4.3'
+__version_tuple__ = version_tuple = (0, 4, 3)

anemoi_datasets-0.4.3/src/anemoi/datasets/commands/compare.py

@@ -0,0 +1,105 @@
+#!/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 numpy as np
+import tqdm
+import zarr
+
+from anemoi.datasets import open_dataset
+
+from . import Command
+
+
+class Compare(Command):
+    """Compare two datasets. This command compares the variables in two datasets and prints the mean of the common variables. It does not compare the data itself (yet)."""
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument("dataset1")
+        command_parser.add_argument("dataset2")
+        command_parser.add_argument("--data", action="store_true", help="Compare the data.")
+        command_parser.add_argument("--statistics", action="store_true", help="Compare the statistics.")
+
+    def run(self, args):
+        ds1 = open_dataset(args.dataset1)
+        ds2 = open_dataset(args.dataset2)
+
+        v1 = set(ds1.variables)
+        v2 = set(ds2.variables)
+
+        print("Only in dataset 1:", ", ".join(sorted(v1 - v2)))
+        print("Only in dataset 2:", ", ".join(sorted(v2 - v1)))
+        print()
+        common = sorted(v1 & v2)
+        print("Common:")
+        print("-------")
+        print()
+
+        for v in common:
+            print(
+                f"{v:14}",
+                f"{ds1.statistics['mean'][ds1.name_to_index[v]]:14g}",
+                f"{ds2.statistics['mean'][ds2.name_to_index[v]]:14g}",
+            )
+
+        if args.data:
+            print()
+            print("Data:")
+            print("-----")
+            print()
+
+            diff = 0
+            for a, b in tqdm.tqdm(zip(ds1, ds2)):
+                if not np.array_equal(a, b, equal_nan=True):
+                    diff += 1
+
+            print(f"Number of different rows: {diff}/{len(ds1)}")
+
+        if args.data:
+            print()
+            print("Data 2:")
+            print("-----")
+            print()
+
+            ds1 = zarr.open(args.dataset1, mode="r")
+            ds2 = zarr.open(args.dataset2, mode="r")
+
+            for name in (
+                "data",
+                "count",
+                "sums",
+                "squares",
+                "mean",
+                "stdev",
+                "minimum",
+                "maximum",
+                "latitudes",
+                "longitudes",
+            ):
+                a1 = ds1[name]
+                a2 = ds2[name]
+
+                if len(a1) != len(a2):
+                    print(f"{name}: lengths mismatch {len(a1)} != {len(a2)}")
+                    continue
+
+                diff = 0
+                for a, b in tqdm.tqdm(zip(a1, a2), leave=False):
+                    if not np.array_equal(a, b, equal_nan=True):
+                        if diff == 0:
+                            print(f"\n{name}: first different row:")
+                            print(a[a != b])
+                            print(b[a != b])
+
+                        diff += 1
+
+                print(f"{name}: {diff} different rows out of {len(a1)}")
+
+
+command = Compare

anemoi_datasets-0.4.3/src/anemoi/datasets/commands/create.py

@@ -0,0 +1,114 @@
+import datetime
+import logging
+import time
+from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import as_completed
+
+import tqdm
+from anemoi.utils.humanize import seconds_to_human
+
+from anemoi.datasets.create.trace import enable_trace
+
+from . import Command
+
+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.
+    """
+
+    now = datetime.datetime.now()
+    LOG.debug(f"Task {what}({args},{kwargs}) starting")
+
+    from anemoi.datasets.create import Creator
+
+    if "trace" in options:
+        enable_trace(options["trace"])
+
+    c = Creator(**options)
+    result = getattr(c, what)(*args, **kwargs)
+
+    LOG.debug(f"Task {what}({args},{kwargs}) completed ({datetime.datetime.now()-now})")
+    return result
+
+
+class Create(Command):
+    """Create a dataset."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument(
+            "--overwrite",
+            action="store_true",
+            help="Overwrite existing files. This will delete the target dataset if it already exists.",
+        )
+        command_parser.add_argument(
+            "--test",
+            action="store_true",
+            help="Build a small dataset, using only the first dates. And, when possible, using low resolution and less ensemble members.",
+        )
+        command_parser.add_argument("config", help="Configuration yaml file defining the recipe to create the dataset.")
+        command_parser.add_argument("path", help="Path to store the created data.")
+        group = command_parser.add_mutually_exclusive_group()
+        group.add_argument("--threads", help="Use `n` parallel thread workers.", type=int, default=0)
+        group.add_argument("--processes", help="Use `n` parallel process workers.", type=int, default=0)
+        command_parser.add_argument("--trace", action="store_true")
+
+    def run(self, args):
+        now = time.time()
+        if args.threads + args.processes:
+            self.parallel_create(args)
+        else:
+            self.serial_create(args)
+        LOG.info(f"Create completed in {seconds_to_human(time.time()-now)}")
+
+    def serial_create(self, args):
+        from anemoi.datasets.create import Creator
+
+        options = vars(args)
+        c = Creator(**options)
+        c.create()
+
+    def parallel_create(self, args):
+        """Some modules, like fsspec do not work well with fork()
+        Other modules may not be thread safe. So we implement
+        parallel loadining using multiprocessing before any
+        of the modules are imported.
+        """
+
+        options = vars(args)
+        parallel = args.threads + args.processes
+        args.use_threads = args.threads > 0
+
+        if args.use_threads:
+            ExecutorClass = ThreadPoolExecutor
+        else:
+            ExecutorClass = ProcessPoolExecutor
+
+        with ExecutorClass(max_workers=1) as executor:
+            total = executor.submit(task, "init", options).result()
+
+        futures = []
+
+        with ExecutorClass(max_workers=parallel) as executor:
+            for n in range(total):
+                futures.append(executor.submit(task, "load", options, parts=f"{n+1}/{total}"))
+
+            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:
+            executor.submit(task, "statistics", options).result()
+            executor.submit(task, "additions", options).result()
+            executor.submit(task, "cleanup", options).result()
+            executor.submit(task, "verify", options).result()
+
+
+command = Create

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/src/anemoi/datasets/commands/inspect.py

@@ -16,7 +16,7 @@ import numpy as np
 import semantic_version
 import tqdm
 from anemoi.utils.humanize import bytes
-from anemoi.utils.humanize import number
+from anemoi.utils.humanize import bytes_to_human
 from anemoi.utils.humanize import when
 from anemoi.utils.text import dotted_line
 from anemoi.utils.text import progress
@@ -215,9 +215,9 @@ class Version:
             total_size, n = compute_directory_size(self.path)
 
         if total_size is not None:
-            print(f"💽 Size       : {bytes(total_size)} ({number(total_size)})")
+            print(f"💽 Size       : {bytes(total_size)} ({bytes_to_human(total_size)})")
         if n is not None:
-            print(f"📁 Files      : {number(n)}")
+            print(f"📁 Files      : {n:,}")
 
     @property
     def statistics(self):

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/src/anemoi/datasets/create/__init__.py

@@ -7,8 +7,15 @@
 # nor does it submit to any jurisdiction.
 #
 
+import logging
 import os
 
+LOG = logging.getLogger(__name__)
+
+
+def _ignore(*args, **kwargs):
+    pass
+
 
 class Creator:
     def __init__(
@@ -16,19 +23,21 @@ class Creator:
         path,
         config=None,
         cache=None,
-        print=print,
+        use_threads=False,
         statistics_tmp=None,
         overwrite=False,
         test=None,
+        progress=None,
         **kwargs,
     ):
         self.path = path  # Output path
         self.config = config
         self.cache = cache
-        self.print = print
+        self.use_threads = use_threads
         self.statistics_tmp = statistics_tmp
         self.overwrite = overwrite
         self.test = test
+        self.progress = progress if progress is not None else _ignore
 
     def init(self, check_name=False):
         # check path
@@ -44,10 +53,11 @@ class Creator:
                 path=self.path,
                 config=self.config,
                 statistics_tmp=self.statistics_tmp,
-                print=self.print,
+                use_threads=self.use_threads,
+                progress=self.progress,
                 test=self.test,
             )
-            obj.initialise(check_name=check_name)
+            return obj.initialise(check_name=check_name)
 
     def load(self, parts=None):
         from .loaders import ContentLoader
@@ -56,7 +66,8 @@ class Creator:
             loader = ContentLoader.from_dataset_config(
                 path=self.path,
                 statistics_tmp=self.statistics_tmp,
-                print=self.print,
+                use_threads=self.use_threads,
+                progress=self.progress,
                 parts=parts,
             )
             loader.load()
@@ -66,7 +77,8 @@ class Creator:
 
         loader = StatisticsAdder.from_dataset(
             path=self.path,
-            print=self.print,
+            use_threads=self.use_threads,
+            progress=self.progress,
             statistics_tmp=self.statistics_tmp,
             statistics_output=output,
             recompute=False,
@@ -74,20 +86,21 @@ class Creator:
             statistics_end=end,
         )
         loader.run()
+        assert loader.ready()
 
     def size(self):
         from .loaders import DatasetHandler
         from .size import compute_directory_sizes
 
         metadata = compute_directory_sizes(self.path)
-        handle = DatasetHandler.from_dataset(path=self.path, print=self.print)
+        handle = DatasetHandler.from_dataset(path=self.path, use_threads=self.use_threads)
         handle.update_metadata(**metadata)
 
     def cleanup(self):
         from .loaders import DatasetHandlerWithStatistics
 
         cleaner = DatasetHandlerWithStatistics.from_dataset(
-            path=self.path, print=self.print, statistics_tmp=self.statistics_tmp
+            path=self.path, use_threads=self.use_threads, progress=self.progress, statistics_tmp=self.statistics_tmp
         )
         cleaner.tmp_statistics.delete()
         cleaner.registry.clean()
@@ -103,15 +116,17 @@ class Creator:
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
         if statistics:
-            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a = StatisticsAddition.from_dataset(path=self.path, use_threads=self.use_threads)
             a.initialise()
 
         for d in delta:
             try:
-                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a = TendenciesStatisticsAddition.from_dataset(
+                    path=self.path, use_threads=self.use_threads, progress=self.progress, delta=d
+                )
                 a.initialise()
             except TendenciesStatisticsDeltaNotMultipleOfFrequency:
-                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+                LOG.info(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], statistics=True):
         from .loaders import StatisticsAddition
@@ -119,15 +134,17 @@ class Creator:
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
         if statistics:
-            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a = StatisticsAddition.from_dataset(path=self.path, use_threads=self.use_threads)
             a.run(parts)
 
         for d in delta:
             try:
-                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a = TendenciesStatisticsAddition.from_dataset(
+                    path=self.path, use_threads=self.use_threads, progress=self.progress, delta=d
+                )
                 a.run(parts)
             except TendenciesStatisticsDeltaNotMultipleOfFrequency:
-                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+                LOG.debug(f"Skipping delta={d} as it is not a multiple of the frequency.")
 
     def finalise_additions(self, delta=[1, 3, 6, 12, 24], statistics=True):
         from .loaders import StatisticsAddition
@@ -135,15 +152,17 @@ class Creator:
         from .loaders import TendenciesStatisticsDeltaNotMultipleOfFrequency
 
         if statistics:
-            a = StatisticsAddition.from_dataset(path=self.path, print=self.print)
+            a = StatisticsAddition.from_dataset(path=self.path, use_threads=self.use_threads)
             a.finalise()
 
         for d in delta:
             try:
-                a = TendenciesStatisticsAddition.from_dataset(path=self.path, print=self.print, delta=d)
+                a = TendenciesStatisticsAddition.from_dataset(
+                    path=self.path, use_threads=self.use_threads, progress=self.progress, delta=d
+                )
                 a.finalise()
             except TendenciesStatisticsDeltaNotMultipleOfFrequency:
-                self.print(f"Skipping delta={d} as it is not a multiple of the frequency.")
+                LOG.debug(f"Skipping delta={d} as it is not a multiple of the frequency.")
 
     def finalise(self, **kwargs):
         self.statistics(**kwargs)
@@ -174,3 +193,10 @@ class Creator:
             return True
         except zarr.errors.PathNotFoundError:
             return False
+
+    def verify(self):
+        from .loaders import DatasetVerifier
+
+        handle = DatasetVerifier.from_dataset(path=self.path, use_threads=self.use_threads)
+
+        handle.verify()

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/src/anemoi/datasets/create/check.py

@@ -56,7 +56,7 @@ class DatasetName:
             raise ValueError(self.error_message)
 
     def _parse(self, name):
-        pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h)-v(\d+)-?(.*)$"
+        pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h)-v(\d+)-?([a-zA-Z0-9-]+)$"
         match = re.match(pattern, name)
 
         assert match, (name, pattern)
@@ -136,18 +136,19 @@ class StatisticsValueError(ValueError):
     pass
 
 
-def check_data_values(arr, *, name: str, log=[], allow_nan=False):
-    if allow_nan is False:
-        allow_nan = lambda x: False  # noqa: E731
+def check_data_values(arr, *, name: str, log=[], allow_nans=False):
 
-    if allow_nan(name):
+    if (isinstance(allow_nans, (set, list, tuple, dict)) and name in allow_nans) or allow_nans:
         arr = arr[~np.isnan(arr)]
 
+    assert arr.size > 0, (name, *log)
+
     min, max = arr.min(), arr.max()
     assert not (np.isnan(arr).any()), (name, min, max, *log)
 
     if min == 9999.0:
         warnings.warn(f"Min value 9999 for {name}")
+
     if max == 9999.0:
         warnings.warn(f"Max value 9999 for {name}")
 

{anemoi_datasets-0.4.0 → anemoi_datasets-0.4.3}/src/anemoi/datasets/create/chunks.py

@@ -57,7 +57,7 @@ class ChunkFilter:
         if not parts:
             warnings.warn(f"Nothing to do for chunk {i}/{n}.")
 
-        LOG.info(f"Running parts: {parts}")
+        LOG.debug(f"Running parts: {parts}")
 
         self.allowed = parts