Allow regridding of cubes with multiple coordinates per axis

CHANGELOG.md

@@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## [Unreleased]
+
+### Added
+
+- [PR#217](https://github.com/SciTools-incubator/iris-esmf-regrid/pull/217)
+  Changed the behaviour of coordinate fetching to allow Cubes with both
+  1D DimCoords and 2D AuxCoords. In this case the DimCoords are prioritised.
+
 ## [0.5] - 2022-10-14
 
 This release improves the support for features such as Bilinear regridding,

esmf_regrid/experimental/unstructured_scheme.py

@@ -7,7 +7,7 @@
 
 from esmf_regrid.esmf_regridder import Regridder
 from esmf_regrid.experimental.unstructured_regrid import MeshInfo
-from esmf_regrid.schemes import _create_cube, _cube_to_GridInfo
+from esmf_regrid.schemes import _create_cube, _cube_to_GridInfo, _get_coord
 
 
 def _map_complete_blocks(src, func, dims, out_sizes):
@@ -152,8 +152,8 @@ def _regrid_unstructured_to_rectilinear__prepare(
     The 'regrid info' returned can be re-used over many 2d slices.
 
     """
-    grid_x = target_grid_cube.coord(axis="x")
-    grid_y = target_grid_cube.coord(axis="y")
+    grid_x = _get_coord(target_grid_cube, "x")
+    grid_y = _get_coord(target_grid_cube, "y")
     mesh = src_mesh_cube.mesh
     location = src_mesh_cube.location
     if mesh is None:
@@ -480,8 +480,8 @@ def _regrid_rectilinear_to_unstructured__prepare(
     The 'regrid info' returned can be re-used over many 2d slices.
 
     """
-    grid_x = src_grid_cube.coord(axis="x")
-    grid_y = src_grid_cube.coord(axis="y")
+    grid_x = _get_coord(src_grid_cube, "x")
+    grid_y = _get_coord(src_grid_cube, "y")
     mesh = target_mesh_cube.mesh
     location = target_mesh_cube.location
     if mesh is None:
@@ -752,8 +752,8 @@ def __call__(self, cube):
             area-weighted regridding via :mod:`ESMF` generated weights.
 
         """
-        grid_x = cube.coord(axis="x")
-        grid_y = cube.coord(axis="y")
+        grid_x = _get_coord(cube, "x")
+        grid_y = _get_coord(cube, "y")
         # Ignore differences in var_name that might be caused by saving.
         self_grid_x = copy.deepcopy(self.grid_x)
         self_grid_x.var_name = grid_x.var_name

esmf_regrid/schemes.py

@@ -8,6 +8,7 @@
 from iris._lazy_data import map_complete_blocks
 import iris.coords
 import iris.cube
+from iris.exceptions import CoordinateNotFoundError
 import numpy as np
 
 from esmf_regrid.esmf_regridder import GridInfo, RefinedGridInfo, Regridder
@@ -19,13 +20,21 @@
 ]
 
 
+def _get_coord(cube, axis):
+    try:
+        coord = cube.coord(axis=axis, dim_coords=True)
+    except CoordinateNotFoundError:
+        coord = cube.coord(axis=axis)
+    return coord
+
+
 def _cube_to_GridInfo(cube, center=False, resolution=None):
     # This is a simplified version of an equivalent function/method in PR #26.
     # It is anticipated that this function will be replaced by the one in PR #26.
     #
     # Returns a GridInfo object describing the horizontal grid of the cube.
     # This may be inherited from code written for the rectilinear regridding scheme.
-    lon, lat = cube.coord(axis="x"), cube.coord(axis="y")
+    lon, lat = _get_coord(cube, "x"), _get_coord(cube, "y")
     #  Checks may cover units, coord systems (e.g. rotated pole), contiguous bounds.
     if cube.coord_system() is None:
         crs = None
@@ -181,10 +190,10 @@ def copy_coords(src_coords, add_method):
 
 
 def _regrid_rectilinear_to_rectilinear__prepare(src_grid_cube, tgt_grid_cube):
-    tgt_x = tgt_grid_cube.coord(axis="x")
-    tgt_y = tgt_grid_cube.coord(axis="y")
-    src_x = src_grid_cube.coord(axis="x")
-    src_y = src_grid_cube.coord(axis="y")
+    tgt_x = _get_coord(tgt_grid_cube, "x")
+    tgt_y = _get_coord(tgt_grid_cube, "y")
+    src_x = _get_coord(src_grid_cube, "x")
+    src_y = _get_coord(src_grid_cube, "y")
 
     if len(src_x.shape) == 1:
         grid_x_dim = src_grid_cube.coord_dims(src_x)[0]
@@ -372,7 +381,7 @@ def __init__(self, src_grid, tgt_grid, mdtol=0):
         self.regridder = regrid_info.regridder
 
         # Record the source grid.
-        self.src_grid = (src_grid.coord(axis="x"), src_grid.coord(axis="y"))
+        self.src_grid = (_get_coord(src_grid, "x"), _get_coord(src_grid, "y"))
 
     def __call__(self, cube):
         """
@@ -395,7 +404,7 @@ def __call__(self, cube):
             area-weighted regridding via :mod:`ESMF` generated weights.
 
         """
-        src_x, src_y = (cube.coord(axis="x"), cube.coord(axis="y"))
+        src_x, src_y = (_get_coord(cube, "x"), _get_coord(cube, "y"))
 
         # Check the source grid matches that used in initialisation
         if self.src_grid != (src_x, src_y):

esmf_regrid/tests/unit/schemes/test_ESMFAreaWeightedRegridder.py

@@ -118,6 +118,59 @@ def extract_weights(regridder):
         assert np.allclose(extract_weights(grid_to_grid), extract_weights(regridder))
 
 
+def test_curvilinear_and_rectilinear():
+    """
+    Test :func:`esmf_regrid.schemes.ESMFAreaWeightedRegridder`.
+
+    Checks that a cube with both curvilinear and rectilinear coords still works.
+    Checks that the DimCoords have priority over AuxCoords.
+    """
+    n_lons_src = 6
+    n_lons_tgt = 3
+    n_lats_src = 4
+    n_lats_tgt = 2
+    lon_bounds = (-180, 180)
+    lat_bounds = (-90, 90)
+    grid_src = _grid_cube(n_lons_src, n_lats_src, lon_bounds, lat_bounds, circular=True)
+    grid_tgt = _grid_cube(n_lons_tgt, n_lats_tgt, lon_bounds, lat_bounds, circular=True)
+    curv_src = _curvilinear_cube(n_lons_src, n_lats_src, lon_bounds, lat_bounds)
+    curv_tgt = _curvilinear_cube(n_lons_tgt, n_lats_tgt, lon_bounds, lat_bounds)
+
+    src = curv_src.copy()
+    grid_lat_src = grid_src.coord("latitude")
+    grid_lat_src.standard_name = "grid_latitude"
+    src.add_dim_coord(grid_lat_src, 0)
+    grid_lon_src = grid_src.coord("longitude")
+    grid_lon_src.standard_name = "grid_longitude"
+    src.add_dim_coord(grid_lon_src, 1)
+
+    tgt = curv_tgt.copy()
+    grid_lat_tgt = grid_tgt.coord("latitude")
+    grid_lat_tgt.standard_name = "grid_latitude"
+    tgt.add_dim_coord(grid_lat_tgt, 0)
+    grid_lon_tgt = grid_tgt.coord("longitude")
+    grid_lon_tgt.standard_name = "grid_longitude"
+    tgt.add_dim_coord(grid_lon_tgt, 1)
+
+    # Change the AuxCoords to check that the DimCoords have priority.
+    src.coord("latitude").bounds[:] = 0
+    src.coord("longitude").bounds[:] = 0
+    tgt.coord("latitude").bounds[:] = 0
+    tgt.coord("longitude").bounds[:] = 0
+
+    # Ensure the source data is all ones.
+    src.data[:] = 1
+
+    rg = ESMFAreaWeightedRegridder(src, tgt)
+    result = rg(src)
+
+    expected = grid_tgt.copy()
+    # If the aux coords had been prioritised, expected.data would be a fully masked array.
+    expected.data[:] = 1
+    assert expected == result
+    assert not np.ma.is_masked(result)
+
+
 def test_unit_equivalence():
     """
     Test initialisation of :func:`esmf_regrid.schemes.ESMFAreaWeightedRegridder`.