Copy additional coordinates when regridding

esmvalcore/preprocessor/_regrid.py

@@ -17,6 +17,7 @@
 import dask.array as da
 import iris
 import iris.coords
+import iris.exceptions
 import numpy as np
 import stratify
 from geopy.geocoders import Nominatim
@@ -791,17 +792,16 @@ def regrid(
     Note that the target grid can be a :class:`~iris.cube.Cube`, a
     :class:`~esmvalcore.dataset.Dataset`, a path to a cube
     (:class:`~pathlib.Path` or :obj:`str`), a grid spec (:obj:`str`) in the
-    form of `MxN`, or a :obj:`dict` specifying the target grid.
+    form of `MxN` degrees, or a :obj:`dict` specifying the target grid.
 
     For the latter, the `target_grid` should be a :obj:`dict` with the
     following keys:
 
     - ``start_longitude``: longitude at the center of the first grid cell.
     - ``end_longitude``: longitude at the center of the last grid cell.
-    - ``step_longitude``: constant longitude distance between grid cell
-        centers.
+    - ``step_longitude``: constant longitude distance between grid cell centers.
     - ``start_latitude``: latitude at the center of the first grid cell.
-    - ``end_latitude``: longitude at the center of the last grid cell.
+    - ``end_latitude``: latitude at the center of the last grid cell.
     - ``step_latitude``: constant latitude distance between grid cell centers.
 
     Parameters
@@ -910,18 +910,9 @@ def regrid(
 
     # Horizontal grids from source and target (almost) match
     # -> Return source cube with target coordinates
-    if cube.coords("latitude") and cube.coords("longitude"):
-        if _horizontal_grid_is_close(cube, target_grid_cube):
-            for coord in ["latitude", "longitude"]:
-                is_dim_coord = cube.coords(coord, dim_coords=True)
-                coord_dims = cube.coord_dims(coord)
-                cube.remove_coord(coord)
-                target_coord = target_grid_cube.coord(coord).copy()
-                if is_dim_coord:
-                    cube.add_dim_coord(target_coord, coord_dims)
-                else:
-                    cube.add_aux_coord(target_coord, coord_dims)
-            return cube
+    if _horizontal_grid_is_close(cube, target_grid_cube):
+        _update_horizontal_coords(target_grid_cube, cube, overwrite=True)
+        return cube
 
     # Load scheme and reuse existing regridder if possible
     if isinstance(scheme, str):
@@ -930,7 +921,22 @@ def regrid(
 
     # Rechunk and actually perform the regridding
     cube = _rechunk(cube, target_grid_cube)
-    return regridder(cube)
+    result = regridder(cube)
+    # Iris only supports regridding of 1D coordinates and iris-esmf-regrid
+    # uses only the DimCoords if both grid_latitude/grid_longitude or
+    # projection_x_coordinate/projection_y_coordinate DimCoords and latitude
+    # longitude AuxCoords are present, so we need to copy the 2D lat/lon
+    # AuxCoords from the target grid if they are not present on the result.
+    #
+    # Similarly, when the target grid is a 2D latitude/longitude grid,
+    # regridding does not use the DimCoord because it has no meaning. This
+    # is typical in CMIP6 ocean data (e.g. coordinates named i/j, cell index
+    # along first/second dimension), so we need to copy the dummy DimCoords
+    # from the target grid in this case to ensure the resulting cube has the
+    # same coordinates when using the regridded cubes as input to the
+    # multi-model statistics or similar preprocessor functions later on.
+    _update_horizontal_coords(target_grid_cube, result, overwrite=False)
+    return result
 
 
 def _cache_clear():
@@ -1006,20 +1012,78 @@ def _horizontal_grid_is_close(cube1: Cube, cube2: Cube) -> bool:
     bool
         ``True`` if grids are close; ``False`` if not.
     """
-    # Go through the 2 expected horizontal coordinates longitude and latitude.
-    for coord in ["latitude", "longitude"]:
-        coord1 = cube1.coord(coord)
-        coord2 = cube2.coord(coord)
-
-        if coord1.shape != coord2.shape:
+    for axis in ("X", "Y"):
+        # DimCoords are kept in-memory, so prefer those for performance.
+        try:
+            coord1 = cube1.coord(axis=axis, dim_coords=True)
+        except iris.exceptions.CoordinateNotFoundError:
+            try:
+                coord1 = cube1.coord(axis=axis, dim_coords=False)
+            except iris.exceptions.CoordinateNotFoundError:
+                # No horizontal coordinate found.
+                return False
+
+        for coord2 in cube2.coords(axis=axis):
+            if coord2.shape == coord1.shape:
+                break
+        else:
+            # No matching coordinate found.
             return False
 
-        if not np.allclose(coord1.bounds, coord2.bounds):
+        if coord1.has_bounds() and coord2.has_bounds():
+            array1 = coord1.core_bounds()
+            array2 = coord2.core_bounds()
+        else:
+            array1 = coord1.core_points()
+            array2 = coord2.core_points()
+
+        if not np.allclose(array1, array2):
             return False
 
     return True
 
 
+def _update_horizontal_coords(src: Cube, tgt: Cube, overwrite: bool) -> None:
+    """Update horizontal coordinates.
+
+    Parameters
+    ----------
+    src:
+        Source cube to copy the coordinates from.
+    tgt:
+        Target cube to copy the coordinates to.
+    overwrite:
+        If ``True``, overwrite existing coordinates on the target cube. If
+        ``False``, only add the coordinates if they are not already present on
+        the target cube.
+    """
+    src_horizontal_dims = get_dims_along_axes(src, ("X", "Y"))
+    tgt_horizontal_dims = get_dims_along_axes(tgt, ("X", "Y"))
+    if len(src_horizontal_dims) != len(tgt_horizontal_dims):
+        return
+
+    # Copy DimCoords. This assumes they are in the same order in the source
+    # and target cube, which seems to be required for the regridders.
+    for src_dim, tgt_dim in zip(
+        src_horizontal_dims,
+        tgt_horizontal_dims,
+        strict=True,
+    ):
+        for coord in src.coords(dimensions=src_dim, dim_coords=True):
+            if overwrite and tgt.coords(coord.name()):
+                tgt.remove_coord(coord.name())
+            if not tgt.coords(coord.name()):
+                tgt.add_dim_coord(coord.copy(), tgt_dim)
+
+    # Copy 2D latitude and longitude coordinates.
+    if len(src_horizontal_dims) == 2 and len(tgt_horizontal_dims) == 2:
+        for coord in src.coords(dimensions=src_horizontal_dims):
+            if overwrite and tgt.coords(coord.name()):
+                tgt.remove_coord(coord.name())
+            if not tgt.coords(coord.name()):
+                tgt.add_aux_coord(coord.copy(), tgt_horizontal_dims)
+
+
 def _create_cube(
     src_cube: Cube,
     data: np.ndarray | da.Array,

tests/integration/preprocessor/_regrid/test_regrid.py

@@ -42,10 +42,15 @@ def setUp(self):
         )
 
         # Setup cube with multiple horizontal dimensions
-        self.multidim_cube = _make_cube(data, grid="rotated", aux_coord=False)
+        self.multidim_cube = _make_cube(
+            data,
+            grid="rotated",
+            dtype=np.float64,
+            aux_coord=False,
+        )
         lats, lons = np.meshgrid(
-            np.arange(1, data.shape[-2] + 1),
-            np.arange(1, data.shape[-1] + 1),
+            np.arange(1, data.shape[-2] + 1).astype(np.float64),
+            np.arange(1, data.shape[-1] + 1).astype(np.float64),
         )
         self.multidim_cube.add_aux_coord(
             iris.coords.AuxCoord(
@@ -166,6 +171,45 @@ def test_regrid__linear_multidim(self):
         result.data = np.round(result.data, 1)
         assert_array_equal(result.data, expected)
 
+    @pytest.mark.parametrize(
+        "use_src_coords",
+        [
+            ["latitude", "longitude"],
+            ["grid_latitude", "grid_longitude"],
+        ],
+    )
+    @pytest.mark.parametrize("close", [True, False])
+    def test_regrid__nearest_cordex_result_has_all_coords(
+        self,
+        use_src_coords,
+        close,
+    ):
+        """Test that the result of regridding has both 1D and 2D lat and lon."""
+        target_grid = self.multidim_cube.copy()
+        offset = 1e-9 if close else 0.1
+        target_grid.coord("grid_longitude").points = (
+            target_grid.coord("grid_longitude").points + offset
+        )
+        target_grid.coord("grid_longitude").bounds = (
+            target_grid.coord("grid_longitude").bounds + offset
+        )
+        target_grid.coord("longitude").points = (
+            target_grid.coord("longitude").points + offset
+        )
+        result = regrid(
+            self.multidim_cube,
+            target_grid,
+            "nearest",
+            use_src_coords=use_src_coords,
+        )
+        for coord in (
+            "latitude",
+            "longitude",
+            "grid_latitude",
+            "grid_longitude",
+        ):
+            assert result.coord(coord) == target_grid.coord(coord)
+
     @pytest.mark.parametrize("cache_weights", [True, False])
     def test_regrid__linear_file(self, tmp_path, cache_weights):
         file = tmp_path / "file.nc"

tests/unit/preprocessor/_regrid/__init__.py

@@ -3,6 +3,7 @@
 from typing import Literal
 
 import iris
+import iris.coord_systems
 import iris.fileformats
 import numpy as np
 from iris.coords import AuxCoord, CellMethod, DimCoord
@@ -77,9 +78,9 @@ def _make_cube(  # noqa: PLR0915,C901
         cube.add_dim_coord(_make_vcoord(z, dtype=dtype), 0)
 
     if grid == "rotated":
-        # Create a synthetic test latitude coordinate.
+        # Create a synthetic test grid_latitude coordinate.
         data = np.arange(y, dtype=dtype) + 1
-        cs = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS)
+        cs = iris.coord_systems.RotatedGeogCS(6.08, -166.92)
         kwargs = {
             "standard_name": "grid_latitude",
             "long_name": "latitude in rotated pole grid",
@@ -93,7 +94,7 @@ def _make_cube(  # noqa: PLR0915,C901
             ycoord.guess_bounds()
         cube.add_dim_coord(ycoord, 1)
 
-        # Create a synthetic test longitude coordinate.
+        # Create a synthetic test grid_longitude coordinate.
         data = np.arange(x, dtype=dtype) + 1
         kwargs = {
             "standard_name": "grid_longitude",

tests/unit/preprocessor/_regrid/test_regrid.py

@@ -3,6 +3,9 @@
 import dask
 import dask.array as da
 import iris
+import iris.coord_systems
+import iris.coords
+import iris.cube
 import numpy as np
 import pytest
 
@@ -36,7 +39,7 @@ def _make_coord(
     coord = iris.coords.DimCoord(
         np.linspace(start, stop, step),
         standard_name=name,
-        units="degrees",
+        units="degrees_north" if name == "latitude" else "degrees_east",
     )
     coord.guess_bounds()
     return coord
@@ -238,6 +241,72 @@ def test_horizontal_grid_is_close(cube2_spec: dict, expected: bool) -> None:
     assert _horizontal_grid_is_close(cube1, cube2) == expected
 
 
+def test_horizontal_grid_is_close_aux_coords():
+    """Test for `_horizontal_grid_is_close`."""
+    cube1 = _make_cube(lat=LAT_SPEC1, lon=LON_SPEC1)
+    cube2 = _make_cube(lat=LAT_SPEC1, lon=LON_SPEC1)
+    lat = cube1.coord("latitude")
+    lon = cube1.coord("longitude")
+    cube1.remove_coord(lat)
+    cube1.remove_coord(lon)
+    cube1.add_aux_coord(lat, 0)
+    cube1.add_aux_coord(lon, 1)
+
+    lat = cube2.coord("latitude")
+    lon = cube2.coord("longitude")
+    cube2.remove_coord(lat)
+    cube2.remove_coord(lon)
+    cube2.add_aux_coord(lat, 0)
+    cube2.add_aux_coord(lon, 1)
+
+    assert _horizontal_grid_is_close(cube1, cube2) is True
+
+
+@pytest.mark.parametrize("close", [True, False])
+def test_horizontal_grid_is_close_rotated(close):
+    """Test for `_horizontal_grid_is_close`."""
+    cs = iris.coord_systems.RotatedGeogCS(6.08, -166.92)
+    cube1 = iris.cube.Cube(
+        np.zeros((2, 3)),
+        dim_coords_and_dims=[
+            (
+                iris.coords.DimCoord(
+                    [1, 2],
+                    standard_name="grid_latitude",
+                    units="degrees",
+                    coord_system=cs,
+                ),
+                0,
+            ),
+            (
+                iris.coords.DimCoord(
+                    [1, 2, 3],
+                    standard_name="grid_longitude",
+                    units="degrees",
+                    coord_system=cs,
+                ),
+                1,
+            ),
+        ],
+    )
+    cube2 = cube1.copy()
+    if not close:
+        cube2.coord("grid_latitude").points = (
+            cube2.coord("grid_latitude").points + 1.0
+        )
+
+    assert _horizontal_grid_is_close(cube1, cube2) is close
+
+
+def test_horizontal_grid_is_close_no_coords():
+    """Test for `_horizontal_grid_is_close`."""
+    cube1 = _make_cube(lat=LAT_SPEC1, lon=LON_SPEC1)
+    cube2 = _make_cube(lat=LAT_SPEC1, lon=LON_SPEC1)
+    cube1.remove_coord("latitude")
+
+    assert _horizontal_grid_is_close(cube1, cube2) is False
+
+
 def test_regrid_is_skipped_if_grids_are_the_same_dim_coord(mocker):
     """Test that regridding is skipped if the grids are the same."""
     mock_get_regridder = mocker.patch(