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

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.0}/.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.0}/.gitignore

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-datasets
-Version: 0.3.10
+Version: 0.4.0
 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.4
+Requires-Dist: anemoi-utils[provenance]>=0.3.5
 Requires-Dist: numpy
 Requires-Dist: pyyaml
 Requires-Dist: semantic-version
@@ -231,7 +231,8 @@ Requires-Dist: tqdm
 Requires-Dist: zarr<=2.17
 Provides-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"
@@ -239,14 +240,16 @@ 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: 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"
@@ -257,13 +260,11 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: requests; extra == "dev"
 Requires-Dist: s3fs; extra == "dev"
 Requires-Dist: sphinx; extra == "dev"
-Requires-Dist: sphinx-argparse; extra == "dev"
 Requires-Dist: sphinx-rtd-theme; extra == "dev"
 Provides-Extra: docs
 Requires-Dist: nbsphinx; extra == "docs"
 Requires-Dist: pandoc; extra == "docs"
 Requires-Dist: sphinx; extra == "docs"
-Requires-Dist: sphinx-argparse; extra == "docs"
 Requires-Dist: sphinx-rtd-theme; extra == "docs"
 Provides-Extra: remote
 Requires-Dist: boto3; extra == "remote"

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.0}/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.0}/pyproject.toml

@@ -50,17 +50,19 @@ dynamic = [
   "version",
 ]
 dependencies = [
-  "anemoi-utils[provenance]>=0.3.4",
+  "anemoi-utils[provenance]>=0.3.5",
   "numpy",
   "pyyaml",
   "semantic-version",
   "tqdm",
   "zarr<=2.17",
+
 ]
 
 optional-dependencies.all = [
   "boto3",
-  "climetlab>=0.23.2",
+  "earthkit-data[mars]>=0.9",
+  "earthkit-geo>=0.2",
   "earthkit-meteo",
   "ecmwflibs>=0.6.3",
   "entrypoints",
@@ -68,8 +70,10 @@ optional-dependencies.all = [
   "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",
@@ -78,7 +82,8 @@ optional-dependencies.create = [
 
 optional-dependencies.dev = [
   "boto3",
-  "climetlab>=0.23.2",
+  "earthkit-data[mars]>=0.9",
+  "earthkit-geo>=0.2",
   "earthkit-meteo",
   "ecmwflibs>=0.6.3",
   "entrypoints",
@@ -89,16 +94,16 @@ optional-dependencies.dev = [
   "requests",
   "s3fs",
   "sphinx",
-  "sphinx-argparse",
   "sphinx-rtd-theme",
 ]
+
 optional-dependencies.docs = [
   "nbsphinx",
   "pandoc",
   "sphinx",
-  "sphinx-argparse",
   "sphinx-rtd-theme",
 ]
+
 optional-dependencies.remote = [
   "boto3",
   "requests",
@@ -110,8 +115,9 @@ optional-dependencies.tests = [
 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.0}/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.0'
+__version_tuple__ = version_tuple = (0, 4, 0)

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

@@ -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.0}/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.0}/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

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.0}/src/anemoi/datasets/create/config.py

@@ -12,7 +12,7 @@ import os
 from copy import deepcopy
 
 import yaml
-from climetlab.core.order import normalize_order_by
+from earthkit.data.core.order import normalize_order_by
 
 from .utils import load_json_or_yaml
 

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

@@ -13,10 +13,10 @@ import importlib
 import entrypoints
 
 
-def assert_is_fieldset(obj):
-    from climetlab.readers.grib.index import FieldSet
+def assert_is_fieldlist(obj):
+    from earthkit.data.indexing.fieldlist import FieldList
 
-    assert isinstance(obj, FieldSet), type(obj)
+    assert isinstance(obj, FieldList), type(obj)
 
 
 def import_function(name, kind):

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.0}/src/anemoi/datasets/create/functions/filters/empty.py

@@ -7,10 +7,10 @@
 # nor does it submit to any jurisdiction.
 #
 
-import climetlab as cml
+import earthkit.data as ekd
 
 
 def execute(context, input, **kwargs):
-    # Usefull to create a pipeline that returns an empty result
-    # So we can reference an earlier step in a function like 'contants'
-    return cml.load_source("empty")
+    # Useful to create a pipeline that returns an empty result
+    # So we can reference an earlier step in a function like 'constants'
+    return ekd.from_source("empty")

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.0}/src/anemoi/datasets/create/functions/filters/rename.py

@@ -9,7 +9,7 @@
 
 import re
 
-from climetlab.indexing.fieldset import FieldArray
+from earthkit.data.indexing.fieldlist import FieldArray
 
 
 class RenamedFieldMapping:
