Add xESMF regridder

doc/conf.py

@@ -433,22 +433,25 @@
 
 # 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}/',
      None),
     '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 -------------------------------------------------------

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)

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)