anemoi-datasets 0.5.18__py3-none-any.whl → 0.5.20__py3-none-any.whl

anemoi/datasets/create/sources/anemoi_dataset.py

@@ -0,0 +1,73 @@
+# (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.
+
+import numpy as np
+
+from .legacy import legacy_source
+
+
+@legacy_source(__file__)
+def execute(context, dates, params=None, **kwargs):
+    import earthkit.data as ekd
+
+    from anemoi.datasets import open_dataset
+
+    ds = open_dataset(**kwargs)
+    # dates_to_index = {date: i for i, date in enumerate(ds.dates)}
+
+    indices = []
+    for date in dates:
+        idx = np.where(ds.dates == date)[0]
+        if len(idx) == 0:
+            continue
+        indices.append((int(idx[0]), date))
+
+    vars = ds.variables
+    if params is None:
+        params = vars
+
+    if not isinstance(params, (list, tuple, set)):
+        params = [params]
+
+    params = set(params)
+    results = []
+
+    ensemble = ds.shape[2] > 1
+    latitudes = ds.latitudes
+    longitudes = ds.longitudes
+
+    for idx, date in indices:
+
+        metadata = dict(valid_datetime=date, latitudes=latitudes, longitudes=longitudes)
+
+        for j, y in enumerate(ds[idx]):
+
+            param = vars[j]
+            if param not in params:
+                continue
+
+            # metadata['name'] = param
+            # metadata['param_level'] = param
+            metadata["param"] = param
+
+            for k, e in enumerate(y):
+                if ensemble:
+                    metadata["number"] = k + 1
+
+                metadata["values"] = e
+
+                results.append(metadata.copy())
+
+    print(results[0].keys())
+
+    # "list-of-dicts" does support resolution
+    results = ekd.from_source("list-of-dicts", results)
+
+    # return new_fieldlist_from_list([new_field_from_latitudes_longitudes(x, latitudes, longitudes) for x in results])
+    return results

anemoi/datasets/create/sources/grib.py

@@ -17,6 +17,7 @@ from typing import Optional
 from typing import Union
 
 import earthkit.data as ekd
+from anemoi.transform.flavour import RuleBasedFlavour
 from earthkit.data import from_source
 from earthkit.data.indexing.fieldlist import FieldArray
 from earthkit.data.utils.patterns import Pattern
@@ -244,6 +245,7 @@ def execute(
     path: Union[str, List[str]],
     latitudes: Optional[Dict[str, Any]] = None,
     longitudes: Optional[Dict[str, Any]] = None,
+    flavour: Optional[Union[str, Dict[str, Any]]] = None,
     *args: Any,
     **kwargs: Any,
 ) -> ekd.FieldList:
@@ -255,6 +257,7 @@ def execute(
         path (Union[str, List[str]]): Path or list of paths to the GRIB files.
         latitudes (Optional[Dict[str, Any]], optional): Latitude information. Defaults to None.
         longitudes (Optional[Dict[str, Any]], optional): Longitude information. Defaults to None.
+        flavour (Optional[Union[str, Dict[str, Any]]], optional): Flavour information. Defaults to None.
         *args (Any): Additional arguments.
         **kwargs (Any): Additional keyword arguments.
 
@@ -264,6 +267,8 @@ def execute(
         The loaded dataset.
     """
     given_paths = path if isinstance(path, list) else [path]
+    if flavour is not None:
+        flavour = RuleBasedFlavour(flavour)
 
     geography = None
     if latitudes is not None and longitudes is not None:
@@ -282,6 +287,8 @@ def execute(
         for path in _expand(paths):
             context.trace("📁", "PATH", path)
             s = from_source("file", path)
+            if flavour is not None:
+                s = flavour.map(s)
             s = s.sel(valid_datetime=dates, **kwargs)
             ds = ds + s