Add 2D guess bounds

lib/iris/analysis/_grid_angles.py

@@ -14,6 +14,7 @@
 import numpy as np
 
 import iris
+from iris.coord_systems import GeogCS, RotatedGeogCS
 
 
 def _3d_xyz_from_latlon(lon, lat):
@@ -453,3 +454,125 @@ def rotate_grid_vectors(u_cube, v_cube, grid_angles_cube=None, grid_angles_kwarg
     v_out.data = np.ma.masked_array(vv, mask=mask)
 
     return u_out, v_out
+
+
+def _vectorised_matmul(mats, vecs):
+    return np.einsum("ijk,ji->ki", mats, vecs)
+
+
+def _generate_180_mats_from_uvecs(uvecs):
+    mats = np.einsum("ji,ki->ijk", uvecs, uvecs) * 2
+    np.einsum("ijj->ij", mats)[:] -= 1
+    return mats
+
+
+def _2D_guess_bounds_first_pass(array):
+    # average and normalise, boundary buffer represents edges and corners
+    result_array = np.zeros((array.shape[0], array.shape[1] + 1, array.shape[2] + 1))
+    pads = ((0, 1), (1, 0))
+    for pad_i in pads:
+        for pad_j in pads:
+            result_array += np.pad(array, ((0, 0), pad_i, pad_j))
+
+    # normalise
+    result_array /= np.linalg.norm(result_array, ord=2, axis=0)[np.newaxis, ...]
+    return result_array
+
+
+def _2D_gb_buffer_outer(array_shape):
+    # return appropriate numpy slice for outer halo
+    _, x, y = array_shape
+    x_i = list(range(x)) + ([x - 1] * (y - 2)) + list(range(x))[::-1] + ([0] * (y - 2))
+    y_i = (
+        ([0] * (x - 1)) + list(range(y)) + ([y - 1] * (x - 2)) + list(range(1, y))[::-1]
+    )
+    return np.s_[:, x_i, y_i]
+
+
+def _2D_gb_buffer_inner(array_shape):
+    # return appropriate numpy slice for inner halo
+    _, x, y = array_shape
+    x_i = (
+        [1]
+        + list(range(1, x - 1))
+        + ([x - 2] * y)
+        + list(range(1, x - 1))[::-1]
+        + ([1] * (y - 1))
+    )
+    y_i = (
+        ([1] * x) + list(range(1, y - 1)) + ([y - 2] * x) + list(range(1, y - 1))[::-1]
+    )
+    return np.s_[:, x_i, y_i]
+
+
+def _2D_guess_bounds_in_place(lons, lats, extrapolate=True):
+    lon_array = lons.points
+    lat_array = lats.points
+    xyz_array = _3d_xyz_from_latlon(lon_array, lat_array)
+
+    result_xyz = _2D_guess_bounds_first_pass(xyz_array)
+    if extrapolate:
+        outer_inds = _2D_gb_buffer_outer(result_xyz.shape)
+        inner_inds = _2D_gb_buffer_inner(result_xyz.shape)
+        mats = _generate_180_mats_from_uvecs(result_xyz[outer_inds])
+        result_xyz[outer_inds] = _vectorised_matmul(mats, result_xyz[inner_inds])
+
+    result_lon_bounds, result_lat_bounds = _latlon_from_xyz(result_xyz)
+
+    # add these bounds cf style
+    lons.bounds = np.stack(
+        [
+            result_lon_bounds[:-1, :-1],
+            result_lon_bounds[:-1, 1:],
+            result_lon_bounds[1:, 1:],
+            result_lon_bounds[1:, :-1],
+        ],
+        axis=2,
+    )
+    lats.bounds = np.stack(
+        [
+            result_lat_bounds[:-1, :-1],
+            result_lat_bounds[:-1, 1:],
+            result_lat_bounds[1:, 1:],
+            result_lat_bounds[1:, :-1],
+        ],
+        axis=2,
+    )
+
+
+def guess_2D_bounds(x, y, extrapolate=True, in_place=False):
+    """Guess the bounds of a pair of 2D coords.
+
+    Parameters
+    ----------
+    x : class:`~iris.coords.AuxCoord`
+        A "longitude" or "grid_longitude" coordinate.
+    y : class:`~iris.coords.AuxCoord`
+        A "latitude" or "grid_latitude" coordinate.
+    extrapolate : bool, default=True
+        If True, extend the edge bounds beyond the limits of the edge points.
+    in_place : bool, default=False
+        If True, modify the coordinate arguments in place.
+    """
+    assert len(x.shape) == len(y.shape) == 2
+    assert x.standard_name in ("longitude", "grid_longitude")
+    assert y.standard_name in ("latitude", "grid_latitude")
+
+    if x.units != "degrees" or y.units != "degrees":
+        msg = "Coordinate units are expected to be degrees."
+        raise ValueError(msg)
+    if not all(
+        isinstance(coord.coord_system, GeogCS | RotatedGeogCS | None)
+        for coord in [x, y]
+    ):
+        msg = "Coordinate systems are expected geodetic."
+        raise ValueError(msg)
+
+    if in_place:
+        new_x = x
+        new_y = y
+    else:
+        new_x = x.copy()
+        new_y = y.copy()
+    _2D_guess_bounds_in_place(new_x, new_y, extrapolate=extrapolate)
+    return new_x, new_y

lib/iris/analysis/cartography.py

@@ -27,7 +27,7 @@
 from iris.util import _meshgrid
 import iris.warnings
 
-from ._grid_angles import gridcell_angles, rotate_grid_vectors
+from ._grid_angles import gridcell_angles, guess_2D_bounds, rotate_grid_vectors
 
 # List of contents to control Sphinx autodocs.
 # Unfortunately essential to get docs for the grid_angles functions.
@@ -39,6 +39,7 @@
     "get_xy_contiguous_bounded_grids",
     "get_xy_grids",
     "gridcell_angles",
+    "guess_2D_bounds",
     "project",
     "rotate_grid_vectors",
     "rotate_pole",

lib/iris/tests/unit/analysis/cartography/test_guess_2D_bounds.py

@@ -0,0 +1,159 @@
+# Copyright Iris contributors
+#
+# This file is part of Iris and is released under the BSD license.
+# See LICENSE in the root of the repository for full licensing details.
+"""Unit tests for the function :func:`iris.analysis.cartography.guess_2D_bounds`."""
+
+import numpy as np
+import pytest
+
+from iris.analysis.cartography import _transform_xy, guess_2D_bounds
+from iris.coord_systems import Mercator, RotatedGeogCS
+from iris.coords import AuxCoord
+from iris.tests.unit.analysis.cartography.test_gridcell_angles import (
+    _2d_multicells_testcube,
+)
+
+
+def _2D_guess_bounds(cube, extrapolate=True, in_place=False):
+    lon = cube.coord(axis="X")
+    lat = cube.coord(axis="Y")
+    new_lon, new_lat = guess_2D_bounds(
+        lon, lat, extrapolate=extrapolate, in_place=in_place
+    )
+    if not in_place:
+        for old, new in [[lon, new_lon], [lat, new_lat]]:
+            dims = cube.coord_dims(old)
+            cube.remove_coord(old)
+            cube.add_aux_coord(new, dims)
+    return cube
+
+
+def test_2D_guess_bounds_contiguity():
+    cube = _2d_multicells_testcube()
+    assert not cube.coord("latitude").is_contiguous()
+    assert not cube.coord("longitude").is_contiguous()
+
+    result_extrap = _2D_guess_bounds(cube)
+    assert result_extrap.coord("latitude").is_contiguous()
+    assert result_extrap.coord("longitude").is_contiguous()
+
+    result_clipped = _2D_guess_bounds(cube, extrapolate=False)
+    assert result_clipped.coord("latitude").is_contiguous()
+    assert result_clipped.coord("longitude").is_contiguous()
+
+
+def test_2D_guess_bounds_rotational_equivalence():
+    # Check that _2D_guess_bounds is rotationally equivalent.
+    cube = _2d_multicells_testcube()
+
+    # Define a rotation with a pair of coordinate systems.
+    rotated_cs = RotatedGeogCS(20, 30).as_cartopy_crs()
+    unrotated_cs = RotatedGeogCS(0, 0).as_cartopy_crs()
+
+    # Guess the bounds before rotating the lat-lon points.
+    _2D_guess_bounds(cube, extrapolate=True, in_place=True)
+    lon_bounds_unrotated = cube.coord("longitude").bounds
+    lat_bounds_unrotated = cube.coord("latitude").bounds
+
+    # Rotate the guessed lat-lon bounds.
+    rotated_lon_bounds, rotated_lat_bounds = _transform_xy(
+        unrotated_cs,
+        lon_bounds_unrotated.flatten(),
+        lat_bounds_unrotated.flatten(),
+        rotated_cs,
+    )
+
+    # Rotate the lat-lon points.
+    lat = cube.coord("latitude")
+    lon = cube.coord("longitude")
+    lon.points, lat.points = _transform_xy(
+        unrotated_cs, lon.points, lat.points, rotated_cs
+    )
+
+    # guess the bounds after rotating the lat-lon points.
+    _2D_guess_bounds(cube, extrapolate=True, in_place=True)
+    lon_bounds_from_rotated_points = cube.coord("longitude").bounds
+    lat_bounds_from_rotated_points = cube.coord("latitude").bounds
+
+    # Check that the results are equivalent.
+    assert np.allclose(rotated_lon_bounds, lon_bounds_from_rotated_points.flatten())
+    assert np.allclose(rotated_lat_bounds, lat_bounds_from_rotated_points.flatten())
+
+
+def test_2D_guess_bounds_transpose_equivalence():
+    # Check that _2D_guess_bounds is transpose equivalent.
+    cube = _2d_multicells_testcube()
+    cube_transposed = _2d_multicells_testcube()
+
+    def transpose_2D_coord(coord):
+        new_points = coord.points.transpose()
+        new_bounds = coord.bounds.transpose((1, 0, 2))[:, :, (0, 3, 2, 1)]
+        new_coord = AuxCoord(
+            new_points,
+            bounds=new_bounds,
+            standard_name=coord.standard_name,
+            units=coord.units,
+        )
+        return new_coord
+
+    cube_transposed.transpose()
+
+    new_lat = transpose_2D_coord(cube_transposed.coord("latitude"))
+    new_lon = transpose_2D_coord(cube_transposed.coord("longitude"))
+    cube_transposed.remove_coord("latitude")
+    cube_transposed.remove_coord("longitude")
+    cube_transposed.add_aux_coord(new_lat, (0, 1))
+    cube_transposed.add_aux_coord(new_lon, (0, 1))
+
+    _2D_guess_bounds(cube, extrapolate=True, in_place=True)
+    _2D_guess_bounds(cube_transposed, extrapolate=True, in_place=True)
+
+    cube_transposed.transpose()
+
+    untransposed_lat = transpose_2D_coord(cube_transposed.coord("latitude"))
+    untransposed_lon = transpose_2D_coord(cube_transposed.coord("longitude"))
+
+    assert np.allclose(untransposed_lat.bounds, cube.coord("latitude").bounds)
+    assert np.allclose(untransposed_lon.bounds, cube.coord("longitude").bounds)
+
+
+def test_2D_guess_bounds_slice_equivalence():
+    # Extrapolation should approximate expected values from an extended regular grid when points are
+    # close enough together.
+    shrink_factor = 1000
+    cube = _2d_multicells_testcube()
+    for coord in cube.coords():
+        coord.points = coord.points / shrink_factor
+    sub_cube = cube[1:, 1:-1]
+    _2D_guess_bounds(cube)
+    _2D_guess_bounds(sub_cube)
+    assert np.allclose(
+        cube[1:, 1:-1].coord("latitude").bounds, sub_cube.coord("latitude").bounds
+    )
+    assert np.allclose(
+        cube[1:, 1:-1].coord("longitude").bounds, sub_cube.coord("longitude").bounds
+    )
+
+
+def test_2D_guess_bounds_coord_systems():
+    rotated_cs = RotatedGeogCS(20, 30)
+    mercator_cs = Mercator()
+
+    rotated_cube = _2d_multicells_testcube()
+    rotated_cube.coord("latitude").coord_system = rotated_cs
+    rotated_cube.coord("latitude").standard_name = "grid_latitude"
+    rotated_cube.coord("longitude").coord_system = rotated_cs
+    rotated_cube.coord("longitude").standard_name = "grid_longitude"
+
+    _2D_guess_bounds(rotated_cube, in_place=True)
+
+    assert rotated_cube.coord("grid_latitude").is_contiguous()
+    assert rotated_cube.coord("grid_longitude").is_contiguous()
+
+    mercator_cube = _2d_multicells_testcube()
+    mercator_cube.coord("latitude").coord_system = mercator_cs
+    mercator_cube.coord("longitude").coord_system = mercator_cs
+
+    with pytest.raises(ValueError, match="Coordinate systems are expected geodetic."):
+        _2D_guess_bounds(mercator_cube)