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,119 @@ 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(cube, extrapolate=True, in_place=False):
+    lons = cube.coord(axis="X")
+    lats = cube.coord(axis="Y")
+    h_dims = cube.coord_dims(lons)
+    assert h_dims == cube.coord_dims(lats)
+    assert len(h_dims) == 2
+
+    if lons.units != "degrees" or lats.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 [lats, lons]
+    ):
+        msg = "Coordinate systems are expected geodetic."
+        raise ValueError(msg)
+
+    if in_place:
+        _2D_guess_bounds_in_place(lons, lats, extrapolate=extrapolate)
+
+    else:
+        new_lons = lons.copy()
+        new_lats = lats.copy()
+        _2D_guess_bounds_in_place(new_lons, new_lats, extrapolate=extrapolate)
+        cube.remove_coord(lons)
+        cube.remove_coord(lats)
+        cube.add_aux_coord(new_lons, h_dims)
+        cube.add_aux_coord(new_lats, h_dims)
+
+        return cube
+
+
+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,
+    )

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

@@ -11,7 +11,9 @@
 import numpy as np
 import pytest
 
-from iris.analysis.cartography import gridcell_angles
+from iris.analysis._grid_angles import _2D_guess_bounds
+from iris.analysis.cartography import _transform_xy, gridcell_angles
+from iris.coord_systems import Mercator, RotatedGeogCS
 from iris.coords import AuxCoord
 from iris.cube import Cube
 from iris.tests import _shared_utils
@@ -301,3 +303,115 @@ def test_fail_noncoord_coord(self):
     def test_fail_bad_method(self):
         with pytest.raises(ValueError, match="unrecognised cell_angle_boundpoints"):
             self._check_multiple_orientations_and_latitudes(method="something_unknown")
+
+
+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_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)