From 7b58507805b7ddb3aa16957c93f0fa56d5d2e90d Mon Sep 17 00:00:00 2001 From: Alessandro Amici Date: Thu, 7 Mar 2024 11:33:44 +0100 Subject: [PATCH] Add support for healpix and Gaussian datasets and for regrid --- xarray_ecmwf/client_cdsapi.py | 4 ++-- xarray_ecmwf/engine_ecmwf.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/xarray_ecmwf/client_cdsapi.py b/xarray_ecmwf/client_cdsapi.py index 6fc3c7b..490a506 100644 --- 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 index f0422e0..9ae0dad 100644 --- 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())