anemoi-datasets 0.5.12__py3-none-any.whl → 0.5.14__py3-none-any.whl

anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.12'
-__version_tuple__ = version_tuple = (0, 5, 12)
+__version__ = version = '0.5.14'
+__version_tuple__ = version_tuple = (0, 5, 14)

anemoi/datasets/commands/create.py

@@ -83,11 +83,12 @@ class Create(Command):
         task("load", options)
         task("finalise", options)
 
-        task("patch", options)
-
         task("init_additions", options)
         task("run_additions", options)
         task("finalise_additions", options)
+
+        task("patch", options)
+
         task("cleanup", options)
         task("verify", options)
 
@@ -153,6 +154,7 @@ class Create(Command):
 
         with ExecutorClass(max_workers=1) as executor:
             executor.submit(task, "finalise-additions", options).result()
+            executor.submit(task, "patch", options).result()
             executor.submit(task, "cleanup", options).result()
             executor.submit(task, "verify", options).result()
 

anemoi/datasets/create/__init__.py

@@ -79,7 +79,10 @@ def json_tidy(o):
         )
         return o.isoformat()
 
-    raise TypeError(repr(o) + " is not JSON serializable")
+    if isinstance(o, (np.float32, np.float64)):
+        return float(o)
+
+    raise TypeError(f"{repr(o)} is not JSON serializable {type(o)}")
 
 
 def build_statistics_dates(dates, start, end):
@@ -596,6 +599,8 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         # There is one cube to load for each result.
         dates = list(result.group_of_dates)
 
+        LOG.debug(f"Loading cube for {len(dates)} dates")
+
         cube = result.get_cube()
         shape = cube.extended_user_shape
         dates_in_data = cube.user_coords["valid_datetime"]
@@ -622,10 +627,14 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
 
         check_shape(cube, dates, dates_in_data)
 
-        def check_dates_in_data(lst, lst2):
-            lst2 = [np.datetime64(_) for _ in lst2]
-            lst = [np.datetime64(_) for _ in lst]
-            assert lst == lst2, ("Dates in data are not the requested ones:", lst, lst2)
+        def check_dates_in_data(dates_in_data, requested_dates):
+            requested_dates = [np.datetime64(_) for _ in requested_dates]
+            dates_in_data = [np.datetime64(_) for _ in dates_in_data]
+            assert dates_in_data == requested_dates, (
+                "Dates in data are not the requested ones:",
+                dates_in_data,
+                requested_dates,
+            )
 
         check_dates_in_data(dates_in_data, dates)
 
@@ -638,12 +647,14 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         indexes = dates_to_indexes(self.dates, dates_in_data)
 
         array = ViewCacheArray(self.data_array, shape=shape, indexes=indexes)
+        LOG.info(f"Loading array shape={shape}, indexes={len(indexes)}")
         self.load_cube(cube, array)
 
         stats = compute_statistics(array.cache, self.variables_names, allow_nans=self._get_allow_nans())
         self.tmp_statistics.write(indexes, stats, dates=dates_in_data)
-
+        LOG.info("Flush data array")
         array.flush()
+        LOG.info("Flushed data array")
 
     def _get_allow_nans(self):
         config = self.main_config
@@ -736,6 +747,11 @@ class AdditionsMixin:
         if not self.delta.total_seconds() % frequency.total_seconds() == 0:
             LOG.debug(f"Delta {self.delta} is not a multiple of frequency {frequency}. Skipping.")
             return True
+
+        if self.dataset.zarr_metadata.get("build", {}).get("additions", None) is False:
+            LOG.warning(f"Additions are disabled for {self.path} in the recipe.")
+            return True
+
         return False
 
     @cached_property

anemoi/datasets/create/check.py

