First pass at mesh to mesh

esmf_regrid/experimental/unstructured_scheme.py

@@ -9,6 +9,8 @@
     _regrid_rectilinear_to_unstructured__prepare,
     _regrid_unstructured_to_rectilinear__perform,
     _regrid_unstructured_to_rectilinear__prepare,
+    _regrid_unstructured_to_unstructured__perform,
+    _regrid_unstructured_to_unstructured__prepare,
 )
 
 
@@ -353,3 +355,77 @@ def __init__(
         self.resolution = src_resolution
         self.mesh, self.location = self._tgt
         self.grid_x, self.grid_y = self._src
+
+
+def regrid_unstructured_to_unstructured(
+    src_mesh_cube,
+    tgt_mesh_cube,
+    mdtol=0,
+    method="conservative",
+    use_src_mask=False,
+    use_tgt_mask=False,
+):
+    r"""
+    Regrid rectilinear :class:`~iris.cube.Cube` onto unstructured mesh.
+
+    Return a new :class:`~iris.cube.Cube` with :attr:`~iris.cube.Cube.data`
+    values calculated using weights generated by :mod:`esmpy` to give the weighted
+    mean of :attr:`~iris.cube.Cube.data` values from ``src_mesh_cube`` regridded onto the
+    horizontal mesh of ``tgt_mesh_cube``. The resulting cube will have the same
+    ``mesh_dim`` as ``src_mesh_cube``.
+
+    Parameters
+    ----------
+    src_mesh_cube : :class:`iris.cube.Cube`
+        A unstructured instance of :class:`~iris.cube.Cube` that supplies the data,
+        metadata and coordinates.
+    tgt_mesh_cube : :class:`iris.cube.Cube`
+        An unstructured instance of :class:`~iris.cube.Cube` that supplies the desired
+        horizontal mesh definition.
+    mdtol : float, default=0
+        Tolerance of missing data. The value returned in each element of the
+        returned :class:`~iris.cube.Cube`\\ 's :attr:`~iris.cube.Cube.data` array
+        will be masked if the fraction of masked
+        data in the overlapping cells of the source cube exceeds ``mdtol``. This
+        fraction is calculated based on the area of masked cells within each
+        target cell. ``mdtol=0`` means no missing data is tolerated while ``mdtol=1``
+        will mean the resulting element will be masked if and only if all the
+        overlapping cells of the ``src_cube`` are masked.
+    method : str, default="conservative"
+        Either "conservative", "bilinear" or "nearest". Corresponds to the :mod:`esmpy` methods
+        :attr:`~esmpy.api.constants.RegridMethod.CONSERVE` or
+        :attr:`~esmpy.api.constants.RegridMethod.BILINEAR` or
+        :attr:`~esmpy.api.constants.RegridMethod.NEAREST` used to calculate weights.
+    use_src_mask : :obj:`~numpy.typing.ArrayLike` or bool, default=False
+        Either an array representing the cells in the source to ignore, or else
+        a boolean value. If True, this array is taken from the mask on the data
+        in ``src_cube``. If False, no mask will be taken and all points will
+        be used in weights calculation.
+    use_tgt_mask : :obj:`~numpy.typing.ArrayLike` or bool, default=False
+        Either an array representing the cells in the target to ignore, or else
+        a boolean value. If True, this array is taken from the mask on the data
+        in ``grid_cube``. If False, no mask will be taken and all points
+        will be used in weights calculation.
+
+    Returns
+    -------
+    :class:`iris.cube.Cube`
+        A new :class:`~iris.cube.Cube` instance.
+
+    """
+    if tgt_mesh_cube.mesh is None:
+        raise ValueError("mesh_cube has no mesh.")
+    src_mask = _get_mask(src_mesh_cube, use_src_mask)
+    tgt_mask = _get_mask(tgt_mesh_cube, use_tgt_mask)
+
+    regrid_info = _regrid_unstructured_to_unstructured__prepare(
+        src_mesh_cube,
+        tgt_mesh_cube,
+        method=method,
+        src_mask=src_mask,
+        tgt_mask=tgt_mask,
+    )
+    result = _regrid_unstructured_to_unstructured__perform(
+        src_mesh_cube, regrid_info, mdtol
+    )
+    return result

esmf_regrid/schemes.py

@@ -516,6 +516,13 @@ def _regrid_rectilinear_to_rectilinear__prepare(
     src_mask=None,
     tgt_mask=None,
 ):
+    """
+    First (setup) part of 'regrid_rectilinear_to_rectilinear'.
+
+    Check inputs and calculate the sparse regrid matrix and related info.
+    The 'regrid info' returned can be re-used over many 2d slices.
+
+    """
     tgt_x = _get_coord(tgt_grid_cube, "x")
     tgt_y = _get_coord(tgt_grid_cube, "y")
     src_x = _get_coord(src_grid_cube, "x")
@@ -544,6 +551,12 @@ def _regrid_rectilinear_to_rectilinear__prepare(
 
 
 def _regrid_rectilinear_to_rectilinear__perform(src_cube, regrid_info, mdtol):
+    """
+    Second (regrid) part of 'regrid_rectilinear_to_rectilinear'.
+
+    Perform the prepared regrid calculation on a single cube.
+
+    """
     grid_x_dim, grid_y_dim = regrid_info.dims
     grid_x, grid_y = regrid_info.target
     regridder = regrid_info.regridder
@@ -769,17 +782,96 @@ def _regrid_rectilinear_to_unstructured__perform(src_cube, regrid_info, mdtol):
 
 
 def _regrid_unstructured_to_unstructured__prepare(
-    src_grid_cube,
-    target_mesh_cube,
+    src_mesh_cube,
+    tgt_cube_or_mesh,
     method,
     precomputed_weights=None,
+    src_mask=None,
+    tgt_mask=None,
+    src_location=None,
     tgt_location=None,
 ):
-    raise NotImplementedError
+    """
+    First (setup) part of 'regrid_unstructured_to_unstructured'.
+
+    Check inputs and calculate the sparse regrid matrix and related info.
+    The 'regrid info' returned can be re-used over many 2d slices.
+
+    """
+    if isinstance(tgt_cube_or_mesh, Mesh):
+        mesh = tgt_cube_or_mesh
+        location = tgt_location
+    else:
+        mesh = tgt_cube_or_mesh.mesh
+        location = tgt_cube_or_mesh.location
+
+    mesh_dim = src_mesh_cube.mesh_dim()
+
+    src_meshinfo = _make_meshinfo(
+        src_mesh_cube, method, src_mask, "source", location=src_location
+    )
+    tgt_meshinfo = _make_meshinfo(
+        tgt_cube_or_mesh, method, tgt_mask, "target", location=tgt_location
+    )
+
+    regridder = Regridder(
+        src_meshinfo,
+        tgt_meshinfo,
+        method=method,
+        precomputed_weights=precomputed_weights,
+    )
+
+    regrid_info = RegridInfo(
+        dims=[mesh_dim],
+        target=MeshRecord(mesh, location),
+        regridder=regridder,
+    )
+
+    return regrid_info
 
 
 def _regrid_unstructured_to_unstructured__perform(src_cube, regrid_info, mdtol):
-    raise NotImplementedError
+    """
+    Second (regrid) part of 'regrid_unstructured_to_unstructured'.
+
+    Perform the prepared regrid calculation on a single cube.
+
+    """
+    (mesh_dim,) = regrid_info.dims
+    mesh, location = regrid_info.target
+    regridder = regrid_info.regridder
+
+    regrid = functools.partial(
+        _regrid_along_dims,
+        regridder,
+        dims=[mesh_dim],
+        num_out_dims=1,
+        mdtol=mdtol,
+    )
+    if location == "face":
+        face_node = mesh.face_node_connectivity
+        chunk_shape = (face_node.shape[face_node.location_axis],)
+    elif location == "node":
+        chunk_shape = mesh.node_coords[0].shape
+    else:
+        raise NotImplementedError(f"Unrecognised location {location}.")
+
+    new_data = _map_complete_blocks(
+        src_cube,
+        regrid,
+        (mesh_dim,),
+        chunk_shape,
+    )
+
+    new_cube = _create_cube(
+        new_data,
+        src_cube,
+        (mesh_dim,),
+        mesh.to_MeshCoords(location),
+        1,
+    )
+
+    return new_cube
 
 
 def regrid_rectilinear_to_rectilinear(

esmf_regrid/tests/unit/experimental/unstructured_scheme/test_regrid_unstructured_to_unstructured.py

@@ -0,0 +1,174 @@
+"""Unit tests for :func:`esmf_regrid.experimental.unstructured_scheme.regrid_unstructured_to_unstructured`."""
+
+from iris.coords import AuxCoord, DimCoord
+from iris.cube import Cube
+import numpy as np
+import pytest
+
+from esmf_regrid.experimental.unstructured_scheme import (
+    regrid_unstructured_to_unstructured,
+)
+from esmf_regrid.tests.unit.schemes.test__mesh_to_MeshInfo import (
+    _gridlike_mesh_cube,
+)
+from esmf_regrid.tests.unit.schemes.test__regrid_unstructured_to_rectilinear__prepare import (
+    _flat_mesh_cube,
+    _full_mesh,
+)
+
+
+def _add_metadata(cube):
+    result = cube.copy()
+    result.units = "K"
+    result.attributes = {"a": 1}
+    result.standard_name = "air_temperature"
+    scalar_height = AuxCoord([5], units="m", standard_name="height")
+    scalar_time = DimCoord([10], units="s", standard_name="time")
+    result.add_aux_coord(scalar_height)
+    result.add_aux_coord(scalar_time)
+    return result
+
+
+def test_flat_cubes():
+    """
+    Basic test for :func:`esmf_regrid.experimental.unstructured_scheme.regrid_unstructured_to_unstructured`.
+
+    Tests with flat cubes as input (a 1D mesh cube and a 2D grid cube).
+    """
+    src = _flat_mesh_cube()
+
+    n_lons = 6
+    n_lats = 5
+    tgt = _flat_mesh_cube()
+    # Ensure data in the target grid is different to the expected data.
+    # i.e. target grid data is all zero, expected data is all one
+    tgt.data[:] = 0
+
+    src = _add_metadata(src)
+    src.data[:] = 1  # Ensure all data in the source is one.
+    result = regrid_unstructured_to_unstructured(src, tgt)
+
+    expected_data = np.ones([n_lats, n_lons])
+    expected_cube = _add_metadata(tgt)
+
+    # Lenient check for data.
+    assert np.allclose(expected_data, result.data)
+
+    # Check metadata and scalar coords.
+    expected_cube.data = result.data
+    assert expected_cube == result
+
+
+@pytest.mark.parametrize("method", ("bilinear", "nearest"))
+def test_node_friendly_methods(method):
+    """
+    Basic test for :func:`esmf_regrid.experimental.unstructured_scheme.regrid_unstructured_to_unstructured`.
+
+    Tests with the bilinear and nearest method.
+    """
+    n_lons = 6
+    n_lats = 5
+    src = _gridlike_mesh_cube(n_lons, n_lats, location="node")
+    tgt = _gridlike_mesh_cube(n_lons, n_lats, location="node")
+    # Ensure data in the target mesh is different to the expected data.
+    # i.e. target mesh data is all zero, expected data is all one
+    tgt.data[:] = 0
+
+    src = _add_metadata(src)
+    src.data[:] = 1  # Ensure all data in the source is one.
+    result = regrid_unstructured_to_unstructured(src, tgt, method=method)
+
+    expected_data = np.ones_like(tgt.data)
+    expected_cube = _add_metadata(tgt)
+
+    # Lenient check for data.
+    assert np.allclose(expected_data, result.data)
+
+    # Check metadata and scalar coords.
+    expected_cube.data = result.data
+    assert expected_cube == result
+
+
+def test_invalid_args():
+    """
+    Test for :func:`esmf_regrid.experimental.unstructured_scheme.regrid_unstructured_to_unstructured`.
+
+    Tests that an appropriate error is raised when arguments are invalid.
+    """
+    n_lons = 6
+    n_lats = 5
+    node_src = _gridlike_mesh_cube(n_lons, n_lats, location="node")
+    edge_src = _gridlike_mesh_cube(n_lons, n_lats, location="edge")
+    face_src = _gridlike_mesh_cube(n_lons, n_lats, location="face")
+    tgt = _gridlike_mesh_cube(n_lons, n_lats)
+
+    with pytest.raises(NotImplementedError):
+        _ = regrid_unstructured_to_unstructured(face_src, tgt, method="other")
+    with pytest.raises(ValueError) as excinfo:
+        _ = regrid_unstructured_to_unstructured(node_src, tgt, method="conservative")
+    expected_message = (
+        "Conservative regridding requires a source cube located on "
+        "the face of a cube, target cube had the node location."
+    )
+    assert expected_message in str(excinfo.value)
+    with pytest.raises(ValueError) as excinfo:
+        _ = regrid_unstructured_to_unstructured(edge_src, tgt, method="bilinear")
+    expected_message = (
+        "bilinear regridding requires a source cube with a node "
+        "or face location, target cube had the edge location."
+    )
+    assert expected_message in str(excinfo.value)
+    with pytest.raises(ValueError) as excinfo:
+        _ = regrid_unstructured_to_unstructured(edge_src, tgt, method="nearest")
+    expected_message = (
+        "nearest regridding requires a source cube with a node "
+        "or face location, target cube had the edge location."
+    )
+    assert expected_message in str(excinfo.value)
+
+
+def test_multidim_cubes():
+    """
+    Test for :func:`esmf_regrid.experimental.unstructured_scheme.regrid_unstructured_to_unstructured`.
+
+    Tests with multidimensional cubes. The source cube contains
+    coordinates on the dimensions before and after the mesh dimension.
+    """
+
+    mesh = _full_mesh()
+    mesh_length = mesh.connectivity(contains_face=True).shape[0]
+
+    h = 2
+    t = 3
+    height = DimCoord(np.arange(h), standard_name="height")
+    time = DimCoord(np.arange(t), standard_name="time")
+
+    src_data = np.empty([t, mesh_length, h])
+    src_data[:] = np.arange(t * h).reshape([t, h])[:, np.newaxis, :]
+    cube = Cube(src_data)
+    mesh_coord_x, mesh_coord_y = mesh.to_MeshCoords("face")
+    cube.add_aux_coord(mesh_coord_x, 1)
+    cube.add_aux_coord(mesh_coord_y, 1)
+    cube.add_dim_coord(time, 0)
+    cube.add_dim_coord(height, 2)
+
+    n_lons = 6
+    n_lats = 5
+    tgt = _gridlike_mesh_cube(n_lons, n_lats)
+
+    result = regrid_unstructured_to_unstructured(cube, tgt)
+
+    # Lenient check for data.
+    expected_data = np.empty([t, n_lats * n_lons, h])
+    expected_data[:] = np.arange(t * h).reshape(t, h)[:, np.newaxis, :]
+    assert np.allclose(expected_data, result.data)
+
+    expected_cube = Cube(expected_data)
+    expected_cube.add_dim_coord(time, 0)
+    expected_cube.add_aux_coord(tgt.coord("latitude"), 1)
+    expected_cube.add_aux_coord(tgt.coord("longitude"), 1)
+    expected_cube.add_dim_coord(height, 2)
+
+    # Check metadata and scalar coords.
+    result.data = expected_data
+    assert expected_cube == result