From 0b1cb1d7d37d1bf2947045d876145adf5577d947 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 20 Oct 2022 16:06:35 +0100 Subject: [PATCH 01/30] support building regridders with masks --- esmf_regrid/_esmf_sdo.py | 12 +++++- esmf_regrid/esmf_regridder.py | 2 + esmf_regrid/schemes.py | 35 +++++++++++++----- .../schemes/test_ESMFAreaWeightedRegridder.py | 37 +++++++++++++++++++ 4 files changed, 75 insertions(+), 11 deletions(-) diff --git a/esmf_regrid/_esmf_sdo.py b/esmf_regrid/_esmf_sdo.py index 8c93f166..cb929ec1 100644 --- a/esmf_regrid/_esmf_sdo.py +++ b/esmf_regrid/_esmf_sdo.py @@ -17,10 +17,11 @@ class SDO(ABC): objects supported by ESMPy, Grids, Meshes, and LocStreams. """ - def __init__(self, shape, index_offset, field_kwargs): + def __init__(self, shape, index_offset, field_kwargs, mask=None): self._shape = shape self._index_offset = index_offset self._field_kwargs = field_kwargs + self._mask = mask @abstractmethod def _make_esmf_sdo(self): @@ -112,6 +113,7 @@ def __init__( circular=False, areas=None, center=False, + mask=None, ): """ Create a :class:`GridInfo` object describing the grid. @@ -196,6 +198,7 @@ def __init__( shape=shape, index_offset=1, field_kwargs={"staggerloc": esmpy.StaggerLoc.CENTER}, + mask=mask, ) def _as_esmf_info(self): @@ -283,6 +286,13 @@ def _make_esmf_sdo(self): grid_center_y = grid.get_coords(1, staggerloc=esmpy.StaggerLoc.CENTER) grid_center_y[:] = truecenterlats + if self._mask is not None: + grid.add_item(esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER) + grid_mask = grid.get_item( + esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER + ) + grid_mask[:] = self._mask + if areas is not None: grid.add_item(esmpy.GridItem.AREA, staggerloc=esmpy.StaggerLoc.CENTER) grid_areas = grid.get_item( diff --git a/esmf_regrid/esmf_regridder.py b/esmf_regrid/esmf_regridder.py index adf5a544..0b1afb28 100644 --- a/esmf_regrid/esmf_regridder.py +++ b/esmf_regrid/esmf_regridder.py @@ -24,6 +24,8 @@ def _get_regrid_weights_dict(src_field, tgt_field, regrid_method): # Choosing the norm_type DSTAREA allows for mdtol type operations # to be performed using the weights information later on. norm_type=esmpy.NormType.DSTAREA, + src_mask_values=np.array([True]), + dst_mask_values=np.array([True]), factors=True, ) # Without specifying deep_copy=true, the information in weights_dict diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 94b064a6..4eeb1af1 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -28,7 +28,7 @@ def _get_coord(cube, axis): return coord -def _cube_to_GridInfo(cube, center=False, resolution=None): +def _cube_to_GridInfo(cube, center=False, resolution=None, mask=None): # This is a simplified version of an equivalent function/method in PR #26. # It is anticipated that this function will be replaced by the one in PR #26. # @@ -51,10 +51,12 @@ def _cube_to_GridInfo(cube, center=False, resolution=None): # TODO: perform checks on lat/lon. elif londim == 2: assert cube.coord_dims(lon) == cube.coord_dims(lat) - assert lon.is_contiguous() - assert lat.is_contiguous() + if mask is None: + assert lon.is_contiguous() + assert lat.is_contiguous() # 2D coords must be AuxCoords, which do not have a circular attribute. circular = False + # TODO: This should be replaced by another method when there is a mask. lon_bound_array = lon.contiguous_bounds() lon_bound_array = lon.units.convert(lon_bound_array, Unit("degrees")) lat_bound_array = lat.contiguous_bounds() @@ -68,6 +70,7 @@ def _cube_to_GridInfo(cube, center=False, resolution=None): crs=crs, circular=circular, center=center, + mask=mask, ) else: grid_info = RefinedGridInfo( @@ -189,7 +192,9 @@ def copy_coords(src_coords, add_method): ) -def _regrid_rectilinear_to_rectilinear__prepare(src_grid_cube, tgt_grid_cube): +def _regrid_rectilinear_to_rectilinear__prepare( + src_grid_cube, tgt_grid_cube, src_mask=None, tgt_mask=None +): tgt_x = _get_coord(tgt_grid_cube, "x") tgt_y = _get_coord(tgt_grid_cube, "y") src_x = _get_coord(src_grid_cube, "x") @@ -201,8 +206,8 @@ def _regrid_rectilinear_to_rectilinear__prepare(src_grid_cube, tgt_grid_cube): else: grid_y_dim, grid_x_dim = src_grid_cube.coord_dims(src_x) - srcinfo = _cube_to_GridInfo(src_grid_cube) - tgtinfo = _cube_to_GridInfo(tgt_grid_cube) + srcinfo = _cube_to_GridInfo(src_grid_cube, mask=src_mask) + tgtinfo = _cube_to_GridInfo(tgt_grid_cube, mask=tgt_mask) regridder = Regridder(srcinfo, tgtinfo) @@ -301,7 +306,7 @@ class ESMFAreaWeighted: :mod:`esmpy` to be able to handle grids in different coordinate systems. """ - def __init__(self, mdtol=0): + def __init__(self, mdtol=0, src_mask=None, tgt_mask=None): """ Area-weighted scheme for regridding between rectilinear grids. @@ -321,6 +326,8 @@ def __init__(self, mdtol=0): 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 def __repr__(self): """Return a representation of the class.""" @@ -345,13 +352,19 @@ def regridder(self, src_grid, tgt_grid): grid as ``src_grid`` that is to be regridded to the grid of ``tgt_grid``. """ - return ESMFAreaWeightedRegridder(src_grid, tgt_grid, mdtol=self.mdtol) + return ESMFAreaWeightedRegridder( + src_grid, + tgt_grid, + mdtol=self.mdtol, + src_mask=self.src_mask, + tgt_mask=self.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): + def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=None, tgt_mask=None): """ Create regridder for conversions between ``src_grid`` and ``tgt_grid``. @@ -374,7 +387,9 @@ def __init__(self, src_grid, tgt_grid, mdtol=0): raise ValueError(msg.format(mdtol)) self.mdtol = mdtol - regrid_info = _regrid_rectilinear_to_rectilinear__prepare(src_grid, tgt_grid) + regrid_info = _regrid_rectilinear_to_rectilinear__prepare( + src_grid, tgt_grid, src_mask=src_mask, tgt_mask=tgt_mask + ) # Store regrid info. self.grid_x = regrid_info.x_coord diff --git a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py index 1cfcb2e1..d2a85f29 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py @@ -229,3 +229,40 @@ def extract_weights(regridder): curv_to_curv_rad, ]: assert np.allclose(extract_weights(grid_to_grid), extract_weights(regridder)) + + +def test_masks(): + """ + Test initialisation of :func:`esmf_regrid.schemes.ESMFAreaWeightedRegridder`. + + Checks that the `src_mask` and `tgt_mask` keywords work properly. + """ + # TODO: check that this handles discontiguities appropriately. + src = _curvilinear_cube(7, 6, [-180, 180], [-90, 90]) + tgt = _curvilinear_cube(6, 7, [-180, 180], [-90, 90]) + + src_mask = np.zeros([6, 7], dtype=bool) + src_mask[0, 0] = True + + tgt_mask = np.zeros([7, 6], dtype=bool) + tgt_mask[0, 0] = True + + rg_src_masked = ESMFAreaWeightedRegridder(src, tgt, src_mask=src_mask) + rg_tgt_masked = ESMFAreaWeightedRegridder(src, tgt, tgt_mask=tgt_mask) + rg_unmasked = ESMFAreaWeightedRegridder(src, tgt) + + weights_src_masked = rg_src_masked.regridder.weight_matrix + weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix + weights_unmasked = rg_unmasked.regridder.weight_matrix + + # Check there are no weights associated with the masked point. + assert weights_src_masked[:, 0].nnz == 0 + assert weights_tgt_masked[0].nnz == 0 + + # Check all other weights are correct. + assert np.allclose( + weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() + ) + assert np.allclose( + weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense() + ) From 6163ff83842137c3b86c8be264071847ba7b8106 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 20 Oct 2022 15:18:33 +0000 Subject: [PATCH 02/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../tests/unit/schemes/test_ESMFAreaWeightedRegridder.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py index d2a85f29..cf282e1b 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py @@ -263,6 +263,4 @@ def test_masks(): assert np.allclose( weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() ) - assert np.allclose( - weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense() - ) + assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) From 6dda5c2cbd213941141f19b7899d180090d3471d Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 20 Oct 2022 21:33:53 +0100 Subject: [PATCH 03/30] handle contiguous bounds for masked data --- esmf_regrid/schemes.py | 33 +++++++++++++++++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 4eeb1af1..cbdf5543 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -28,6 +28,31 @@ def _get_coord(cube, axis): return coord +def _contiguous_masked(bounds, mask): + mask = np.array(mask, dtype=int) + + blank_filter = np.zeros([size + 1 for size in mask.shape], dtype=int) + filter_0 = blank_filter.copy() + filter_0[:-1, :-1] = 1 - mask + filter_0[0, 0] = 1 + filter_1 = blank_filter.copy() + filter_1[:-1, 1:] = 1 - mask + filter_1[0, 1:] = 1 + filter_1 = filter_1 * (1 - filter_0) + filter_3 = blank_filter.copy() + filter_3[1:, :-1] = 1 - mask + filter_3[1:, 0] = 1 + filter_3 = filter_3 * (1 - filter_0 - filter_1) + filter_2 = 1 - filter_0 - filter_1 - filter_3 + + contiguous_bounds = blank_filter.astype(bounds.dtype) + contiguous_bounds[:-1, :-1] += bounds[:, :, 0] * filter_0[:-1, :-1] + contiguous_bounds[:-1, 1:] += bounds[:, :, 1] * filter_1[:-1, 1:] + contiguous_bounds[1:, 1:] += bounds[:, :, 2] * filter_2[1:, 1:] + contiguous_bounds[1:, :-1] += bounds[:, :, 3] * filter_3[1:, :-1] + return contiguous_bounds + + def _cube_to_GridInfo(cube, center=False, resolution=None, mask=None): # This is a simplified version of an equivalent function/method in PR #26. # It is anticipated that this function will be replaced by the one in PR #26. @@ -57,9 +82,13 @@ def _cube_to_GridInfo(cube, center=False, resolution=None, mask=None): # 2D coords must be AuxCoords, which do not have a circular attribute. circular = False # TODO: This should be replaced by another method when there is a mask. - lon_bound_array = lon.contiguous_bounds() + if mask is None: + lon_bound_array = lon.contiguous_bounds() + lat_bound_array = lat.contiguous_bounds() + else: + lon_bound_array = _contiguous_masked(lon.bounds, mask) + lat_bound_array = _contiguous_masked(lat.bounds, mask) lon_bound_array = lon.units.convert(lon_bound_array, Unit("degrees")) - lat_bound_array = lat.contiguous_bounds() lat_bound_array = lat.units.convert(lat_bound_array, Unit("degrees")) if resolution is None: grid_info = GridInfo( From 73a30aacd90b8d5dc24c422fa60ff6f02203b58c Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 21 Oct 2022 10:31:34 +0100 Subject: [PATCH 04/30] improve documentation --- esmf_regrid/schemes.py | 50 +++++++++++++++++++++++++++++++++++++----- 1 file changed, 45 insertions(+), 5 deletions(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index cbdf5543..b2c04af2 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -29,23 +29,63 @@ def _get_coord(cube, axis): def _contiguous_masked(bounds, mask): + """ + Returns the (N+1, M+1) bound values for bounds of 2D coordinate of shape (N,M). + + Assumes the bounds are contiguous outside of the mask. The returned bounds + fully describe the unmasked bounds when this is the case. This function is + designed to replicate the behaviour of coord.contiguous_bounds() for unmasked + bounds. + + For example, suppose we have a masked array: + + # 0 0 0 0 + # 0 - - 0 + # 0 - - 0 + # 0 0 0 0 + + The indices of the bounds which the final array will derive from are as follows: + + # (0, 0, 0) (0, 1, 0) (0, 2, 0) (0, 3, 0) (0, 3, 1) + # (1, 0, 0) (1, 0, 1) (0, 2, 3) (1, 3, 0) (1, 3, 1) + # (2, 0, 0) (2, 0, 1) (1, 1, 2) (2, 3, 0) (2, 3, 1) + # (3, 0, 0) (3, 1, 0) (3, 2, 0) (3, 3, 0) (3, 3, 1) + # (3, 0, 3) (3, 1, 3) (3, 3, 3) (3, 3, 3) (3, 3, 2) + + Note that only the center bound derives from a masked cell. + """ mask = np.array(mask, dtype=int) - blank_filter = np.zeros([size + 1 for size in mask.shape], dtype=int) - filter_0 = blank_filter.copy() + # Construct a blank template array in the shape of the output. + blank_template = np.zeros([size + 1 for size in mask.shape], dtype=int) + + # Define the bounds which will derive from the bounds in index 0. + filter_0 = blank_template.copy() filter_0[:-1, :-1] = 1 - mask + # Ensure the corner points are covered appropriately. filter_0[0, 0] = 1 - filter_1 = blank_filter.copy() + + # Define the bounds which will derive from the bounds in index 1. + filter_1 = blank_template.copy() filter_1[:-1, 1:] = 1 - mask + # Ensure the corner and edge bounds are covered appropriately. filter_1[0, 1:] = 1 + # Do not cover any points already covered. filter_1 = filter_1 * (1 - filter_0) - filter_3 = blank_filter.copy() + + # Define the bounds which will derive from the bounds in index 3. + filter_3 = blank_template.copy() filter_3[1:, :-1] = 1 - mask + # Ensure the corner and edge bounds are covered appropriately. filter_3[1:, 0] = 1 + # Do not cover any points already covered. filter_3 = filter_3 * (1 - filter_0 - filter_1) + + # The bounds deriving from index 2 will be all those not already covered. filter_2 = 1 - filter_0 - filter_1 - filter_3 - contiguous_bounds = blank_filter.astype(bounds.dtype) + # Copy the bounds information over into the appropriate places. + contiguous_bounds = blank_template.astype(bounds.dtype) contiguous_bounds[:-1, :-1] += bounds[:, :, 0] * filter_0[:-1, :-1] contiguous_bounds[:-1, 1:] += bounds[:, :, 1] * filter_1[:-1, 1:] contiguous_bounds[1:, 1:] += bounds[:, :, 2] * filter_2[1:, 1:] From 3035301ea523898818b160e076fd3d1be2b0bd75 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 21 Oct 2022 11:03:25 +0100 Subject: [PATCH 05/30] test for discontiguities --- esmf_regrid/schemes.py | 2 +- .../tests/unit/schemes/test_ESMFAreaWeightedRegridder.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index b2c04af2..14236394 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -30,7 +30,7 @@ def _get_coord(cube, axis): def _contiguous_masked(bounds, mask): """ - Returns the (N+1, M+1) bound values for bounds of 2D coordinate of shape (N,M). + Return the (N+1, M+1) bound values for bounds of 2D coordinate of shape (N,M). Assumes the bounds are contiguous outside of the mask. The returned bounds fully describe the unmasked bounds when this is the case. This function is diff --git a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py index cf282e1b..748d87a5 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py @@ -237,15 +237,19 @@ def test_masks(): Checks that the `src_mask` and `tgt_mask` keywords work properly. """ - # TODO: check that this handles discontiguities appropriately. src = _curvilinear_cube(7, 6, [-180, 180], [-90, 90]) tgt = _curvilinear_cube(6, 7, [-180, 180], [-90, 90]) + # Make src and tgt discontiguous at (0, 0) src_mask = np.zeros([6, 7], dtype=bool) src_mask[0, 0] = True + src.coord("latitude").bounds[0, 0] = 0 + src.coord("longitude").bounds[0, 0] = 0 tgt_mask = np.zeros([7, 6], dtype=bool) tgt_mask[0, 0] = True + tgt.coord("latitude").bounds[0, 0] = 0 + tgt.coord("longitude").bounds[0, 0] = 0 rg_src_masked = ESMFAreaWeightedRegridder(src, tgt, src_mask=src_mask) rg_tgt_masked = ESMFAreaWeightedRegridder(src, tgt, tgt_mask=tgt_mask) From 07daffd0f463a93f6f6b23419de274eb7a6b3907 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 21 Oct 2022 14:35:58 +0100 Subject: [PATCH 06/30] support mask fetching --- esmf_regrid/_esmf_sdo.py | 13 +++-- .../experimental/unstructured_regrid.py | 7 +++ esmf_regrid/schemes.py | 50 ++++++++++++++++--- .../schemes/test_ESMFAreaWeightedRegridder.py | 16 +++--- 4 files changed, 70 insertions(+), 16 deletions(-) diff --git a/esmf_regrid/_esmf_sdo.py b/esmf_regrid/_esmf_sdo.py index cb929ec1..d1d12728 100644 --- a/esmf_regrid/_esmf_sdo.py +++ b/esmf_regrid/_esmf_sdo.py @@ -63,6 +63,11 @@ def index_offset(self): """Return the index offset.""" return self._index_offset + @property + def mask(self): + """Return the index offset.""" + return self._mask + def _array_to_matrix(self, array): """ Reshape data to a form that is compatible with weight matrices. @@ -112,8 +117,8 @@ def __init__( crs=None, circular=False, areas=None, - center=False, mask=None, + center=False, ): """ Create a :class:`GridInfo` object describing the grid. @@ -142,6 +147,8 @@ def __init__( Array describing the areas associated with each face. If ``None``, then :mod:`esmpy` will use its own calculated areas. + mask: :obj:`~numpy.typing.ArrayLike`, optional + Array describing which elements :mod:`esmpy` will ignore. center : bool, default=False Describes if the center points of the grid cells are used in regridding calculations. @@ -286,12 +293,12 @@ def _make_esmf_sdo(self): grid_center_y = grid.get_coords(1, staggerloc=esmpy.StaggerLoc.CENTER) grid_center_y[:] = truecenterlats - if self._mask is not None: + if self.mask is not None: grid.add_item(esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER) grid_mask = grid.get_item( esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER ) - grid_mask[:] = self._mask + grid_mask[:] = self.mask if areas is not None: grid.add_item(esmpy.GridItem.AREA, staggerloc=esmpy.StaggerLoc.CENTER) diff --git a/esmf_regrid/experimental/unstructured_regrid.py b/esmf_regrid/experimental/unstructured_regrid.py index a74e3ac1..e03d3d46 100644 --- a/esmf_regrid/experimental/unstructured_regrid.py +++ b/esmf_regrid/experimental/unstructured_regrid.py @@ -26,6 +26,7 @@ def __init__( node_start_index, elem_start_index=0, areas=None, + mask=None, elem_coords=None, location="face", ): @@ -56,6 +57,8 @@ def __init__( areas: :obj:`~numpy.typing.ArrayLike`, optional Array describing the areas associated with each face. If ``None``, then :mod:`esmpy` will use its own calculated areas. + mask: :obj:`~numpy.typing.ArrayLike`, optional + Array describing which elements :mod:`esmpy` will ignore. elem_coords : :obj:`~numpy.typing.ArrayLike`, optional An ``Nx2`` array describing the location of the face centers of the mesh. ``elem_coords[:,0]`` describes the longitudes in degrees and @@ -84,6 +87,7 @@ def __init__( shape=shape, index_offset=self.esi, field_kwargs=field_kwargs, + mask=mask, ) def _as_esmf_info(self): @@ -108,6 +112,7 @@ def _as_esmf_info(self): elemId, elemType, elemConn, + self.mask, self.areas, elemCoord, ) @@ -124,6 +129,7 @@ def _make_esmf_sdo(self): elemId, elemType, elemConn, + mask, areas, elemCoord, ) = info @@ -139,6 +145,7 @@ def _make_esmf_sdo(self): elemId, elemType, elemConn, + element_mask=mask, element_area=areas, element_coords=elemCoord, ) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 14236394..273e7aaa 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -28,6 +28,23 @@ def _get_coord(cube, axis): return coord +def _get_mask(cube): + 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): + mask = data.mask + if cube.coord_dims(src_x) != cube.coord_dims(src_y): + mask = mask.T + else: + mask = None + return mask + + def _contiguous_masked(bounds, mask): """ Return the (N+1, M+1) bound values for bounds of 2D coordinate of shape (N,M). @@ -113,21 +130,21 @@ def _cube_to_GridInfo(cube, center=False, resolution=None, mask=None): assert isinstance(lon, iris.coords.DimCoord) assert isinstance(lat, iris.coords.DimCoord) circular = lon.circular + lon_bound_array = lon.contiguous_bounds() + lat_bound_array = lat.contiguous_bounds() # TODO: perform checks on lat/lon. elif londim == 2: assert cube.coord_dims(lon) == cube.coord_dims(lat) if mask is None: assert lon.is_contiguous() assert lat.is_contiguous() + lon_bound_array = lon.contiguous_bounds() + lat_bound_array = lat.contiguous_bounds() + else: + lon_bound_array = _contiguous_masked(lon.bounds, mask) + lat_bound_array = _contiguous_masked(lat.bounds, mask) # 2D coords must be AuxCoords, which do not have a circular attribute. circular = False - # TODO: This should be replaced by another method when there is a mask. - if mask is None: - lon_bound_array = lon.contiguous_bounds() - lat_bound_array = lat.contiguous_bounds() - else: - lon_bound_array = _contiguous_masked(lon.bounds, mask) - lat_bound_array = _contiguous_masked(lat.bounds, mask) lon_bound_array = lon.units.convert(lon_bound_array, Unit("degrees")) lat_bound_array = lat.units.convert(lat_bound_array, Unit("degrees")) if resolution is None: @@ -389,6 +406,12 @@ def __init__(self, mdtol=0, src_mask=None, tgt_mask=None): 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 : :obj:`~numpy.typing.ArrayLike`, bool, optional + Array describing which elements :mod:`ESMF` will ignore on the source grid. + If True, the mask will be derived from the cube. + tgt_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional + Array describing which elements :mod:`ESMF` will ignore on the target grid. + If True, the mask will be derived from the cube. """ if not (0 <= mdtol <= 1): @@ -449,6 +472,12 @@ def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=None, tgt_mask=None): 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, optional + Array describing which elements :mod:`ESMF` will ignore on the source grid. + If True, the mask will be derived from the cube. + tgt_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional + Array describing which elements :mod:`ESMF` will ignore on the target grid. + If True, the mask will be derived from the cube. """ if not (0 <= mdtol <= 1): @@ -456,6 +485,11 @@ def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=None, tgt_mask=None): raise ValueError(msg.format(mdtol)) self.mdtol = mdtol + if src_mask is True: + src_mask = _get_mask(src_grid) + if tgt_mask is True: + tgt_mask = _get_mask(tgt_grid) + regrid_info = _regrid_rectilinear_to_rectilinear__prepare( src_grid, tgt_grid, src_mask=src_mask, tgt_mask=tgt_mask ) @@ -464,6 +498,8 @@ def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=None, tgt_mask=None): 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")) diff --git a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py index 748d87a5..b79912df 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py @@ -243,16 +243,20 @@ def test_masks(): # Make src and tgt discontiguous at (0, 0) src_mask = np.zeros([6, 7], dtype=bool) src_mask[0, 0] = True - src.coord("latitude").bounds[0, 0] = 0 - src.coord("longitude").bounds[0, 0] = 0 + src.data = np.ma.array(src.data, mask=src_mask) + src_discontiguous = src.copy() + src_discontiguous.coord("latitude").bounds[0, 0] = 0 + src_discontiguous.coord("longitude").bounds[0, 0] = 0 tgt_mask = np.zeros([7, 6], dtype=bool) tgt_mask[0, 0] = True - tgt.coord("latitude").bounds[0, 0] = 0 - tgt.coord("longitude").bounds[0, 0] = 0 + tgt.data = np.ma.array(tgt.data, mask=tgt_mask) + tgt_discontiguous = tgt.copy() + tgt_discontiguous.coord("latitude").bounds[0, 0] = 0 + tgt_discontiguous.coord("longitude").bounds[0, 0] = 0 - rg_src_masked = ESMFAreaWeightedRegridder(src, tgt, src_mask=src_mask) - rg_tgt_masked = ESMFAreaWeightedRegridder(src, tgt, tgt_mask=tgt_mask) + rg_src_masked = ESMFAreaWeightedRegridder(src_discontiguous, tgt, src_mask=True) + rg_tgt_masked = ESMFAreaWeightedRegridder(src, tgt_discontiguous, tgt_mask=True) rg_unmasked = ESMFAreaWeightedRegridder(src, tgt) weights_src_masked = rg_src_masked.regridder.weight_matrix From 0c76b427e286b959b98181015302a08883432a02 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 21 Oct 2022 15:16:14 +0100 Subject: [PATCH 07/30] clarify boolean arguments --- esmf_regrid/schemes.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 273e7aaa..68671d18 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -392,7 +392,7 @@ class ESMFAreaWeighted: :mod:`esmpy` to be able to handle grids in different coordinate systems. """ - def __init__(self, mdtol=0, src_mask=None, tgt_mask=None): + def __init__(self, mdtol=0, src_mask=False, tgt_mask=False): """ Area-weighted scheme for regridding between rectilinear grids. @@ -406,12 +406,12 @@ def __init__(self, mdtol=0, src_mask=None, tgt_mask=None): 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 : :obj:`~numpy.typing.ArrayLike`, bool, optional - Array describing which elements :mod:`ESMF` will ignore on the source grid. - If True, the mask will be derived from the cube. - tgt_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional - Array describing which elements :mod:`ESMF` will ignore on the target grid. - If True, the mask will be derived from the cube. + 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 + If True, derive a mask from target cube which will tell :mod:`ESMF` + which points to ignore. """ if not (0 <= mdtol <= 1): @@ -456,7 +456,7 @@ def regridder(self, src_grid, tgt_grid): 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=None, tgt_mask=None): + def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=False, tgt_mask=False): """ Create regridder for conversions between ``src_grid`` and ``tgt_grid``. @@ -473,11 +473,11 @@ def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=None, tgt_mask=None): ``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, optional - Array describing which elements :mod:`ESMF` will ignore on the source grid. - If True, the mask will be derived from the cube. + 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 - Array describing which elements :mod:`ESMF` will ignore on the target grid. - If True, the mask will be derived from the cube. + Array describing which elements :mod:`ESMF` will ignore on the tgt_grid. + If True, the mask will be derived from tgt_grid. """ if not (0 <= mdtol <= 1): @@ -487,8 +487,12 @@ def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=None, tgt_mask=None): 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 regrid_info = _regrid_rectilinear_to_rectilinear__prepare( src_grid, tgt_grid, src_mask=src_mask, tgt_mask=tgt_mask From 79ba39ce9db71456c6aa2370556610c72b0b1f9e Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 1 Mar 2023 11:13:01 +0000 Subject: [PATCH 08/30] extend mask support --- esmf_regrid/_esmf_sdo.py | 27 ++++++++- .../experimental/unstructured_scheme.py | 55 +++++++++++++++++-- esmf_regrid/schemes.py | 1 + .../test_GridToMeshESMFRegridder.py | 39 +++++++++++++ .../test_MeshToGridESMFRegridder.py | 39 +++++++++++++ 5 files changed, 152 insertions(+), 9 deletions(-) diff --git a/esmf_regrid/_esmf_sdo.py b/esmf_regrid/_esmf_sdo.py index d1d12728..af432501 100644 --- a/esmf_regrid/_esmf_sdo.py +++ b/esmf_regrid/_esmf_sdo.py @@ -43,6 +43,11 @@ def _refined_shape(self): """Return shape passed to ESMF.""" return self._shape + @property + def _refined_mask(self): + """Return mask passed to ESMF.""" + return self._mask + @property def dims(self): """Return number of dimensions.""" @@ -65,7 +70,7 @@ def index_offset(self): @property def mask(self): - """Return the index offset.""" + """Return the mask.""" return self._mask def _array_to_matrix(self, array): @@ -298,7 +303,7 @@ def _make_esmf_sdo(self): grid_mask = grid.get_item( esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER ) - grid_mask[:] = self.mask + grid_mask[:] = self._refined_mask if areas is not None: grid.add_item(esmpy.GridItem.AREA, staggerloc=esmpy.StaggerLoc.CENTER) @@ -331,6 +336,7 @@ def __init__( latbounds, resolution=3, crs=None, + mask=None, ): """ Create a :class:`RefinedGridInfo` object describing the grid. @@ -371,7 +377,7 @@ def __init__( # Create dummy lat/lon values lons = np.zeros(self.n_lons_orig) lats = np.zeros(self.n_lats_orig) - super().__init__(lons, lats, lonbounds, latbounds, crs=crs) + super().__init__(lons, lats, lonbounds, latbounds, crs=crs, mask=mask) if self.n_lats_orig == 1 and np.allclose(latbounds, [-90, 90]): self._refined_latbounds = np.array([-90, 0, 90]) @@ -403,6 +409,21 @@ def _refined_shape(self): self.n_lons_orig * self.lon_expansion, ) + @property + def _refined_mask(self): + """Return mask passed to ESMF.""" + new_mask = np.broadcast_to( + self.mask[:, np.newaxis, :, np.newaxis], + [ + self.n_lats_orig, + self.lat_expansion, + self.n_lons_orig, + self.lon_expansion, + ], + ) + new_mask = new_mask.reshape(self._refined_shape) + return new_mask + def _collapse_weights(self, is_tgt): """ Return a matrix to collapse the weight matrix. diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index f90d1823..557f77db 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -7,7 +7,7 @@ from esmf_regrid.esmf_regridder import Regridder from esmf_regrid.experimental.unstructured_regrid import MeshInfo -from esmf_regrid.schemes import _create_cube, _cube_to_GridInfo, _get_coord +from esmf_regrid.schemes import _create_cube, _cube_to_GridInfo, _get_coord, _get_mask def _map_complete_blocks(src, func, dims, out_sizes): @@ -106,7 +106,7 @@ def _map_complete_blocks(src, func, dims, out_sizes): ) -def _mesh_to_MeshInfo(mesh, location): +def _mesh_to_MeshInfo(mesh, location, mask=None): # Returns a MeshInfo object describing the mesh of the cube. assert mesh.topology_dimension == 2 if None in mesh.face_coords: @@ -119,6 +119,7 @@ def _mesh_to_MeshInfo(mesh, location): mesh.face_node_connectivity.start_index, elem_coords=elem_coords, location=location, + mask=mask, ) return meshinfo @@ -144,6 +145,8 @@ def _regrid_unstructured_to_rectilinear__prepare( method, precomputed_weights=None, resolution=None, + src_mask=None, + tgt_mask=None, ): """ First (setup) part of 'regrid_unstructured_to_rectilinear'. @@ -185,8 +188,10 @@ def _regrid_unstructured_to_rectilinear__prepare( # mesh belongs to. mesh_dim = src_mesh_cube.mesh_dim() - meshinfo = _mesh_to_MeshInfo(mesh, location) - gridinfo = _cube_to_GridInfo(target_grid_cube, center=center, resolution=resolution) + meshinfo = _mesh_to_MeshInfo(mesh, location, mask=src_mask) + gridinfo = _cube_to_GridInfo( + target_grid_cube, center=center, resolution=resolution, mask=tgt_mask + ) regridder = Regridder( meshinfo, gridinfo, method=method, precomputed_weights=precomputed_weights @@ -248,6 +253,8 @@ def regrid_unstructured_to_rectilinear( mdtol=0, method="conservative", resolution=None, + src_mask=None, + tgt_mask=None, ): r""" Regrid unstructured :class:`~iris.cube.Cube` onto rectilinear grid. @@ -304,6 +311,8 @@ def regrid_unstructured_to_rectilinear( grid_cube, method=method, resolution=resolution, + src_mask=src_mask, + tgt_mask=tgt_mask, ) result = _regrid_unstructured_to_rectilinear__perform(src_cube, regrid_info, mdtol) return result @@ -320,6 +329,8 @@ def __init__( method="conservative", precomputed_weights=None, resolution=None, + src_mask=False, + tgt_mask=False, ): """ Create regridder for conversions between source mesh and target grid. @@ -383,12 +394,23 @@ 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 + if tgt_mask is True: + tgt_mask = _get_mask(target_grid_cube) + elif tgt_mask is False: + tgt_mask = None + 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, ) # Record source mesh. @@ -473,6 +495,8 @@ def _regrid_rectilinear_to_unstructured__prepare( method, precomputed_weights=None, resolution=None, + src_mask=None, + tgt_mask=None, ): """ First (setup) part of 'regrid_rectilinear_to_unstructured'. @@ -517,8 +541,10 @@ def _regrid_rectilinear_to_unstructured__prepare( else: grid_y_dim, grid_x_dim = src_grid_cube.coord_dims(grid_x) - meshinfo = _mesh_to_MeshInfo(mesh, location) - gridinfo = _cube_to_GridInfo(src_grid_cube, center=center, resolution=resolution) + meshinfo = _mesh_to_MeshInfo(mesh, location, mask=tgt_mask) + gridinfo = _cube_to_GridInfo( + src_grid_cube, center=center, resolution=resolution, mask=src_mask + ) regridder = Regridder( gridinfo, meshinfo, method=method, precomputed_weights=precomputed_weights @@ -578,6 +604,8 @@ def regrid_rectilinear_to_unstructured( mdtol=0, method="conservative", resolution=None, + src_mask=None, + tgt_mask=None, ): r""" Regrid rectilinear :class:`~iris.cube.Cube` onto unstructured mesh. @@ -638,6 +666,8 @@ def regrid_rectilinear_to_unstructured( mesh_cube, method=method, resolution=resolution, + src_mask=src_mask, + tgt_mask=tgt_mask, ) result = _regrid_rectilinear_to_unstructured__perform(src_cube, regrid_info, mdtol) return result @@ -654,6 +684,8 @@ def __init__( method="conservative", precomputed_weights=None, resolution=None, + src_mask=False, + tgt_mask=False, ): """ Create regridder for conversions between source grid and target mesh. @@ -712,12 +744,23 @@ 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 + if tgt_mask is True: + tgt_mask = _get_mask(target_mesh_cube) + elif tgt_mask is False: + tgt_mask = None + 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, ) # Store regrid info. diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 68671d18..22d51e22 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -164,6 +164,7 @@ def _cube_to_GridInfo(cube, center=False, resolution=None, mask=None): lat_bound_array, crs=crs, resolution=resolution, + mask=mask, ) return grid_info diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py index 1290f409..ea1c8f82 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py @@ -424,3 +424,42 @@ def test_curvilinear(): result.data = expected_data assert expected_cube == result assert result_lazy == result + + +def test_masks(): + """ + Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`. + Checks that the `src_mask` and `tgt_mask` keywords work properly. + """ + src = _curvilinear_cube(7, 6, [-180, 180], [-90, 90]) + tgt = _gridlike_mesh_cube(6, 7) + + # Make src and tgt discontiguous at (0, 0) + src_mask = np.zeros([6, 7], dtype=bool) + src_mask[0, 0] = True + src.data = np.ma.array(src.data, mask=src_mask) + src_discontiguous = src.copy() + src_discontiguous.coord("latitude").bounds[0, 0] = 0 + src_discontiguous.coord("longitude").bounds[0, 0] = 0 + + tgt_mask = np.zeros([7 * 6], dtype=bool) + tgt_mask[0] = True + tgt.data = np.ma.array(tgt.data, mask=tgt_mask) + + rg_src_masked = GridToMeshESMFRegridder(src_discontiguous, tgt, src_mask=True) + rg_tgt_masked = GridToMeshESMFRegridder(src, tgt, tgt_mask=True) + rg_unmasked = GridToMeshESMFRegridder(src, tgt) + + weights_src_masked = rg_src_masked.regridder.weight_matrix + weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix + weights_unmasked = rg_unmasked.regridder.weight_matrix + + # Check there are no weights associated with the masked point. + assert weights_src_masked[:, 0].nnz == 0 + assert weights_tgt_masked[0].nnz == 0 + + # Check all other weights are correct. + assert np.allclose( + weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() + ) + assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py index 5b0d1a68..ae95c94a 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py @@ -400,3 +400,42 @@ def test_curvilinear(): result.data = expected_data assert expected_cube == result assert result_lazy == result + + +def test_masks(): + """ + Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`. + Checks that the `src_mask` and `tgt_mask` keywords work properly. + """ + src = _gridlike_mesh_cube(7, 6) + tgt = _curvilinear_cube(6, 7, [-180, 180], [-90, 90]) + + # Make src and tgt discontiguous at (0, 0) + src_mask = np.zeros([6 * 7], dtype=bool) + src_mask[0] = True + src.data = np.ma.array(src.data, mask=src_mask) + + tgt_mask = np.zeros([7, 6], dtype=bool) + tgt_mask[0, 0] = True + tgt.data = np.ma.array(tgt.data, mask=tgt_mask) + tgt_discontiguous = tgt.copy() + tgt_discontiguous.coord("latitude").bounds[0, 0] = 0 + tgt_discontiguous.coord("longitude").bounds[0, 0] = 0 + + rg_src_masked = MeshToGridESMFRegridder(src, tgt, src_mask=True) + rg_tgt_masked = MeshToGridESMFRegridder(src, tgt_discontiguous, tgt_mask=True) + rg_unmasked = MeshToGridESMFRegridder(src, tgt) + + weights_src_masked = rg_src_masked.regridder.weight_matrix + weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix + weights_unmasked = rg_unmasked.regridder.weight_matrix + + # Check there are no weights associated with the masked point. + assert weights_src_masked[:, 0].nnz == 0 + assert weights_tgt_masked[0].nnz == 0 + + # Check all other weights are correct. + assert np.allclose( + weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() + ) + assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) From de746180655cecd4aec23f6830fd7bd1b008fa81 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 1 Mar 2023 16:36:34 +0000 Subject: [PATCH 09/30] add tests --- .../benchmarks/esmf_regridder/__init__.py | 36 +++++++++++- benchmarks/benchmarks/generate_data.py | 55 +++++++++++++++++++ .../test_GridToMeshESMFRegridder.py | 39 ++++++++++++- .../test_MeshToGridESMFRegridder.py | 39 ++++++++++++- 4 files changed, 166 insertions(+), 3 deletions(-) diff --git a/benchmarks/benchmarks/esmf_regridder/__init__.py b/benchmarks/benchmarks/esmf_regridder/__init__.py index aa399e2e..c213364a 100644 --- a/benchmarks/benchmarks/esmf_regridder/__init__.py +++ b/benchmarks/benchmarks/esmf_regridder/__init__.py @@ -15,7 +15,7 @@ MeshToGridESMFRegridder, ) from esmf_regrid.schemes import ESMFAreaWeightedRegridder -from ..generate_data import _grid_cube, _gridlike_mesh_cube +from ..generate_data import _curvilinear_cube, _grid_cube, _gridlike_mesh_cube def _make_small_grid_args(): @@ -452,3 +452,37 @@ def time_save(self, _, tp, rgt): def time_load(self, _, tp, rgt): """Benchmark the loading time.""" _ = self.load_regridder(self.source_file) + + +class TimeMaskedRegridding: + """Benchmarks for :class:`~esmf_regrid.esmf_regrid.schemes.ESMFAreaWeightedRegridder`.""" + + def setup(self): + """ASV setup method.""" + src = _curvilinear_cube(250, 251, [-180, 180], [-90, 90]) + tgt = _curvilinear_cube(251, 250, [-180, 180], [-90, 90]) + + # Make src and tgt discontiguous at (0, 0) + src_mask = np.zeros([250, 251], dtype=bool) + src_mask[0, :] = True + src.data = np.ma.array(src.data, mask=src_mask) + src.coord("latitude").bounds[0, :, :2] = 0 + src.coord("longitude").bounds[0, :, :2] = 0 + + tgt_mask = np.zeros([251, 250], dtype=bool) + tgt_mask[:, 0] = True + tgt.data = np.ma.array(tgt.data, mask=tgt_mask) + tgt.coord("latitude").bounds[:, 0, ::3] = 0 + tgt.coord("longitude").bounds[:, 0, ::3] = 0 + + self.regrid_class = ESMFAreaWeightedRegridder + self.src = src + self.tgt = tgt + + def time_prepare_without_masks(self): + """Benchmark the prepare time with discontiguities but no mask.""" + _ = self.regrid_class(self.src, self.tgt) + + def time_prepare_with_masks(self): + """Benchmark the prepare time with discontiguities and masks.""" + _ = self.regrid_class(self.src, self.tgt, src_mask=True, tgt_mask=True) diff --git a/benchmarks/benchmarks/generate_data.py b/benchmarks/benchmarks/generate_data.py index bcabac8b..0a787454 100644 --- a/benchmarks/benchmarks/generate_data.py +++ b/benchmarks/benchmarks/generate_data.py @@ -161,6 +161,61 @@ def external(*args, **kwargs): return return_cube +def _curvilinear_cube( + n_lons, + n_lats, + lon_outer_bounds, + lat_outer_bounds, +): + """Call _curvilinear_cube via :func:`run_function_elsewhere`.""" + + def external(*args, **kwargs): + """ + Prep and call _curvilinear_cube, saving to a NetCDF file. + + Saving to a file allows the original python executable to pick back up. + + Remember that all arguments must work as strings, hence the fresh + construction of a ``coord_system`` within the function. + + """ + from iris import save + + from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import ( + _curvilinear_cube as original, + ) + + save_path = kwargs.pop("save_path") + + cube = original(*args, **kwargs) + save(cube, save_path) + + file_name_sections = [ + "_curvilinear_cube", + n_lons, + n_lats, + lon_outer_bounds, + lat_outer_bounds, + ] + file_name = "_".join(str(section) for section in file_name_sections) + # Remove 'unsafe' characters. + file_name = re.sub(r"\W+", "", file_name) + save_path = (BENCHMARK_DATA / file_name).with_suffix(".nc") + + if not REUSE_DATA or not save_path.is_file(): + _ = run_function_elsewhere( + external, + n_lons, + n_lats, + lon_outer_bounds, + lat_outer_bounds, + save_path=str(save_path), + ) + + return_cube = load_cube(str(save_path)) + return return_cube + + def _gridlike_mesh_cube(n_lons, n_lats): """Call _gridlike_mesh via :func:`run_function_elsewhere`.""" diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py index ea1c8f82..fd49b15b 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py @@ -429,12 +429,13 @@ def test_curvilinear(): def test_masks(): """ Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`. + Checks that the `src_mask` and `tgt_mask` keywords work properly. """ src = _curvilinear_cube(7, 6, [-180, 180], [-90, 90]) tgt = _gridlike_mesh_cube(6, 7) - # Make src and tgt discontiguous at (0, 0) + # Make src discontiguous at (0, 0) src_mask = np.zeros([6, 7], dtype=bool) src_mask[0, 0] = True src.data = np.ma.array(src.data, mask=src_mask) @@ -463,3 +464,39 @@ def test_masks(): weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() ) assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) + + +def test_masks_with_resolution(): + """ + Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`. + + Checks that the `src_mask` and `tgt_mask` keywords work properly. + """ + src = _grid_cube(7, 6, [-180, 180], [-90, 90]) + tgt = _gridlike_mesh_cube(6, 7) + + src_mask = np.zeros([6, 7], dtype=bool) + src_mask[0, 0] = True + src.data = np.ma.array(src.data, mask=src_mask) + + tgt_mask = np.zeros([7 * 6], dtype=bool) + tgt_mask[0] = True + tgt.data = np.ma.array(tgt.data, mask=tgt_mask) + + rg_src_masked = GridToMeshESMFRegridder(src, tgt, src_mask=True, resolution=2) + rg_tgt_masked = GridToMeshESMFRegridder(src, tgt, tgt_mask=True, resolution=2) + rg_unmasked = GridToMeshESMFRegridder(src, tgt, resolution=2) + + weights_src_masked = rg_src_masked.regridder.weight_matrix + weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix + weights_unmasked = rg_unmasked.regridder.weight_matrix + + # Check there are no weights associated with the masked point. + assert weights_src_masked[:, 0].nnz == 0 + assert weights_tgt_masked[0].nnz == 0 + + # Check all other weights are correct. + assert np.allclose( + weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() + ) + assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py index ae95c94a..f7314a49 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py @@ -405,16 +405,17 @@ def test_curvilinear(): def test_masks(): """ Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`. + Checks that the `src_mask` and `tgt_mask` keywords work properly. """ src = _gridlike_mesh_cube(7, 6) tgt = _curvilinear_cube(6, 7, [-180, 180], [-90, 90]) - # Make src and tgt discontiguous at (0, 0) src_mask = np.zeros([6 * 7], dtype=bool) src_mask[0] = True src.data = np.ma.array(src.data, mask=src_mask) + # Make tgt discontiguous at (0, 0) tgt_mask = np.zeros([7, 6], dtype=bool) tgt_mask[0, 0] = True tgt.data = np.ma.array(tgt.data, mask=tgt_mask) @@ -439,3 +440,39 @@ def test_masks(): weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() ) assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) + + +def test_masks_with_resolution(): + """ + Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`. + + Checks that the `src_mask` and `tgt_mask` keywords work properly. + """ + src = _gridlike_mesh_cube(7, 6) + tgt = _grid_cube(6, 7, [-180, 180], [-90, 90]) + + src_mask = np.zeros([6 * 7], dtype=bool) + src_mask[0] = True + src.data = np.ma.array(src.data, mask=src_mask) + + tgt_mask = np.zeros([7, 6], dtype=bool) + tgt_mask[0, 0] = True + tgt.data = np.ma.array(tgt.data, mask=tgt_mask) + + rg_src_masked = MeshToGridESMFRegridder(src, tgt, src_mask=True, resolution=2) + rg_tgt_masked = MeshToGridESMFRegridder(src, tgt, tgt_mask=True, resolution=2) + rg_unmasked = MeshToGridESMFRegridder(src, tgt, resolution=2) + + weights_src_masked = rg_src_masked.regridder.weight_matrix + weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix + weights_unmasked = rg_unmasked.regridder.weight_matrix + + # Check there are no weights associated with the masked point. + assert weights_src_masked[:, 0].nnz == 0 + assert weights_tgt_masked[0].nnz == 0 + + # Check all other weights are correct. + assert np.allclose( + weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() + ) + assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) From 894790380c56953f20e0c3e89f8a959a9033d92c Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 9 Mar 2023 13:45:52 +0000 Subject: [PATCH 10/30] add load/save support and improve docstrings --- esmf_regrid/experimental/io.py | 38 ++++++++ .../experimental/unstructured_scheme.py | 25 ++++++ esmf_regrid/schemes.py | 22 +++-- .../experimental/io/test_round_tripping.py | 88 ++++++++++++++++++- 4 files changed, 163 insertions(+), 10 deletions(-) diff --git a/esmf_regrid/experimental/io.py b/esmf_regrid/experimental/io.py index 7a07733e..3f4b8458 100644 --- a/esmf_regrid/experimental/io.py +++ b/esmf_regrid/experimental/io.py @@ -20,7 +20,9 @@ ] REGRIDDER_NAME_MAP = {rg_class.__name__: rg_class for rg_class in SUPPORTED_REGRIDDERS} SOURCE_NAME = "regridder_source_field" +SOURCE_MASK_NAME = "regridder_source_mask" TARGET_NAME = "regridder_target_field" +TARGET_MASK_NAME = "regridder_target_mask" WEIGHTS_NAME = "regridder_weights" WEIGHTS_SHAPE_NAME = "weights_shape" WEIGHTS_ROW_NAME = "weight_matrix_rows" @@ -69,6 +71,12 @@ def _standard_grid_cube(grid, name): if regridder_type == "GridToMeshESMFRegridder": src_grid = (rg.grid_y, rg.grid_x) src_cube = _standard_grid_cube(src_grid, SOURCE_NAME) + if rg.src_mask is not False: + src_mask = rg.src_mask.astype(int) + src_mask_coord = AuxCoord( + src_mask, var_name=SOURCE_MASK_NAME, long_name=SOURCE_MASK_NAME + ) + src_cube.add_aux_coord(src_mask_coord, [0, 1]) tgt_mesh = rg.mesh tgt_location = rg.location @@ -77,6 +85,13 @@ def _standard_grid_cube(grid, name): tgt_cube = Cube(tgt_data, var_name=TARGET_NAME, long_name=TARGET_NAME) for coord in tgt_mesh_coords: tgt_cube.add_aux_coord(coord, 0) + if rg.tgt_mask is not False: + tgt_mask = rg.tgt_mask.astype(int) + tgt_mask_coord = AuxCoord( + tgt_mask, var_name=TARGET_MASK_NAME, long_name=TARGET_MASK_NAME + ) + tgt_cube.add_aux_coord(tgt_mask_coord, 0) + elif regridder_type == "MeshToGridESMFRegridder": src_mesh = rg.mesh src_location = rg.location @@ -85,9 +100,21 @@ def _standard_grid_cube(grid, name): src_cube = Cube(src_data, var_name=SOURCE_NAME, long_name=SOURCE_NAME) for coord in src_mesh_coords: src_cube.add_aux_coord(coord, 0) + if rg.src_mask is not False: + src_mask = rg.src_mask.astype(int) + src_mask_coord = AuxCoord( + src_mask, var_name=SOURCE_MASK_NAME, long_name=SOURCE_MASK_NAME + ) + src_cube.add_aux_coord(src_mask_coord, 0) tgt_grid = (rg.grid_y, rg.grid_x) tgt_cube = _standard_grid_cube(tgt_grid, TARGET_NAME) + if rg.tgt_mask is not False: + tgt_mask = rg.tgt_mask.astype(int) + tgt_mask_coord = AuxCoord( + tgt_mask, var_name=TARGET_MASK_NAME, long_name=TARGET_MASK_NAME + ) + tgt_cube.add_aux_coord(tgt_mask_coord, [0, 1]) else: msg = ( f"Expected a regridder of type `GridToMeshESMFRegridder` or " @@ -205,6 +232,15 @@ def load_regridder(filename): mdtol = weights_cube.attributes[MDTOL] + if src_cube.coords(SOURCE_MASK_NAME): + src_mask = src_cube.coord(SOURCE_MASK_NAME).points + else: + src_mask = False + if tgt_cube.coords(TARGET_MASK_NAME): + tgt_mask = tgt_cube.coord(TARGET_MASK_NAME).points + else: + tgt_mask = False + regridder = scheme( src_cube, tgt_cube, @@ -212,6 +248,8 @@ def load_regridder(filename): method=method, precomputed_weights=weight_matrix, resolution=resolution, + src_mask=src_mask, + tgt_mask=tgt_mask, ) esmf_version = weights_cube.attributes[VERSION_ESMF] diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index 557f77db..636b8b40 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -361,6 +361,17 @@ 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 + 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 + 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. + """ # TODO: Record information about the identity of the mesh. This would @@ -398,10 +409,12 @@ def __init__( 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 partial_regrid_info = _regrid_unstructured_to_rectilinear__prepare( src_mesh_cube, @@ -716,6 +729,16 @@ 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 + 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 + 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 + will be used in weights calculation. """ if method not in ["conservative", "bilinear"]: @@ -748,10 +771,12 @@ def __init__( 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 partial_regrid_info = _regrid_rectilinear_to_unstructured__prepare( src_grid_cube, diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 22d51e22..2afb87c7 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -37,7 +37,7 @@ def _get_mask(cube): slices = cube.slices([src_x, src_y]) data = next(slices).data if np.ma.is_masked(data): - mask = data.mask + mask = np.ma.getmaskarray(data) if cube.coord_dims(src_x) != cube.coord_dims(src_y): mask = mask.T else: @@ -426,7 +426,7 @@ def __repr__(self): """Return a representation of the class.""" return "ESMFAreaWeighted(mdtol={})".format(self.mdtol) - def regridder(self, src_grid, tgt_grid): + def regridder(self, src_grid, tgt_grid, src_mask=None, tgt_mask=None): """ Create regridder to perform regridding from ``src_grid`` to ``tgt_grid``. @@ -436,6 +436,12 @@ def regridder(self, src_grid, tgt_grid): 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 + 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 + Array describing which elements :mod:`ESMF` will ignore on the tgt_grid. + If True, the mask will be derived from tgt_grid. Returns ------- @@ -445,12 +451,16 @@ def regridder(self, src_grid, tgt_grid): 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 return ESMFAreaWeightedRegridder( src_grid, tgt_grid, mdtol=self.mdtol, - src_mask=self.src_mask, - tgt_mask=self.tgt_mask, + src_mask=src_mask, + tgt_mask=tgt_mask, ) @@ -473,10 +483,10 @@ 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, optional + 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, optional + 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. diff --git a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py index 27fac6b5..0c562a69 100644 --- a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py +++ b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py @@ -18,7 +18,11 @@ def _make_grid_to_mesh_regridder( - method="conservative", resolution=None, grid_dims=1, circular=True + method="conservative", + resolution=None, + grid_dims=1, + circular=True, + masks=False, ): src_lons = 3 src_lats = 4 @@ -41,14 +45,33 @@ def _make_grid_to_mesh_regridder( location = "face" tgt = _gridlike_mesh_cube(tgt_lons, tgt_lats, location=location) + if masks: + src_data = ma.array(src.data) + src_data[0, 0] = ma.masked + src.data = src_data + src_mask = True + tgt_data = ma.array(tgt.data) + tgt_data[0] = ma.masked + tgt.data = tgt_data + tgt_mask = True + else: + src_mask = False + tgt_mask = False + rg = GridToMeshESMFRegridder( - src, tgt, method=method, mdtol=0.5, resolution=resolution + src, + tgt, + method=method, + mdtol=0.5, + resolution=resolution, + src_mask=src_mask, + tgt_mask=tgt_mask, ) return rg, src def _make_mesh_to_grid_regridder( - method="conservative", resolution=None, grid_dims=1, circular=True + method="conservative", resolution=None, grid_dims=1, circular=True, masks=False ): src_lons = 3 src_lats = 4 @@ -68,8 +91,27 @@ def _make_mesh_to_grid_regridder( location = "face" src = _gridlike_mesh_cube(src_lons, src_lats, location=location) + if masks: + src_data = ma.array(src.data) + src_data[0] = ma.masked + src.data = src_data + src_mask = True + tgt_data = ma.array(tgt.data) + tgt_data[0, 0] = ma.masked + tgt.data = tgt_data + tgt_mask = True + else: + src_mask = False + tgt_mask = False + rg = MeshToGridESMFRegridder( - src, tgt, method=method, mdtol=0.5, resolution=resolution + src, + tgt, + method=method, + mdtol=0.5, + resolution=resolution, + src_mask=src_mask, + tgt_mask=tgt_mask, ) return rg, src @@ -195,6 +237,25 @@ def test_GridToMeshESMFRegridder_curvilinear_round_trip(tmp_path): assert np.array_equal(original_result.mask, loaded_result.mask) +def test_GridToMeshESMFRegridder_masked_round_trip(tmp_path): + """Test save/load round tripping for `GridToMeshESMFRegridder`.""" + original_rg, src = _make_grid_to_mesh_regridder(masks=True) + filename = tmp_path / "regridder.nc" + save_regridder(original_rg, filename) + loaded_rg = load_regridder(str(filename)) + + # Compare the weight matrices. + original_matrix = original_rg.regridder.weight_matrix + loaded_matrix = loaded_rg.regridder.weight_matrix + # Ensure the original and loaded weight matrix have identical type. + assert type(original_matrix) is type(loaded_matrix) # noqa E721 + assert np.array_equal(original_matrix.todense(), loaded_matrix.todense()) + + # Ensure the masks are preserved + assert np.array_equal(loaded_rg.src_mask, original_rg.src_mask) + assert np.array_equal(loaded_rg.tgt_mask, original_rg.tgt_mask) + + def test_MeshToGridESMFRegridder_round_trip(tmp_path): """Test save/load round tripping for `MeshToGridESMFRegridder`.""" original_rg, src = _make_mesh_to_grid_regridder(circular=True) @@ -311,3 +372,22 @@ def test_MeshToGridESMFRegridder_curvilinear_round_trip(tmp_path): loaded_result = loaded_rg(src).data assert np.array_equal(original_result, loaded_result) assert np.array_equal(original_result.mask, loaded_result.mask) + + +def test_MeshToGridESMFRegridder_masked_round_trip(tmp_path): + """Test save/load round tripping for `MeshToGridESMFRegridder`.""" + original_rg, src = _make_mesh_to_grid_regridder(masks=True) + filename = tmp_path / "regridder.nc" + save_regridder(original_rg, filename) + loaded_rg = load_regridder(str(filename)) + + # Compare the weight matrices. + original_matrix = original_rg.regridder.weight_matrix + loaded_matrix = loaded_rg.regridder.weight_matrix + # Ensure the original and loaded weight matrix have identical type. + assert type(original_matrix) is type(loaded_matrix) # noqa E721 + assert np.array_equal(original_matrix.todense(), loaded_matrix.todense()) + + # Ensure the masks are preserved + assert np.array_equal(loaded_rg.src_mask, original_rg.src_mask) + assert np.array_equal(loaded_rg.tgt_mask, original_rg.tgt_mask) From d3f543af8ec7374b829f5b55a3412773321cb08a Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 10 Mar 2023 11:52:01 +0000 Subject: [PATCH 11/30] additional fixes --- esmf_regrid/experimental/io.py | 8 ++++---- esmf_regrid/schemes.py | 12 ++++++++++++ 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/esmf_regrid/experimental/io.py b/esmf_regrid/experimental/io.py index 3f4b8458..6f62e699 100644 --- a/esmf_regrid/experimental/io.py +++ b/esmf_regrid/experimental/io.py @@ -71,7 +71,7 @@ def _standard_grid_cube(grid, name): if regridder_type == "GridToMeshESMFRegridder": src_grid = (rg.grid_y, rg.grid_x) src_cube = _standard_grid_cube(src_grid, SOURCE_NAME) - if rg.src_mask is not False: + if isinstance(rg.src_mask, np.ndarray): src_mask = rg.src_mask.astype(int) src_mask_coord = AuxCoord( src_mask, var_name=SOURCE_MASK_NAME, long_name=SOURCE_MASK_NAME @@ -85,7 +85,7 @@ def _standard_grid_cube(grid, name): tgt_cube = Cube(tgt_data, var_name=TARGET_NAME, long_name=TARGET_NAME) for coord in tgt_mesh_coords: tgt_cube.add_aux_coord(coord, 0) - if rg.tgt_mask is not False: + if isinstance(rg.tgt_mask, np.ndarray): tgt_mask = rg.tgt_mask.astype(int) tgt_mask_coord = AuxCoord( tgt_mask, var_name=TARGET_MASK_NAME, long_name=TARGET_MASK_NAME @@ -100,7 +100,7 @@ def _standard_grid_cube(grid, name): src_cube = Cube(src_data, var_name=SOURCE_NAME, long_name=SOURCE_NAME) for coord in src_mesh_coords: src_cube.add_aux_coord(coord, 0) - if rg.src_mask is not False: + if isinstance(rg.src_mask, np.ndarray): src_mask = rg.src_mask.astype(int) src_mask_coord = AuxCoord( src_mask, var_name=SOURCE_MASK_NAME, long_name=SOURCE_MASK_NAME @@ -109,7 +109,7 @@ def _standard_grid_cube(grid, name): tgt_grid = (rg.grid_y, rg.grid_x) tgt_cube = _standard_grid_cube(tgt_grid, TARGET_NAME) - if rg.tgt_mask is not False: + if isinstance(rg.tgt_mask, np.ndarray): tgt_mask = rg.tgt_mask.astype(int) tgt_mask_coord = AuxCoord( tgt_mask, var_name=TARGET_MASK_NAME, long_name=TARGET_MASK_NAME diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 2afb87c7..a72b7eb1 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -31,12 +31,24 @@ def _get_coord(cube, axis): def _get_mask(cube): src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y")) + horizontal_dims = set(cube.coord_dims(src_x)) | set(cube.coord_dims(src_y)) + other_dims = tuple(set(range(cube.ndim)) - horizontal_dims) + 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 From 2b248f0976e92d8cc71a97526358b284d74ec90d Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 13 Mar 2023 00:01:29 +0000 Subject: [PATCH 12/30] address review comments --- esmf_regrid/_esmf_sdo.py | 14 ++++++++------ esmf_regrid/esmf_regridder.py | 6 ++++-- esmf_regrid/experimental/io.py | 35 +++++++++++----------------------- 3 files changed, 23 insertions(+), 32 deletions(-) diff --git a/esmf_regrid/_esmf_sdo.py b/esmf_regrid/_esmf_sdo.py index af432501..ee657908 100644 --- a/esmf_regrid/_esmf_sdo.py +++ b/esmf_regrid/_esmf_sdo.py @@ -298,17 +298,19 @@ def _make_esmf_sdo(self): grid_center_y = grid.get_coords(1, staggerloc=esmpy.StaggerLoc.CENTER) grid_center_y[:] = truecenterlats + def add_get_item(grid, **kwargs): + grid.add_item(**kwargs) + return grid.get_item(**kwargs) + if self.mask is not None: - grid.add_item(esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER) - grid_mask = grid.get_item( - esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER + grid_mask = add_get_item( + grid, item=esmpy.GridItem.MASK, staggerloc=esmpy.StaggerLoc.CENTER ) grid_mask[:] = self._refined_mask if areas is not None: - grid.add_item(esmpy.GridItem.AREA, staggerloc=esmpy.StaggerLoc.CENTER) - grid_areas = grid.get_item( - esmpy.GridItem.AREA, staggerloc=esmpy.StaggerLoc.CENTER + grid_areas = add_get_item( + grid, item=esmpy.GridItem.AREA, staggerloc=esmpy.StaggerLoc.CENTER ) grid_areas[:] = areas.T diff --git a/esmf_regrid/esmf_regridder.py b/esmf_regrid/esmf_regridder.py index 0b1afb28..cbb87349 100644 --- a/esmf_regrid/esmf_regridder.py +++ b/esmf_regrid/esmf_regridder.py @@ -15,6 +15,8 @@ def _get_regrid_weights_dict(src_field, tgt_field, regrid_method): + # The value, in array form, that ESMF should treat as an affirmative mask. + expected_mask = np.array([True]) regridder = esmpy.Regrid( src_field, tgt_field, @@ -24,8 +26,8 @@ def _get_regrid_weights_dict(src_field, tgt_field, regrid_method): # Choosing the norm_type DSTAREA allows for mdtol type operations # to be performed using the weights information later on. norm_type=esmpy.NormType.DSTAREA, - src_mask_values=np.array([True]), - dst_mask_values=np.array([True]), + src_mask_values=expected_mask, + dst_mask_values=expected_mask, factors=True, ) # Without specifying deep_copy=true, the information in weights_dict diff --git a/esmf_regrid/experimental/io.py b/esmf_regrid/experimental/io.py index 6f62e699..ea5df746 100644 --- a/esmf_regrid/experimental/io.py +++ b/esmf_regrid/experimental/io.py @@ -35,6 +35,13 @@ RESOLUTION = "resolution" +def _add_mask_to_cube(mask, cube, name): + if isinstance(mask, np.ndarray): + mask = mask.astype(int) + mask_coord = AuxCoord(mask, var_name=name, long_name=name) + cube.add_aux_coord(mask_coord, list(range(cube.ndim))) + + def save_regridder(rg, filename): """ Save a regridder scheme instance. @@ -71,12 +78,7 @@ def _standard_grid_cube(grid, name): if regridder_type == "GridToMeshESMFRegridder": src_grid = (rg.grid_y, rg.grid_x) src_cube = _standard_grid_cube(src_grid, SOURCE_NAME) - if isinstance(rg.src_mask, np.ndarray): - src_mask = rg.src_mask.astype(int) - src_mask_coord = AuxCoord( - src_mask, var_name=SOURCE_MASK_NAME, long_name=SOURCE_MASK_NAME - ) - src_cube.add_aux_coord(src_mask_coord, [0, 1]) + _add_mask_to_cube(rg.src_mask, src_cube, SOURCE_MASK_NAME) tgt_mesh = rg.mesh tgt_location = rg.location @@ -85,12 +87,7 @@ def _standard_grid_cube(grid, name): tgt_cube = Cube(tgt_data, var_name=TARGET_NAME, long_name=TARGET_NAME) for coord in tgt_mesh_coords: tgt_cube.add_aux_coord(coord, 0) - if isinstance(rg.tgt_mask, np.ndarray): - tgt_mask = rg.tgt_mask.astype(int) - tgt_mask_coord = AuxCoord( - tgt_mask, var_name=TARGET_MASK_NAME, long_name=TARGET_MASK_NAME - ) - tgt_cube.add_aux_coord(tgt_mask_coord, 0) + _add_mask_to_cube(rg.tgt_mask, tgt_cube, TARGET_MASK_NAME) elif regridder_type == "MeshToGridESMFRegridder": src_mesh = rg.mesh @@ -100,21 +97,11 @@ def _standard_grid_cube(grid, name): src_cube = Cube(src_data, var_name=SOURCE_NAME, long_name=SOURCE_NAME) for coord in src_mesh_coords: src_cube.add_aux_coord(coord, 0) - if isinstance(rg.src_mask, np.ndarray): - src_mask = rg.src_mask.astype(int) - src_mask_coord = AuxCoord( - src_mask, var_name=SOURCE_MASK_NAME, long_name=SOURCE_MASK_NAME - ) - src_cube.add_aux_coord(src_mask_coord, 0) + _add_mask_to_cube(rg.src_mask, src_cube, SOURCE_MASK_NAME) tgt_grid = (rg.grid_y, rg.grid_x) tgt_cube = _standard_grid_cube(tgt_grid, TARGET_NAME) - if isinstance(rg.tgt_mask, np.ndarray): - tgt_mask = rg.tgt_mask.astype(int) - tgt_mask_coord = AuxCoord( - tgt_mask, var_name=TARGET_MASK_NAME, long_name=TARGET_MASK_NAME - ) - tgt_cube.add_aux_coord(tgt_mask_coord, [0, 1]) + _add_mask_to_cube(rg.tgt_mask, tgt_cube, TARGET_MASK_NAME) else: msg = ( f"Expected a regridder of type `GridToMeshESMFRegridder` or " From 5566ecf7308fdaaff2f00b35963736fd6d3250eb Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 16 Mar 2023 14:51:59 +0000 Subject: [PATCH 13/30] address review comments --- .../experimental/unstructured_scheme.py | 82 ++++++------ esmf_regrid/schemes.py | 118 +++++++++--------- 2 files changed, 108 insertions(+), 92 deletions(-) diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index 636b8b40..1c1381b5 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -253,8 +253,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. @@ -299,6 +299,16 @@ 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 ------- @@ -306,6 +316,9 @@ def regrid_unstructured_to_rectilinear( 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, @@ -329,8 +342,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. @@ -361,12 +374,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 @@ -405,16 +418,8 @@ 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, @@ -422,8 +427,8 @@ def __init__( 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. @@ -617,8 +622,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. @@ -667,6 +672,16 @@ 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 ------- @@ -674,6 +689,9 @@ def regrid_rectilinear_to_unstructured( 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, @@ -697,8 +715,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. @@ -729,12 +747,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 @@ -767,16 +785,8 @@ 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, @@ -784,8 +794,8 @@ def __init__( 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. diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index a72b7eb1..a6a9860b 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -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): @@ -405,7 +413,7 @@ class ESMFAreaWeighted: :mod:`esmpy` 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. @@ -419,10 +427,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. @@ -431,14 +439,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``. @@ -448,10 +456,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. @@ -463,23 +471,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``. @@ -495,12 +505,16 @@ 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): @@ -508,25 +522,17 @@ def __init__(self, src_grid, tgt_grid, mdtol=0, src_mask=False, tgt_mask=False): 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")) From bded83b396ea1ade3309cc0e2cb6e0d39a455365 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Mar 2023 14:59:36 +0000 Subject: [PATCH 14/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- esmf_regrid/schemes.py | 1 - 1 file changed, 1 deletion(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index a6a9860b..f6495459 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -32,7 +32,6 @@ def _get_mask(cube, use_mask=True): if use_mask == False: return None elif use_mask == True: - src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y")) horizontal_dims = set(cube.coord_dims(src_x)) | set(cube.coord_dims(src_y)) From d029ae0f0abf3dfab51908b4dd41185f4e466615 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 16 Mar 2023 15:36:58 +0000 Subject: [PATCH 15/30] fix errors --- benchmarks/benchmarks/esmf_regridder/__init__.py | 13 +++++++------ esmf_regrid/experimental/io.py | 12 ++++++------ esmf_regrid/schemes.py | 4 ++-- .../test_GridToMeshESMFRegridder.py | 12 ++++++------ .../test_MeshToGridESMFRegridder.py | 12 ++++++------ .../unit/schemes/test_ESMFAreaWeightedRegridder.py | 6 +++--- 6 files changed, 30 insertions(+), 29 deletions(-) diff --git a/benchmarks/benchmarks/esmf_regridder/__init__.py b/benchmarks/benchmarks/esmf_regridder/__init__.py index c213364a..0a71f0eb 100644 --- a/benchmarks/benchmarks/esmf_regridder/__init__.py +++ b/benchmarks/benchmarks/esmf_regridder/__init__.py @@ -479,10 +479,11 @@ def setup(self): self.src = src self.tgt = tgt - def time_prepare_without_masks(self): - """Benchmark the prepare time with discontiguities but no mask.""" - _ = self.regrid_class(self.src, self.tgt) - def time_prepare_with_masks(self): - """Benchmark the prepare time with discontiguities and masks.""" - _ = self.regrid_class(self.src, self.tgt, src_mask=True, tgt_mask=True) + """Benchmark prepare time with discontiguities and masks.""" + try: + _ = self.regrid_class( + self.src, self.tgt, use_src_mask=True, use_tgt_mask=True + ) + except TypeError: + _ = self.regrid_class(self.src, self.tgt) diff --git a/esmf_regrid/experimental/io.py b/esmf_regrid/experimental/io.py index ea5df746..08777652 100644 --- a/esmf_regrid/experimental/io.py +++ b/esmf_regrid/experimental/io.py @@ -220,13 +220,13 @@ def load_regridder(filename): mdtol = weights_cube.attributes[MDTOL] if src_cube.coords(SOURCE_MASK_NAME): - src_mask = src_cube.coord(SOURCE_MASK_NAME).points + use_src_mask = src_cube.coord(SOURCE_MASK_NAME).points else: - src_mask = False + use_src_mask = False if tgt_cube.coords(TARGET_MASK_NAME): - tgt_mask = tgt_cube.coord(TARGET_MASK_NAME).points + use_tgt_mask = tgt_cube.coord(TARGET_MASK_NAME).points else: - tgt_mask = False + use_tgt_mask = False regridder = scheme( src_cube, @@ -235,8 +235,8 @@ def load_regridder(filename): method=method, precomputed_weights=weight_matrix, resolution=resolution, - src_mask=src_mask, - tgt_mask=tgt_mask, + use_src_mask=use_src_mask, + use_tgt_mask=use_tgt_mask, ) esmf_version = weights_cube.attributes[VERSION_ESMF] diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index f6495459..7ea0b8e7 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -29,9 +29,9 @@ def _get_coord(cube, axis): def _get_mask(cube, use_mask=True): - if use_mask == False: + if use_mask is False: return None - elif use_mask == True: + elif use_mask is True: src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y")) horizontal_dims = set(cube.coord_dims(src_x)) | set(cube.coord_dims(src_y)) diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py index fd49b15b..ecd4c19e 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py @@ -430,7 +430,7 @@ def test_masks(): """ Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`. - Checks that the `src_mask` and `tgt_mask` keywords work properly. + Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. """ src = _curvilinear_cube(7, 6, [-180, 180], [-90, 90]) tgt = _gridlike_mesh_cube(6, 7) @@ -447,8 +447,8 @@ def test_masks(): tgt_mask[0] = True tgt.data = np.ma.array(tgt.data, mask=tgt_mask) - rg_src_masked = GridToMeshESMFRegridder(src_discontiguous, tgt, src_mask=True) - rg_tgt_masked = GridToMeshESMFRegridder(src, tgt, tgt_mask=True) + rg_src_masked = GridToMeshESMFRegridder(src_discontiguous, tgt, use_src_mask=True) + rg_tgt_masked = GridToMeshESMFRegridder(src, tgt, use_tgt_mask=True) rg_unmasked = GridToMeshESMFRegridder(src, tgt) weights_src_masked = rg_src_masked.regridder.weight_matrix @@ -470,7 +470,7 @@ def test_masks_with_resolution(): """ Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`. - Checks that the `src_mask` and `tgt_mask` keywords work properly. + Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. """ src = _grid_cube(7, 6, [-180, 180], [-90, 90]) tgt = _gridlike_mesh_cube(6, 7) @@ -483,8 +483,8 @@ def test_masks_with_resolution(): tgt_mask[0] = True tgt.data = np.ma.array(tgt.data, mask=tgt_mask) - rg_src_masked = GridToMeshESMFRegridder(src, tgt, src_mask=True, resolution=2) - rg_tgt_masked = GridToMeshESMFRegridder(src, tgt, tgt_mask=True, resolution=2) + rg_src_masked = GridToMeshESMFRegridder(src, tgt, use_src_mask=True, resolution=2) + rg_tgt_masked = GridToMeshESMFRegridder(src, tgt, use_tgt_mask=True, resolution=2) rg_unmasked = GridToMeshESMFRegridder(src, tgt, resolution=2) weights_src_masked = rg_src_masked.regridder.weight_matrix diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py index f7314a49..5a67509b 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py @@ -406,7 +406,7 @@ def test_masks(): """ Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`. - Checks that the `src_mask` and `tgt_mask` keywords work properly. + Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. """ src = _gridlike_mesh_cube(7, 6) tgt = _curvilinear_cube(6, 7, [-180, 180], [-90, 90]) @@ -423,8 +423,8 @@ def test_masks(): tgt_discontiguous.coord("latitude").bounds[0, 0] = 0 tgt_discontiguous.coord("longitude").bounds[0, 0] = 0 - rg_src_masked = MeshToGridESMFRegridder(src, tgt, src_mask=True) - rg_tgt_masked = MeshToGridESMFRegridder(src, tgt_discontiguous, tgt_mask=True) + rg_src_masked = MeshToGridESMFRegridder(src, tgt, use_src_mask=True) + rg_tgt_masked = MeshToGridESMFRegridder(src, tgt_discontiguous, use_tgt_mask=True) rg_unmasked = MeshToGridESMFRegridder(src, tgt) weights_src_masked = rg_src_masked.regridder.weight_matrix @@ -446,7 +446,7 @@ def test_masks_with_resolution(): """ Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`. - Checks that the `src_mask` and `tgt_mask` keywords work properly. + Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. """ src = _gridlike_mesh_cube(7, 6) tgt = _grid_cube(6, 7, [-180, 180], [-90, 90]) @@ -459,8 +459,8 @@ def test_masks_with_resolution(): tgt_mask[0, 0] = True tgt.data = np.ma.array(tgt.data, mask=tgt_mask) - rg_src_masked = MeshToGridESMFRegridder(src, tgt, src_mask=True, resolution=2) - rg_tgt_masked = MeshToGridESMFRegridder(src, tgt, tgt_mask=True, resolution=2) + rg_src_masked = MeshToGridESMFRegridder(src, tgt, use_src_mask=True, resolution=2) + rg_tgt_masked = MeshToGridESMFRegridder(src, tgt, use_tgt_mask=True, resolution=2) rg_unmasked = MeshToGridESMFRegridder(src, tgt, resolution=2) weights_src_masked = rg_src_masked.regridder.weight_matrix diff --git a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py index b79912df..8e737689 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py @@ -235,7 +235,7 @@ def test_masks(): """ Test initialisation of :func:`esmf_regrid.schemes.ESMFAreaWeightedRegridder`. - Checks that the `src_mask` and `tgt_mask` keywords work properly. + Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. """ src = _curvilinear_cube(7, 6, [-180, 180], [-90, 90]) tgt = _curvilinear_cube(6, 7, [-180, 180], [-90, 90]) @@ -255,8 +255,8 @@ def test_masks(): tgt_discontiguous.coord("latitude").bounds[0, 0] = 0 tgt_discontiguous.coord("longitude").bounds[0, 0] = 0 - rg_src_masked = ESMFAreaWeightedRegridder(src_discontiguous, tgt, src_mask=True) - rg_tgt_masked = ESMFAreaWeightedRegridder(src, tgt_discontiguous, tgt_mask=True) + rg_src_masked = ESMFAreaWeightedRegridder(src_discontiguous, tgt, use_src_mask=True) + rg_tgt_masked = ESMFAreaWeightedRegridder(src, tgt_discontiguous, use_tgt_mask=True) rg_unmasked = ESMFAreaWeightedRegridder(src, tgt) weights_src_masked = rg_src_masked.regridder.weight_matrix From 748d560b2568ebcc67064aebc24f2786956d7317 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 16 Mar 2023 15:47:07 +0000 Subject: [PATCH 16/30] fix errors --- .../experimental/io/test_round_tripping.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py index 0c562a69..843a8f8c 100644 --- a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py +++ b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py @@ -49,14 +49,14 @@ def _make_grid_to_mesh_regridder( src_data = ma.array(src.data) src_data[0, 0] = ma.masked src.data = src_data - src_mask = True + use_src_mask = True tgt_data = ma.array(tgt.data) tgt_data[0] = ma.masked tgt.data = tgt_data - tgt_mask = True + use_tgt_mask = True else: - src_mask = False - tgt_mask = False + use_src_mask = False + use_tgt_mask = False rg = GridToMeshESMFRegridder( src, @@ -64,8 +64,8 @@ def _make_grid_to_mesh_regridder( method=method, mdtol=0.5, resolution=resolution, - src_mask=src_mask, - tgt_mask=tgt_mask, + use_src_mask=use_src_mask, + use_tgt_mask=use_tgt_mask, ) return rg, src @@ -95,14 +95,14 @@ def _make_mesh_to_grid_regridder( src_data = ma.array(src.data) src_data[0] = ma.masked src.data = src_data - src_mask = True + use_src_mask = True tgt_data = ma.array(tgt.data) tgt_data[0, 0] = ma.masked tgt.data = tgt_data - tgt_mask = True + use_tgt_mask = True else: - src_mask = False - tgt_mask = False + use_src_mask = False + use_tgt_mask = False rg = MeshToGridESMFRegridder( src, @@ -110,8 +110,8 @@ def _make_mesh_to_grid_regridder( method=method, mdtol=0.5, resolution=resolution, - src_mask=src_mask, - tgt_mask=tgt_mask, + use_src_mask=use_src_mask, + use_tgt_mask=use_tgt_mask, ) return rg, src From 0bf03c82a4f625a6a5597b3b9c9835c9a474c49c Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Fri, 17 Mar 2023 16:05:18 +0000 Subject: [PATCH 17/30] parameterise test --- .../test_MeshToGridESMFRegridder.py | 60 ++++++------------- 1 file changed, 18 insertions(+), 42 deletions(-) diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py index 5a67509b..34c9b8eb 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py @@ -402,14 +402,21 @@ def test_curvilinear(): assert result_lazy == result -def test_masks(): +@pytest.mark.parametrize( + "resolution", (None, 2), ids=("no resolution", "with resolution") +) +def test_masks(resolution): """ Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`. Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. """ src = _gridlike_mesh_cube(7, 6) - tgt = _curvilinear_cube(6, 7, [-180, 180], [-90, 90]) + if resolution is None: + tgt = _curvilinear_cube(6, 7, [-180, 180], [-90, 90]) + else: + # The resolution keyword is only valid for rectilinear grids. + tgt = _grid_cube(6, 7, [-180, 180], [-90, 90]) src_mask = np.zeros([6 * 7], dtype=bool) src_mask[0] = True @@ -420,48 +427,17 @@ def test_masks(): tgt_mask[0, 0] = True tgt.data = np.ma.array(tgt.data, mask=tgt_mask) tgt_discontiguous = tgt.copy() - tgt_discontiguous.coord("latitude").bounds[0, 0] = 0 - tgt_discontiguous.coord("longitude").bounds[0, 0] = 0 - - rg_src_masked = MeshToGridESMFRegridder(src, tgt, use_src_mask=True) - rg_tgt_masked = MeshToGridESMFRegridder(src, tgt_discontiguous, use_tgt_mask=True) - rg_unmasked = MeshToGridESMFRegridder(src, tgt) - - weights_src_masked = rg_src_masked.regridder.weight_matrix - weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix - weights_unmasked = rg_unmasked.regridder.weight_matrix - - # Check there are no weights associated with the masked point. - assert weights_src_masked[:, 0].nnz == 0 - assert weights_tgt_masked[0].nnz == 0 + if resolution is None: + tgt_discontiguous.coord("latitude").bounds[0, 0] = 0 + tgt_discontiguous.coord("longitude").bounds[0, 0] = 0 - # Check all other weights are correct. - assert np.allclose( - weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() + rg_src_masked = MeshToGridESMFRegridder( + src, tgt, use_src_mask=True, resolution=resolution ) - assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) - - -def test_masks_with_resolution(): - """ - Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.MeshToGridESMFRegridder`. - - Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. - """ - src = _gridlike_mesh_cube(7, 6) - tgt = _grid_cube(6, 7, [-180, 180], [-90, 90]) - - src_mask = np.zeros([6 * 7], dtype=bool) - src_mask[0] = True - src.data = np.ma.array(src.data, mask=src_mask) - - tgt_mask = np.zeros([7, 6], dtype=bool) - tgt_mask[0, 0] = True - tgt.data = np.ma.array(tgt.data, mask=tgt_mask) - - rg_src_masked = MeshToGridESMFRegridder(src, tgt, use_src_mask=True, resolution=2) - rg_tgt_masked = MeshToGridESMFRegridder(src, tgt, use_tgt_mask=True, resolution=2) - rg_unmasked = MeshToGridESMFRegridder(src, tgt, resolution=2) + rg_tgt_masked = MeshToGridESMFRegridder( + src, tgt_discontiguous, use_tgt_mask=True, resolution=resolution + ) + rg_unmasked = MeshToGridESMFRegridder(src, tgt, resolution=resolution) weights_src_masked = rg_src_masked.regridder.weight_matrix weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix From e6a4ed8bb7110f88e313bcbc1ab441f2e29b0ee2 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Tue, 21 Mar 2023 09:36:18 +0000 Subject: [PATCH 18/30] address review comments --- esmf_regrid/schemes.py | 43 ++++++++++++++++++++++++++++-------------- 1 file changed, 29 insertions(+), 14 deletions(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 7ea0b8e7..0d4370ae 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -66,12 +66,15 @@ def _get_mask(cube, use_mask=True): def _contiguous_masked(bounds, mask): """ - Return the (N+1, M+1) bound values for bounds of 2D coordinate of shape (N,M). + Return the (N+1, M+1) vertices for 2D coordinate bounds of shape (N, M, 4). - Assumes the bounds are contiguous outside of the mask. The returned bounds - fully describe the unmasked bounds when this is the case. This function is - designed to replicate the behaviour of coord.contiguous_bounds() for unmasked - bounds. + Assumes the bounds are contiguous outside of the mask. As long as the only + discontiguities are associated with masked points, the returned vertices will + represent every bound associated with an unmasked point. This function is + designed to replicate the behaviour of iris.Coord.contiguous_bounds() for such + bounds. For each vertex in the returned array there are up to four choices of + bounds to derive from. Bounds associated with umasked points will be prioritised + in this choice. For example, suppose we have a masked array: @@ -80,6 +83,10 @@ def _contiguous_masked(bounds, mask): # 0 - - 0 # 0 0 0 0 + Now the indices of the final bounds dimension correspond to positions on the + vertex array. For a bound whose first indixes are (i, j) the corresponding + position in the vertex array of the four final indices are as follows: + 0=(j, i), 1=(j, i+1), 2=(j+1, i+1), 3=(j+1, i) The indices of the bounds which the final array will derive from are as follows: # (0, 0, 0) (0, 1, 0) (0, 2, 0) (0, 3, 0) (0, 3, 1) @@ -88,22 +95,30 @@ def _contiguous_masked(bounds, mask): # (3, 0, 0) (3, 1, 0) (3, 2, 0) (3, 3, 0) (3, 3, 1) # (3, 0, 3) (3, 1, 3) (3, 3, 3) (3, 3, 3) (3, 3, 2) - Note that only the center bound derives from a masked cell. + Note that only the center bound derives from a masked cell as there are no + unmasked points to derive from. """ mask = np.array(mask, dtype=int) # Construct a blank template array in the shape of the output. blank_template = np.zeros([size + 1 for size in mask.shape], dtype=int) + # Define the slices of the resulting vertex array which derive from the + # bounds in index 0 to 3. + slice_0 = np.s_[:-1, :-1] + slice_1 = np.s_[:-1, 1:] + slice_2 = np.s_[1:, 1:] + slice_3 = np.s_[1:, :-1] + # Define the bounds which will derive from the bounds in index 0. filter_0 = blank_template.copy() - filter_0[:-1, :-1] = 1 - mask + filter_0[slice_0] = 1 - mask # Ensure the corner points are covered appropriately. filter_0[0, 0] = 1 # Define the bounds which will derive from the bounds in index 1. filter_1 = blank_template.copy() - filter_1[:-1, 1:] = 1 - mask + filter_1[slice_1] = 1 - mask # Ensure the corner and edge bounds are covered appropriately. filter_1[0, 1:] = 1 # Do not cover any points already covered. @@ -111,7 +126,7 @@ def _contiguous_masked(bounds, mask): # Define the bounds which will derive from the bounds in index 3. filter_3 = blank_template.copy() - filter_3[1:, :-1] = 1 - mask + filter_3[slice_3] = 1 - mask # Ensure the corner and edge bounds are covered appropriately. filter_3[1:, 0] = 1 # Do not cover any points already covered. @@ -122,10 +137,10 @@ def _contiguous_masked(bounds, mask): # Copy the bounds information over into the appropriate places. contiguous_bounds = blank_template.astype(bounds.dtype) - contiguous_bounds[:-1, :-1] += bounds[:, :, 0] * filter_0[:-1, :-1] - contiguous_bounds[:-1, 1:] += bounds[:, :, 1] * filter_1[:-1, 1:] - contiguous_bounds[1:, 1:] += bounds[:, :, 2] * filter_2[1:, 1:] - contiguous_bounds[1:, :-1] += bounds[:, :, 3] * filter_3[1:, :-1] + contiguous_bounds[slice_0] += bounds[:, :, 0] * filter_0[slice_0] + contiguous_bounds[slice_1] += bounds[:, :, 1] * filter_1[slice_1] + contiguous_bounds[slice_2] += bounds[:, :, 2] * filter_2[slice_2] + contiguous_bounds[slice_3] += bounds[:, :, 3] * filter_3[slice_3] return contiguous_bounds @@ -154,7 +169,7 @@ def _cube_to_GridInfo(cube, center=False, resolution=None, mask=None): # TODO: perform checks on lat/lon. elif londim == 2: assert cube.coord_dims(lon) == cube.coord_dims(lat) - if mask is None: + if np.any(mask): assert lon.is_contiguous() assert lat.is_contiguous() lon_bound_array = lon.contiguous_bounds() From 4287b5fda27414a907de7a3d9cf611b1284db730 Mon Sep 17 00:00:00 2001 From: stephenworsley <49274989+stephenworsley@users.noreply.github.com> Date: Tue, 21 Mar 2023 09:38:26 +0000 Subject: [PATCH 19/30] Apply suggestions from code review Co-authored-by: Martin Yeo <40734014+trexfeathers@users.noreply.github.com> --- benchmarks/benchmarks/generate_data.py | 3 +-- .../unit/experimental/io/test_round_tripping.py | 12 +++++++++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/benchmarks/benchmarks/generate_data.py b/benchmarks/benchmarks/generate_data.py index 0a787454..17b3c7ae 100644 --- a/benchmarks/benchmarks/generate_data.py +++ b/benchmarks/benchmarks/generate_data.py @@ -175,8 +175,7 @@ def external(*args, **kwargs): Saving to a file allows the original python executable to pick back up. - Remember that all arguments must work as strings, hence the fresh - construction of a ``coord_system`` within the function. + Remember that all arguments must work as strings. """ from iris import save diff --git a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py index 843a8f8c..ec7cec3c 100644 --- a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py +++ b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py @@ -237,9 +237,15 @@ def test_GridToMeshESMFRegridder_curvilinear_round_trip(tmp_path): assert np.array_equal(original_result.mask, loaded_result.mask) -def test_GridToMeshESMFRegridder_masked_round_trip(tmp_path): - """Test save/load round tripping for `GridToMeshESMFRegridder`.""" - original_rg, src = _make_grid_to_mesh_regridder(masks=True) +# TODO: parametrize the rest of the tests in this module. +@pytest.mark.parametrize( + "rg_maker", + [_make_grid_to_mesh_regridder, _make_mesh_to_grid_regridder], + ids=["grid_to_mesh", "mesh_to_grid"], +) +def test_MeshESMFRegridder_masked_round_trip(tmp_path, rg_maker): + """Test save/load round tripping for the Mesh regridder classes.""" + original_rg, src = rg_maker(masks=True) filename = tmp_path / "regridder.nc" save_regridder(original_rg, filename) loaded_rg = load_regridder(str(filename)) From 4987e2601786d6fa4f96b977146f3ff5ff939bd6 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Tue, 21 Mar 2023 10:00:25 +0000 Subject: [PATCH 20/30] fix test --- esmf_regrid/tests/unit/experimental/io/test_round_tripping.py | 1 + 1 file changed, 1 insertion(+) diff --git a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py index ec7cec3c..210dba16 100644 --- a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py +++ b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py @@ -2,6 +2,7 @@ import numpy as np from numpy import ma +import pytest from esmf_regrid.experimental.io import load_regridder, save_regridder from esmf_regrid.experimental.unstructured_scheme import ( From 891bc8072441ba972d3e88df408682fa7fc0a2d0 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Tue, 21 Mar 2023 10:03:41 +0000 Subject: [PATCH 21/30] simplify tests --- .../experimental/io/test_round_tripping.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py index 210dba16..157a92c2 100644 --- a/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py +++ b/esmf_regrid/tests/unit/experimental/io/test_round_tripping.py @@ -379,22 +379,3 @@ def test_MeshToGridESMFRegridder_curvilinear_round_trip(tmp_path): loaded_result = loaded_rg(src).data assert np.array_equal(original_result, loaded_result) assert np.array_equal(original_result.mask, loaded_result.mask) - - -def test_MeshToGridESMFRegridder_masked_round_trip(tmp_path): - """Test save/load round tripping for `MeshToGridESMFRegridder`.""" - original_rg, src = _make_mesh_to_grid_regridder(masks=True) - filename = tmp_path / "regridder.nc" - save_regridder(original_rg, filename) - loaded_rg = load_regridder(str(filename)) - - # Compare the weight matrices. - original_matrix = original_rg.regridder.weight_matrix - loaded_matrix = loaded_rg.regridder.weight_matrix - # Ensure the original and loaded weight matrix have identical type. - assert type(original_matrix) is type(loaded_matrix) # noqa E721 - assert np.array_equal(original_matrix.todense(), loaded_matrix.todense()) - - # Ensure the masks are preserved - assert np.array_equal(loaded_rg.src_mask, original_rg.src_mask) - assert np.array_equal(loaded_rg.tgt_mask, original_rg.tgt_mask) From 62baf48ddad82e9f14127dc03b21a5b935f32a88 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Tue, 21 Mar 2023 10:18:36 +0000 Subject: [PATCH 22/30] fix tests --- esmf_regrid/schemes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 0d4370ae..b55f8bef 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -169,7 +169,7 @@ def _cube_to_GridInfo(cube, center=False, resolution=None, mask=None): # TODO: perform checks on lat/lon. elif londim == 2: assert cube.coord_dims(lon) == cube.coord_dims(lat) - if np.any(mask): + if not np.any(mask): assert lon.is_contiguous() assert lat.is_contiguous() lon_bound_array = lon.contiguous_bounds() From c8db5991f15be2ccc366534e1fcda3e062675ac4 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 22 Mar 2023 11:18:45 +0000 Subject: [PATCH 23/30] add test --- .../unit/schemes/test__cube_to_GridInfo.py | 45 ++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py index 900d14ee..13680e4a 100644 --- a/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py +++ b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py @@ -8,7 +8,7 @@ import scipy.sparse from esmf_regrid.esmf_regridder import Regridder -from esmf_regrid.schemes import _cube_to_GridInfo +from esmf_regrid.schemes import _cube_to_GridInfo, _contiguous_masked def _generate_points_and_bounds(n, outer_bounds): @@ -195,3 +195,46 @@ def test_curvilinear_grid(): circular_gridinfo.circular = True rg_circular = Regridder(circular_gridinfo, gridinfo) assert np.allclose(expected_weights.todense(), rg_circular.weight_matrix.todense()) + + +def test__contiguous_bounds(): + """Test generation of contiguous bounds.""" + # Generate a CF style bounds array with unique values for each bound. + # The values will represent the index in the bounds array. + cf_bds = np.zeros([5, 6, 4]) + cf_bds += np.arange(5)[:, np.newaxis, np.newaxis] * 100 + cf_bds += np.arange(6)[np.newaxis, :, np.newaxis] * 10 + cf_bds += np.arange(4)[np.newaxis, np.newaxis, :] + + # Define a mask such that each possible 2x2 sub array is represented. + # e.g. + # [[1, 1], + # [0, 1]] + mask = np.array( + [ + [0, 1, 1, 0, 0, 0], + [0, 1, 1, 0, 0, 1], + [1, 0, 1, 1, 1, 1], + [1, 0, 0, 0, 1, 0], + [0, 1, 1, 0, 0, 0], + ] + ) + + # Calculate the values on the vertices. + vertices = _contiguous_masked(cf_bds, mask) + # Note, the only values deriving from masked points are: + # 11, 12, 152, 203, 412 + # These are the only vertices not connected to any unmasked points. + # fmt: off + expected_vertices = np.array( + [ + [0, 1, 11, 30, 40, 50, 51], + [100, 101, 12, 130, 140, 141, 52], + [103, 210, 211, 133, 143, 142, 152], + [203, 310, 320, 330, 331, 350, 351], + [400, 401, 323, 430, 440, 450, 451], + [403, 402, 412, 433, 443, 453, 452], + ] + ) + # fmt: on + assert np.array_equal(expected_vertices, vertices) From 6f700b6757548f69fe6e2a28658e89f453c1af76 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 22 Mar 2023 12:05:18 +0000 Subject: [PATCH 24/30] parameterize test --- .../test_GridToMeshESMFRegridder.py | 60 ++++++------------- 1 file changed, 18 insertions(+), 42 deletions(-) diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py index ecd4c19e..c8e1106c 100644 --- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py +++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py @@ -426,13 +426,20 @@ def test_curvilinear(): assert result_lazy == result -def test_masks(): +@pytest.mark.parametrize( + "resolution", (None, 2), ids=("no resolution", "with resolution") +) +def test_masks(resolution): """ Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`. Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. """ - src = _curvilinear_cube(7, 6, [-180, 180], [-90, 90]) + if resolution is None: + src = _curvilinear_cube(7, 6, [-180, 180], [-90, 90]) + else: + # The resolution keyword is only valid for rectilinear grids. + src = _grid_cube(7, 6, [-180, 180], [-90, 90]) tgt = _gridlike_mesh_cube(6, 7) # Make src discontiguous at (0, 0) @@ -440,52 +447,21 @@ def test_masks(): src_mask[0, 0] = True src.data = np.ma.array(src.data, mask=src_mask) src_discontiguous = src.copy() - src_discontiguous.coord("latitude").bounds[0, 0] = 0 - src_discontiguous.coord("longitude").bounds[0, 0] = 0 + if resolution is None: + src_discontiguous.coord("latitude").bounds[0, 0] = 0 + src_discontiguous.coord("longitude").bounds[0, 0] = 0 tgt_mask = np.zeros([7 * 6], dtype=bool) tgt_mask[0] = True tgt.data = np.ma.array(tgt.data, mask=tgt_mask) - rg_src_masked = GridToMeshESMFRegridder(src_discontiguous, tgt, use_src_mask=True) - rg_tgt_masked = GridToMeshESMFRegridder(src, tgt, use_tgt_mask=True) - rg_unmasked = GridToMeshESMFRegridder(src, tgt) - - weights_src_masked = rg_src_masked.regridder.weight_matrix - weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix - weights_unmasked = rg_unmasked.regridder.weight_matrix - - # Check there are no weights associated with the masked point. - assert weights_src_masked[:, 0].nnz == 0 - assert weights_tgt_masked[0].nnz == 0 - - # Check all other weights are correct. - assert np.allclose( - weights_src_masked[:, 1:].todense(), weights_unmasked[:, 1:].todense() + rg_src_masked = GridToMeshESMFRegridder( + src_discontiguous, tgt, use_src_mask=True, resolution=resolution ) - assert np.allclose(weights_tgt_masked[1:].todense(), weights_unmasked[1:].todense()) - - -def test_masks_with_resolution(): - """ - Test initialisation of :func:`esmf_regrid.experimental.unstructured_scheme.GridToMeshESMFRegridder`. - - Checks that the `use_src_mask` and `use_tgt_mask` keywords work properly. - """ - src = _grid_cube(7, 6, [-180, 180], [-90, 90]) - tgt = _gridlike_mesh_cube(6, 7) - - src_mask = np.zeros([6, 7], dtype=bool) - src_mask[0, 0] = True - src.data = np.ma.array(src.data, mask=src_mask) - - tgt_mask = np.zeros([7 * 6], dtype=bool) - tgt_mask[0] = True - tgt.data = np.ma.array(tgt.data, mask=tgt_mask) - - rg_src_masked = GridToMeshESMFRegridder(src, tgt, use_src_mask=True, resolution=2) - rg_tgt_masked = GridToMeshESMFRegridder(src, tgt, use_tgt_mask=True, resolution=2) - rg_unmasked = GridToMeshESMFRegridder(src, tgt, resolution=2) + rg_tgt_masked = GridToMeshESMFRegridder( + src, tgt, use_tgt_mask=True, resolution=resolution + ) + rg_unmasked = GridToMeshESMFRegridder(src, tgt, resolution=resolution) weights_src_masked = rg_src_masked.regridder.weight_matrix weights_tgt_masked = rg_tgt_masked.regridder.weight_matrix From be718287854a726860862950938b0f1ff4919cc0 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 22 Mar 2023 13:15:28 +0000 Subject: [PATCH 25/30] flake8 fix --- esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py index 13680e4a..cd1fba17 100644 --- a/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py +++ b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py @@ -8,7 +8,7 @@ import scipy.sparse from esmf_regrid.esmf_regridder import Regridder -from esmf_regrid.schemes import _cube_to_GridInfo, _contiguous_masked +from esmf_regrid.schemes import _contiguous_masked, _cube_to_GridInfo def _generate_points_and_bounds(n, outer_bounds): @@ -228,8 +228,8 @@ def test__contiguous_bounds(): # fmt: off expected_vertices = np.array( [ - [0, 1, 11, 30, 40, 50, 51], - [100, 101, 12, 130, 140, 141, 52], + [0, 1, 11, 30, 40, 50, 51], # flake8: E241 + [100, 101, 12, 130, 140, 141, 52], # flake8: E241 [103, 210, 211, 133, 143, 142, 152], [203, 310, 320, 330, 331, 350, 351], [400, 401, 323, 430, 440, 450, 451], From 9659cd082aa7bb45f30164f91e5cfa862a6c0352 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 22 Mar 2023 13:21:30 +0000 Subject: [PATCH 26/30] flake8 fix --- esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py index cd1fba17..b83eef45 100644 --- a/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py +++ b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py @@ -228,8 +228,8 @@ def test__contiguous_bounds(): # fmt: off expected_vertices = np.array( [ - [0, 1, 11, 30, 40, 50, 51], # flake8: E241 - [100, 101, 12, 130, 140, 141, 52], # flake8: E241 + [0, 1, 11, 30, 40, 50, 51], # noqa: E241 + [100, 101, 12, 130, 140, 141, 52], # noqa: E241 [103, 210, 211, 133, 143, 142, 152], [203, 310, 320, 330, 331, 350, 351], [400, 401, 323, 430, 440, 450, 451], From eb1da3e330a7797ae86b3bad0006168dbf3ae928 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 22 Mar 2023 14:38:14 +0000 Subject: [PATCH 27/30] lazy mask calculations --- esmf_regrid/schemes.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index b55f8bef..5fd9d205 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -5,6 +5,7 @@ import functools from cf_units import Unit +import dask.array as da from iris._lazy_data import map_complete_blocks import iris.coords import iris.cube @@ -45,9 +46,10 @@ def _get_mask(cube, use_mask=True): 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) + full_mask = da.ma.getmaskarray(cube.core_data()) if not np.array_equal( - np.all(full_mask, axis=other_dims), np.any(full_mask, axis=other_dims) + da.all(full_mask, axis=other_dims).compute(), + da.any(full_mask, axis=other_dims).compute(), ): raise ValueError( "The mask derived from the cube is not constant over non-horizontal dimensions." From b309b4a882ddd20b608437ff81cf93edea84e58e Mon Sep 17 00:00:00 2001 From: stephenworsley <49274989+stephenworsley@users.noreply.github.com> Date: Wed, 22 Mar 2023 15:03:18 +0000 Subject: [PATCH 28/30] Apply suggestions from code review Co-authored-by: Martin Yeo <40734014+trexfeathers@users.noreply.github.com> --- esmf_regrid/experimental/unstructured_scheme.py | 16 ++++++++-------- esmf_regrid/schemes.py | 8 ++++---- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index 1c1381b5..a19bef60 100644 --- a/esmf_regrid/experimental/unstructured_scheme.py +++ b/esmf_regrid/experimental/unstructured_scheme.py @@ -299,12 +299,12 @@ 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 + 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_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 + 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 ``target_grid_cube``. If False, no mask will be taken and all points @@ -374,12 +374,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). - use_src_mask : :obj:`~numpy.typing.ArrayLike`, bool, default=False + 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_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 + 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 ``target_grid_cube``. If False, no mask will be taken and all points @@ -672,12 +672,12 @@ 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 + 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_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 + 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 ``target_grid_cube``. If False, no mask will be taken and all points @@ -747,12 +747,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). - use_src_mask : :obj:`~numpy.typing.ArrayLike`, bool, default=False + 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_grid_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 + 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 ``target_mesh_cube``. If False, no mask will be taken and all points diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 5fd9d205..18c015b1 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -472,10 +472,10 @@ def regridder(self, src_grid, tgt_grid, use_src_mask=None, use_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. - use_src_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional + use_src_mask : :obj:`~numpy.typing.ArrayLike` or bool, optional Array describing which elements :mod:`ESMF` will ignore on the src_grid. If True, the mask will be derived from src_grid. - use_tgt_mask : :obj:`~numpy.typing.ArrayLike`, bool, optional + use_tgt_mask : :obj:`~numpy.typing.ArrayLike` or bool, optional Array describing which elements :mod:`ESMF` will ignore on the tgt_grid. If True, the mask will be derived from tgt_grid. @@ -521,12 +521,12 @@ def __init__( 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. - use_src_mask : :obj:`~numpy.typing.ArrayLike`, bool, default=False + 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_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 + use_tgt_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 ``tgt_grid``. If False, no mask will be taken and all points will From 441e9c6e8cd529a2ead559157a03cdda4403f5de Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 22 Mar 2023 16:49:36 +0000 Subject: [PATCH 29/30] address review comments --- esmf_regrid/schemes.py | 7 +- .../unit/schemes/test__cube_to_GridInfo.py | 77 ++++++++++++------- 2 files changed, 52 insertions(+), 32 deletions(-) diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 18c015b1..1b26bef6 100644 --- a/esmf_regrid/schemes.py +++ b/esmf_regrid/schemes.py @@ -31,7 +31,7 @@ def _get_coord(cube, axis): def _get_mask(cube, use_mask=True): if use_mask is False: - return None + result = None elif use_mask is True: src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y")) @@ -61,9 +61,10 @@ def _get_mask(cube, use_mask=True): mask = mask.T else: mask = None - return mask + result = mask else: - return use_mask + result = use_mask + return result def _contiguous_masked(bounds, mask): diff --git a/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py index b83eef45..6d2ef88b 100644 --- a/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py +++ b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py @@ -5,6 +5,7 @@ from iris.cube import Cube from iris.fileformats.pp import EARTH_RADIUS import numpy as np +import pytest import scipy.sparse from esmf_regrid.esmf_regridder import Regridder @@ -197,7 +198,8 @@ def test_curvilinear_grid(): assert np.allclose(expected_weights.todense(), rg_circular.weight_matrix.todense()) -def test__contiguous_bounds(): +@pytest.mark.parametrize("masked", (True, False), ids=("masked", "unmasked")) +def test__contiguous_bounds(masked): """Test generation of contiguous bounds.""" # Generate a CF style bounds array with unique values for each bound. # The values will represent the index in the bounds array. @@ -206,35 +208,52 @@ def test__contiguous_bounds(): cf_bds += np.arange(6)[np.newaxis, :, np.newaxis] * 10 cf_bds += np.arange(4)[np.newaxis, np.newaxis, :] - # Define a mask such that each possible 2x2 sub array is represented. - # e.g. - # [[1, 1], - # [0, 1]] - mask = np.array( - [ - [0, 1, 1, 0, 0, 0], - [0, 1, 1, 0, 0, 1], - [1, 0, 1, 1, 1, 1], - [1, 0, 0, 0, 1, 0], - [0, 1, 1, 0, 0, 0], - ] - ) + if masked: + # Define a mask such that each possible 2x2 sub array is represented. + # e.g. + # [[1, 1], + # [0, 1]] + mask = np.array( + [ + [0, 1, 1, 0, 0, 0], + [0, 1, 1, 0, 0, 1], + [1, 0, 1, 1, 1, 1], + [1, 0, 0, 0, 1, 0], + [0, 1, 1, 0, 0, 0], + ] + ) + else: + mask = np.zeros([5, 6]) # Calculate the values on the vertices. vertices = _contiguous_masked(cf_bds, mask) - # Note, the only values deriving from masked points are: - # 11, 12, 152, 203, 412 - # These are the only vertices not connected to any unmasked points. - # fmt: off - expected_vertices = np.array( - [ - [0, 1, 11, 30, 40, 50, 51], # noqa: E241 - [100, 101, 12, 130, 140, 141, 52], # noqa: E241 - [103, 210, 211, 133, 143, 142, 152], - [203, 310, 320, 330, 331, 350, 351], - [400, 401, 323, 430, 440, 450, 451], - [403, 402, 412, 433, 443, 453, 452], - ] - ) - # fmt: on + if masked: + # Note, the only values deriving from masked points are: + # 11, 12, 152, 203, 412 + # These are the only vertices not connected to any unmasked points. + # fmt: off + expected_vertices = np.array( + [ + [0, 1, 11, 30, 40, 50, 51], # noqa: E241 + [100, 101, 12, 130, 140, 141, 52], # noqa: E241 + [103, 210, 211, 133, 143, 142, 152], + [203, 310, 320, 330, 331, 350, 351], + [400, 401, 323, 430, 440, 450, 451], + [403, 402, 412, 433, 443, 453, 452], + ] + ) + # fmt: on + else: + # fmt: off + expected_vertices = np.array( + [ + [0, 10, 20, 30, 40, 50, 51], # noqa: E241 + [100, 110, 120, 130, 140, 150, 151], + [200, 210, 220, 230, 240, 250, 251], + [300, 310, 320, 330, 340, 350, 351], + [400, 410, 420, 430, 440, 450, 451], + [403, 413, 423, 433, 443, 453, 452], + ] + ) + # fmt: on assert np.array_equal(expected_vertices, vertices) From e06c20eaf80d1122a919b8157b71eae498777efb Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Wed, 22 Mar 2023 16:59:48 +0000 Subject: [PATCH 30/30] add to changelog --- CHANGELOG.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b7d89e9..bc92a08b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,11 @@ and this project adheres to [Semantic Versioning](http://semver.org/). - [PR#220](https://github.com/SciTools-incubator/iris-esmf-regrid/pull/220) Matured the benchmarking architecture in line with the latest setup in SciTools/iris. +- [PR#241](https://github.com/SciTools-incubator/iris-esmf-regrid/pull/241) + Fixed compatibility with esmpy 8.4. +- [PR#219](https://github.com/SciTools-incubator/iris-esmf-regrid/pull/219) + Added support for 2D AuxCoords with discontiguities under masked values + with the use_src_mask and use_tgt_mask keywords. ## [0.5] - 2022-10-14