Add support for healpix and Gaussian datasets and for regrid

bopen · Mar 7, 2024 · 7b58507 · 7b58507
1 parent 5ddc781
commit 7b58507
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/xarray_ecmwf/client_cdsapi.py b/xarray_ecmwf/client_cdsapi.py
@@ -11,7 +11,7 @@
 
 LOGGER = logging.getLogger(__name__)
 
-COORDINATES_ORDER = ("valid_time", "time", "step", "isobaricInhPa", "number")
+DIMS_ORDER = ("valid_time", "time", "step", "isobaricInhPa", "number", "values")
 
 
 @attrs.define
@@ -154,7 +154,7 @@ def get_coords_attrs_and_dtype(
             da = list(sample_ds.data_vars.values())[0]
             coords: dict[str, Any] = {}
             # ensure order
-            for name in COORDINATES_ORDER:
+            for name in DIMS_ORDER:
                 if name in chunked_request_coords:
                     coords[name] = chunked_request_coords[name]
                 elif name in da.dims:

diff --git a/xarray_ecmwf/engine_ecmwf.py b/xarray_ecmwf/engine_ecmwf.py
@@ -127,6 +127,7 @@ def open_dataset(  # type:ignore
         cfgrib_kwargs: dict[str, Any] = {},
         request_chunker_kwargs: dict[str, Any] = {},
         request_client_class: type[client_common.RequestClientProtocol] | None = None,
+        open_dataset: Callable[[str, ...], xr.Dataset] = xr.open_dataset,
     ) -> xr.Dataset:
         if not isinstance(filename_or_obj, dict):
             raise TypeError("argument must be a valid request dictionary")
@@ -141,7 +142,7 @@ def open_dataset(  # type:ignore
         if not cache_kwargs.get("cache_file", True):
             cfgrib_kwargs = cfgrib_kwargs | {"indexpath": ""}
 
-        open_dataset = functools.partial(xr.open_dataset, **cfgrib_kwargs)
+        open_dataset = functools.partial(open_dataset, **cfgrib_kwargs)
         dataset_cacher = DatasetCacher(request_client, open_dataset, **cache_kwargs)
         LOGGER.info(request_chunker.get_request_dimensions())