Skip to content

Commit

Permalink
address review comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@stephenworsley
stephenworsley committed Mar 16, 2023
1 parent f42ad74 commit 52ddfc2
Show file tree
Hide file tree
Showing 2 changed files with 108 additions and 92 deletions.
82 changes: 46 additions & 36 deletions esmf_regrid/experimental/unstructured_scheme.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -252,8 +252,8 @@ def regrid_unstructured_to_rectilinear(
mdtol=0,
method="conservative",
resolution=None,
src_mask=None,
tgt_mask=None,
use_src_mask=False,
use_tgt_mask=False,
):
r"""
Regrid unstructured :class:`~iris.cube.Cube` onto rectilinear grid.
Expand Down Expand Up @@ -298,13 +298,26 @@ def regrid_unstructured_to_rectilinear(
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`, 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_mesh_cube``. If False, no mask will be taken and all points will
be used in weights calculation.
use_tgt_mask : :obj:`~numpy.typing.ArrayLike`, 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 ``target_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.
"""
src_mask = _get_mask(src_cube, use_src_mask)
tgt_mask = _get_mask(grid_cube, use_tgt_mask)

regrid_info = _regrid_unstructured_to_rectilinear__prepare(
src_cube,
grid_cube,
Expand All @@ -328,8 +341,8 @@ def __init__(
method="conservative",
precomputed_weights=None,
resolution=None,
src_mask=False,
tgt_mask=False,
use_src_mask=False,
use_tgt_mask=False,
):
"""
Create regridder for conversions between source mesh and target grid.
Expand Down Expand Up @@ -360,12 +373,12 @@ def __init__(
given to ESMF for calculation. If resolution is set, target_grid_cube
must have strictly increasing bounds (bounds may be transposed plus or
minus 360 degrees to make the bounds strictly increasing).
src_mask : bool, array, default=False
use_src_mask : :obj:`~numpy.typing.ArrayLike`, 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_mesh_cube``. If False, no mask will be taken and all points will
be used in weights calculation.
tgt_mask : bool, array, default=False
use_tgt_mask : :obj:`~numpy.typing.ArrayLike`, 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 ``target_grid_cube``. If False, no mask will be taken and all points
Expand Down Expand Up @@ -404,25 +417,17 @@ def __init__(
)
self.resolution = resolution

if src_mask is True:
src_mask = _get_mask(src_mesh_cube)
elif src_mask is False:
src_mask = None
self.src_mask = src_mask
if tgt_mask is True:
tgt_mask = _get_mask(target_grid_cube)
elif tgt_mask is False:
tgt_mask = None
self.tgt_mask = tgt_mask
self.src_mask = _get_mask(src_mesh_cube, use_src_mask)
self.tgt_mask = _get_mask(target_grid_cube, use_tgt_mask)

partial_regrid_info = _regrid_unstructured_to_rectilinear__prepare(
src_mesh_cube,
target_grid_cube,
method=self.method,
precomputed_weights=precomputed_weights,
resolution=resolution,
src_mask=src_mask,
tgt_mask=tgt_mask,
src_mask=self.src_mask,
tgt_mask=self.tgt_mask,
)

# Record source mesh.
Expand Down Expand Up @@ -616,8 +621,8 @@ def regrid_rectilinear_to_unstructured(
mdtol=0,
method="conservative",
resolution=None,
src_mask=None,
tgt_mask=None,
use_src_mask=False,
use_tgt_mask=False,
):
r"""
Regrid rectilinear :class:`~iris.cube.Cube` onto unstructured mesh.
Expand Down Expand Up @@ -666,13 +671,26 @@ def regrid_rectilinear_to_unstructured(
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`, 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_mesh_cube``. If False, no mask will be taken and all points will
be used in weights calculation.
use_tgt_mask : :obj:`~numpy.typing.ArrayLike`, 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 ``target_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.
"""
src_mask = _get_mask(src_cube, use_src_mask)
tgt_mask = _get_mask(mesh_cube, use_tgt_mask)

regrid_info = _regrid_rectilinear_to_unstructured__prepare(
src_cube,
mesh_cube,
Expand All @@ -696,8 +714,8 @@ def __init__(
method="conservative",
precomputed_weights=None,
resolution=None,
src_mask=False,
tgt_mask=False,
use_src_mask=False,
use_tgt_mask=False,
):
"""
Create regridder for conversions between source grid and target mesh.
Expand Down Expand Up @@ -728,12 +746,12 @@ def __init__(
given to ESMF for calculation. If resolution is set, src_grid_cube
must have strictly increasing bounds (bounds may be transposed plus or
minus 360 degrees to make the bounds strictly increasing).
src_mask : :obj:`~numpy.typing.ArrayLike`, bool, default=False
use_src_mask : :obj:`~numpy.typing.ArrayLike`, 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_grid_cube``. If False, no mask will be taken and all points will
be used in weights calculation.
tgt_mask : :obj:`~numpy.typing.ArrayLike`, bool, default=False
use_tgt_mask : :obj:`~numpy.typing.ArrayLike`, 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 ``target_mesh_cube``. If False, no mask will be taken and all points
Expand Down Expand Up @@ -766,25 +784,17 @@ def __init__(
self.method = method
self.resolution = resolution

if src_mask is True:
src_mask = _get_mask(src_grid_cube)
elif src_mask is False:
src_mask = None
self.src_mask = src_mask
if tgt_mask is True:
tgt_mask = _get_mask(target_mesh_cube)
elif tgt_mask is False:
tgt_mask = None
self.tgt_mask = tgt_mask
self.src_mask = _get_mask(src_grid_cube, use_src_mask)
self.tgt_mask = _get_mask(target_mesh_cube, use_tgt_mask)

partial_regrid_info = _regrid_rectilinear_to_unstructured__prepare(
src_grid_cube,
target_mesh_cube,
method=self.method,
precomputed_weights=precomputed_weights,
resolution=self.resolution,
src_mask=src_mask,
tgt_mask=tgt_mask,
src_mask=self.src_mask,
tgt_mask=self.tgt_mask,
)

# Store regrid info.
Expand Down
118 changes: 62 additions & 56 deletions esmf_regrid/schemes.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,33 +28,41 @@ def _get_coord(cube, axis):
return coord


def _get_mask(cube):
src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y"))
def _get_mask(cube, use_mask=True):
if use_mask == False:
return None
elif use_mask == True:

horizontal_dims = set(cube.coord_dims(src_x)) | set(cube.coord_dims(src_y))
other_dims = tuple(set(range(cube.ndim)) - horizontal_dims)
src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y"))

if cube.coord_dims(src_x) == cube.coord_dims(src_y):
slices = cube.slices([src_x])
else:
slices = cube.slices([src_x, src_y])
data = next(slices).data
if np.ma.is_masked(data):
# Check that the mask is constant along all other dimensions.
full_mask = np.ma.getmaskarray(cube.data)
if not np.array_equal(
np.all(full_mask, axis=other_dims), np.any(full_mask, axis=other_dims)
):
raise ValueError(
"The mask derived from the cube is not constant over non-horizontal dimensions."
"Consider passing in an explicit mask instead."
)
mask = np.ma.getmaskarray(data)
if cube.coord_dims(src_x) != cube.coord_dims(src_y):
mask = mask.T
horizontal_dims = set(cube.coord_dims(src_x)) | set(cube.coord_dims(src_y))
other_dims = tuple(set(range(cube.ndim)) - horizontal_dims)

# Find a representative slice of data that spans both horizontal coords.
if cube.coord_dims(src_x) == cube.coord_dims(src_y):
slices = cube.slices([src_x])
else:
slices = cube.slices([src_x, src_y])
data = next(slices).data
if np.ma.is_masked(data):
# Check that the mask is constant along all other dimensions.
full_mask = np.ma.getmaskarray(cube.data)
if not np.array_equal(
np.all(full_mask, axis=other_dims), np.any(full_mask, axis=other_dims)
):
raise ValueError(
"The mask derived from the cube is not constant over non-horizontal dimensions."
"Consider passing in an explicit mask instead."
)
mask = np.ma.getmaskarray(data)
# Due to structural reasons, the mask should be transposed for curvilinear grids.
if cube.coord_dims(src_x) != cube.coord_dims(src_y):
mask = mask.T
else:
mask = None
return mask
else:
mask = None
return mask
return use_mask


def _contiguous_masked(bounds, mask):
Expand Down Expand Up @@ -404,7 +412,7 @@ class ESMFAreaWeighted:
:mod:`ESMF` to be able to handle grids in different coordinate systems.
"""

def __init__(self, mdtol=0, src_mask=False, tgt_mask=False):
def __init__(self, mdtol=0, use_src_mask=False, use_tgt_mask=False):
"""
Area-weighted scheme for regridding between rectilinear grids.
Expand All @@ -418,10 +426,10 @@ def __init__(self, mdtol=0, src_mask=False, tgt_mask=False):
data is tolerated while ``mdtol=1`` will mean the resulting element
will be masked if and only if all the overlapping elements of the
source grid are masked.
src_mask : bool, default=False
use_src_mask : bool, default=False
If True, derive a mask from source cube which will tell :mod:`ESMF`
which points to ignore.
tgt_mask : bool, default=False
use_tgt_mask : bool, default=False
If True, derive a mask from target cube which will tell :mod:`ESMF`
which points to ignore.
Expand All @@ -430,14 +438,14 @@ def __init__(self, mdtol=0, src_mask=False, tgt_mask=False):
msg = "Value for mdtol must be in range 0 - 1, got {}."
raise ValueError(msg.format(mdtol))
self.mdtol = mdtol
self.src_mask = src_mask
self.tgt_mask = tgt_mask
self.use_src_mask = use_src_mask
self.use_tgt_mask = use_tgt_mask

def __repr__(self):
"""Return a representation of the class."""
return "ESMFAreaWeighted(mdtol={})".format(self.mdtol)

def regridder(self, src_grid, tgt_grid, src_mask=None, tgt_mask=None):
def regridder(self, src_grid, tgt_grid, use_src_mask=None, use_tgt_mask=None):
"""
Create regridder to perform regridding from ``src_grid`` to ``tgt_grid``.
Expand All @@ -447,10 +455,10 @@ def regridder(self, src_grid, tgt_grid, src_mask=None, tgt_mask=None):
The :class:`~iris.cube.Cube` defining the source grid.
tgt_grid : :class:`iris.cube.Cube`
The :class:`~iris.cube.Cube` defining the target grid.
src_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional
use_src_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional
Array describing which elements :mod:`ESMF` will ignore on the src_grid.
If True, the mask will be derived from src_grid.
tgt_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional
use_tgt_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional
Array describing which elements :mod:`ESMF` will ignore on the tgt_grid.
If True, the mask will be derived from tgt_grid.
Expand All @@ -462,23 +470,25 @@ def regridder(self, src_grid, tgt_grid, src_mask=None, tgt_mask=None):
grid as ``src_grid`` that is to be regridded to the grid of
``tgt_grid``.
"""
if src_mask is None:
src_mask = self.src_mask
if tgt_mask is None:
tgt_mask = self.tgt_mask
if use_src_mask is None:
use_src_mask = self.use_src_mask
if use_tgt_mask is None:
use_tgt_mask = self.use_tgt_mask
return ESMFAreaWeightedRegridder(
src_grid,
tgt_grid,
mdtol=self.mdtol,
src_mask=src_mask,
tgt_mask=tgt_mask,
use_src_mask=use_src_mask,
use_tgt_mask=use_tgt_mask,
)


class ESMFAreaWeightedRegridder:
r"""Regridder class for unstructured to rectilinear :class:`~iris.cube.Cube`\\ s."""

def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=False, tgt_mask=False):
def __init__(
self, src_grid, tgt_grid, mdtol=0, use_src_mask=False, use_tgt_mask=False
):
"""
Create regridder for conversions between ``src_grid`` and ``tgt_grid``.
Expand All @@ -494,38 +504,34 @@ def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=False, tgt_mask=False):
exceeds ``mdtol``. ``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 contributing elements of data are masked.
src_mask : :obj:`~numpy.typing.ArrayLike`, bool, default=False
Array describing which elements :mod:`ESMF` will ignore on the src_grid.
If True, the mask will be derived from src_grid.
tgt_mask : :obj:`~numpy.typing.ArrayLike`, bool, default=False
Array describing which elements :mod:`ESMF` will ignore on the tgt_grid.
If True, the mask will be derived from tgt_grid.
use_src_mask : :obj:`~numpy.typing.ArrayLike`, 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_grid``. If False, no mask will be taken and all points will
be used in weights calculation.
use_tgt_mask : :obj:`~numpy.typing.ArrayLike`, 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 ``tgt_grid``. If False, no mask will be taken and all points will
be used in weights calculation.
"""
if not (0 <= mdtol <= 1):
msg = "Value for mdtol must be in range 0 - 1, got {}."
raise ValueError(msg.format(mdtol))
self.mdtol = mdtol

if src_mask is True:
src_mask = _get_mask(src_grid)
elif src_mask is False:
src_mask = None
if tgt_mask is True:
tgt_mask = _get_mask(tgt_grid)
elif tgt_mask is False:
tgt_mask = None
self.src_mask = _get_mask(src_grid, use_src_mask)
self.tgt_mask = _get_mask(tgt_grid, use_tgt_mask)

regrid_info = _regrid_rectilinear_to_rectilinear__prepare(
src_grid, tgt_grid, src_mask=src_mask, tgt_mask=tgt_mask
src_grid, tgt_grid, src_mask=self.src_mask, tgt_mask=self.tgt_mask
)

# Store regrid info.
self.grid_x = regrid_info.x_coord
self.grid_y = regrid_info.y_coord
self.regridder = regrid_info.regridder
self.src_mask = src_mask
self.tgt_mask = tgt_mask

# Record the source grid.
self.src_grid = (_get_coord(src_grid, "x"), _get_coord(src_grid, "y"))
Expand Down

0 comments on commit 52ddfc2

Please sign in to comment.