diff --git a/doc/conf.py b/doc/conf.py index 3f443ced03..294b60491d 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -433,7 +433,7 @@ # Configuration for intersphinx intersphinx_mapping = { - 'cf_units': ('https://cf-units.readthedocs.io/en/latest/', None), + 'cf_units': ('https://cf-units.readthedocs.io/en/stable/', None), 'cftime': ('https://unidata.github.io/cftime/', None), 'esmvalcore': (f'https://docs.esmvaltool.org/projects/ESMValCore/en/{rtd_version}/', @@ -441,14 +441,17 @@ 'esmvaltool': (f'https://docs.esmvaltool.org/en/{rtd_version}/', None), 'dask': ('https://docs.dask.org/en/stable/', None), 'distributed': ('https://distributed.dask.org/en/stable/', None), - 'iris': ('https://scitools-iris.readthedocs.io/en/latest/', None), - 'iris-esmf-regrid': ('https://iris-esmf-regrid.readthedocs.io/en/latest', + 'iris': ('https://scitools-iris.readthedocs.io/en/stable/', None), + 'iris-esmf-regrid': ('https://iris-esmf-regrid.readthedocs.io/en/stable', None), 'matplotlib': ('https://matplotlib.org/stable/', None), + 'ncdata': ('https://ncdata.readthedocs.io/en/stable/', None), 'numpy': ('https://numpy.org/doc/stable/', None), - 'pyesgf': ('https://esgf-pyclient.readthedocs.io/en/latest/', None), + 'pyesgf': ('https://esgf-pyclient.readthedocs.io/en/stable/', None), 'python': ('https://docs.python.org/3/', None), 'scipy': ('https://docs.scipy.org/doc/scipy/', None), + 'xarray': ('https://docs.xarray.dev/en/stable/', None), + 'xesmf': ('https://xesmf.readthedocs.io/en/stable/', None), } # -- Extlinks extension ------------------------------------------------------- diff --git a/esmvalcore/preprocessor/_regrid_xesmf.py b/esmvalcore/preprocessor/_regrid_xesmf.py new file mode 100644 index 0000000000..05063a44aa --- /dev/null +++ b/esmvalcore/preprocessor/_regrid_xesmf.py @@ -0,0 +1,165 @@ +"""xESMF regridding. + +To use this, install xesmf and ncdata, e.g. ``mamba install xesmf ncdata``. +""" + +import inspect + +import dask.array as da +import iris.cube +import numpy as np + + +class xESMFRegridder: # noqa + """xESMF regridding function. + + This is a wrapper around :class:`xesmf.frontend.Regridder` so it can be + used in :meth:`iris.cube.Cube.regrid`. + + Supports lazy regridding. + + Parameters + ---------- + src_cube: + Cube describing the source grid. + tgt_cube: + Cube describing the target grid. + **kwargs: + Any keyword argument to :class:`xesmf.frontend.Regridder` or + :meth:`xesmf.frontend.Regridder.__call__` can be provided. + + Attributes + ---------- + kwargs: + Keyword arguments to :class:`xesmf.frontend.Regridder`. + default_call_kwargs: + Default keyword arguments to :meth:`xesmf.frontend.Regridder.__call__`. + """ + + def __init__( + self, + src_cube: iris.cube.Cube, + tgt_cube: iris.cube.Cube, + **kwargs, + ) -> None: + import ncdata.iris_xarray + import xesmf + + call_arg_names = list( + inspect.signature(xesmf.Regridder.__call__).parameters)[2:] + self.kwargs = { + k: v + for k, v in kwargs.items() if k not in call_arg_names + } + self.default_call_kwargs = { + k: v + for k, v in kwargs.items() if k in call_arg_names + } + + src_cube = src_cube.copy(da.ma.filled(src_cube.core_data(), np.nan)) + tgt_cube = tgt_cube.copy(da.ma.filled(tgt_cube.core_data(), np.nan)) + src_ds = ncdata.iris_xarray.cubes_to_xarray([src_cube]) + tgt_ds = ncdata.iris_xarray.cubes_to_xarray([tgt_cube]) + + self._regridder = xesmf.Regridder(src_ds, tgt_ds, **self.kwargs) + + def __repr__(self) -> str: + """Return a string representation of the class.""" + kwargs = self.kwargs | self.default_call_kwargs + return f"{self.__class__.__name__}(**{kwargs})" + + def __call__(self, src_cube: iris.cube.Cube, **kwargs) -> iris.cube.Cube: + """Run the regridder. + + Parameters + ---------- + src_cube: + The cube to regrid. + **kwargs: + Keyword arguments to :meth:`xesmf.frontend.Regridder.__call__`. + + Returns + ------- + iris.cube.Cube + The regridded cube. + """ + import ncdata.iris_xarray + + call_args = dict(self.default_call_kwargs) + call_args.update(kwargs) + + src_cube = src_cube.copy(da.ma.filled(src_cube.core_data(), np.nan)) + src_ds = ncdata.iris_xarray.cubes_to_xarray([src_cube]) + + tgt_ds = self._regridder(src_ds, **call_args) + + cube = ncdata.iris_xarray.cubes_from_xarray( + tgt_ds, + iris_load_kwargs={'constraints': src_cube.standard_name}, + )[0] + cube.data = da.ma.masked_where(da.isnan(cube.core_data()), + cube.core_data()) + return cube + + +class xESMF: # noqa + """xESMF regridding scheme. + + This is a wrapper around :class:`xesmf.frontend.Regridder` so it can be + used in :meth:`iris.cube.Cube.regrid`. It uses the :mod:`ncdata` package to + convert the :class:`iris.cube.Cube` to an :class:`xarray.Dataset` before + regridding and back after regridding. + + Supports lazy regridding. + + Masks are converted to :obj:`numpy.nan` before regridding and converted + back to masks after regridding. + + Parameters + ---------- + **kwargs: + Any keyword argument to :class:`xesmf.frontend.Regridder` or + :meth:`xesmf.frontend.Regridder.__call__` can be provided. By default, + the arguments ``ignore_degenerate=True``, ``keep_attrs=True``, + ``skipna=True``, and ``unmapped_to_nan=True`` will be used. + + Attributes + ---------- + kwargs: + Keyword arguments that will be provided to the regridder. + """ + + def __init__(self, **kwargs) -> None: + args = { + 'ignore_degenerate': True, + 'skipna': True, + 'keep_attrs': True, + 'unmapped_to_nan': True, + } + args.update(kwargs) + self.kwargs = args + + def __repr__(self) -> str: + """Return string representation of class.""" + return f'{self.__class__.__name__}(**{self.kwargs})' + + def regridder( + self, + src_cube: iris.cube.Cube, + tgt_cube: iris.cube.Cube, + ) -> xESMFRegridder: + """Create xESMF regridding function. + + Parameters + ---------- + src_cube: + Cube defining the source grid. + tgt_cube: + Cube defining the target grid. + + Returns + ------- + xESMFRegridder + xESMF regridding function. + """ + return xESMFRegridder(src_cube, tgt_cube, **self.kwargs) diff --git a/esmvalcore/preprocessor/regrid_schemes.py b/esmvalcore/preprocessor/regrid_schemes.py index 91af9e3fdf..c8a6f2d13c 100644 --- a/esmvalcore/preprocessor/regrid_schemes.py +++ b/esmvalcore/preprocessor/regrid_schemes.py @@ -17,10 +17,10 @@ UnstructuredLinearRegridder, UnstructuredNearest, ) +from esmvalcore.preprocessor._regrid_xesmf import xESMF, xESMFRegridder logger = logging.getLogger(__name__) - __all__ = [ 'ESMPyAreaWeighted', 'ESMPyLinear', @@ -31,6 +31,8 @@ 'UnstructuredLinear', 'UnstructuredLinearRegridder', 'UnstructuredNearest', + 'xESMF', + 'xESMFRegridder', ] @@ -51,7 +53,6 @@ class GenericRegridder: Cube, \*\*kwargs) -> Cube. **kwargs: Keyword arguments for the generic regridding function. - """ def __init__( @@ -79,7 +80,6 @@ def __call__(self, cube: Cube) -> Cube: ------- Cube Regridded cube. - """ return self.func(cube, self.tgt_cube, **self.kwargs) @@ -98,7 +98,6 @@ class GenericFuncScheme: Cube, \*\*kwargs) -> Cube. **kwargs: Keyword arguments for the generic regridding function. - """ def __init__(self, func: Callable, **kwargs): @@ -125,6 +124,5 @@ def regridder(self, src_cube: Cube, tgt_cube: Cube) -> GenericRegridder: ------- GenericRegridder Regridder instance. - """ return GenericRegridder(src_cube, tgt_cube, self.func, **self.kwargs)