@@ -58,7 +58,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+)-?([a-zA-Z0-9-]+)?$"
+        pattern = r"^(\w+)-([\w-]+)-(\w+)-(\w+)-(\d\d\d\d)-(\d\d\d\d)-(\d+h|\d+m)-v(\d+)-?([a-zA-Z0-9-]+)?$"
         match = re.match(pattern, name)
 
         if not match:

anemoi/datasets/create/functions/__init__.py

@@ -22,6 +22,7 @@ def assert_is_fieldlist(obj):
 def import_function(name, kind):
 
     from anemoi.transform.filters import filter_registry
+    from anemoi.transform.sources import source_registry
 
     name = name.replace("-", "_")
 
@@ -45,7 +46,20 @@ def import_function(name, kind):
         if filter_registry.lookup(name, return_none=True):
 
             def proc(context, data, *args, **kwargs):
-                return filter_registry.create(name, *args, **kwargs)(data)
+                filter = filter_registry.create(name, *args, **kwargs)
+                filter.context = context
+                # filter = filter_registry.create(context, name, *args, **kwargs)
+                return filter.forward(data)
+
+            return proc
+
+    if kind == "sources":
+        if source_registry.lookup(name, return_none=True):
+
+            def proc(context, data, *args, **kwargs):
+                source = source_registry.create(name, *args, **kwargs)
+                # source = source_registry.create(context, name, *args, **kwargs)
+                return source.forward(data)
 
             return proc
 

anemoi/datasets/create/functions/filters/orog_to_z.py