@@ -26,8 +26,8 @@ class RenamedFieldMapping:
         self.what = what
         self.renaming = renaming
 
-    def metadata(self, key):
-        value = self.field.metadata(key)
+    def metadata(self, key, **kwargs):
+        value = self.field.metadata(key, **kwargs)
         if key == self.what:
             return self.renaming.get(value, value)
         return value
@@ -48,10 +48,10 @@ class RenamedFieldFormat:
         self.format = format
         self.bits = re.findall(r"{(\w+)}", format)
 
-    def metadata(self, key):
-        value = self.field.metadata(key)
+    def metadata(self, key, **kwargs):
+        value = self.field.metadata(key, **kwargs)
         if "{" + key + "}" in self.format:
-            bits = {b: self.field.metadata(b) for b in self.bits}
+            bits = {b: self.field.metadata(b, **kwargs) for b in self.bits}
             return self.format.format(**bits)
         return value
 

anemoi_datasets-0.4.0/src/anemoi/datasets/create/functions/filters/rotate_winds.py

@@ -0,0 +1,80 @@
+# (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.
+#
+
+from collections import defaultdict
+
+from earthkit.data.indexing.fieldlist import FieldArray
+from earthkit.geo.rotate import rotate_vector
+
+
+class NewDataField:
+    def __init__(self, field, data):
+        self.field = field
+        self.data = data
+
+    def to_numpy(self, *args, **kwargs):
+        return self.data
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(
+    context,
+    input,
+    x_wind,
+    y_wind,
+    source_projection=None,
+    target_projection="+proj=longlat",
+):
+    from pyproj import CRS
+
+    result = FieldArray()
+
+    wind_params = (x_wind, y_wind)
+    wind_pairs = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+
+        if param not in wind_params:
+            result.append(f)
+            continue
+
+        key = tuple(key.items())
+
+        if param in wind_pairs[key]:
+            raise ValueError(f"Duplicate wind component {param} for {key}")
+
+        wind_pairs[key][param] = f
+
+    for _, pairs in wind_pairs.items():
+        if len(pairs) != 2:
+            raise ValueError("Missing wind component")
+
+        x = pairs[x_wind]
+        y = pairs[y_wind]
+
+        assert x.grid_mapping == y.grid_mapping
+
+        lats, lons = x.grid_points()
+        x_new, y_new = rotate_vector(
+            lats,
+            lons,
+            x.to_numpy(flatten=True),
+            y.to_numpy(flatten=True),
+            (source_projection if source_projection is not None else CRS.from_cf(x.grid_mapping)),
+            target_projection,
+        )
+
+        result.append(NewDataField(x, x_new))
+        result.append(NewDataField(y, y_new))
+
+    return result

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.0}/src/anemoi/datasets/create/functions/filters/unrotate_winds.py

@@ -9,60 +9,9 @@
 
 from collections import defaultdict
 
-import numpy as np
-from climetlab.indexing.fieldset import FieldArray
-
-
-def normalise(x):
-    return max(min(x, 1.0), -1.0)
-
-
-def normalise_longitude(lon, minimum):
-    while lon < minimum:
-        lon += 360
-
-    while lon >= minimum + 360:
-        lon -= 360
-
-    return lon
-
-
-def rotate_winds(
-    lats,
-    lons,
-    raw_lats,
-    raw_lons,
-    x_wind,
-    y_wind,
-    south_pole_latitude,
-    south_pole_longitude,
-    south_pole_rotation_angle=0,
-):
-    # Code from MIR
-    assert south_pole_rotation_angle == 0
-    C = np.deg2rad(90 - south_pole_latitude)
-    cos_C = np.cos(C)
-    sin_C = np.sin(C)
-
-    new_x = np.zeros_like(x_wind)
-    new_y = np.zeros_like(y_wind)
-
-    for i, (vx, vy, lat, lon, raw_lat, raw_lon) in enumerate(zip(x_wind, y_wind, lats, lons, raw_lats, raw_lons)):
-        lonRotated = south_pole_longitude - lon
-        lon_rotated = normalise_longitude(lonRotated, -180)
-        lon_unrotated = raw_lon
-
-        a = np.deg2rad(lon_rotated)
-        b = np.deg2rad(lon_unrotated)
-        q = 1 if (sin_C * lon_rotated < 0.0) else -1.0  # correct quadrant
-
-        cos_c = normalise(np.cos(a) * np.cos(b) + np.sin(a) * np.sin(b) * cos_C)
-        sin_c = q * np.sqrt(1.0 - cos_c * cos_c)
-
-        new_x[i] = cos_c * vx + sin_c * vy
-        new_y[i] = -sin_c * vx + cos_c * vy
-
-    return new_x, new_y
+# import numpy as np
+from earthkit.data.indexing.fieldlist import FieldArray
+from earthkit.geo.rotate import unrotate_vector
 
 
 class NewDataField:
