diff --git a/lib/iris/experimental/regrid.py b/lib/iris/experimental/regrid.py index afbf87dee8..0838220f3c 100644 --- a/lib/iris/experimental/regrid.py +++ b/lib/iris/experimental/regrid.py @@ -491,14 +491,30 @@ def _regrid_area_weighted_array( cached_x_bounds.append(x_bounds) cached_x_indices.append(x_indices) + # Ensure we have x_dim and y_dim. + x_dim_orig = copy.copy(x_dim) + y_dim_orig = copy.copy(y_dim) + if y_dim is None: + src_data = np.expand_dims(src_data, axis=src_data.ndim) + y_dim = src_data.ndim - 1 + if x_dim is None: + src_data = np.expand_dims(src_data, axis=src_data.ndim) + x_dim = src_data.ndim - 1 + # Move y_dim and x_dim to last dimensions + src_data = np.moveaxis(src_data, x_dim, -1) + if x_dim < y_dim: + src_data = np.moveaxis(src_data, y_dim - 1, -2) + elif x_dim > y_dim: + src_data = np.moveaxis(src_data, y_dim, -2) + x_dim = src_data.ndim - 1 + y_dim = src_data.ndim - 2 + # Create empty data array to match the new grid. # Note that dtype is not preserved and that the array is # masked to allow for regions that do not overlap. new_shape = list(src_data.shape) - if x_dim is not None: - new_shape[x_dim] = grid_x_bounds.shape[0] - if y_dim is not None: - new_shape[y_dim] = grid_y_bounds.shape[0] + new_shape[x_dim] = grid_x_bounds.shape[0] + new_shape[y_dim] = grid_y_bounds.shape[0] # Use input cube dtype or convert values to the smallest possible float # dtype when necessary. @@ -516,15 +532,9 @@ def _regrid_area_weighted_array( new_data.mask = False # Axes of data over which the weighted mean is calculated. - axes = [] - if y_dim is not None: - axes.append(y_dim) - if x_dim is not None: - axes.append(x_dim) - axis = tuple(axes) + axis = (y_dim, x_dim) # Simple for loop approach. - indices = [slice(None)] * new_data.ndim for j, (y_0, y_1) in enumerate(grid_y_bounds): # Reverse lower and upper if dest grid is decreasing. if grid_y_decreasing: @@ -550,11 +560,7 @@ def _regrid_area_weighted_array( or not x_within_bounds[i] ): # Mask out element(s) in new_data - if x_dim is not None: - indices[x_dim] = i - if y_dim is not None: - indices[y_dim] = j - new_data[tuple(indices)] = ma.masked + new_data[..., j, i] = ma.masked else: # Calculate weighted mean of data points. # Slice out relevant data (this may or may not be a view() @@ -568,24 +574,16 @@ def _regrid_area_weighted_array( # Calculate weights based on areas of cropped bounds. weights = area_func(y_bounds, x_bounds) - if x_dim is not None: - indices[x_dim] = x_indices - if y_dim is not None: - indices[y_dim] = y_indices - data = src_data[tuple(indices)] + data = src_data[..., y_indices, x_indices] # Transpose weights to match dim ordering in data. weights_shape_y = weights.shape[0] weights_shape_x = weights.shape[1] - if x_dim is not None and y_dim is not None and x_dim < y_dim: - weights = weights.T # Broadcast the weights array to allow numpy's ma.average # to be called. weights_padded_shape = [1] * data.ndim - if y_dim is not None: - weights_padded_shape[y_dim] = weights_shape_y - if x_dim is not None: - weights_padded_shape[x_dim] = weights_shape_x + weights_padded_shape[y_dim] = weights_shape_y + weights_padded_shape[x_dim] = weights_shape_x # Assign new shape to raise error on copy. weights.shape = weights_padded_shape # Broadcast weights to match shape of data. @@ -597,17 +595,30 @@ def _regrid_area_weighted_array( ) # Insert data (and mask) values into new array. - if x_dim is not None: - indices[x_dim] = i - if y_dim is not None: - indices[y_dim] = j - new_data[tuple(indices)] = new_data_pt + new_data[..., j, i] = new_data_pt # Remove new mask if original data was not masked # and no values in the new array are masked. if not src_masked and not new_data.mask.any(): new_data = new_data.data + # Restore axis to original order + if x_dim_orig is None and y_dim_orig is None: + new_data = np.squeeze(new_data, axis=x_dim) + new_data = np.squeeze(new_data, axis=y_dim) + elif y_dim_orig is None: + new_data = np.squeeze(new_data, axis=y_dim) + new_data = np.moveaxis(new_data, -1, x_dim_orig) + elif x_dim_orig is None: + new_data = np.squeeze(new_data, axis=x_dim) + new_data = np.moveaxis(new_data, -1, y_dim_orig) + elif x_dim_orig < y_dim_orig: + new_data = np.moveaxis(new_data, -1, x_dim_orig) + new_data = np.moveaxis(new_data, -1, y_dim_orig) + else: + new_data = np.moveaxis(new_data, -2, y_dim_orig) + new_data = np.moveaxis(new_data, -1, x_dim_orig) + return new_data diff --git a/lib/iris/tests/experimental/regrid/test_regrid_area_weighted_rectilinear_src_and_grid.py b/lib/iris/tests/experimental/regrid/test_regrid_area_weighted_rectilinear_src_and_grid.py index 97127cb0b4..289db98560 100644 --- a/lib/iris/tests/experimental/regrid/test_regrid_area_weighted_rectilinear_src_and_grid.py +++ b/lib/iris/tests/experimental/regrid/test_regrid_area_weighted_rectilinear_src_and_grid.py @@ -322,19 +322,36 @@ def test_regrid_latlon_reduced_res(self): res = regrid_area_weighted(src, dest) self.assertCMLApproxData(res, RESULT_DIR + ("latlonreduced.cml",)) - def test_regrid_transposed(self): - src = self.simple_cube.copy() - dest = _subsampled_grid(src, 2, 3) - # Transpose src so that the coords are not y, x ordered. - src.transpose() + def test_regrid_reorder_axis(self): + src = self.realistic_cube[0, :4, :3, :2] + z = src.coord("model_level_number") + lat = src.coord("grid_latitude") + lon = src.coord("grid_longitude") + dest = _resampled_grid(self.realistic_cube[0, 0, :3, :2], 3, 3) res = regrid_area_weighted(src, dest) - self.assertCMLApproxData(res, RESULT_DIR + ("trasposed.cml",)) - # Using original and transposing the result should give the - # same answer. - src = self.simple_cube.copy() + self.assertArrayShapeStats(src, (4, 3, 2), 288.08868, 0.008262919) + self.assertArrayShapeStats(res, (4, 9, 6), 288.08865, 0.00826281) + # Reshape src so that the coords are ordered [x, z, y], + # the mean and std statistics should be the same + data = np.moveaxis(src.data.copy(), 2, 0) + src = iris.cube.Cube(data) + src.add_dim_coord(lat, 2) + src.add_dim_coord(z, 1) + src.add_dim_coord(lon, 0) res = regrid_area_weighted(src, dest) - res.transpose() - self.assertCMLApproxData(res, RESULT_DIR + ("trasposed.cml",)) + self.assertArrayShapeStats(src, (2, 4, 3), 288.08868, 0.008262919) + self.assertArrayShapeStats(res, (6, 4, 9), 288.08865, 0.00826281) + # Reshape src so that the coords are ordered [y, x, z], + # the mean and std statistics should be the same + data = np.moveaxis(src.data.copy(), 2, 0) + src = iris.cube.Cube(data) + src.add_dim_coord(z, 2) + src.add_dim_coord(lon, 1) + src.add_dim_coord(lat, 0) + dest = _resampled_grid(self.realistic_cube[0, 0, :3, :2], 3, 3) + res = regrid_area_weighted(src, dest) + self.assertArrayShapeStats(src, (3, 2, 4), 288.08868, 0.008262919) + self.assertArrayShapeStats(res, (9, 6, 4), 288.08865, 0.00826281) def test_regrid_lon_to_half_res(self): src = self.simple_cube @@ -446,6 +463,16 @@ def test_cross_section(self): self.assertCMLApproxData( res, RESULT_DIR + ("const_lat_cross_section.cml",) ) + # Constant latitude, data order [x, z] + # Using original and transposing the result should give the + # same answer. + src.transpose() + dest.transpose() + res = regrid_area_weighted(src, dest) + res.transpose() + self.assertCMLApproxData( + res, RESULT_DIR + ("const_lat_cross_section.cml",) + ) # Constant longitude src = self.realistic_cube[0, :, :, 10] @@ -460,6 +487,16 @@ def test_cross_section(self): self.assertCMLApproxData( res, RESULT_DIR + ("const_lon_cross_section.cml",) ) + # Constant longitude, data order [y, z] + # Using original and transposing the result should give the + # same answer. + src.transpose() + dest.transpose() + res = regrid_area_weighted(src, dest) + res.transpose() + self.assertCMLApproxData( + res, RESULT_DIR + ("const_lon_cross_section.cml",) + ) def test_scalar_source_cube(self): src = self.simple_cube[1, 2] diff --git a/lib/iris/tests/results/experimental/regrid/regrid_area_weighted_rectilinear_src_and_grid/trasposed.cml b/lib/iris/tests/results/experimental/regrid/regrid_area_weighted_rectilinear_src_and_grid/trasposed.cml deleted file mode 100644 index 9aecbfc388..0000000000 --- a/lib/iris/tests/results/experimental/regrid/regrid_area_weighted_rectilinear_src_and_grid/trasposed.cml +++ /dev/null @@ -1,20 +0,0 @@ - - - - - - - - - - - - - - - - - - - diff --git a/lib/iris/tests/results/experimental/regrid/regrid_area_weighted_rectilinear_src_and_grid/trasposed.data.0.json b/lib/iris/tests/results/experimental/regrid/regrid_area_weighted_rectilinear_src_and_grid/trasposed.data.0.json deleted file mode 100644 index 0525283f0a..0000000000 --- a/lib/iris/tests/results/experimental/regrid/regrid_area_weighted_rectilinear_src_and_grid/trasposed.data.0.json +++ /dev/null @@ -1 +0,0 @@ -{"std": 1.0, "min": 4.499593812507495, "max": 6.499593812507495, "shape": [2], "masked": false, "mean": 5.499593812507495} \ No newline at end of file