diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index f17aadcd..ccda7214 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/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 diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 21731830..59e723d2 100644 --- a/esmf_regrid/schemes.py +++ b/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( diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_regrid_unstructured_to_unstructured.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_regrid_unstructured_to_unstructured.py new file mode 100644 index 00000000..8be55fe9 --- /dev/null +++ b/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 diff --git a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeighted.py b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeighted.py index e9986a74..e73bee5d 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeighted.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeighted.py @@ -14,7 +14,13 @@ @pytest.mark.parametrize( "src_type,tgt_type", - [("grid", "grid"), ("grid", "mesh"), ("grid", "just_mesh"), ("mesh", "grid")], + [ + ("grid", "grid"), + ("grid", "mesh"), + ("grid", "just_mesh"), + ("mesh", "grid"), + ("mesh", "mesh"), + ], ) def test_cube_regrid(src_type, tgt_type): """ diff --git a/esmf_regrid/tests/unit/schemes/test_ESMFBilinear.py b/esmf_regrid/tests/unit/schemes/test_ESMFBilinear.py index d92350bf..393bbe0a 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFBilinear.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFBilinear.py @@ -14,7 +14,13 @@ @pytest.mark.parametrize( "src_type,tgt_type", - [("grid", "grid"), ("grid", "mesh"), ("grid", "just_mesh"), ("mesh", "grid")], + [ + ("grid", "grid"), + ("grid", "mesh"), + ("grid", "just_mesh"), + ("mesh", "grid"), + ("mesh", "mesh"), + ], ) def test_cube_regrid(src_type, tgt_type): """ diff --git a/esmf_regrid/tests/unit/schemes/test_ESMFNearest.py b/esmf_regrid/tests/unit/schemes/test_ESMFNearest.py index 85f8ec49..b5812b5a 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFNearest.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFNearest.py @@ -19,7 +19,13 @@ @pytest.mark.parametrize( "src_type,tgt_type", - [("grid", "grid"), ("grid", "mesh"), ("grid", "just_mesh"), ("mesh", "grid")], + [ + ("grid", "grid"), + ("grid", "mesh"), + ("grid", "just_mesh"), + ("mesh", "grid"), + ("mesh", "mesh"), + ], ) def test_cube_regrid(src_type, tgt_type): """