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 diff --git a/benchmarks/benchmarks/esmf_regridder/__init__.py b/benchmarks/benchmarks/esmf_regridder/__init__.py index aa399e2e..0a71f0eb 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,38 @@ 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_with_masks(self): + """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/benchmarks/benchmarks/generate_data.py b/benchmarks/benchmarks/generate_data.py index bcabac8b..17b3c7ae 100644 --- a/benchmarks/benchmarks/generate_data.py +++ b/benchmarks/benchmarks/generate_data.py @@ -161,6 +161,60 @@ 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. + + """ + 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/_esmf_sdo.py b/esmf_regrid/_esmf_sdo.py index 8c93f166..ee657908 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): @@ -42,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.""" @@ -62,6 +68,11 @@ def index_offset(self): """Return the index offset.""" return self._index_offset + @property + def mask(self): + """Return the mask.""" + return self._mask + def _array_to_matrix(self, array): """ Reshape data to a form that is compatible with weight matrices. @@ -111,6 +122,7 @@ def __init__( crs=None, circular=False, areas=None, + mask=None, center=False, ): """ @@ -140,6 +152,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. @@ -196,6 +210,7 @@ def __init__( shape=shape, index_offset=1, field_kwargs={"staggerloc": esmpy.StaggerLoc.CENTER}, + mask=mask, ) def _as_esmf_info(self): @@ -283,10 +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_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 @@ -314,6 +338,7 @@ def __init__( latbounds, resolution=3, crs=None, + mask=None, ): """ Create a :class:`RefinedGridInfo` object describing the grid. @@ -354,7 +379,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]) @@ -386,6 +411,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/esmf_regridder.py b/esmf_regrid/esmf_regridder.py index adf5a544..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,6 +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=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 7a07733e..08777652 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" @@ -33,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. @@ -69,6 +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) + _add_mask_to_cube(rg.src_mask, src_cube, SOURCE_MASK_NAME) tgt_mesh = rg.mesh tgt_location = rg.location @@ -77,6 +87,8 @@ 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) + _add_mask_to_cube(rg.tgt_mask, tgt_cube, TARGET_MASK_NAME) + elif regridder_type == "MeshToGridESMFRegridder": src_mesh = rg.mesh src_location = rg.location @@ -85,9 +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) + _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) + _add_mask_to_cube(rg.tgt_mask, tgt_cube, TARGET_MASK_NAME) else: msg = ( f"Expected a regridder of type `GridToMeshESMFRegridder` or " @@ -205,6 +219,15 @@ def load_regridder(filename): mdtol = weights_cube.attributes[MDTOL] + if src_cube.coords(SOURCE_MASK_NAME): + use_src_mask = src_cube.coord(SOURCE_MASK_NAME).points + else: + use_src_mask = False + if tgt_cube.coords(TARGET_MASK_NAME): + use_tgt_mask = tgt_cube.coord(TARGET_MASK_NAME).points + else: + use_tgt_mask = False + regridder = scheme( src_cube, tgt_cube, @@ -212,6 +235,8 @@ def load_regridder(filename): method=method, precomputed_weights=weight_matrix, resolution=resolution, + use_src_mask=use_src_mask, + use_tgt_mask=use_tgt_mask, ) esmf_version = weights_cube.attributes[VERSION_ESMF] 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/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py index f90d1823..a19bef60 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, + use_src_mask=False, + use_tgt_mask=False, ): r""" Regrid unstructured :class:`~iris.cube.Cube` onto rectilinear grid. @@ -292,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` 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` 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 + will be used in weights calculation. Returns ------- @@ -299,11 +316,16 @@ 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, 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 +342,8 @@ def __init__( method="conservative", precomputed_weights=None, resolution=None, + use_src_mask=False, + use_tgt_mask=False, ): """ Create regridder for conversions between source mesh and target grid. @@ -350,6 +374,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). + 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` 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 + will be used in weights calculation. + """ # TODO: Record information about the identity of the mesh. This would @@ -383,12 +418,17 @@ def __init__( ) self.resolution = resolution + self.src_mask = _get_mask(src_mesh_cube, use_src_mask) + self.tgt_mask = _get_mask(target_grid_cube, use_tgt_mask) + partial_regrid_info = _regrid_unstructured_to_rectilinear__prepare( src_mesh_cube, target_grid_cube, method=self.method, precomputed_weights=precomputed_weights, resolution=resolution, + src_mask=self.src_mask, + tgt_mask=self.tgt_mask, ) # Record source mesh. @@ -473,6 +513,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 +559,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 +622,8 @@ def regrid_rectilinear_to_unstructured( mdtol=0, method="conservative", resolution=None, + use_src_mask=False, + use_tgt_mask=False, ): r""" Regrid rectilinear :class:`~iris.cube.Cube` onto unstructured mesh. @@ -626,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` 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` 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 + will be used in weights calculation. Returns ------- @@ -633,11 +689,16 @@ 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, 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 +715,8 @@ def __init__( method="conservative", precomputed_weights=None, resolution=None, + use_src_mask=False, + use_tgt_mask=False, ): """ Create regridder for conversions between source grid and target mesh. @@ -684,6 +747,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). + 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` 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 + will be used in weights calculation. """ if method not in ["conservative", "bilinear"]: @@ -712,12 +785,17 @@ def __init__( self.method = method self.resolution = resolution + self.src_mask = _get_mask(src_grid_cube, use_src_mask) + self.tgt_mask = _get_mask(target_mesh_cube, use_tgt_mask) + partial_regrid_info = _regrid_rectilinear_to_unstructured__prepare( src_grid_cube, target_mesh_cube, method=self.method, precomputed_weights=precomputed_weights, resolution=self.resolution, + src_mask=self.src_mask, + tgt_mask=self.tgt_mask, ) # Store regrid info. diff --git a/esmf_regrid/schemes.py b/esmf_regrid/schemes.py index 94b064a6..1b26bef6 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 @@ -28,7 +29,125 @@ def _get_coord(cube, axis): return coord -def _cube_to_GridInfo(cube, center=False, resolution=None): +def _get_mask(cube, use_mask=True): + if use_mask is False: + result = None + 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)) + 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 = da.ma.getmaskarray(cube.core_data()) + if not np.array_equal( + 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." + "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 + result = mask + else: + result = use_mask + return result + + +def _contiguous_masked(bounds, mask): + """ + 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. 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: + + # 0 0 0 0 + # 0 - - 0 + # 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) + # (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 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[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[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. + filter_1 = filter_1 * (1 - filter_0) + + # Define the bounds which will derive from the bounds in index 3. + filter_3 = blank_template.copy() + 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. + 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 + + # Copy the bounds information over into the appropriate places. + contiguous_bounds = blank_template.astype(bounds.dtype) + 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 + + +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. # @@ -48,16 +167,22 @@ def _cube_to_GridInfo(cube, center=False, resolution=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) - assert lon.is_contiguous() - assert lat.is_contiguous() + if not np.any(mask): + 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 - lon_bound_array = lon.contiguous_bounds() 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( @@ -68,6 +193,7 @@ def _cube_to_GridInfo(cube, center=False, resolution=None): crs=crs, circular=circular, center=center, + mask=mask, ) else: grid_info = RefinedGridInfo( @@ -75,6 +201,7 @@ def _cube_to_GridInfo(cube, center=False, resolution=None): lat_bound_array, crs=crs, resolution=resolution, + mask=mask, ) return grid_info @@ -189,7 +316,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 +330,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 +430,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, use_src_mask=False, use_tgt_mask=False): """ Area-weighted scheme for regridding between rectilinear grids. @@ -315,18 +444,26 @@ def __init__(self, mdtol=0): 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. + use_src_mask : bool, default=False + If True, derive a mask from source cube which will tell :mod:`ESMF` + which points to ignore. + use_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): msg = "Value for mdtol must be in range 0 - 1, got {}." raise ValueError(msg.format(mdtol)) self.mdtol = mdtol + 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): + 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``. @@ -336,6 +473,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. + 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` 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. Returns ------- @@ -345,13 +488,25 @@ 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) + 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, + 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): + 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``. @@ -367,6 +522,16 @@ def __init__(self, src_grid, tgt_grid, mdtol=0): 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` 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` 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 + be used in weights calculation. """ if not (0 <= mdtol <= 1): @@ -374,7 +539,12 @@ 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) + 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=self.src_mask, tgt_mask=self.tgt_mask + ) # Store regrid info. self.grid_x = regrid_info.x_coord 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..157a92c2 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 ( @@ -18,7 +19,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 +46,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 + use_src_mask = True + tgt_data = ma.array(tgt.data) + tgt_data[0] = ma.masked + tgt.data = tgt_data + use_tgt_mask = True + else: + use_src_mask = False + use_tgt_mask = False + rg = GridToMeshESMFRegridder( - src, tgt, method=method, mdtol=0.5, resolution=resolution + src, + tgt, + method=method, + mdtol=0.5, + resolution=resolution, + use_src_mask=use_src_mask, + use_tgt_mask=use_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 +92,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 + use_src_mask = True + tgt_data = ma.array(tgt.data) + tgt_data[0, 0] = ma.masked + tgt.data = tgt_data + use_tgt_mask = True + else: + use_src_mask = False + use_tgt_mask = False + rg = MeshToGridESMFRegridder( - src, tgt, method=method, mdtol=0.5, resolution=resolution + src, + tgt, + method=method, + mdtol=0.5, + resolution=resolution, + use_src_mask=use_src_mask, + use_tgt_mask=use_tgt_mask, ) return rg, src @@ -195,6 +238,31 @@ def test_GridToMeshESMFRegridder_curvilinear_round_trip(tmp_path): assert np.array_equal(original_result.mask, loaded_result.mask) +# 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)) + + # 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) 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..c8e1106c 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,55 @@ def test_curvilinear(): result.data = expected_data assert expected_cube == result assert result_lazy == result + + +@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. + """ + 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) + 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() + 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, resolution=resolution + ) + 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 + 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..34c9b8eb 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,55 @@ def test_curvilinear(): result.data = expected_data assert expected_cube == result assert result_lazy == result + + +@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) + 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 + 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) + tgt_discontiguous = tgt.copy() + if resolution is None: + 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, resolution=resolution + ) + 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 + 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/schemes/test_ESMFAreaWeightedRegridder.py b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py index 1cfcb2e1..8e737689 100644 --- a/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py +++ b/esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py @@ -229,3 +229,46 @@ 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 `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]) + + # 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, 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 = 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 + 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/schemes/test__cube_to_GridInfo.py b/esmf_regrid/tests/unit/schemes/test__cube_to_GridInfo.py index 900d14ee..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,10 +5,11 @@ 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 -from esmf_regrid.schemes import _cube_to_GridInfo +from esmf_regrid.schemes import _contiguous_masked, _cube_to_GridInfo def _generate_points_and_bounds(n, outer_bounds): @@ -195,3 +196,64 @@ 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()) + + +@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. + 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, :] + + 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) + 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)