diff --git a/lib/iris/analysis/_grid_angles.py b/lib/iris/analysis/_grid_angles.py index e54365e248..f77dd8cbfe 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,233 @@ 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): + # We interpret the 3D array `mats` as if it were a list of matrices varying over + # its last dimension so that mats[:,:,i] represents a single matrix. We consider + # the 2D array `vecs` to represent a list of vectors varying over its last dimension + # so that vecs[:,i] represents a single vector. The output will be such that for: + # result[:,i] == mats[:,:,i] @ vecs[:,i]. + return np.einsum("jki,ji->ki", mats, vecs) + + +def _generate_180_mats_from_uvecs(uvecs): + # Generates a 3D array representing a list of matrices which can be used by the + # function _vectorised_matmul. Both the input `uvecs` and the `result` vary in their + # last dimension so that uvecs[:,i] is a length 3 vector corresponding to the + # rotation matrix result[:,:,i], where this is a rotation of 180 degrees about the + # axis passing through uvecs[:,i] and the origin. + # If the vector uvecs[:,i] = (x,y,z), the equivalent rotation matrix result[:,:,i] + # will be: + # | 2x^2 - 1 2xy 2xz | + # | 2xy 2y^2 - 1 2yz | + # | 2xz 2yz 2z^2 - 1 | + mats = np.einsum("ji,ki->jki", uvecs, uvecs) * 2 + + # At this point the matrix mats[:,:,i] will be: + # | 2x^2 2xy 2xz | + # | 2xy 2y^2 2yz | + # | 2xz 2yz 2z^2 | + # to achieve the desired result, we take one from the diagonal. + np.einsum("jji->ji", mats)[:] -= 1 + return mats + + +def _2D_guess_bounds_first_pass(array): + # Average and normalise sets of neighbouring points. This calculation actually + # normalises the sum rather than the average, since this is equivalent. + # Corners of the resulting array will be the sum of only a single corner of the + # source array. Edges of the resulting array will be the midpoint between edge + # points in the source. Internal points in the resulting array will correspond + # to internal bounds in the final result of `guess_2D_bounds`. + # Note: if `extrapolate` is False, this array will contain all the bounds in the + # final result of `guess_2D_bounds`. + 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. + # This slice preserves the first dimension, which captures the 3D vector points and + # lists a series of indices in the next 2 dimensions. + # Each index in this slice corresponds to a corner or edge bound. + + # This halo starts at the index [:, 0, 0], the next set of indices increase in the + # second dimension until the index [:, -1, 0], then the last dimension increases + # until it reaches the index [:, -1, -1], the remaining indices follow the edge of + # the array anit-clockwise until it reaches the last index at [:, 0, 1]. + _, 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. + # This slice preserves the first dimension, which captures the 3D vector points and + # lists a series of indices in the next 2 dimensions. + # Each index in this slice corresponds to an internal bound neighbouring a corner or + # edge bound. + # For every index in the outer halo, this gives the nearest index not in the outer halo. + # Note: the internal bounds which are nearest to the corner bounds ar each the nearest + # bound for at least two other edge or corner bounds. Therefore, these indices will + # occur at least 3 times in the returned slice. + + # This halo starts at the index [:, 1, 1], the next index is also [:, 1, 1]. The + # next set of indices increase in the second dimension until the index [:, -2, 1], + # this index is repeated twice. In the edge case where the index [:, 1, 1] is + # equivalent to [:, -2, 1] the first two copies of the index will be followed + # immediately by two more copies. Then the last dimension increases until it reaches + # the index [:, -2, -2], repeating twice again, the remaining indices follow this + # pattern until it reaches the last index at [:, 1, 1]. + _, 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 + # Convert from lat-lon to 3D space. + xyz_array = _3d_xyz_from_latlon(lon_array, lat_array) + + # Create an array with internal-bounds, edge-midpoints, and corner-points. + result_xyz = _2D_guess_bounds_first_pass(xyz_array) + if extrapolate: + # Generate slice of edge-midpoints/edge-bounds and corner-points/corner-bounds. + outer_inds = _2D_gb_buffer_outer(result_xyz.shape) + # Generate slice of internal-bounds which neighbour the above bounds. + inner_inds = _2D_gb_buffer_inner(result_xyz.shape) + # Generate rotation matrices about corner-points and edge-midpoints + mats = _generate_180_mats_from_uvecs(result_xyz[outer_inds]) + # Rotate internal-bounds about their corresponding corner-point/edge-midpoint. + # Replace the corner-point/edge-midpoint with this result. + result_xyz[outer_inds] = _vectorised_matmul(mats, result_xyz[inner_inds]) + + # Convert back from 3D to lat-lon. + result_lon_bounds, result_lat_bounds = _latlon_from_xyz(result_xyz) + + # Reformat these bounds as CF style bounds. + 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. + + Bounds on a 2D coordinate are structured so that each point of the coordinate + has 4 associated bounds, with each pair of neighbouring points sharing 2 of their + bounds in common. We can categorise these bounds by how many points they are + associated with: + + - Internal bounds, belonging to 4 surrounding points. + - Edge bounds, belonging only to 2 points on the edge of the grid. + - Corner bounds, belonging only to 1 point on the corner of the grid. + + Since each of these categories of bounds has a different number of points associated + with them, they must also be calculated differently: + + - Each internal bound is calculated from the 4 surrounding points. The surrounding + points are first transformed from 2D lat-lon space to 3D space on the surface of a + unit sphere. The average of these points in 3D space is then taken. This point is then + projected onto that unit sphere and transformed back to 2D lat-lon space. Note that in + the edge case where the 3D average is precisely the origin, it will not be possible to + project onto the unit sphere and an error will be raised. + - The calculation for edge bounds depends on if the `extrapolate` keyword is True or + False. When `extrapolate` is False, the edge bound is calculated as the midpoint of + the 2 neighbouring bounds. This calculation is done, as above, by converting into 3D + space, averaging, and then converting back to lat-lon space. When `extrapolate` is True, + edge bounds are calculated using the neighbouring internal bound. The edge bound and the + neighbouring internal bound will be precisely the 2 bounds which the edge points share + in common. The extrapolated edge bound is calculated by rotating the internal bound 180 + degrees about the midpoint between the 2 neighbouring points calculated previously. This + rotation is done in 3D space about the axis defined by the line which passes through the + midpoint and the origin. This is equivalent to defining the edge bound so that the + midpoint between the two bounds is the same as the midpoint between the two points. + - The calculation for the corner bounds similarly depends on `extrapolate`. When + `extrapolate` is False, each corner bound is precisely equal to its associated corner + point. When `extrapolate` is True, we calculate the corner bound by taking the only + internal bound associated with the corner point and, as above, rotating it by 180 + degrees about the corner point. + + + Parameters + ---------- + x : class:`~iris.coords.AuxCoord` + A "longitude" or "grid_longitude" coordinate. Coordinate must be 2D. + y : class:`~iris.coords.AuxCoord` + A "latitude" or "grid_latitude" coordinate. Coordinate must be 2D. + 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. + """ + if x.shape != y.shape: + msg = "Coordinates do not have the same shape." + raise ValueError(msg) + if len(x.shape) != 2: + msg = "Coordinates are not 2D." + raise ValueError(msg) + if x.shape[0] < 2 or x.shape[1] < 2: + msg = "Coordinates must have length at least 2 in each dimension." + raise ValueError(msg) + if x.standard_name not in ("longitude", "grid_longitude"): + msg = "X coordinate is not 'longitude' or 'grid_longitude'." + raise ValueError(msg) + if y.standard_name not in ("latitude", "grid_latitude"): + msg = "Y coordinate is not 'latitude' or 'grid_latitude'." + raise ValueError(msg) + + 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 diff --git a/lib/iris/analysis/cartography.py b/lib/iris/analysis/cartography.py index e88711d51f..c3f77cbb11 100644 --- a/lib/iris/analysis/cartography.py +++ b/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", diff --git a/lib/iris/tests/unit/analysis/cartography/test_guess_2D_bounds.py b/lib/iris/tests/unit/analysis/cartography/test_guess_2D_bounds.py new file mode 100644 index 0000000000..45cde05a3c --- /dev/null +++ b/lib/iris/tests/unit/analysis/cartography/test_guess_2D_bounds.py @@ -0,0 +1,206 @@ +# 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) + + +def test_invalid_coords_1D(): + lat_1D = AuxCoord(np.arange(5), standard_name="latitude") + lon_1D = AuxCoord(np.arange(5), standard_name="longitude") + with pytest.raises(ValueError, match="Coordinates are not 2D."): + guess_2D_bounds(lon_1D, lat_1D) + + +def test_invalid_coords_shape(): + lat = AuxCoord(np.ones([2, 3]), standard_name="latitude") + lon = AuxCoord(np.ones([3, 2]), standard_name="longitude") + with pytest.raises(ValueError, match="Coordinates do not have the same shape."): + guess_2D_bounds(lon, lat) + + +def test_invalid_coords_size(): + expected_msg = "Coordinates must have length at least 2 in each dimension." + + lat = AuxCoord(np.ones([1, 3]), standard_name="latitude") + lon = AuxCoord(np.ones([1, 3]), standard_name="longitude") + with pytest.raises(ValueError, match=expected_msg): + guess_2D_bounds(lon, lat) + + lat = AuxCoord(np.ones([3, 1]), standard_name="latitude") + lon = AuxCoord(np.ones([3, 1]), standard_name="longitude") + with pytest.raises(ValueError, match=expected_msg): + guess_2D_bounds(lon, lat) + + +def test_invalid_coords_name(): + lat_valid = AuxCoord(np.ones([2, 3]), standard_name="latitude") + lon_valid = AuxCoord(np.ones([2, 3]), standard_name="longitude") + lat_invalid = AuxCoord(np.ones([2, 3]), long_name="lat") + lon_invalid = AuxCoord(np.ones([2, 3]), long_name="long") + with pytest.raises( + ValueError, match="X coordinate is not 'longitude' or 'grid_longitude'." + ): + guess_2D_bounds(lon_invalid, lat_valid) + with pytest.raises( + ValueError, match="X coordinate is not 'longitude' or 'grid_longitude'." + ): + guess_2D_bounds(lat_valid, lon_valid) + with pytest.raises( + ValueError, match="Y coordinate is not 'latitude' or 'grid_latitude'." + ): + guess_2D_bounds(lon_valid, lat_invalid)