def _regrid_area_weighted_rectilinear_src_and_grid__perform(
src_cube, regrid_info, mdtol):
"""
Second (regrid) part of 'regrid_area_weighted_rectilinear_src_and_grid'.
Perform the prepared regrid calculation on a single 2d cube.
"""
(
src_x,
src_y,
src_x_dim,
src_y_dim,
grid_x,
grid_y,
meshgrid_x,
meshgrid_y,
weights_info,
index_info,
) = regrid_info
# Calculate new data array for regridded cube.
regrid = functools.partial(
_regrid_area_weighted_array,
x_dim=src_x_dim,
y_dim=src_y_dim,
weights_info=weights_info,
index_info=index_info,
mdtol=mdtol,
)
new_data = map_complete_blocks(src_cube, regrid, (src_y_dim, src_x_dim),
meshgrid_x.shape)
# Wrap up the data as a Cube.
regrid_callback = RectilinearRegridder._regrid
new_cube = RectilinearRegridder._create_cube(
new_data,
src_cube,
src_x_dim,
src_y_dim,
src_x,
src_y,
grid_x,
grid_y,
meshgrid_x,
meshgrid_y,
regrid_callback,
)
# Slice out any length 1 dimensions.
indices = [slice(None, None)] * new_data.ndim
if src_x_dim is not None and new_cube.shape[src_x_dim] == 1:
indices[src_x_dim] = 0
if src_y_dim is not None and new_cube.shape[src_y_dim] == 1:
indices[src_y_dim] = 0
if 0 in indices:
new_cube = new_cube[tuple(indices)]
return new_cube
def _regrid_rectilinear_to_rectilinear__perform(src_cube, regrid_info, mdtol):
grid_x_dim, grid_y_dim, grid_x, grid_y, regridder = regrid_info
# Set up a function which can accept just chunk of data as an argument.
regrid = functools.partial(
_regrid_along_grid_dims,
regridder,
grid_x_dim=grid_x_dim,
grid_y_dim=grid_y_dim,
mdtol=mdtol,
)
# Apply regrid to all the chunks of src_cube, ensuring first that all
# chunks cover the entire horizontal plane (otherwise they would break
# the regrid function).
if len(grid_x.shape) == 1:
chunk_shape = (len(grid_x.points), len(grid_y.points))
else:
chunk_shape = grid_x.shape
new_data = map_complete_blocks(
src_cube,
regrid,
(grid_x_dim, grid_y_dim),
chunk_shape,
)
new_cube = _create_cube(
new_data,
src_cube,
(grid_x_dim, grid_y_dim),
(grid_x, grid_y),
2,
)
return new_cube
def test_rechunk(self):
lazy_array = da.asarray(self.array, chunks=((1, 1), (2, 2)))
cube, _ = create_mock_cube(lazy_array)
result = map_complete_blocks(cube,
self.func,
dims=(1, ),
out_sizes=(4, ))
self.assertTrue(is_lazy_data(result))
self.assertArrayEqual(result.compute(), self.func_result)
def test_multidimensional_input(self):
array = np.arange(2 * 3 * 4).reshape(2, 3, 4)
lazy_array = da.asarray(array, chunks=((1, 1), (1, 2), (4, )))
cube, _ = create_mock_cube(lazy_array)
result = map_complete_blocks(cube,
self.func,
dims=(1, 2),
out_sizes=(3, 4))
self.assertTrue(is_lazy_data(result))
self.assertArrayEqual(result.compute(), array + 1)
def test_different_out_shape(self):
lazy_array = da.asarray(self.array, chunks=((1, 1), (4, )))
cube, _ = create_mock_cube(lazy_array)
def func(_):
return np.arange(2).reshape(1, 2)
func_result = [[0, 1], [0, 1]]
result = map_complete_blocks(cube, func, dims=(1, ), out_sizes=(2, ))
self.assertTrue(is_lazy_data(result))
self.assertArrayEqual(result.compute(), func_result)
def test_lazy_input(self):
lazy_array = da.asarray(self.array, chunks=((1, 1), (4, )))
cube, cube_data = create_mock_cube(lazy_array)
result = map_complete_blocks(cube,
self.func,
dims=(1, ),
out_sizes=(4, ))
self.assertTrue(is_lazy_data(result))
self.assertArrayEqual(result.compute(), self.func_result)
# check correct data was accessed
cube.lazy_data.assert_called_once()
cube_data.assert_not_called()
def test_non_lazy_input(self):
# Check that a non-lazy input doesn't trip up the functionality.
cube, cube_data = create_mock_cube(self.array)
result = map_complete_blocks(cube,
self.func,
dims=(1, ),
out_sizes=(4, ))
self.assertFalse(is_lazy_data(result))
self.assertArrayEqual(result, self.func_result)
# check correct data was accessed
cube.lazy_data.assert_not_called()
cube_data.assert_called_once()
def __call__(self, src):
"""
Regrid this :class:`~iris.cube.Cube` on to the target grid of
this :class:`RectilinearRegridder`.
The given cube must be defined with the same grid as the source
grid used to create this :class:`RectilinearRegridder`.
If the source cube has lazy data, the returned cube will also
have lazy data.
Args:
* src:
A :class:`~iris.cube.Cube` to be regridded.
Returns:
A cube defined with the horizontal dimensions of the target
and the other dimensions from this cube. The data values of
this cube will be converted to values on the new grid using
either nearest-neighbour or linear interpolation.
.. note::
If the source cube has lazy data,
`chunks <https://docs.dask.org/en/latest/array-chunks.html>`__
in the horizontal dimensions will be combined before regridding.
"""
from iris.cube import Cube
# Validity checks.
if not isinstance(src, Cube):
raise TypeError("'src' must be a Cube")
if get_xy_dim_coords(src) != self._src_grid:
raise ValueError("The given cube is not defined on the same "
"source grid as this regridder.")
src_x_coord, src_y_coord = get_xy_dim_coords(src)
grid_x_coord, grid_y_coord = self._tgt_grid
src_cs = src_x_coord.coord_system
grid_cs = grid_x_coord.coord_system
if src_cs is None and grid_cs is None:
if not (src_x_coord.is_compatible(grid_x_coord)
and src_y_coord.is_compatible(grid_y_coord)):
raise ValueError("The rectilinear grid coordinates of the "
"given cube and target grid have no "
"coordinate system but they do not have "
"matching coordinate metadata.")
elif src_cs is None or grid_cs is None:
raise ValueError("The rectilinear grid coordinates of the given "
"cube and target grid must either both have "
"coordinate systems or both have no coordinate "
"system but with matching coordinate metadata.")
# Check the source grid units.
for coord in (src_x_coord, src_y_coord):
self._check_units(coord)
# Convert the grid to a 2D sample grid in the src CRS.
sample_grid = self._sample_grid(src_cs, grid_x_coord, grid_y_coord)
sample_grid_x, sample_grid_y = sample_grid
# Compute the interpolated data values.
x_dim = src.coord_dims(src_x_coord)[0]
y_dim = src.coord_dims(src_y_coord)[0]
# Define regrid function
regrid = functools.partial(
self._regrid,
x_dim=x_dim,
y_dim=y_dim,
src_x_coord=src_x_coord,
src_y_coord=src_y_coord,
sample_grid_x=sample_grid_x,
sample_grid_y=sample_grid_y,
method=self._method,
extrapolation_mode=self._extrapolation_mode,
)
data = map_complete_blocks(src, regrid, (y_dim, x_dim),
sample_grid_x.shape)
# Wrap up the data as a Cube.
regrid_callback = functools.partial(self._regrid,
method=self._method,
extrapolation_mode="nan")
result = self._create_cube(
data,
src,
x_dim,
y_dim,
src_x_coord,
src_y_coord,
grid_x_coord,
grid_y_coord,
sample_grid_x,
sample_grid_y,
regrid_callback,
)
return result