PI-2472: Tweak area weighting regrid enforce xdim ydim

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
 
 

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]