anemoi-datasets 0.3.10__tar.gz → 0.4.2__tar.gz

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/.github/workflows/python-publish.yml

@@ -6,7 +6,7 @@ name: Upload Python Package
 on:
 
   push: {}
-
+  pull_request:
   release:
     types: [created]
 
@@ -45,7 +45,7 @@ jobs:
         pip freeze
 
     - name: Tests
-      run: pytest -k data
+      run: pytest
 
   deploy:
 

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/.gitignore

@@ -187,3 +187,4 @@ _build/
 ~*
 *.sync
 *.dot
+_dev/

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/.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.3.10 → anemoi_datasets-0.4.2}/.readthedocs.yaml

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

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-datasets
-Version: 0.3.10
+Version: 0.4.2
 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,49 +223,65 @@ 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.4
+Requires-Dist: anemoi-utils[provenance]>=0.3.11
 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: climetlab>=0.23.2; 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"
 Provides-Extra: create
-Requires-Dist: climetlab>=0.23.2; extra == "create"
+Requires-Dist: earthkit-data[mars]>=0.9; extra == "create"
+Requires-Dist: earthkit-geo>=0.2; extra == "create"
 Requires-Dist: earthkit-meteo; extra == "create"
 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: climetlab>=0.23.2; 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; 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; 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.2/docs/building/advanced-options.rst

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

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

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

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

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

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/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.3.10 → anemoi_datasets-0.4.2}/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.3.10 → anemoi_datasets-0.4.2}/docs/using/subsetting.rst

@@ -40,9 +40,9 @@ while ``end="2020"`` is equivalent to ``end="2020-12-31"``.
 
 Note also how the ``frequency`` of the dataset will change how the
 ``end`` option is interpreted: - ``end="2020"`` with a ``frequency`` of
-one hour is equivalent to ``end="2020-12-31 23:00"`` - ``end="2020"``
+one hour is equivalent to ``end="2020-12-31 23:00:00"`` - ``end="2020"``
 with a ``frequency`` of 6 hours is equivalent to ``end="2020-12-31
-18:00"``
+18:00:00"``
 
 .. _frequency:
 

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/pyproject.toml

@@ -50,26 +50,32 @@ dynamic = [
   "version",
 ]
 dependencies = [
-  "anemoi-utils[provenance]>=0.3.4",
+  "anemoi-utils[provenance]>=0.3.11",
   "numpy",
   "pyyaml",
   "semantic-version",
   "tqdm",
-  "zarr<=2.17",
+  "zarr",
 ]
 
 optional-dependencies.all = [
+  "aiohttp",
   "boto3",
-  "climetlab>=0.23.2",
+  "earthkit-data[mars]>=0.9",
+  "earthkit-geo>=0.2",
   "earthkit-meteo",
   "ecmwflibs>=0.6.3",
   "entrypoints",
+  "gcsfs",
+  "kerchunk",
   "pyproj",
   "requests",
   "s3fs",
 ]
+
 optional-dependencies.create = [
-  "climetlab>=0.23.2",
+  "earthkit-data[mars]>=0.9",
+  "earthkit-geo>=0.2",
   "earthkit-meteo",
   "ecmwflibs>=0.6.3",
   "entrypoints",
@@ -77,41 +83,59 @@ optional-dependencies.create = [
 ]
 
 optional-dependencies.dev = [
+  "aiohttp",
   "boto3",
-  "climetlab>=0.23.2",
+  "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",
+  "sphinx-argparse<0.5",
   "sphinx-rtd-theme",
 ]
+
 optional-dependencies.docs = [
   "nbsphinx",
   "pandoc",
+  "rstfmt",
   "sphinx",
-  "sphinx-argparse",
+  "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"
-# Changelog = "https://github.com/ecmwf/anemoi-datasets/CHANGELOG.md"
 urls.Repository = "https://github.com/ecmwf/anemoi-datasets/"
+# Changelog = "https://github.com/ecmwf/anemoi-datasets/CHANGELOG.md"
+
 scripts.anemoi-datasets = "anemoi.datasets.__main__:main"
 
 [tool.setuptools.package-data]

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/src/anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.3.10'
-__version_tuple__ = version_tuple = (0, 3, 10)
+__version__ = version = '0.4.2'
+__version_tuple__ = version_tuple = (0, 4, 2)

anemoi_datasets-0.4.2/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.2/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.3.10 → anemoi_datasets-0.4.2}/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):
@@ -382,7 +382,7 @@ class NoVersion(Version):
     @property
     def last_date(self):
         monthly = find(self.metadata, "monthly")
