Enable the vector_regrid flag via the backend (#436)

Vector regridding was recently added to ESMF's C and Python APIs
(esmf-org/esmf@c153f29350). This commit
exposes this option via the xesmf backend to ESMPy. (See
https://earthsystemmodeling.org/docs/release/latest/ESMF_refdoc/node5.html#SECTION050121700000000000000
for general information on vector regridding in ESMF.)

I looked into making changes to expose this via the xesmf frontend as
well, but it seems like this is going to be more involved, so I held off
on this: For most regridding operations, the regridding matrix only
considers the spatial dimensions, and then this matrix can be applied
separately across all non-spatial dimensions. However, for vector
regridding, the regridding matrix accounts for the vector dimension
(which must be of size 2) in addition to the spatial dimensions, and
then the regridding operation acts on all of these dimensions. My sense
is that xesmf is set up to assume the typical operation, and some work
will be needed to add flexibility to handle this new vector regridding
configuration.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: raphael dussin <[email protected]>

Author: 3 people

xesmf/backend.py

@@ -17,6 +17,7 @@
 
 import os
 import warnings
+from collections.abc import Sequence
 
 try:
     import esmpy as ESMF
@@ -313,7 +314,7 @@ def get_shape(self, loc=ESMF.MeshLoc.ELEMENT):
         return (self.size[loc], 1)
 
 
-def esmf_regrid_build(
+def esmf_regrid_build(  # noqa: C901
     sourcegrid,
     destgrid,
     method,
@@ -323,6 +324,7 @@ def esmf_regrid_build(
     extrap_dist_exponent=None,
     extrap_num_src_pnts=None,
     ignore_degenerate=None,
+    vector_regrid=None,
 ):
     """
     Create an ESMF.Regrid object, containing regridding weights.
@@ -381,9 +383,22 @@ def esmf_regrid_build(
         If False (default), raise error if grids contain degenerated cells
         (i.e. triangles or lines, instead of quadrilaterals)
 
+    vector_regrid : bool, optional
+        If True, treat a single extra (non-spatial) dimension in the source and
+        destination data fields as the components of a vector. (If True and
+        there is more than one extra dimension in either the source or
+        destination data fields, an error will be raised.) If not specified,
+        defaults to False.
+
+        Only vector dimensions of size 2 are supported. The first entry is
+        interpreted as the east component and the second as the north component.
+        i.e., ``extra_dims`` must be ``[2]``.
+
+        Requires ESMPy 8.9.0 or newer.
+
     Returns
     -------
-    grid : ESMF.Grid object
+    regrid : ESMF.Regrid object
 
     """
 
@@ -430,6 +445,12 @@ def esmf_regrid_build(
                 'destination grid has no corner information. ' 'cannot use conservative regridding.'
             )
 
+    if vector_regrid:
+        # Check this ESMPy requirement in order to give a more helpful error message if it
+        # isn't met
+        if not (isinstance(extra_dims, Sequence) and len(extra_dims) == 1 and extra_dims[0] == 2):
+            raise ValueError('`vector_regrid` currently requires `extra_dims` to be `[2]`')
+
     # ESMF.Regrid requires Field (Grid+data) as input, not just Grid.
     # Extra dimensions are specified when constructing the Field objects,
     # not when constructing the Regrid object later on.
@@ -480,6 +501,12 @@ def esmf_regrid_build(
     )
     if allow_masked_values:
         kwargs.update(dict(src_mask_values=[0], dst_mask_values=[0]))
+    # Only add the vector_regrid argument if it is given and true; this supports backwards
+    # compatibility with versions of ESMPy prior to 8.9.0 that do not have this option.
+    # (If the user explicitly sets the vector_regrid argument, however, that will still be
+    # passed through; this will lead to an exception with older ESMPy versions.)
+    if vector_regrid:
+        kwargs['vector_regrid'] = vector_regrid
 
     regrid = ESMF.Regrid(sourcefield, destfield, **kwargs)
 

xesmf/tests/test_backend.py

@@ -295,6 +295,36 @@ def test_read_weights(tmp_path):
         read_weights(ds.drop_vars('col'), lon_in.size, lon_out.size)
 
 
+def test_vector_regrid():
+    grid_in = Grid.from_xarray(lon_in.T, lat_in.T)
+    grid_out = Grid.from_xarray(lon_out.T, lat_out.T)
+    data_in_vec = np.stack(
+        [np.cos(data_in) * np.sign(lat_in), np.sin(data_in) * np.sign(lon_in)], axis=0
+    )
+
+    try:
+        regrid_vec = esmf_regrid_build(
+            grid_in, grid_out, 'bilinear', extra_dims=[2], vector_regrid=True
+        )
+    except TypeError:
+        pytest.skip('vector_regrid argument not supported by ESMPy versions < 8.9')
+    assert regrid_vec.vector_regrid is True
+    data_out_vec = esmf_regrid_apply(regrid_vec, data_in_vec.T).T
+
+    regrid_nonvec = esmf_regrid_build(grid_in, grid_out, 'bilinear', extra_dims=[2])
+    assert regrid_nonvec.vector_regrid in (None, False)
+    data_out_nonvec = esmf_regrid_apply(regrid_nonvec, data_in_vec.T).T
+
+    # The result of vector regridding should differ from the result of non-vector
+    # regridding, but not by much. (This threshold of 0.1 was determined empirically and
+    # may need to be adjusted if details of the test inputs change.)
+    assert np.array_equal(data_out_vec, data_out_nonvec) is False
+    assert np.max(np.abs(data_out_vec - data_out_nonvec)) < 0.1
+
+    esmf_regrid_finalize(regrid_vec)
+    esmf_regrid_finalize(regrid_nonvec)
+
+
 def test_deprecated():
     from xesmf.backend import esmf_grid, esmf_locstream