@@ -0,0 +1,58 @@
+# (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 collections import defaultdict
+
+from earthkit.data.indexing.fieldlist import FieldArray
+
+
+class NewDataField:
+    def __init__(self, field, data, new_name):
+        self.field = field
+        self.data = data
+        self.new_name = new_name
+
+    def to_numpy(self, *args, **kwargs):
+        return self.data
+
+    def metadata(self, key=None, **kwargs):
+        if key is None:
+            return self.field.metadata(**kwargs)
+
+        value = self.field.metadata(key, **kwargs)
+        if key == "param":
+            return self.new_name
+        return value
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(context, input, orog, z="z"):
+    """Convert orography [m] to z (geopotential height)"""
+    result = FieldArray()
+
+    processed_fields = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param == orog:
+            key = tuple(key.items())
+
+            if param in processed_fields[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            output = f.to_numpy(flatten=True) * 9.80665
+            result.append(NewDataField(f, output, z))
+        else:
+            result.append(f)
+
+    return result

anemoi/datasets/create/functions/filters/sum.py

@@ -0,0 +1,71 @@
+# (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 collections import defaultdict
+
+from earthkit.data.indexing.fieldlist import FieldArray
+
+
+class NewDataField:
+    def __init__(self, field, data, new_name):
+        self.field = field
+        self.data = data
+        self.new_name = new_name
+
+    def to_numpy(self, *args, **kwargs):
+        return self.data
+
+    def metadata(self, key=None, **kwargs):
+        if key is None:
+            return self.field.metadata(**kwargs)
+
+        value = self.field.metadata(key, **kwargs)
+        if key == "param":
+            return self.new_name
+        return value
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(context, input, params, output):
+    """Computes the sum over a set of variables"""
+    result = FieldArray()
+
+    needed_fields = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param in params:
+            key = tuple(key.items())
+
+            if param in needed_fields[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            needed_fields[key][param] = f
+        else:
+            result.append(f)
+
+    for keys, values in needed_fields.items():
+
+        if len(values) != len(params):
+            raise ValueError("Missing fields")
+
+        s = None
+        for k, v in values.items():
+            c = v.to_numpy(flatten=True)
+            if s is None:
+                s = c
+            else:
+                s += c
+        result.append(NewDataField(values[list(values.keys())[0]], s, output))
+
+    return result

anemoi/datasets/create/functions/filters/wz_to_w.py

@@ -0,0 +1,79 @@
+# (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 collections import defaultdict
+
+from earthkit.data.indexing.fieldlist import FieldArray
+
+
+class NewDataField:
+    def __init__(self, field, data, new_name):
+        self.field = field
+        self.data = data
+        self.new_name = new_name
+
+    def to_numpy(self, *args, **kwargs):
+        return self.data
+
+    def metadata(self, key=None, **kwargs):
+        if key is None:
+            return self.field.metadata(**kwargs)
+
+        value = self.field.metadata(key, **kwargs)
+        if key == "param":
+            return self.new_name
+        return value
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(context, input, wz, t, w="w"):
+    """Convert geometric vertical velocity (m/s) to vertical velocity (Pa / s)"""
+    result = FieldArray()
+
+    params = (wz, t)
+    pairs = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param in params:
+            key = tuple(key.items())
+
+            if param in pairs[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            pairs[key][param] = f
+            if param == t:
+                result.append(f)
+        else:
+            result.append(f)
+
+    for keys, values in pairs.items():
+
+        if len(values) != 2:
+            raise ValueError("Missing fields")
+
+        wz_pl = values[wz].to_numpy(flatten=True)
+        t_pl = values[t].to_numpy(flatten=True)
+        pressure = keys[4][1] * 100  # TODO: REMOVE HARDCODED INDICES
+
+        w_pl = wz_to_w(wz_pl, t_pl, pressure)
+        result.append(NewDataField(values[wz], w_pl, w))
+
+    return result
+
+
+def wz_to_w(wz, t, pressure):
+    g = 9.81
+    Rd = 287.058
+
+    return -wz * g * pressure / (t * Rd)

anemoi/datasets/create/functions/sources/accumulations.py

@@ -253,6 +253,7 @@ def _compute_accumulations(
     data_accumulation_period=None,
     patch=_identity,
     base_times=None,
+    use_cdsapi_dataset=None,
 ):
     adjust_step = isinstance(user_accumulation_period, int)
 
@@ -311,7 +312,9 @@ def _compute_accumulations(
 
                 requests.append(patch(r))
 
-    ds = mars(context, dates, *requests, request_already_using_valid_datetime=True)
+    ds = mars(
+        context, dates, *requests, request_already_using_valid_datetime=True, use_cdsapi_dataset=use_cdsapi_dataset
+    )
 
     accumulations = {}
     for a in [AccumulationClass(out, frequency=frequency, **r) for r in requests]:
@@ -366,7 +369,7 @@ def _scda(request):
     return request
 
 
-def accumulations(context, dates, **request):
+def accumulations(context, dates, use_cdsapi_dataset=None, **request):
     _to_list(request["param"])
     class_ = request.get("class", "od")
     stream = request.get("stream", "oper")
@@ -379,6 +382,7 @@ def accumulations(context, dates, **request):
     KWARGS = {
         ("od", "oper"): dict(patch=_scda),
         ("od", "elda"): dict(base_times=(6, 18)),
+        ("od", "enfo"): dict(base_times=(0, 6, 12, 18)),
         ("ea", "oper"): dict(data_accumulation_period=1, base_times=(6, 18)),
         ("ea", "enda"): dict(data_accumulation_period=3, base_times=(6, 18)),
         ("rr", "oper"): dict(base_times=(0, 3, 6, 9, 12, 15, 18, 21)),
@@ -394,6 +398,7 @@ def accumulations(context, dates, **request):
         dates,
         request,
         user_accumulation_period=user_accumulation_period,
+        use_cdsapi_dataset=use_cdsapi_dataset,
         **kwargs,
     )
 

anemoi/datasets/create/functions/sources/eccc_fstd.py

@@ -0,0 +1,16 @@
+# (C) Copyright 2025 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 .xarray import load_many
+
+
+def execute(context, dates, path, *args, **kwargs):
+    options = {"engine": "fstd"}
+    return load_many("🍁", context, dates, path, *args, options=options, **kwargs)

anemoi/datasets/create/functions/sources/mars.py

@@ -246,6 +246,7 @@ def mars(
     *requests,
     request_already_using_valid_datetime=False,
     date_key="date",
+    use_cdsapi_dataset=None,
     **kwargs,
 ):
 
@@ -305,7 +306,10 @@ def mars(
                     f"⚠️ Unknown key {k}={v} in MARS request. Did you mean '{did_you_mean(k, MARS_KEYS)}' ?"
                 )
         try:
-            ds = ds + from_source("mars", **r)
+            if use_cdsapi_dataset:
+                ds = ds + from_source("cds", use_cdsapi_dataset, r)
+            else:
+                ds = ds + from_source("mars", **r)
         except Exception as e:
             if "File is empty:" not in str(e):
                 raise

anemoi/datasets/create/functions/sources/xarray/__init__.py

@@ -29,7 +29,7 @@ def check(what, ds, paths, **kwargs):
         raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, {what}s={paths})")
 
 
-def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs):
+def load_one(emoji, context, dates, dataset, *, options={}, flavour=None, patch=None, **kwargs):
     import xarray as xr
 
     """
@@ -54,10 +54,10 @@ def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs)
     else:
         data = xr.open_dataset(dataset, **options)
 
-    fs = XarrayFieldList.from_xarray(data, flavour)
+    fs = XarrayFieldList.from_xarray(data, flavour=flavour, patch=patch)
 
     if len(dates) == 0:
-        return fs.sel(**kwargs)
+        result = fs.sel(**kwargs)
     else:
         result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
 

anemoi/datasets/create/functions/sources/xarray/field.py

@@ -92,6 +92,10 @@ class XArrayField(Field):
     def grid_points(self):
         return self.owner.grid_points()
 
+    def to_latlon(self, flatten=True):
+        assert flatten
+        return dict(lat=self.latitudes, lon=self.longitudes)
+
     @property
     def resolution(self):
         return None
@@ -120,6 +124,6 @@ class XArrayField(Field):
     def __repr__(self):
         return repr(self._metadata)
 
-    def _values(self):
+    def _values(self, dtype=None):
         # we don't use .values as this will download the data
         return self.selection

anemoi/datasets/create/functions/sources/xarray/fieldlist.py

@@ -16,6 +16,7 @@ from earthkit.data.core.fieldlist import FieldList
 
 from .field import EmptyFieldList
 from .flavour import CoordinateGuesser
+from .patch import patch_dataset
 from .time import Time
 from .variable import FilteredVariable
 from .variable import Variable
@@ -49,7 +50,11 @@ class XarrayFieldList(FieldList):
         raise IndexError(k)
 
     @classmethod
-    def from_xarray(cls, ds, flavour=None):
+    def from_xarray(cls, ds, *, flavour=None, patch=None):
+
+        if patch is not None:
+            ds = patch_dataset(ds, patch)
+
         variables = []
 
         if isinstance(flavour, str):
@@ -83,6 +88,8 @@ class XarrayFieldList(FieldList):
             _skip_attr(variable, "bounds")
             _skip_attr(variable, "grid_mapping")
 
+        LOG.debug("Xarray data_vars: %s", ds.data_vars)
+
         # Select only geographical variables
         for name in ds.data_vars:
 
@@ -97,6 +104,7 @@ class XarrayFieldList(FieldList):
                 c = guess.guess(ds[coord], coord)
                 assert c, f"Could not guess coordinate for {coord}"
                 if coord not in variable.dims:
+                    LOG.debug("%s: coord=%s (not a dimension): dims=%s", variable, coord, variable.dims)
                     c.is_dim = False
                 coordinates.append(c)
 
@@ -104,6 +112,7 @@ class XarrayFieldList(FieldList):
             assert grid_coords <= 2
 
             if grid_coords < 2:
+                LOG.debug("Skipping %s (not 2D): %s", variable, [(c, c.is_grid, c.is_dim) for c in coordinates])
                 continue
 
             v = Variable(

anemoi/datasets/create/functions/sources/xarray/metadata.py

@@ -24,6 +24,7 @@ class _MDMapping:
     def __init__(self, variable):
         self.variable = variable
         self.time = variable.time
+        # Aliases
         self.mapping = dict(param="variable")
         for c in variable.coordinates:
             for v in c.mars_names:
@@ -34,7 +35,6 @@ class _MDMapping:
         return self.mapping.get(key, key)
 
     def from_user(self, kwargs):
-        print("from_user", kwargs, self)
         return {self._from_user(k): v for k, v in kwargs.items()}
 
     def __repr__(self):
@@ -81,22 +81,16 @@ class XArrayMetadata(RawMetadata):
     def _valid_datetime(self):
         return self._get("valid_datetime")
 
-    def _get(self, key, **kwargs):
+    def get(self, key, astype=None, **kwargs):
 
         if key in self._d:
+            if astype is not None:
+                return astype(self._d[key])
             return self._d[key]
 
-        if key.startswith("mars."):
-            key = key[5:]
-            if key not in self.MARS_KEYS:
-                if kwargs.get("raise_on_missing", False):
-                    raise KeyError(f"Invalid key '{key}' in namespace='mars'")
-                else:
-                    return kwargs.get("default", None)
-
         key = self._mapping._from_user(key)
 
-        return super()._get(key, **kwargs)
+        return super().get(key, astype=astype, **kwargs)
 
 
 class XArrayFieldGeography(Geography):

anemoi/datasets/create/functions/sources/xarray/patch.py

@@ -0,0 +1,44 @@
+# (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
+
+LOG = logging.getLogger(__name__)
+
+
+def patch_attributes(ds, attributes):
+    for name, value in attributes.items():
+        variable = ds[name]
+        variable.attrs.update(value)
+
+    return ds
+
+
+def patch_coordinates(ds, coordinates):
+    for name in coordinates:
+        ds = ds.assign_coords({name: ds[name]})
+
+    return ds
+
+
+PATCHES = {
+    "attributes": patch_attributes,
+    "coordinates": patch_coordinates,
+}
+
+
+def patch_dataset(ds, patch):
+    for what, values in patch.items():
+        if what not in PATCHES:
+            raise ValueError(f"Unknown patch type {what!r}")
+
+        ds = PATCHES[what](ds, values)
+
+    return ds

anemoi/datasets/create/functions/sources/xarray/time.py

@@ -62,12 +62,18 @@ class Time:
 
         raise NotImplementedError(f"{len(date_coordinate)=} {len(time_coordinate)=} {len(step_coordinate)=}")
 
+    def select_valid_datetime(self, variable):
+        raise NotImplementedError(f"{self.__class__.__name__}.select_valid_datetime()")
+
 
 class Constant(Time):
 
     def fill_time_metadata(self, coords_values, metadata):
         return None
 
+    def select_valid_datetime(self, variable):
+        return None
+
 
 class Analysis(Time):
 
@@ -83,6 +89,9 @@ class Analysis(Time):
 
         return valid_datetime
 
+    def select_valid_datetime(self, variable):
+        return self.time_coordinate_name
+
 
 class ForecastFromValidTimeAndStep(Time):
 
@@ -116,6 +125,9 @@ class ForecastFromValidTimeAndStep(Time):
 
         return valid_datetime
 
+    def select_valid_datetime(self, variable):
+        return self.time_coordinate_name
+
 
 class ForecastFromValidTimeAndBaseTime(Time):
 
@@ -138,6 +150,9 @@ class ForecastFromValidTimeAndBaseTime(Time):
 
         return valid_datetime
 
+    def select_valid_datetime(self, variable):
+        return self.time_coordinate_name
+
 
 class ForecastFromBaseTimeAndDate(Time):