def test_flatten_spatial_dimensions(request, fixture_name):
"""Test the flatten_spatial_dimensions function"""
cube = request.getfixturevalue(fixture_name)
flat, lat_index, lon_index = flatten_spatial_dimensions(cube)
if cube.data.ndim == 2:
np.testing.assert_equal(flat, cube.data.flatten())
else:
assert flat.shape == (20, 3)
np.testing.assert_equal(flat[0:2, :], [[1, 21, 41], [2, 22, 42]])
def process(self, cube_in: Cube, cube_in_mask: Cube, cube_out_mask: Cube) -> Cube:
"""
Regridding considering land_sea mask. please note cube_in must use
lats/lons rectlinear system(GeogCS). cube_in_mask and cube_in could be
different resolution. cube_out could be either in lats/lons rectlinear
system or LambertAzimuthalEqualArea system. Grid points in cube_out
domain but not in cube_in domain will be masked.
Args:
cube_in:
Cube of data to be regridded.
cube_in_mask:
Cube of land_binary_mask data ((land:1, sea:0). used to determine
where the input model data is representing land and sea points.
cube_out_mask:
Cube of land_binary_mask data on target grid (land:1, sea:0).
Returns:
Regridded result cube.
"""
# if cube_in's coordinate descending, make it assending.
# if mask considered, reverse mask cube's coordinate if descending
cube_in = ensure_ascending_coord(cube_in)
if WITH_MASK in self.regrid_mode:
cube_in_mask = ensure_ascending_coord(cube_in_mask)
# check if input source grid is on even-spacing, ascending lat/lon system
# return grid spacing for latitude and logitude
lat_spacing, lon_spacing = calculate_input_grid_spacing(cube_in)
# Gather output latitude/longitudes from output template cube
if (
cube_out_mask.coord(axis="x").standard_name == "projection_x_coordinate"
and cube_out_mask.coord(axis="y").standard_name == "projection_y_coordinate"
):
out_latlons = np.dstack(transform_grid_to_lat_lon(cube_out_mask)).reshape(
(-1, 2)
)
else:
out_latlons = latlon_from_cube(cube_out_mask)
# Subset the input cube so that extra spatial area beyond the output is removed
# This is a performance optimisation to reduce the size of the dataset being processed
total_out_point_num = out_latlons.shape[0]
lat_max, lon_max = out_latlons.max(axis=0)
lat_min, lon_min = out_latlons.min(axis=0)
if WITH_MASK in self.regrid_mode:
cube_in, cube_in_mask = slice_mask_cube_by_domain(
cube_in, cube_in_mask, (lat_max, lon_max, lat_min, lon_min)
)
else: # not WITH_MASK
cube_in = slice_cube_by_domain(
cube_in, (lat_max, lon_max, lat_min, lon_min)
)
# group cube_out's grid points into outside or inside cube_in's domain
(
outside_input_domain_index,
inside_input_domain_index,
) = group_target_points_with_source_domain(cube_in, out_latlons)
# exclude out-of-input-domain target point here
if len(outside_input_domain_index) > 0:
out_latlons = out_latlons[inside_input_domain_index]
# Gather input latitude/longitudes from input cube
in_latlons = latlon_from_cube(cube_in)
# Number of grid points in X dimension is used to work out length of flattened array
# stripes for finding surrounding points for bilinear interpolation
in_lons_size = cube_in.coord(axis="x").shape[0] # longitude
# Reshape input data so that spatial dimensions can be handled as one
in_values, lats_index, lons_index = flatten_spatial_dimensions(cube_in)
# Locate nearby input points for output points
indexes = basic_indexes(
out_latlons, in_latlons, in_lons_size, lat_spacing, lon_spacing
)
if WITH_MASK in self.regrid_mode:
in_classified = classify_input_surface_type(cube_in_mask, in_latlons)
out_classified = classify_output_surface_type(cube_out_mask)
if len(outside_input_domain_index) > 0:
out_classified = out_classified[inside_input_domain_index]
# Identify mismatched surface types from input and output classifications
surface_type_mask = similar_surface_classify(
in_classified, out_classified, indexes
)
# Initialise distances and weights to zero. Weights are only used for the bilinear case
distances = np.zeros((out_latlons.shape[0], NUM_NEIGHBOURS), dtype=np.float32)
weights = np.zeros((out_latlons.shape[0], NUM_NEIGHBOURS), dtype=np.float32)
# handle nearest option
if NEAREST in self.regrid_mode:
for i in range(NUM_NEIGHBOURS):
distances[:, i] = np.square(
in_latlons[indexes[:, i], 0] - out_latlons[:, 0]
) + np.square(in_latlons[indexes[:, i], 1] - out_latlons[:, 1])
# for nearest-with-mask-2,adjust indexes and distance for mismatched
# surface type location
if WITH_MASK in self.regrid_mode:
distances, indexes = nearest_with_mask_regrid(
distances,
indexes,
surface_type_mask,
in_latlons,
out_latlons,
in_classified,
out_classified,
self.vicinity,
)
# apply nearest distance rule
output_flat = nearest_regrid(distances, indexes, in_values)
elif BILINEAR in self.regrid_mode:
# Assume all four nearby points are same surface type and calculate default weights
# These will be updated for mask/mismatched surface type further below
index_range = np.arange(weights.shape[0])
weights[index_range] = basic_weights(
index_range, indexes, out_latlons, in_latlons, lat_spacing, lon_spacing,
)
if WITH_MASK in self.regrid_mode:
# For bilinear-with-mask-2, adjust weights and indexes for mismatched
# surface type locations
weights, indexes = adjust_for_surface_mismatch(
in_latlons,
out_latlons,
in_classified,
out_classified,
weights,
indexes,
surface_type_mask,
in_lons_size,
self.vicinity,
lat_spacing,
lon_spacing,
)
# apply bilinear rule
output_flat = apply_weights(indexes, in_values, weights)
# check if we need mask cube_out grid points which are out of cube_in range
if len(outside_input_domain_index) > 0:
output_flat = mask_target_points_outside_source_domain(
total_out_point_num,
outside_input_domain_index,
inside_input_domain_index,
output_flat,
)
# Un-flatten spatial dimensions and put into output cube
output_array = unflatten_spatial_dimensions(
output_flat, cube_out_mask, in_values, lats_index, lons_index
)
output_cube = create_regrid_cube(output_array, cube_in, cube_out_mask)
return output_cube