diff --git a/lib/iris/analysis/_grid_angles.py b/lib/iris/analysis/_grid_angles.py index e54365e248..874fa5e751 100644 --- a/lib/iris/analysis/_grid_angles.py +++ b/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, + ) diff --git a/lib/iris/tests/unit/analysis/cartography/test_gridcell_angles.py b/lib/iris/tests/unit/analysis/cartography/test_gridcell_angles.py index 32af84317c..b5d1386854 100644 --- a/lib/iris/tests/unit/analysis/cartography/test_gridcell_angles.py +++ b/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)