diff --git a/lib/iris/experimental/regrid.py b/lib/iris/experimental/regrid.py index afbf87dee8..ba4aff035a 100644 --- a/lib/iris/experimental/regrid.py +++ b/lib/iris/experimental/regrid.py @@ -491,6 +491,31 @@ def _regrid_area_weighted_array( cached_x_bounds.append(x_bounds) cached_x_indices.append(x_indices) + # Move y_dim and x_dim to last dimensions + x_dim_orig = copy.copy(x_dim) + y_dim_orig = copy.copy(y_dim) + if x_dim is None and y_dim is None: + # e.g. a scalar point such as a vertical profile + pass + elif x_dim is not None and y_dim is None: + # test cross_section along line latitude + src_data = np.moveaxis(src_data, x_dim, -1) + x_dim = src_data.ndim - 1 + elif y_dim is not None and x_dim is None: + # test cross_section along line longitude + src_data = np.moveaxis(src_data, y_dim, -1) + y_dim = src_data.ndim - 1 + elif x_dim < y_dim: + src_data = np.moveaxis(src_data, x_dim, -1) + src_data = np.moveaxis(src_data, y_dim - 1, -2) + x_dim = src_data.ndim - 1 + y_dim = src_data.ndim - 2 + else: + src_data = np.moveaxis(src_data, x_dim, -1) + 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. @@ -577,8 +602,6 @@ def _regrid_area_weighted_array( # 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 @@ -608,6 +631,20 @@ def _regrid_area_weighted_array( 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: + pass + elif x_dim_orig is not None and y_dim_orig is None: + new_data = np.moveaxis(new_data, -1, x_dim_orig) + elif y_dim_orig is not None and x_dim_orig is None: + 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