PI-2472: Tweak area weighting regrid move xdim and ydim axes (#3594)

* _regrid_area_weighted_array: Set axis order to y_dim, x_dim last dimensions

* _regrid_area_weighted_array: Extra tests for axes ordering

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
 
 

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]