Add an iris-esmf-regrid based regridding scheme (#2457)

Co-authored-by: Manuel Schlund <32543114+schlunma@users.noreply.github.com>

Author: bouweandela

doc/conf.py

@@ -438,20 +438,21 @@
 
 # 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),
+    'esmpy': ('https://earthsystemmodeling.org/esmpy_doc/release/latest/html/',
+              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',
-                         None),
+    'iris': ('https://scitools-iris.readthedocs.io/en/stable/', None),
+    'esmf_regrid': ('https://iris-esmf-regrid.readthedocs.io/en/stable/', None),
     'matplotlib': ('https://matplotlib.org/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),
 }

doc/quickstart/find_data.rst

@@ -398,32 +398,17 @@ ESMValCore can automatically make native ICON data `UGRID
 loading the data.
 The UGRID conventions provide a standardized format to store data on
 unstructured grids, which is required by many software packages or tools to
-work correctly.
+work correctly and specifically by Iris to interpret the grid as a
+:ref:`mesh <iris:ugrid>`.
 An example is the horizontal regridding of native ICON data to a regular grid.
-While the built-in :ref:`nearest scheme <built-in regridding
-schemes>` can handle unstructured grids not in UGRID format, using more complex
-regridding algorithms (for example provided by the
-:doc:`iris-esmf-regrid:index` package through :ref:`generic regridding
-schemes`) requires the input data in UGRID format.
-The following code snippet provides a preprocessor that regrids native ICON
-data to a 1°x1° grid using `ESMF's first-order conservative regridding
-algorithm <https://earthsystemmodeling.org/regrid/#regridding-methods>`__:
-
-.. code-block:: yaml
-
-   preprocessors:
-     regrid_icon:
-       regrid:
-         target_grid: 1x1
-         scheme:
-           reference: esmf_regrid.schemes:ESMFAreaWeighted
-
+While the :ref:`built-in regridding schemes <built-in regridding schemes>`
+`linear` and `nearest`  can handle unstructured grids (i.e., not UGRID-compliant) and meshes (i.e., UGRID-compliant),
+the `area_weighted` scheme requires the input data in UGRID format.
 This automatic UGRIDization is enabled by default, but can be switched off with
 the facet ``ugrid: false`` in the recipe or the extra facets (see below).
-This is useful for diagnostics that do not support input data in UGRID format
-(yet) like the :ref:`Psyplot diagnostic <esmvaltool:recipes_psyplot_diag>` or
-if you want to use the built-in :ref:`nearest scheme <built-in
-regridding schemes>` regridding scheme.
+This is useful for diagnostics that act on the native ICON grid and do not
+support input data in UGRID format (yet), like the
+:ref:`Psyplot diagnostic <esmvaltool:recipes_psyplot_diag>`.
 
 For 3D ICON variables, ESMValCore tries to add the pressure level information
 (from the variables `pfull` and `phalf`) and/or altitude information (from the

doc/recipe/preprocessor.rst

@@ -890,15 +890,15 @@ The arguments are defined below:
 Regridding (interpolation, extrapolation) schemes
 -------------------------------------------------
 
-ESMValCore has a number of built-in regridding schemes, which are presented in
-:ref:`built-in regridding schemes`. Additionally, it is also possible to use
-third party regridding schemes designed for use with :doc:`Iris
-<iris:index>`. This is explained in :ref:`generic regridding schemes`.
+ESMValCore provides three default regridding schemes, which are presented in
+:ref:`default regridding schemes`. Additionally, it is also possible to use
+third party regridding schemes designed for use with :meth:`iris.cube.Cube.regrid`.
+This is explained in :ref:`generic regridding schemes`.
 
 Grid types
 ~~~~~~~~~~
 
-In ESMValCore, we distinguish between three grid types (note that these might
+In ESMValCore, we distinguish between various grid types (note that these might
 differ from other definitions):
 
 * **Regular grid**: A rectilinear grid with 1D latitude and 1D longitude
@@ -907,30 +907,34 @@ differ from other definitions):
   longitude coordinates with common dimensions.
 * **Unstructured grid**: A grid with 1D latitude and 1D longitude coordinates
   with common dimensions (i.e., a simple list of points).
+* **Mesh**: A mesh as supported by Iris and described in :ref:`iris:ugrid`.
 
-.. _built-in regridding schemes:
+.. _default regridding schemes:
 
-Built-in regridding schemes
-~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Default regridding schemes
+~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 * ``linear``: Bilinear regridding.
   For source data on a regular grid, uses :obj:`~iris.analysis.Linear` with
   `extrapolation_mode='mask'`.
-  For source data on an irregular grid, uses
-  :class:`~esmvalcore.preprocessor.regrid_schemes.ESMPyLinear`.
+  For source and/or target data on an irregular grid or mesh, uses
+  :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with
+  `method='bilinear'`.
   For source data on an unstructured grid, uses
   :class:`~esmvalcore.preprocessor.regrid_schemes.UnstructuredLinear`.
 * ``nearest``: Nearest-neighbor regridding.
   For source data on a regular grid, uses :obj:`~iris.analysis.Nearest` with
   `extrapolation_mode='mask'`.
-  For source data on an irregular grid, uses
-  :class:`~esmvalcore.preprocessor.regrid_schemes.ESMPyNearest`.
+  For source and/or target data on an irregular grid or mesh, uses
+  :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with
+  `method='nearest'`.
   For source data on an unstructured grid, uses
   :class:`~esmvalcore.preprocessor.regrid_schemes.UnstructuredNearest`.
 * ``area_weighted``: First-order conservative (area-weighted) regridding.
   For source data on a regular grid, uses :obj:`~iris.analysis.AreaWeighted`.
-  For source data on an irregular grid, uses
-  :class:`~esmvalcore.preprocessor.regrid_schemes.ESMPyAreaWeighted`.
+  For source and/or target data on an irregular grid or mesh, uses
+  :class:`~esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` with
+  `method='conservative'`.
   Source data on an unstructured grid is not supported.
 
 .. _generic regridding schemes:
@@ -950,7 +954,9 @@ afforded by the built-in schemes described above.
 
 To facilitate this, the :func:`~esmvalcore.preprocessor.regrid` preprocessor
 allows the use of any scheme designed for Iris. The scheme must be installed
-and importable. To use this feature, the ``scheme`` key passed to the
+and importable. Several such schemes are provided by :mod:`iris.analysis` and
+:mod:`esmvalcore.preprocessor.regrid_schemes`.
+To use this feature, the ``scheme`` key passed to the
 preprocessor must be a dictionary instead of a simple string that contains all
 necessary information. That includes a ``reference`` to the desired scheme
 itself, as well as any arguments that should be passed through to the
@@ -996,10 +1002,13 @@ module, the second refers to the scheme, i.e. some callable that will be called
 with the remaining entries of the ``scheme`` dictionary passed as keyword
 arguments.
 
-One package that aims to capitalize on the :ref:`support for unstructured grids
-introduced in Iris 3.2 <iris:ugrid>` is :doc:`iris-esmf-regrid:index`.
+One package that aims to capitalize on the :ref:`support for meshes
+introduced in Iris 3.2 <iris:ugrid>` is :doc:`esmf_regrid:index`.
 It aims to provide lazy regridding for structured regular and irregular grids,
-as well as unstructured grids.
+as well as meshes. It is recommended to use these schemes through
+the :obj:`esmvalcore.preprocessor.regrid_schemes.IrisESMFRegrid` scheme though,
+as that provides more efficient handling of masks.
+
 An example of its usage in a preprocessor is:
 
 .. code-block:: yaml
@@ -1009,16 +1018,19 @@ An example of its usage in a preprocessor is:
         regrid:
           target_grid: 2.5x2.5
           scheme:
-            reference: esmf_regrid.schemes:ESMFAreaWeighted
+            reference: esmvalcore.preprocessor.regrid_schemes:IrisESMFRegrid
+            method: conservative
             mdtol: 0.7
+            use_src_mask: true
+            collapse_src_mask_along: ZT
 
 Additionally, the use of generic schemes that take source and target grid cubes as
 arguments is also supported. The call function for such schemes must be defined as
 `(src_cube, grid_cube, **kwargs)` and they must return `iris.cube.Cube` objects.
 The `regrid` module will automatically pass the source and grid cubes as inputs
 of the scheme. An example of this usage is
 the :func:`~esmf_regrid.schemes.regrid_rectilinear_to_rectilinear`
-scheme available in :doc:`iris-esmf-regrid:index`:
+scheme available in :doc:`esmf_regrid:index`:
 
 .. code-block:: yaml
 

environment.yml

@@ -19,7 +19,7 @@ dependencies:
   - geopy
   - humanfriendly
   - iris >=3.10.0
-  - iris-esmf-regrid >=0.10.0  # github.com/SciTools-incubator/iris-esmf-regrid/pull/342
+  - iris-esmf-regrid >=0.11.0
   - iris-grib
   - isodate
   - jinja2

esmvalcore/preprocessor/_regrid.py

@@ -30,6 +30,9 @@
 from esmvalcore.cmor.table import CMOR_TABLES
 from esmvalcore.exceptions import ESMValCoreDeprecationWarning
 from esmvalcore.iris_helpers import has_irregular_grid, has_unstructured_grid
+from esmvalcore.preprocessor._shared import (
+    get_dims_along_axes,
+)
 from esmvalcore.preprocessor._shared import (
     get_array_module,
     preserve_float_dtype,
@@ -39,10 +42,8 @@
     add_cell_measure,
 )
 from esmvalcore.preprocessor.regrid_schemes import (
-    ESMPyAreaWeighted,
-    ESMPyLinear,
-    ESMPyNearest,
     GenericFuncScheme,
+    IrisESMFRegrid,
     UnstructuredLinear,
     UnstructuredNearest,
 )
@@ -91,9 +92,17 @@
 # curvilinear grids; i.e., grids that can be described with 2D latitude and 2D
 # longitude coordinates with common dimensions)
 HORIZONTAL_SCHEMES_IRREGULAR = {
-    'area_weighted': ESMPyAreaWeighted(),
-    'linear': ESMPyLinear(),
-    'nearest': ESMPyNearest(),
+    'area_weighted': IrisESMFRegrid(method='conservative'),
+    'linear': IrisESMFRegrid(method='bilinear'),
+    'nearest': IrisESMFRegrid(method='nearest'),
+}
+
+# Supported horizontal regridding schemes for meshes
+# https://scitools-iris.readthedocs.io/en/stable/further_topics/ugrid/index.html
+HORIZONTAL_SCHEMES_MESH = {
+    'area_weighted': IrisESMFRegrid(method='conservative'),
+    'linear': IrisESMFRegrid(method='bilinear'),
+    'nearest': IrisESMFRegrid(method='nearest'),
 }
 
 # Supported horizontal regridding schemes for unstructured grids (i.e., grids,
@@ -533,29 +542,7 @@ def _get_target_grid_cube(
     return target_grid_cube
 
 
-def _attempt_irregular_regridding(cube: Cube, scheme: str) -> bool:
-    """Check if irregular regridding with ESMF should be used."""
-    if not has_irregular_grid(cube):
-        return False
-    if scheme not in HORIZONTAL_SCHEMES_IRREGULAR:
-        raise ValueError(
-            f"Regridding scheme '{scheme}' does not support irregular data, "
-            f"expected one of {list(HORIZONTAL_SCHEMES_IRREGULAR)}")
-    return True
-
-
-def _attempt_unstructured_regridding(cube: Cube, scheme: str) -> bool:
-    """Check if unstructured regridding should be used."""
-    if not has_unstructured_grid(cube):
-        return False
-    if scheme not in HORIZONTAL_SCHEMES_UNSTRUCTURED:
-        raise ValueError(
-            f"Regridding scheme '{scheme}' does not support unstructured "
-            f"data, expected one of {list(HORIZONTAL_SCHEMES_UNSTRUCTURED)}")
-    return True
-
-
-def _load_scheme(src_cube: Cube, scheme: str | dict):
+def _load_scheme(src_cube: Cube, tgt_cube: Cube, scheme: str | dict):
     """Return scheme that can be used in :meth:`iris.cube.Cube.regrid`."""
     loaded_scheme: Any = None
 
@@ -586,23 +573,27 @@ def _load_scheme(src_cube: Cube, scheme: str | dict):
         logger.debug("Loaded regridding scheme %s", loaded_scheme)
         return loaded_scheme
 
-    # Scheme is a dict -> assume this describes a generic regridding scheme
     if isinstance(scheme, dict):
+        # Scheme is a dict -> assume this describes a generic regridding scheme
         loaded_scheme = _load_generic_scheme(scheme)
-
-    # Scheme is a str -> load appropriate regridding scheme depending on the
-    # type of input data
-    elif _attempt_irregular_regridding(src_cube, scheme):
-        loaded_scheme = HORIZONTAL_SCHEMES_IRREGULAR[scheme]
-    elif _attempt_unstructured_regridding(src_cube, scheme):
-        loaded_scheme = HORIZONTAL_SCHEMES_UNSTRUCTURED[scheme]
     else:
-        loaded_scheme = HORIZONTAL_SCHEMES_REGULAR.get(scheme)
-
-    if loaded_scheme is None:
-        raise ValueError(
-            f"Got invalid regridding scheme string '{scheme}', expected one "
-            f"of {list(HORIZONTAL_SCHEMES_REGULAR)}")
+        # Scheme is a str -> load appropriate regridding scheme depending on
+        # the type of input data
+        if has_irregular_grid(src_cube) or has_irregular_grid(tgt_cube):
+            grid_type = 'irregular'
+        elif src_cube.mesh is not None or tgt_cube.mesh is not None:
+            grid_type = 'mesh'
+        elif has_unstructured_grid(src_cube):
+            grid_type = 'unstructured'
+        else:
+            grid_type = 'regular'
+
+        schemes = globals()[f"HORIZONTAL_SCHEMES_{grid_type.upper()}"]
+        if scheme not in schemes:
+            raise ValueError(
+                f"Regridding scheme '{scheme}' not available for {grid_type} "
+                f"data, expected one of: {', '.join(schemes)}")
+        loaded_scheme = schemes[scheme]
 
     logger.debug("Loaded regridding scheme %s", loaded_scheme)
 
@@ -676,14 +667,14 @@ def _get_regridder(
                     return regridder
 
         # Regridder is not in cached -> return a new one and cache it
-        loaded_scheme = _load_scheme(src_cube, scheme)
+        loaded_scheme = _load_scheme(src_cube, tgt_cube, scheme)
         regridder = loaded_scheme.regridder(src_cube, tgt_cube)
         _CACHED_REGRIDDERS.setdefault(name_shape_key, {})
         _CACHED_REGRIDDERS[name_shape_key][coord_key] = regridder
 
     # (2) Weights caching disabled
     else:
-        loaded_scheme = _load_scheme(src_cube, scheme)
+        loaded_scheme = _load_scheme(src_cube, tgt_cube, scheme)
         regridder = loaded_scheme.regridder(src_cube, tgt_cube)
 
     return regridder
@@ -860,36 +851,40 @@ def _cache_clear():
 
 def _rechunk(cube: Cube, target_grid: Cube) -> Cube:
     """Re-chunk cube with optimal chunk sizes for target grid."""
-    if not cube.has_lazy_data() or cube.ndim < 3:
-        # Only rechunk lazy multidimensional data
+    if not cube.has_lazy_data():
+        # Only rechunk lazy data
         return cube
 
-    lon_coord = target_grid.coord(axis='X')
-    lat_coord = target_grid.coord(axis='Y')
-    if lon_coord.ndim != 1 or lat_coord.ndim != 1:
-        # This function only supports 1D lat/lon coordinates.
-        return cube
+    # Extract grid dimension information from source cube
+    src_grid_indices = get_dims_along_axes(cube, ["X", "Y"])
+    src_grid_shape = tuple(cube.shape[i] for i in src_grid_indices)
+    src_grid_ndims = len(src_grid_indices)
 
-    lon_dim, = target_grid.coord_dims(lon_coord)
-    lat_dim, = target_grid.coord_dims(lat_coord)
-    grid_indices = sorted((lon_dim, lat_dim))
-    target_grid_shape = tuple(target_grid.shape[i] for i in grid_indices)
+    # Extract grid dimension information from target cube.
+    tgt_grid_indices = get_dims_along_axes(target_grid, ["X", "Y"])
+    tgt_grid_shape = tuple(target_grid.shape[i] for i in tgt_grid_indices)
+    tgt_grid_ndims = len(tgt_grid_indices)
 
-    if 2 * np.prod(cube.shape[-2:]) > np.prod(target_grid_shape):
+    if 2 * np.prod(src_grid_shape) > np.prod(tgt_grid_shape):
         # Only rechunk if target grid is more than a factor of 2 larger,
         # because rechunking will keep the original chunk in memory.
         return cube
 
+    # Compute a good chunk size for the target array
+    # This uses the fact that horizontal dimension(s) are the last dimension(s)
+    # of the input cube and also takes into account that iris regridding needs
+    # unchunked data along the grid dimensions.
     data = cube.lazy_data()
+    tgt_shape = data.shape[:-src_grid_ndims] + tgt_grid_shape
+    tgt_chunks = data.chunks[:-src_grid_ndims] + tgt_grid_shape
 
-    # Compute a good chunk size for the target array
-    tgt_shape = data.shape[:-2] + target_grid_shape
-    tgt_chunks = data.chunks[:-2] + target_grid_shape
-    tgt_data = da.empty(tgt_shape, dtype=data.dtype, chunks=tgt_chunks)
-    tgt_data = tgt_data.rechunk({i: "auto" for i in range(cube.ndim - 2)})
+    tgt_data = da.empty(tgt_shape, chunks=tgt_chunks, dtype=data.dtype)
+    tgt_data = tgt_data.rechunk(
+        {i: "auto"
+         for i in range(tgt_data.ndim - tgt_grid_ndims)})
 
     # Adjust chunks to source array and rechunk
-    chunks = tgt_data.chunks[:-2] + data.shape[-2:]
+    chunks = tgt_data.chunks[:-tgt_grid_ndims] + data.shape[-src_grid_ndims:]
     cube.data = data.rechunk(chunks)
 
     return cube