def test_create_cube_2D():
"""Test creation of 2D output grid."""
data = np.ones([2, 3])
# Create a source cube with metadata and scalar coords
src_cube = Cube(np.zeros(5))
src_cube.units = "K"
src_cube.attributes = {"a": 1}
src_cube.standard_name = "air_temperature"
scalar_height = AuxCoord([5], units="m", standard_name="height")
scalar_time = DimCoord([10], units="s", standard_name="time")
src_cube.add_aux_coord(scalar_height)
src_cube.add_aux_coord(scalar_time)
mesh_dim = 0
grid_x = DimCoord(np.arange(3), standard_name="longitude")
grid_y = DimCoord(np.arange(2), standard_name="latitude")
cube = _create_cube(data, src_cube, (mesh_dim,), (grid_x, grid_y), 2)
src_metadata = src_cube.metadata
expected_cube = Cube(data)
expected_cube.metadata = src_metadata
expected_cube.add_dim_coord(grid_x, 1)
expected_cube.add_dim_coord(grid_y, 0)
expected_cube.add_aux_coord(scalar_height)
expected_cube.add_aux_coord(scalar_time)
assert expected_cube == cube
def add_coordinate_to_cube(
cube: Cube,
new_coord: DimCoord,
new_dim_location: int = 0,
copy_metadata: bool = True,
) -> Cube:
"""Create a copy of input cube with an additional dimension coordinate
added to the cube at the specified axis. The data from input cube is broadcast
over this new dimension.
Args:
cube:
cube to add realization dimension to.
new_coord:
new coordinate to add to input cube.
new_dim_location:
position in cube.data to position the new dimension coord. Default is
to add the new coordinate as the leading dimension.
copy_metadata:
flag as to whether to carry metadata over to output cube.
Returns:
A copy of cube broadcast over the new dimension coordinate.
"""
input_dim_count = len(cube.dim_coords)
if (new_dim_location > input_dim_count) or (new_dim_location < 0):
raise ValueError(
f"New dimension location: {new_dim_location} incompatible \
with cube containing {input_dim_count}.")
new_dim_coords = list(cube.dim_coords) + [new_coord]
new_dims = list(range(input_dim_count + 1))
new_dim_coords_and_dims = list(zip(new_dim_coords, new_dims))
aux_coords = cube.aux_coords
aux_coord_dims = [cube.coord_dims(coord.name()) for coord in aux_coords]
new_aux_coords_and_dims = list(zip(aux_coords, aux_coord_dims))
new_coord_size = len(new_coord.points)
new_data = np.broadcast_to(cube.data[..., np.newaxis],
shape=cube.shape + (new_coord_size, )).astype(
cube.data.dtype)
output_cube = Cube(
new_data,
dim_coords_and_dims=new_dim_coords_and_dims,
aux_coords_and_dims=new_aux_coords_and_dims,
)
if copy_metadata:
output_cube.metadata = cube.metadata
final_dim_order = np.insert(np.arange(input_dim_count),
new_dim_location,
values=input_dim_count)
output_cube.transpose(final_dim_order)
return output_cube
def _copy_cube_transformed(src_cube, data, coord_func):
"""
Returns a new cube based on the src_cube, but with the given data,
and with the coordinates transformed via coord_func.
The data must have the same number of dimensions as the source cube.
"""
from iris.cube import Cube
assert src_cube.ndim == data.ndim
# Start with just the metadata and the data...
new_cube = Cube(data)
new_cube.metadata = src_cube.metadata
new_cube.metadata = src_cube.metadata
# ... and then create all the coordinates.
# Record a mapping from old coordinate IDs to new coordinates,
# for subsequent use in creating updated aux_factories.
coord_mapping = {}
def copy_coords(source_coords, add_method):
for coord in source_coords:
new_coord = coord_func(coord)
add_method(new_coord, src_cube.coord_dims(coord))
coord_mapping[id(coord)] = new_coord
copy_coords(src_cube.dim_coords, new_cube.add_dim_coord)
copy_coords(src_cube.aux_coords, new_cube.add_aux_coord)
for factory in src_cube.aux_factories:
new_cube.add_aux_factory(factory.updated(coord_mapping))
return new_cube
def _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 a new Cube for the result of regridding the source Cube onto
the new grid.
All the metadata and coordinates of the result Cube are copied from
the source Cube, with two exceptions:
- Grid dimension coordinates are copied from the grid Cube.
- Auxiliary coordinates which span the grid dimensions are
ignored, except where they provide a reference surface for an
:class:`iris.aux_factory.AuxCoordFactory`.
Args:
* data:
The regridded data as an N-dimensional NumPy array.
* src:
The source Cube.
* x_dim:
The X dimension within the source Cube.
* y_dim:
The Y dimension within the source Cube.
* src_x_coord:
The X :class:`iris.coords.DimCoord`.
* src_y_coord:
The Y :class:`iris.coords.DimCoord`.
* grid_x_coord:
The :class:`iris.coords.DimCoord` for the new grid's X
coordinate.
* grid_y_coord:
The :class:`iris.coords.DimCoord` for the new grid's Y
coordinate.
* sample_grid_x:
A 2-dimensional array of sample X values.
* sample_grid_y:
A 2-dimensional array of sample Y values.
* regrid_callback:
The routine that will be used to calculate the interpolated
values of any reference surfaces.
Returns:
The new, regridded Cube.
"""
from iris.cube import Cube
#
# XXX: At the moment requires to be a static method as used by
# experimental regrid_area_weighted_rectilinear_src_and_grid
#
# Create a result cube with the appropriate metadata
result = Cube(data)
result.metadata = copy.deepcopy(src.metadata)
# Copy across all the coordinates which don't span the grid.
# Record a mapping from old coordinate IDs to new coordinates,
# for subsequent use in creating updated aux_factories.
coord_mapping = {}
def copy_coords(src_coords, add_method):
for coord in src_coords:
dims = src.coord_dims(coord)
if coord == src_x_coord:
coord = grid_x_coord
elif coord == src_y_coord:
coord = grid_y_coord
elif x_dim in dims or y_dim in dims:
continue
result_coord = coord.copy()
add_method(result_coord, dims)
coord_mapping[id(coord)] = result_coord
copy_coords(src.dim_coords, result.add_dim_coord)
copy_coords(src.aux_coords, result.add_aux_coord)
def regrid_reference_surface(
src_surface_coord,
surface_dims,
x_dim,
y_dim,
src_x_coord,
src_y_coord,
sample_grid_x,
sample_grid_y,
regrid_callback,
):
# Determine which of the reference surface's dimensions span the X
# and Y dimensions of the source cube.
surface_x_dim = surface_dims.index(x_dim)
surface_y_dim = surface_dims.index(y_dim)
surface = regrid_callback(
src_surface_coord.points,
surface_x_dim,
surface_y_dim,
src_x_coord,
src_y_coord,
sample_grid_x,
sample_grid_y,
)
surface_coord = src_surface_coord.copy(surface)
return surface_coord
# Copy across any AuxFactory instances, and regrid their reference
# surfaces where required.
for factory in src.aux_factories:
for coord in factory.dependencies.values():
if coord is None:
continue
dims = src.coord_dims(coord)
if x_dim in dims and y_dim in dims:
result_coord = regrid_reference_surface(
coord,
dims,
x_dim,
y_dim,
src_x_coord,
src_y_coord,
sample_grid_x,
sample_grid_y,
regrid_callback,
)
result.add_aux_coord(result_coord, dims)
coord_mapping[id(coord)] = result_coord
try:
result.add_aux_factory(factory.updated(coord_mapping))
except KeyError:
msg = ("Cannot update aux_factory {!r} because of dropped"
" coordinates.".format(factory.name()))
warnings.warn(msg)
return result
def _regrid_weighted_curvilinear_to_rectilinear__perform(
src_cube, regrid_info):
"""
Second (regrid) part of 'regrid_weighted_curvilinear_to_rectilinear'.
Perform the prepared regrid calculation on a single 2d cube.
"""
from iris.cube import Cube
sparse_matrix, sum_weights, rows, grid_cube = regrid_info
# Calculate the numerator of the weighted mean (M, 1).
is_masked = ma.isMaskedArray(src_cube.data)
if not is_masked:
data = src_cube.data
else:
# Use raw data array
data = src_cube.data.data
# Check if there are any masked source points to take account of.
is_masked = np.ma.is_masked(src_cube.data)
if is_masked:
# Zero any masked source points so they add nothing in output sums.
mask = src_cube.data.mask
data[mask] = 0.0
# Calculate a new 'sum_weights' to allow for missing source points.
# N.B. it is more efficient to use the original once-calculated
# sparse matrix, but in this case we can't.
# Hopefully, this post-multiplying by the validities is less costly
# than repeating the whole sparse calculation.
valid_src_cells = ~mask.flat[:]
src_cell_validity_factors = sparse_diags(
np.array(valid_src_cells, dtype=int), 0)
valid_weights = sparse_matrix * src_cell_validity_factors
sum_weights = valid_weights.sum(axis=1).getA()
# Work out where output cells are missing all contributions.
# This allows for where 'rows' contains output cells that have no
# data because of missing input points.
zero_sums = sum_weights == 0.0
# Make sure we can still divide by sum_weights[rows].
sum_weights[zero_sums] = 1.0
# Calculate sum in each target cell, over contributions from each source
# cell.
numerator = sparse_matrix * data.reshape(-1, 1)
# Create a template for the weighted mean result.
weighted_mean = ma.masked_all(numerator.shape, dtype=numerator.dtype)
# Calculate final results in all relevant places.
weighted_mean[rows] = numerator[rows] / sum_weights[rows]
if is_masked:
# Ensure masked points where relevant source cells were all missing.
if np.any(zero_sums):
# Make masked if it wasn't.
weighted_mean = np.ma.asarray(weighted_mean)
# Mask where contributing sums were zero.
weighted_mean[zero_sums] = np.ma.masked
# Construct the final regridded weighted mean cube.
tx = grid_cube.coord(axis="x", dim_coords=True)
ty = grid_cube.coord(axis="y", dim_coords=True)
(tx_dim, ) = grid_cube.coord_dims(tx)
(ty_dim, ) = grid_cube.coord_dims(ty)
dim_coords_and_dims = list(zip((ty.copy(), tx.copy()), (ty_dim, tx_dim)))
cube = Cube(
weighted_mean.reshape(grid_cube.shape),
dim_coords_and_dims=dim_coords_and_dims,
)
cube.metadata = copy.deepcopy(src_cube.metadata)
for coord in src_cube.coords(dimensions=()):
cube.add_aux_coord(coord.copy())
return cube
