diff --git a/docs/examples/index.rst b/docs/examples/index.rst index 96bf3490f..677f5af43 100644 --- a/docs/examples/index.rst +++ b/docs/examples/index.rst @@ -169,6 +169,7 @@ Xarray engine xarray_engine_split.ipynb xarray_engine_squeeze.ipynb xarray_engine_chunks.ipynb + xarray_cupy.ipynb Targets and encoders +++++++++++++++++++++ diff --git a/docs/examples/xarray_cupy.ipynb b/docs/examples/xarray_cupy.ipynb new file mode 100644 index 000000000..884d71136 --- /dev/null +++ b/docs/examples/xarray_cupy.ipynb @@ -0,0 +1,1679 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "47cee8d5-1d5a-4e60-8bcd-fdf6496faadf", + "metadata": {}, + "source": [ + "## Xarray: using CuPy" + ] + }, + { + "cell_type": "markdown", + "id": "ec94158a-eb91-4cb0-a316-748b91f9696c", + "metadata": {}, + "source": [ + "This notebook demonstrates how to use Xarray on a GPU with CuPy. Since CuPy is not a dependency for earthkit-data it has to be installed separately. Also a CUDA-based GPU environment has to be up and running for the notebook to work." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "ebb1cb85-5e14-45dd-b956-432d970eed1c", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + " " + ] + } + ], + "source": [ + "# Get GRIB data on pressure levels\n", + "import earthkit.data as ekd\n", + "ds = ekd.from_source(\"sample\", \"pl.grib\")" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "8bd9660a-9d96-4bed-9c6e-4b3f1fdaeaea", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.Dataset> Size: 176kB\n",
+       "Dimensions:                  (forecast_reference_time: 4, step: 2, level: 2,\n",
+       "                              latitude: 19, longitude: 36)\n",
+       "Coordinates:\n",
+       "  * forecast_reference_time  (forecast_reference_time) datetime64[ns] 32B 202...\n",
+       "  * step                     (step) timedelta64[ns] 16B 00:00:00 06:00:00\n",
+       "  * level                    (level) int64 16B 500 700\n",
+       "  * latitude                 (latitude) float64 152B 90.0 80.0 ... -80.0 -90.0\n",
+       "  * longitude                (longitude) float64 288B 0.0 10.0 ... 340.0 350.0\n",
+       "Data variables:\n",
+       "    r                        (forecast_reference_time, step, level, latitude, longitude) float64 88kB ...\n",
+       "    t                        (forecast_reference_time, step, level, latitude, longitude) float64 88kB ...\n",
+       "Attributes:\n",
+       "    class:        od\n",
+       "    stream:       oper\n",
+       "    levtype:      pl\n",
+       "    type:         fc\n",
+       "    expver:       0001\n",
+       "    date:         20240603\n",
+       "    time:         0\n",
+       "    domain:       g\n",
+       "    number:       0\n",
+       "    Conventions:  CF-1.8\n",
+       "    institution:  ECMWF
" + ], + "text/plain": [ + " Size: 176kB\n", + "Dimensions: (forecast_reference_time: 4, step: 2, level: 2,\n", + " latitude: 19, longitude: 36)\n", + "Coordinates:\n", + " * forecast_reference_time (forecast_reference_time) datetime64[ns] 32B 202...\n", + " * step (step) timedelta64[ns] 16B 00:00:00 06:00:00\n", + " * level (level) int64 16B 500 700\n", + " * latitude (latitude) float64 152B 90.0 80.0 ... -80.0 -90.0\n", + " * longitude (longitude) float64 288B 0.0 10.0 ... 340.0 350.0\n", + "Data variables:\n", + " r (forecast_reference_time, step, level, latitude, longitude) float64 88kB ...\n", + " t (forecast_reference_time, step, level, latitude, longitude) float64 88kB ...\n", + "Attributes:\n", + " class: od\n", + " stream: oper\n", + " levtype: pl\n", + " type: fc\n", + " expver: 0001\n", + " date: 20240603\n", + " time: 0\n", + " domain: g\n", + " number: 0\n", + " Conventions: CF-1.8\n", + " institution: ECMWF" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Create a lazy loaded Xarray with Numpy arrays\n", + "r = ds.to_xarray()\n", + "r" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "debcf314-41e8-4dcc-916b-ac602e936e16", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "numpy.ndarray" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "type(r.t.data)" + ] + }, + { + "cell_type": "markdown", + "id": "28901396-3ded-48c8-9ed1-8ef70ab8df26", + "metadata": {}, + "source": [ + "#### Move to the GPU as CuPy" + ] + }, + { + "cell_type": "markdown", + "id": "b6b9b07b-44bc-4fd4-b312-d3c5bee787a7", + "metadata": {}, + "source": [ + "We use the ``to_device()`` method, which is available on the ``earthkit`` Xarray accessor. The first argument specifies the device. When the device is not \"cpu\" and the ``array_backend`` keyword argument is not specified it is automatically set to \"cupy\"." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "e6e251f2-75ac-4de5-8b9a-98cd12780063", + "metadata": {}, + "outputs": [], + "source": [ + "r_cp = r.earthkit.to_device(\"cuda:0\") \n", + "# equivalent code:\n", + "# r_cp = r.earthkit.to_device(\"cuda:0\", array_backend=\"cupy\") " + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "2716e8e3-4e53-4ba9-ae9d-57fa1f983fa3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "cupy.ndarray" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "type(r_cp.t.data)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "6257255f-59e3-4e05-a63c-7e3e381f1258", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.DataArray 't' ()> Size: 8B\n",
+       "array(261.56490497)
" + ], + "text/plain": [ + " Size: 8B\n", + "array(261.56490497)" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Xarray computations work\n", + "r_cp.t.mean()" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "1106e56a-d9db-4484-88c6-43a7458892f7", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "cupy.ndarray" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Alter the values\n", + "r_cp += 1\n", + "type(r_cp.t.data)" + ] + }, + { + "cell_type": "markdown", + "id": "126a7326-8c24-45bf-b5f2-a195a341a621", + "metadata": {}, + "source": [ + "#### Move back to the CPU as Numpy" + ] + }, + { + "cell_type": "markdown", + "id": "6b08cf88-3d20-4bc3-9589-52f098f502b2", + "metadata": {}, + "source": [ + "We use ``to_device()`` again to move back the dataset to the cpu. When the device is \"cpu\" and the ``array_backend`` keyword argument is not specified it is automatically set to \"numpy\"." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "4e35953f-c075-4118-85a3-2ac108af1bcf", + "metadata": {}, + "outputs": [], + "source": [ + "r_np = r.earthkit.to_device(\"cpu\")\n", + "# equivalent code:\n", + "# r_np = r.earthkit.to_device(\"cpu\", array_backend=\"numpy\") " + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "701838e0-0be8-4066-99db-adc07b79e197", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "numpy.ndarray" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "type(r_np.t.data)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "5e565c0b-f847-40d6-8869-22091780bc98", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.DataArray 't' ()> Size: 8B\n",
+       "array(261.56490497)
" + ], + "text/plain": [ + " Size: 8B\n", + "array(261.56490497)" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# The dataset contains the values altered on the GPU\n", + "r_np.t.mean()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "82ab0386-e111-47e0-9aac-064305f84503", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python (conda-ek-gpu)", + "language": "python", + "name": "earthkit-gpu" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/earthkit/data/indexing/tensor.py b/src/earthkit/data/indexing/tensor.py index 4e6460eb5..006e360eb 100644 --- a/src/earthkit/data/indexing/tensor.py +++ b/src/earthkit/data/indexing/tensor.py @@ -392,6 +392,20 @@ def to_numpy(self, index=None, **kwargs): shape += list(n.shape[1:]) return n.reshape(*shape) if len(shape) > 0 else n + @flatten_arg + def to_array(self, index=None, **kwargs): + if index is not None: + if all(i == slice(None, None, None) for i in index): + index = None + + if index is None: + return self.source.to_array(**kwargs).reshape(*self.full_shape) + else: + n = self.source.to_array(index=index, **kwargs) + shape = list(self._user_shape) + shape += list(n.shape[1:]) + return n.reshape(*shape) if len(shape) > 0 else n + @flatten_arg def latitudes(self, **kwargs): return self.source[0].data("lat", **kwargs) diff --git a/src/earthkit/data/readers/grib/codes.py b/src/earthkit/data/readers/grib/codes.py index 5285d3e05..ccca48f5e 100644 --- a/src/earthkit/data/readers/grib/codes.py +++ b/src/earthkit/data/readers/grib/codes.py @@ -48,6 +48,12 @@ def get(self, handle, dtype=None): else: return v + @staticmethod + def to_numpy_dtype(dtype): + from earthkit.utils.array.dtype import to_numpy_dtype + + return to_numpy_dtype(dtype, default=np.float64) + class GribCodesValueAccessor(GribCodesFloatArrayAccessor): KEY = "values" @@ -56,6 +62,7 @@ def __init__(self): super().__init__() def get(self, handle, dtype=None): + dtype = self.to_numpy_dtype(dtype) if dtype is np.float32 and self.HAS_FLOAT_SUPPORT: return eccodes.codes_get_array(handle, self.KEY, ktype=dtype) else: @@ -68,6 +75,10 @@ class GribCodesLatitudeAccessor(GribCodesFloatArrayAccessor): def __init__(self): super().__init__() + def get(self, handle, dtype=None): + dtype = self.to_numpy_dtype(dtype) + return super().get(handle, dtype=dtype) + class GribCodesLongitudeAccessor(GribCodesFloatArrayAccessor): KEY = "longitudes" @@ -75,6 +86,10 @@ class GribCodesLongitudeAccessor(GribCodesFloatArrayAccessor): def __init__(self): super().__init__() + def get(self, handle, dtype=None): + dtype = self.to_numpy_dtype(dtype) + return super().get(handle, dtype=dtype) + VALUE_ACCESSOR = GribCodesValueAccessor() LATITUDE_ACCESSOR = GribCodesLatitudeAccessor() diff --git a/src/earthkit/data/readers/grib/xarray.py b/src/earthkit/data/readers/grib/xarray.py index 77ea868db..ae8c55cce 100644 --- a/src/earthkit/data/readers/grib/xarray.py +++ b/src/earthkit/data/readers/grib/xarray.py @@ -345,8 +345,9 @@ def to_xarray(self, engine="earthkit", xarray_open_dataset_kwargs=None, **kwargs to False unless the ``profile`` overwrites it. * dtype: str, numpy.dtype or None Typecode or data-type of the array data. - * array_module: module - The module to use for array operations. Default is numpy. + * array_backend: str, array namespace, ArrayBackend, None + The array backend/namespace to use for array operations. The default value (None) is + expanded to "numpy". * direct_backend: bool, None If True, the backend is used directly bypassing :py:meth:`xarray.open_dataset` and ignoring all non-backend related kwargs. If False, the data is read via diff --git a/src/earthkit/data/testing.py b/src/earthkit/data/testing.py index 115a87ced..2ffe43e14 100644 --- a/src/earthkit/data/testing.py +++ b/src/earthkit/data/testing.py @@ -233,6 +233,30 @@ def write_to_file(mode, path, ds, **kwargs): raise ValueError(f"Invalid {mode=}") +def check_array( + v, + shape=None, + first=None, + last=None, + meanv=None, + eps=1e-3, + array_backend=None, +): + if array_backend is None: + from earthkit.utils.array import get_backend + + array_backend = get_backend(v) + + v = array_backend.to_numpy(v) + + import numpy as np + + assert v.shape == shape + assert np.isclose(v[0], first, eps) + assert np.isclose(v[-1], last, eps) + assert np.isclose(v.mean(), meanv, eps) + + def main(path): import sys diff --git a/src/earthkit/data/utils/xarray/builder.py b/src/earthkit/data/utils/xarray/builder.py index 0fb0f14be..b3436d97b 100644 --- a/src/earthkit/data/utils/xarray/builder.py +++ b/src/earthkit/data/utils/xarray/builder.py @@ -11,7 +11,6 @@ from abc import ABCMeta from abc import abstractmethod -import numpy import xarray import xarray.core.indexing as indexing @@ -179,12 +178,9 @@ def __init__(self, tensor, dims, shape, xp, dtype, var_name): self.dims = dims self.shape = shape self._var_name = var_name + self.dtype = dtype + self.xp = xp - # xp and dtype must be set for xarray - self.xp = xp if xp is not None else numpy - if dtype is None: - dtype = numpy.dtype("float64") - self.dtype = xp.dtype(dtype) from dask.utils import SerializableLock self.lock = SerializableLock() @@ -230,19 +226,19 @@ def _raw_indexing_method(self, key: tuple): # LOG.debug(f" {field_index=}") - result = r.to_numpy(index=field_index, dtype=self.dtype) + try: + result = r.to_array(index=field_index, array_backend=self.xp, dtype=self.dtype) + except Exception as e: + LOG.exception("Error in to_array:", e) + raise + + # LOG.debug(f" {result.shape=}" # ensure axes are squeezed when needed singles = [i for i in list(range(len(r.user_shape))) if isinstance(key[i], int)] if singles: result = result.squeeze(axis=tuple(singles)) - # LOG.debug(f" {result.shape=}") - - # Loading as numpy but then converting to the target array module - if self.xp and self.xp != numpy: - result = self.xp.asarray(result) - return result @@ -253,8 +249,21 @@ def __init__(self, ds, profile, dims, grid=None, fixed_local_attrs=None): self.dims = dims self.flatten_values = profile.flatten_values - self.dtype = profile.dtype - self.array_module = profile.array_module + + # Array backend/namespace + array_backend = profile.array_backend + if array_backend is None: + array_backend = "numpy" + + from earthkit.utils.array import get_backend + + self.array_backend = get_backend(array_backend) + assert self.array_backend is not None, f"Unsupported array_backend : {array_backend}" + + dtype = profile.dtype + if dtype is None: + dtype = "float64" + self.dtype = self.array_backend.make_dtype(dtype) # Note: these coords inside the tensor are called user_coords and # the corresponding dims are called user_dims @@ -470,7 +479,7 @@ def build_values(self, tensor, var_dims, name): tensor, var_dims, tensor.full_shape, - self.array_module, + self.array_backend.namespace, self.dtype, name, ) @@ -514,7 +523,7 @@ def build_values(self, tensor, var_dims, name): for f in tensor.source: f.keep = False - return tensor.to_numpy(dtype=self.dtype) + return tensor.to_array(dtype=self.dtype, array_backend=self.array_backend) class DatasetBuilder: diff --git a/src/earthkit/data/utils/xarray/defaults.yaml b/src/earthkit/data/utils/xarray/defaults.yaml index ebfd73028..6633364c0 100644 --- a/src/earthkit/data/utils/xarray/defaults.yaml +++ b/src/earthkit/data/utils/xarray/defaults.yaml @@ -34,7 +34,7 @@ decode_timedelta: # values flatten_values: false dtype: float64 -array_module: numpy +array_backend: numpy # other lazy_load: true diff --git a/src/earthkit/data/utils/xarray/engine.py b/src/earthkit/data/utils/xarray/engine.py index da471cc8b..ad2fc3f31 100644 --- a/src/earthkit/data/utils/xarray/engine.py +++ b/src/earthkit/data/utils/xarray/engine.py @@ -55,6 +55,7 @@ def open_dataset( strict=None, dtype=None, array_module=None, + array_backend=None, errors=None, ): r""" @@ -306,8 +307,9 @@ def open_dataset( to False unless the ``profile`` overwrites it. dtype: str, numpy.dtype or None Typecode or data-type of the array data. - array_module: module - The module to use for array operations. Default is numpy. + array_backend: str, array namespace, ArrayBackend, None + The array backend/namespace to use for array operations. The default value (None) is + expanded to "numpy". """ fieldlist = self._fieldlist(filename_or_obj, source_type) @@ -317,6 +319,13 @@ def open_dataset( else: from .builder import SingleDatasetBuilder + if array_module is not None: + import warnings + + warnings.warn("'array_module' is deprecated. Use 'array_backend' instead", DeprecationWarning) + if array_backend is None: + array_backend = array_module + _kwargs = dict( variable_key=variable_key, drop_variables=drop_variables, @@ -351,7 +360,7 @@ def open_dataset( release_source=release_source, strict=strict, dtype=dtype, - array_module=array_module, + array_backend=array_backend, errors=errors, ) @@ -486,6 +495,15 @@ def to_netcdf(self, *args, **kwargs): return ds.to_netcdf(*args, **kwargs) + def to_device(self, device, *args, array_backend=None, **kwargs): + """Return a **new** DataArray whose data live on *device*.""" + from earthkit.utils.array import to_device + + moved = to_device(self._obj.data, device, *args, array_backend=array_backend, **kwargs) + da = self._obj.copy(deep=False) + da.data = moved + return da + @xarray.register_dataset_accessor("earthkit") class XarrayEarthkitDataSet(XarrayEarthkit): @@ -517,3 +535,12 @@ def to_netcdf(self, *args, **kwargs): break return ds.to_netcdf(*args, **kwargs) + + def to_device(self, device, *args, array_backend=None, **kwargs): + """Return a new Dataset with every data variable on the specified ``device``.""" + from earthkit.utils.array import to_device + + ds = self._obj.copy(deep=False) + for name, var in ds.data_vars.items(): + ds[name].data = to_device(var.data, device, *args, array_backend=array_backend, **kwargs) + return ds diff --git a/src/earthkit/data/utils/xarray/profile.py b/src/earthkit/data/utils/xarray/profile.py index 446326085..3415b3a37 100644 --- a/src/earthkit/data/utils/xarray/profile.py +++ b/src/earthkit/data/utils/xarray/profile.py @@ -279,12 +279,18 @@ def __init__( # values self.flatten_values = kwargs.pop("flatten_values") self.dtype = kwargs.pop("dtype") - self.array_module = kwargs.pop("array_module") + self.array_backend = kwargs.pop("array_backend") - if self.array_module == "numpy": - import numpy as np + if "array_module" in kwargs: + raise ValueError( + "'array_module' is deprecated. Use 'array_backend' instead. " + "If you are using 'array_module', please update your code to use 'array_backend'." + ) - self.array_module = np + # if self.array_backend == "numpy": + # import numpy as np + + # self.array_module = np if kwargs: raise ValueError(f"Unsupported options: {kwargs}") @@ -328,6 +334,19 @@ def from_conf(cls, name, conf, *args, **kwargs): kwargs = copy.deepcopy(kwargs) opt = copy.deepcopy(PROFILE_CONF.defaults) + def _deprec_array_module(data): + """Deprecated: use 'array_backend' instead""" + if "array_module" in data: + import warnings + + warnings.warn("'array_module' is deprecated. Use 'array_backend' instead", DeprecationWarning) + + array_module = kwargs.pop("array_module") + if data.get("array_backend", None) is None: + data["array_backend"] = array_module + + _deprec_array_module(kwargs) + for d in [conf, kwargs]: for k, v in d.items(): if k in PROFILE_CONF.defaults and v is not None: diff --git a/tests/documentation/test_notebooks.py b/tests/documentation/test_notebooks.py index 9262f0b10..8c54be4df 100644 --- a/tests/documentation/test_notebooks.py +++ b/tests/documentation/test_notebooks.py @@ -42,6 +42,7 @@ "grib_to_xarray.ipynb", "grib_to_fdb_target.ipynb", "xarray_engine.ipynb", + "xarray_cupy.ipynb", "netcdf_opendap.ipynb", ] diff --git a/tests/grib/test_grib_geography.py b/tests/grib/test_grib_geography.py index 99351af49..04ecaf8d7 100644 --- a/tests/grib/test_grib_geography.py +++ b/tests/grib/test_grib_geography.py @@ -18,6 +18,7 @@ import earthkit.data from earthkit.data.testing import NO_GEO +from earthkit.data.testing import check_array from earthkit.data.testing import earthkit_examples_file from earthkit.data.testing import earthkit_test_data_file from earthkit.data.utils import projections @@ -29,13 +30,6 @@ from grib_fixtures import load_grib_data # noqa: E402 -def check_array(v, shape=None, first=None, last=None, meanv=None, eps=1e-3): - assert v.shape == shape - assert np.isclose(v[0], first, eps) - assert np.isclose(v[-1], last, eps) - assert np.isclose(v.mean(), meanv, eps) - - @pytest.mark.parametrize("fl_type", FL_TYPES) @pytest.mark.parametrize("index", [0, None]) def test_grib_to_latlon_single(fl_type, index): @@ -48,7 +42,7 @@ def test_grib_to_latlon_single(fl_type, index): check_array_type(v["lon"], array_backend, dtype="float64") check_array_type(v["lat"], array_backend, dtype="float64") check_array( - v["lon"], + array_backend.to_numpy(v["lon"]), (84,), first=0.0, last=330.0, @@ -56,7 +50,7 @@ def test_grib_to_latlon_single(fl_type, index): eps=eps, ) check_array( - v["lat"], + array_backend.to_numpy(v["lat"]), (84,), first=90, last=-90, @@ -79,12 +73,12 @@ def test_grib_to_latlon_single_shape(fl_type, index): # x assert v["lon"].shape == (7, 12) for x in v["lon"]: - assert np.allclose(x, np.linspace(0, 330, 12)) + assert np.allclose(array_backend.to_numpy(x), np.linspace(0, 330, 12)) # y assert v["lat"].shape == (7, 12) for i, y in enumerate(v["lat"]): - assert np.allclose(y, np.ones(12) * (90 - i * 30)) + assert np.allclose(array_backend.to_numpy(y), np.ones(12) * (90 - i * 30)) @pytest.mark.parametrize("fl_type", FL_NUMPY) @@ -126,7 +120,7 @@ def test_grib_to_points_single(fl_type, index): check_array_type(v["x"], array_backend, dtype="float64") check_array_type(v["y"], array_backend, dtype="float64") check_array( - v["x"], + array_backend.to_numpy(v["x"]), (84,), first=0.0, last=330.0, @@ -134,7 +128,7 @@ def test_grib_to_points_single(fl_type, index): eps=eps, ) check_array( - v["y"], + array_backend.to_numpy(v["y"]), (84,), first=90, last=-90, diff --git a/tests/grib/test_grib_values.py b/tests/grib/test_grib_values.py index f8059197f..cda942395 100644 --- a/tests/grib/test_grib_values.py +++ b/tests/grib/test_grib_values.py @@ -16,6 +16,8 @@ import pytest from earthkit.utils.testing import check_array_type +from earthkit.data.testing import check_array + here = os.path.dirname(__file__) sys.path.insert(0, here) from grib_fixtures import FL_FILE # noqa: E402 @@ -24,13 +26,6 @@ from grib_fixtures import load_grib_data # noqa: E402 -def check_array(v, shape=None, first=None, last=None, meanv=None, eps=1e-3): - assert v.shape == shape - assert np.isclose(v[0], first, eps) - assert np.isclose(v[-1], last, eps) - assert np.isclose(v.mean(), meanv, eps) - - @pytest.mark.parametrize("fl_type", FL_TYPES) def test_grib_values_1(fl_type): f, array_backend = load_grib_data("test_single.grib", fl_type, folder="data") @@ -48,6 +43,7 @@ def test_grib_values_1(fl_type): last=227.18560791015625, meanv=274.36566743396577, eps=eps, + array_backend=array_backend, ) # field @@ -55,7 +51,7 @@ def test_grib_values_1(fl_type): check_array_type(v1, array_backend) assert v1.shape == (84,) - assert np.allclose(v, v1, eps) + assert array_backend.allclose(v, v1, eps) @pytest.mark.parametrize("fl_type", FL_FILE) @@ -567,9 +563,10 @@ def test_grib_values_with_missing(fl_type): assert ns.count_nonzero(ns.isnan(v)) == 38 mask = array_backend.from_other([12, 14, 15, 24, 25, 26] + list(range(28, 60))) - assert np.isclose(v[0], 260.4356, eps) - assert np.isclose(v[11], 260.4356, eps) - assert np.isclose(v[-1], 227.1856, eps) + v1 = array_backend.to_numpy(v) + assert np.isclose(v1[0], 260.4356, eps) + assert np.isclose(v1[11], 260.4356, eps) + assert np.isclose(v1[-1], 227.1856, eps) m = v[mask] assert len(m) == 38 assert ns.count_nonzero(ns.isnan(m)) == 38 diff --git a/tests/xr_engine/test_xr_engine.py b/tests/xr_engine/test_xr_engine.py index 73618284c..79ef85bf8 100644 --- a/tests/xr_engine/test_xr_engine.py +++ b/tests/xr_engine/test_xr_engine.py @@ -782,14 +782,21 @@ def test_xr_engine_invalid_kwargs(kwargs): @pytest.mark.cache -def test_xr_engine_dtype(): +@pytest.mark.parametrize( + "dtype,expected_dtype", + [ + (np.float32, np.float32), + ("float32", np.float32), + (np.float64, np.float64), + ("float64", np.float64), + ], +) +def test_xr_engine_dtype(dtype, expected_dtype): ds_ek = from_source("url", earthkit_remote_test_data_file("test-data/xr_engine/level/pl.grib")) - ds = ds_ek.to_xarray(dtype=np.float32) - assert ds["t"].values.dtype == np.float32 - - ds = ds_ek.to_xarray(dtype=np.float64) - assert ds["t"].values.dtype == np.float64 + ds = ds_ek.to_xarray(dtype=dtype) + assert ds["t"].data.dtype == expected_dtype + assert ds["t"].values.dtype == expected_dtype @pytest.mark.cache diff --git a/tests/xr_engine/test_xr_torch.py b/tests/xr_engine/test_xr_torch.py new file mode 100644 index 000000000..e1a6826cc --- /dev/null +++ b/tests/xr_engine/test_xr_torch.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 + +# (C) Copyright 2020 ECMWF. +# +# 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 pytest +from earthkit.utils.array import _TORCH +from earthkit.utils.testing import NO_TORCH +from earthkit.utils.testing import check_array_type + +from earthkit.data import from_source +from earthkit.data.testing import earthkit_remote_test_data_file + + +@pytest.mark.skipif(NO_TORCH, reason="No pytorch installed") +@pytest.mark.cache +def test_xr_engine_torch_core(): + ds_ek = from_source("url", earthkit_remote_test_data_file("test-data/xr_engine/level/pl.grib")) + + ds = ds_ek.to_xarray(array_backend="torch") + check_array_type(ds["t"].data, _TORCH) + + +@pytest.mark.skipif(NO_TORCH, reason="No pytorch installed") +@pytest.mark.cache +def test_xr_engine_torch_core_compat(): + ds_ek = from_source("url", earthkit_remote_test_data_file("test-data/xr_engine/level/pl.grib")) + + ds = ds_ek.to_xarray(array_module="torch") + check_array_type(ds["t"].data, _TORCH) + + +@pytest.mark.skipif(NO_TORCH, reason="No pytorch installed") +@pytest.mark.cache +def test_xr_engine_torch_dtype(): + ds_ek = from_source("url", earthkit_remote_test_data_file("test-data/xr_engine/level/pl.grib")) + + def _check_dtype(dtype, expected_dtype): + ds = ds_ek.to_xarray(array_backend="torch", dtype=dtype) + assert ds["t"].data.dtype == expected_dtype + + dtype = _TORCH.float32 + expected_dtype = _TORCH.float32 + _check_dtype(dtype, expected_dtype) + + dtype = "float32" + expected_dtype = _TORCH.float32 + _check_dtype(dtype, expected_dtype) + + dtype = _TORCH.float64 + expected_dtype = _TORCH.float64 + _check_dtype(dtype, expected_dtype) + + dtype = "float64" + expected_dtype = _TORCH.float64 + _check_dtype(dtype, expected_dtype)