From ed880838184511b84efd63072f137ff25c4e0933 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Thu, 16 Mar 2023 14:51:59 +0000 Subject: [PATCH] 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 cba13456..8bd1c508 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 942144e9..9b41a490 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:`ESMF` to be able to handle grids in different coordinate systems. """ - def __init__(self, mdtol=0, src_mask=False, tgt_mask=False): + def __init__(self, mdtol=0, use_src_mask=False, use_tgt_mask=False): """ Area-weighted scheme for regridding between rectilinear grids. @@ -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"))