-        time = max([int(t) for t in find(self.metadata["climetlab"], "time")])
+        time = max([int(t) for t in find(self.metadata["earthkit-data"], "time")])
         assert isinstance(time, int), (time, type(time))
         if time > 100:
             time = time // 100
@@ -390,7 +390,7 @@ class NoVersion(Version):
 
     @property
     def frequency(self):
-        time = find(self.metadata["climetlab"], "time")
+        time = find(self.metadata["earthkit-data"], "time")
         return 24 // len(time)
 
     @property
@@ -444,9 +444,9 @@ class Version0_4(Version):
         z = self.zarr
 
         # for backward compatibility
-        if "climetlab" in z.attrs:
-            climetlab_version = z.attrs["climetlab"].get("versions", {}).get("climetlab", "unkwown")
-            print(f"climetlab version used to create this zarr: {climetlab_version}. Not supported.")
+        if "earthkit-data" in z.attrs:
+            ekd_version = z.attrs["earthkit-data"].get("versions", {}).get("earthkit-data", "unkwown")
+            print(f"earthkit-data version used to create this zarr: {ekd_version}. Not supported.")
             return
 
         version = z.attrs.get("version")
@@ -455,7 +455,7 @@ class Version0_4(Version):
             print(" Cannot find metadata information about versions.")
         else:
             print(f"Zarr format (version {version})", end="")
-            print(f" created by climetlab={versions.pop('climetlab')}", end="")
+            print(f" created by earthkit-data={versions.pop('earthkit-data')}", end="")
             timestamp = z.attrs.get("creation_timestamp")
             timestamp = datetime.datetime.fromisoformat(timestamp)
             print(f" on {timestamp}", end="")

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/src/anemoi/datasets/commands/scan.py

@@ -3,7 +3,7 @@ import os
 import sys
 from collections import defaultdict
 
-import climetlab as cml
+import earthkit.data as ekd
 import tqdm
 import yaml
 
@@ -50,9 +50,9 @@ class Scan(Command):
         for path in tqdm.tqdm(paths, leave=False):
             if not match(path):
                 continue
-            for field in tqdm.tqdm(cml.load_source("file", path), leave=False):
-                dates.add(field.valid_datetime())
-                mars = field.as_mars()
+            for field in tqdm.tqdm(ekd.from_source("file", path), leave=False):
+                dates.add(field.datetime()["valid_time"])
+                mars = field.metadata(namespace="mars")
                 keys = tuple(mars.get(k) for k in KEYS)
                 gribs[keys].add(path)
                 for k, v in mars.items():

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.2}/src/anemoi/datasets/compute/recentre.py

@@ -10,10 +10,10 @@
 import logging
 
 import numpy as np
-from climetlab.core.temporary import temp_file
-from climetlab.readers.grib.output import new_grib_output
+from earthkit.data.core.temporary import temp_file
+from earthkit.data.readers.grib.output import new_grib_output
 
-from anemoi.datasets.create.functions import assert_is_fieldset
+from anemoi.datasets.create.functions import assert_is_fieldlist
 
 LOG = logging.getLogger(__name__)
 
@@ -96,7 +96,7 @@ def recentre(
 
     for i, centre_field in enumerate(centre):
         param = centre_field.metadata("param")
-        centre_field_as_mars = centre_field.as_mars()
+        centre_field_as_mars = centre_field.metadata(namespace="mars")
 
         # load the centre field
         centre_np = centre_field.to_numpy()
@@ -106,8 +106,13 @@ def recentre(
 
         for j in range(n_numbers):
             ensemble_field = members[i * n_numbers + j]
-            ensemble_field_as_mars = ensemble_field.as_mars()
-            check_compatible(centre_field, ensemble_field, centre_field_as_mars, ensemble_field_as_mars)
+            ensemble_field_as_mars = ensemble_field.metadata(namespace="mars")
+            check_compatible(
+                centre_field,
+                ensemble_field,
+                centre_field_as_mars,
+                ensemble_field_as_mars,
+            )
             members_np[j] = ensemble_field.to_numpy()
 
             ensemble_field_as_mars = tuple(sorted(ensemble_field_as_mars.items()))
@@ -149,10 +154,10 @@ def recentre(
     if output is not None:
         return path
 
-    from climetlab import load_source
+    from earthkit.data import from_source
 
-    ds = load_source("file", path)
-    assert_is_fieldset(ds)
+    ds = from_source("file", path)
+    assert_is_fieldlist(ds)
     # save a reference to the tmp file so it is deleted
     # only when the dataset is not used anymore
     ds._tmp = tmp