Add support for patching xarrays when creating datasets (#160)

* Add support for patching xarrays when creating datasets

* Update CHANGELOG.md

Author: b8raoult

CHANGELOG.md

@@ -16,6 +16,7 @@ Keep it human-readable, your future self will thank you!
 - Fix negative variance for constant variables (#148)
 - Fix cutout slicing of grid dimension (#145)
 - update acumulation (#158)
+- Add ability to patch xarrays (#160)
 
 ### Added
 

src/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])
 

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

@@ -16,6 +16,7 @@
 
 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 @@ def __getitem__(self, i):
         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 @@ def _skip_attr(v, attr_name):
             _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,13 +104,15 @@ def _skip_attr(v, attr_name):
                 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)
 
             grid_coords = sum(1 for c in coordinates if c.is_grid and c.is_dim)
             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(

src/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

src/anemoi/datasets/data/dataset.py

@@ -512,3 +512,50 @@ def _compute_constant_fields_from_statistics(self):
                 result.append(v)
 
         return result
+
+    def plot(self, date, variable, member=0, **kwargs):
+        """For debugging purposes, plot a field.
+
+        Parameters
+        ----------
+        date : int or datetime.datetime or numpy.datetime64 or str
+            The date to plot.
+        variable : int or str
+            The variable to plot.
+        member : int, optional
+            The ensemble member to plot.
+
+        **kwargs:
+            Additional arguments to pass to matplotlib.pyplot.tricontourf
+
+
+        Returns
+        -------
+            matplotlib.pyplot.Axes
+        """
+
+        from anemoi.utils.devtools import plot_values
+        from earthkit.data.utils.dates import to_datetime
+
+        if not isinstance(date, int):
+            date = np.datetime64(to_datetime(date)).astype(self.dates[0].dtype)
+            index = np.where(self.dates == date)[0]
+            if len(index) == 0:
+                raise ValueError(
+                    f"Date {date} not found in the dataset {self.dates[0]} to {self.dates[-1]} by {self.frequency}"
+                )
+            date_index = index[0]
+        else:
+            date_index = date
+
+        if isinstance(variable, int):
+            variable_index = variable
+        else:
+            if variable not in self.variables:
+                raise ValueError(f"Unknown variable {variable} (available: {self.variables})")
+
+            variable_index = self.name_to_index[variable]
+
+        values = self[date_index, variable_index, member]
+
+        return plot_values(values, self.latitudes, self.longitudes, **kwargs)

tests/xarray/test_zarr.py

@@ -73,7 +73,7 @@ def test_weatherbench():
         "levtype": "pl",
     }
 
-    fs = XarrayFieldList.from_xarray(ds, flavour)
+    fs = XarrayFieldList.from_xarray(ds, flavour=flavour)
 
     assert_field_list(
         fs,
@@ -116,7 +116,7 @@ def test_noaa_replay():
         "levtype": "pl",
     }
 
-    fs = XarrayFieldList.from_xarray(ds, flavour)
+    fs = XarrayFieldList.from_xarray(ds, flavour=flavour)
 
     assert_field_list(
         fs,
@@ -141,7 +141,7 @@ def test_planetary_computer_conus404():
         },
     }
 
-    fs = XarrayFieldList.from_xarray(ds, flavour)
+    fs = XarrayFieldList.from_xarray(ds, flavour=flavour)
 
     assert_field_list(
         fs,