@@ -85,7 +34,7 @@ def execute(context, input, u, v):
     wind_pairs = defaultdict(dict)
 
     for f in input:
-        key = f.as_mars()
+        key = f.metadata(namespace="mars")
         param = key.pop("param")
 
         if param not in wind_params:
@@ -107,18 +56,19 @@ def execute(context, input, u, v):
         y = pairs[v]
 
         lats, lons = x.grid_points()
-        raw_lats, raw_longs = x.grid_points_raw()
+        raw_lats, raw_longs = x.grid_points_unrotated()
 
         assert x.rotation == y.rotation
 
-        u_new, v_new = rotate_winds(
+        u_new, v_new = unrotate_vector(
             lats,
             lons,
-            raw_lats,
-            raw_longs,
-            x.to_numpy(reshape=False),
-            y.to_numpy(reshape=False),
-            *x.rotation,
+            x.to_numpy(flatten=True),
+            y.to_numpy(flatten=True),
+            *x.rotation[:2],
+            south_pole_rotation_angle=x.rotation[2],
+            lat_unrotated=raw_lats,
+            lon_unrotated=raw_longs,
         )
 
         result.append(NewDataField(x, u_new))
@@ -128,9 +78,9 @@ def execute(context, input, u, v):
 
 
 if __name__ == "__main__":
-    from climetlab import load_source
+    from earthkit.data import from_source
 
-    source = load_source(
+    source = from_source(
         "mars",
         date=-1,
         param="10u/10v",

{anemoi_datasets-0.3.10 → anemoi_datasets-0.4.0}/src/anemoi/datasets/create/functions/sources/accumulations.py

@@ -11,11 +11,11 @@ import logging
 import warnings
 from copy import deepcopy
 
-import climetlab as cml
+import earthkit.data as ekd
 import numpy as np
-from climetlab.core.temporary import temp_file
-from climetlab.readers.grib.output import new_grib_output
-from climetlab.utils.availability import Availability
+from earthkit.data.core.temporary import temp_file
+from earthkit.data.readers.grib.output import new_grib_output
+from earthkit.data.utils.availability import Availability
 
 from anemoi.datasets.create.utils import to_datetime_list
 
@@ -26,7 +26,7 @@ LOG = logging.getLogger(__name__)
 
 def member(field):
     # Bug in eccodes has number=0 randomly
-    number = field.metadata("number")
+    number = field.metadata("number", default=0)
     if number is None:
         number = 0
     return number
@@ -54,16 +54,25 @@ class Accumulation:
 
     def check(self, field):
         if self._check is None:
-            self._check = field.as_mars()
+            self._check = field.metadata(namespace="mars")
 
-            assert self.param == field.metadata("param"), (self.param, field.metadata("param"))
-            assert self.date == field.metadata("date"), (self.date, field.metadata("date"))
-            assert self.time == field.metadata("time"), (self.time, field.metadata("time"))
+            assert self.param == field.metadata("param"), (
+                self.param,
+                field.metadata("param"),
+            )
+            assert self.date == field.metadata("date"), (
+                self.date,
+                field.metadata("date"),
+            )
+            assert self.time == field.metadata("time"), (
+                self.time,
+                field.metadata("time"),
+            )
             assert self.number == member(field), (self.number, member(field))
 
             return
 
-        mars = field.as_mars()
+        mars = field.metadata(namespace="mars")
         keys1 = sorted(self._check.keys())
         keys2 = sorted(mars.keys())
 
@@ -196,7 +205,11 @@ class AccumulationFromLastStep(Accumulation):
 
     def compute(self, values, startStep, endStep):
 
-        assert endStep - startStep == self.frequency, (startStep, endStep, self.frequency)
+        assert endStep - startStep == self.frequency, (
+            startStep,
+            endStep,
+            self.frequency,
+        )
 
         if self.startStep is None:
             self.startStep = startStep
@@ -307,14 +320,13 @@ def compute_accumulations(
                 )
 
     compressed = Availability(requests)
-    ds = cml.load_source("empty")
+    ds = ekd.from_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)
+        ds = ds + ekd.from_source("mars", **request)
 
     accumulations = {}
     for a in [AccumulationClass(out, frequency=frequency, **r) for r in requests]:
@@ -341,7 +353,7 @@ def compute_accumulations(
 
     out.close()
 
-    ds = cml.load_source("file", path)
+    ds = ekd.from_source("file", path)
 
     assert len(ds) / len(param) / len(number) == len(dates), (
         len(ds),