Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

First pass at mesh to mesh #311

Merged
merged 15 commits into from
Nov 2, 2023
93 changes: 93 additions & 0 deletions esmf_regrid/experimental/unstructured_scheme.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down Expand Up @@ -353,3 +355,94 @@ 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_cube`` regridded onto the
horizontal mesh of ``mesh_cube``. The dimensions on the :class:`~iris.cube.Cube` associated
with the grid will be replaced by a dimension associated with the
:attr:`~iris.cube.Cube.mesh`.
That dimension will be the first of the grid dimensions, whether
it is associated with the ``x`` or ``y`` coordinate. Since two dimensions are
being replaced by one, coordinates associated with dimensions after
the grid will become associated with dimensions one lower.
This function requires that the horizontal dimension of ``mesh_cube`` is
described by a 2D mesh with data located on the faces of that mesh
for conservative regridding and located on either faces or nodes for
bilinear regridding.
This function allows the horizontal grid of ``grid_cube`` to be either
rectilinear or curvilinear (i.e. expressed in terms of two orthogonal
1D coordinates or via a pair of 2D coordinates).
This function also requires that the :class:`~iris.coords.Coord`\\ s describing the
horizontal grid have :attr:`~iris.coords.Coord.bounds`.

Parameters
----------
src_cube : :class:`iris.cube.Cube`
A rectilinear instance of :class:`~iris.cube.Cube` that supplies the data,
metadata and coordinates.
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.
src_resolution : int, optional
If present, represents the amount of latitude slices per cell
given to ESMF for calculation.
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
74 changes: 70 additions & 4 deletions esmf_regrid/schemes.py
Original file line number Diff line number Diff line change
Expand Up @@ -769,17 +769,83 @@ 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,
stephenworsley marked this conversation as resolved.
Show resolved Hide resolved
):
stephenworsley marked this conversation as resolved.
Show resolved Hide resolved
raise NotImplementedError
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):
stephenworsley marked this conversation as resolved.
Show resolved Hide resolved
raise NotImplementedError
(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(
Expand Down
Original file line number Diff line number Diff line change
@@ -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
8 changes: 7 additions & 1 deletion esmf_regrid/tests/unit/schemes/test_ESMFAreaWeighted.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
8 changes: 7 additions & 1 deletion esmf_regrid/tests/unit/schemes/test_ESMFBilinear.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
Loading
Loading