Skip to content

Commit

Permalink
First pass at mesh to mesh (#311)
Browse files Browse the repository at this point in the history
* First pass at mesh to mesh

* fix tests #1

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remve resolution and correct ref's

* remove circular

* new changes to fix tests

* test fix

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix test

* test fixes

* match cube dimension

* fixing last test

* all tests fixed

* adding and updating docstrings

* more docstring changes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
HGWright and pre-commit-ci[bot] authored Nov 2, 2023
1 parent 4f685cf commit b041fdb
Show file tree
Hide file tree
Showing 6 changed files with 367 additions and 7 deletions.
76 changes: 76 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,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
100 changes: 96 additions & 4 deletions esmf_regrid/schemes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
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
Loading

0 comments on commit b041fdb

Please sign in to comment.