diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index 461e75002a..4a037e0c4c 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/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, diff --git a/tests/integration/preprocessor/_regrid/test_regrid.py b/tests/integration/preprocessor/_regrid/test_regrid.py index e374262619..6f0e772d55 100644 --- a/tests/integration/preprocessor/_regrid/test_regrid.py +++ b/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" diff --git a/tests/unit/preprocessor/_regrid/__init__.py b/tests/unit/preprocessor/_regrid/__init__.py index c9938e051d..211f90ddee 100644 --- a/tests/unit/preprocessor/_regrid/__init__.py +++ b/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", diff --git a/tests/unit/preprocessor/_regrid/test_regrid.py b/tests/unit/preprocessor/_regrid/test_regrid.py index 7b41305c7e..9e0610b0db 100644 --- a/tests/unit/preprocessor/_regrid/test_regrid.py +++ b/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(