From f7d84ab41084093ff2d6d7db22f63be3fb800d29 Mon Sep 17 00:00:00 2001 From: Stephen Moseley Date: Fri, 23 Aug 2024 15:44:29 +0100 Subject: [PATCH] Resolves differences with reviewed version --- improver/utilities/spatial.py | 446 +++++++++++++++--- ...st_DifferenceBetweenAdjacentGridSquares.py | 29 +- 2 files changed, 393 insertions(+), 82 deletions(-) diff --git a/improver/utilities/spatial.py b/improver/utilities/spatial.py index f7dc660023..5655e8e05d 100644 --- a/improver/utilities/spatial.py +++ b/improver/utilities/spatial.py @@ -5,7 +5,7 @@ """ Provides support utilities.""" import copy -import warnings +from abc import ABC, abstractmethod from typing import List, Optional, Tuple, Union import cartopy.crs as ccrs @@ -14,8 +14,8 @@ import numpy as np from cartopy.crs import CRS from cf_units import Unit -from iris.coord_systems import GeogCS -from iris.coords import AuxCoord, CellMethod, Coord +from iris.coord_systems import CoordSystem, GeogCS +from iris.coords import AuxCoord, CellMethod, Coord, DimCoord from iris.cube import Cube, CubeList from numpy import ndarray from numpy.ma import MaskedArray @@ -24,9 +24,11 @@ from improver import BasePlugin, PostProcessingPlugin from improver.metadata.amend import update_diagnostic_name from improver.metadata.constants import FLOAT_DTYPE -from improver.metadata.constants.attributes import MANDATORY_ATTRIBUTE_DEFAULTS from improver.metadata.probabilistic import in_vicinity_name_format, is_probability -from improver.metadata.utilities import create_new_diagnostic_cube +from improver.metadata.utilities import ( + create_new_diagnostic_cube, + generate_mandatory_attributes, +) from improver.utilities.cube_checker import check_cube_coordinates, spatial_coords_match from improver.utilities.cube_manipulation import enforce_coordinate_ordering @@ -125,9 +127,8 @@ def distance_to_number_of_grid_cells( d_error = f"Distance of {distance}m" if distance <= 0: raise ValueError(f"Please specify a positive distance in metres. {d_error}") - # calculate grid spacing along chosen axis - grid_spacing_metres = calculate_grid_spacing(cube, "metres", axis=axis) + grid_spacing_metres = calculate_grid_spacing(cube, "m", axis=axis) grid_cells = distance / abs(grid_spacing_metres) if return_int: @@ -159,10 +160,336 @@ def number_of_grid_cells_to_distance(cube: Cube, grid_points: int) -> float: return radius_in_metres +class BaseDistanceCalculator(ABC): + """Base class for distance calculators for cubes with different coordinate systems/axis types""" + + def __init__(self, cube: Cube): + """ + Args: + cube: + Cube for which the distances will be calculated. + """ + self.cube = cube + self.x_separations_axis, self.y_separation_axis = self.get_difference_axes() + + @staticmethod + def build_distances_cube(distances: ndarray, dims: List[Coord], axis: str) -> Cube: + """ + Constructs an output cube with units of metres. + Args: + distances: + Data array containing calculated distances with which to populate the output cube. + dims: + Coordinate axes for the output cube. Must match the shape of distances. + axis: + The axis along which distances have been calculated. + """ + return Cube( + distances, + long_name=f"{axis}_distance_between_grid_points", + units="m", + dim_coords_and_dims=dims, + ) + + @staticmethod + def get_midpoints(axis: Coord) -> np.ndarray: + """ + Returns the midpoints along the supplied axis. If the axis is circular, the difference + between the last and first point is included with the assumption that this is in units of + degrees. + """ + points = axis.points + + if axis.circular: + points = np.hstack((points, 360 + points[0])) + mean_points = (points[1:] + points[:-1]) / 2 + + return mean_points.astype(axis.dtype) + + def get_difference_axes(self) -> Tuple[DimCoord, DimCoord]: + """Derives and returns the x and y coords for a cube of differences along one axis""" + input_cube_x_axis = self.cube.coord(axis="x") + input_cube_y_axis = self.cube.coord(axis="y") + distance_cube_x_axis = input_cube_x_axis.copy( + points=self.get_midpoints(input_cube_x_axis) + ) + distance_cube_y_axis = input_cube_y_axis.copy( + points=self.get_midpoints(input_cube_y_axis) + ) + return distance_cube_x_axis, distance_cube_y_axis + + @abstractmethod + def _get_x_distances(self) -> Cube: + """ + Abstract method for calculating distances along the x axis of the input cube. + The resulting cube shall have two dimensions as the result may be a function of position + along the y axis. + """ + + @abstractmethod + def _get_y_distances(self) -> Cube: + """ + Abstract method for calculating distances along the y axis of the input cube. + The resulting cube shall have two dimensions. + """ + + def get_distances(self) -> Tuple[Cube, Cube]: + """ + Calculates and returns the distances between grid points calculated along the cube's + x and y axis. + + Returns: + - 2D Cube of x-axis distances. + - 2D Cube of y-axis distances. + """ + return self._get_x_distances(), self._get_y_distances() + + +class LatLonCubeDistanceCalculator(BaseDistanceCalculator): + """ + Distance calculator for cubes using a Geographic Coordinate system. + Assumes that latitude and longitude are given in degrees, and that the origin is at the + intersection of the equator and the prime meridian. + Distances are calculated assuming a spherical earth, resulting in a < 0.15% error when compared + with the full haversine formula. + """ + + def __init__(self, cube: Cube): + super().__init__(cube) + self.lats, self.longs = self._get_cube_latlon_points() + self.sphere_radius = cube.coord(axis="x").coord_system.semi_major_axis + + def _get_cube_latlon_points(self) -> Tuple[ndarray, ndarray]: + """ + Extracts the y-axis and x-axis grid points used by a cube + with a geographic coordinate system. + + Returns: + - latitude points used by the cube's grid (in degrees). + - longitude points used by the cube's grid (in degrees). + Raises: + ValueError: Input cube does not use geographic coordinates, and/or + uses units other than degrees. + """ + if ( + self.cube.coord(axis="x").units == "degrees" + and self.cube.coord(axis="y").units == "degrees" + ): + longs = self.cube.coord(axis="x").points + lats = self.cube.coord(axis="y").points + return lats, longs + + raise ValueError( + "Cannot parse spatial axes of the cube provided. " + "Expected lat-long cube with units of degrees." + ) + + def _get_x_distances(self) -> Cube: + """ + Calculates the x-axis distances between adjacent grid points of a cube which uses + Geographic coordinates. + + Returns: + A 2D cube containing the x-axis distances between adjacent grid points of the input + cube in metres. As the earth is an oblate spheroid, the x-axis distances vary as + a function of the y-axis. + If the x-axis is marked as being circular, the distance between the last and first + points is included in the output. + x-axis coord positions are shifted to the mid-point of each pair. + """ + lats_as_col = np.expand_dims(self.lats, axis=1) + + if self.cube.coord(axis="x").circular: + longs = np.hstack([self.longs, 360 + self.longs[0]]) + else: + longs = self.longs + lon_diffs = np.diff(longs) + + x_distances = ( + self.sphere_radius * np.cos(np.deg2rad(lats_as_col)) * np.deg2rad(lon_diffs) + ) + + dims = [(self.cube.coord(axis="y"), 0), (self.x_separations_axis, 1)] + return self.build_distances_cube(x_distances, dims, "x") + + def _get_y_distances(self) -> Cube: + """ + Calculates the y-axis distances between adjacent grid points of a cube which uses + Geographic coordinates. + + Returns: + A 2D cube containing the y-axis distances between adjacent grid points of the input + cube in metres. + y-axis coord positions are shifted to the mid-point of each pair. + """ + lat_diffs = np.diff(self.lats) + + y_distances = self.sphere_radius * np.deg2rad(lat_diffs) + + y_distances_grid = np.tile(np.expand_dims(y_distances, axis=1), len(self.longs)) + dims = [(self.y_separation_axis, 0), (self.cube.coord(axis="x"), 1)] + return self.build_distances_cube(y_distances_grid, dims, "y") + + +class ProjectionCubeDistanceCalculator(BaseDistanceCalculator): + """ + Distance calculator for cubes using a projected coordinate system. + Assumes that x and y coordinates can be expressed in metres. + Distances are calculated assuming an equal-area projection. + """ + + def __init__(self, cube: Cube): + """ + Args: + cube: + Cube for which the distances will be calculated. + Raises: + NotImplementedError: + If the x-axis is marked as being circular. + """ + if cube.coord(axis="x").circular: + raise NotImplementedError( + "Cannot calculate distances between bounding points of a circular projected " + "coordinate." + ) + super().__init__(cube) + + def _get_x_distances(self) -> Cube: + """ + Calculates the x-axis distances between adjacent grid points of a cube which uses + Equal Area coordinates. + + Returns: + A 2D cube containing the x-axis distances between the grid points of the input + cube in metres. + x-axis coord positions are shifted to the mid-point of each pair. + """ + x_distances = calculate_grid_spacing(self.cube, axis="x", units="m") + data = np.full( + (self.cube.shape[0], len(self.x_separations_axis.points)), x_distances + ) + dims = [ + (self.cube.coord("projection_y_coordinate"), 0), + (self.x_separations_axis, 1), + ] + return self.build_distances_cube(data, dims, "x") + + def _get_y_distances(self) -> Cube: + """ + Calculates the y-axis distances between adjacent grid points of a cube which uses + Equal Area coordinates. + + Returns: + A 2D cube containing the y-axis distances between the grid points of the input + cube in metres. + y-axis coord positions are shifted to the mid-point of each pair. + """ + y_grid_spacing = calculate_grid_spacing(self.cube, axis="y", units="m") + data = np.full( + (len(self.y_separation_axis.points), self.cube.data.shape[1]), + y_grid_spacing, + ) + dims = [ + (self.y_separation_axis, 0), + (self.cube.coord("projection_x_coordinate"), 1), + ] + return self.build_distances_cube(data, dims, "y") + + +class DistanceBetweenGridSquares(BasePlugin): + """ + Calculates the distances between adjacent grid squares within a cube. + The distances are calculated along the x and y axes individually. + Returned distances are in metres. + The class can handle cubes with either Geographic (lat-long) or Equal Area projections. + For lat-lon cubes, the distances are calculated assuming a spherical earth. + This causes a < 0.15% error compared with the full haversine formula. + """ + + def _select_distance_calculator(self, cube: Cube): + """ + Chooses which distance calculator class to apply based on the cube's spatial coordinates. + + Args: + cube: + Cube for which the distances will be calculated. + Raises: + ValueError: Cube does not have enough information from which to calculate distances + or uses an unsupported coordinate system. + """ + if self._cube_xy_dimensions_are_distances(cube): + self.distance_calculator = ProjectionCubeDistanceCalculator(cube) + elif self._get_cube_spatial_type(cube) == GeogCS: + self.distance_calculator = LatLonCubeDistanceCalculator(cube) + else: + raise ValueError( + "Unsupported cube coordinate system or insufficent information to " + "calculate cube distances. Cube must either have coordinates for the " + "x and y axis with distance units, or use the Geographic (GeogCS) " + "coordinate system. For cubes with x and y dimensions expressed as angles, " + "distance between points cannot be calculated without a coordinate system." + ) + + @staticmethod + def _get_cube_spatial_type(cube: Cube) -> CoordSystem: + """ + Finds the coordinate system used by a cube. + + Args: + cube: + Cube to find the coordinate system of. + + Returns: + The coordinate system of the cube as an Iris Coordinate System. + """ + coord_system = cube.coord_system() + return type(coord_system) + + @staticmethod + def _cube_xy_dimensions_are_distances(cube: Cube) -> bool: + """ + Returns true if the given cube has coordinates mapping to the x and y axes with units + measuring distance (as opposed to angular separation) and false otherwise. + Args: + cube: + The iris cube to evaluate. + + Returns: + Boolean representing whether the cube has x and y axes defined in a distance unit. + """ + try: + cube.coord(axis="x").convert_units("m") + cube.coord(axis="y").convert_units("m") + return True + except ( + TypeError, + ValueError, + iris.exceptions.UnitConversionError, + iris.exceptions.CoordinateNotFoundError, + ): + return False + + def process(self, cube: Cube) -> Tuple[Cube, Cube]: + """ + Calculate the distances between grid points along the x and y axes + and return the result in separate cubes. + + Args: + cube: + Cube for which the distances will be calculated. + + Returns: + - Cube of x-axis distances. + - Cube of y-axis distances. + """ + self._select_distance_calculator(cube) + return self.distance_calculator.get_distances() + + class DifferenceBetweenAdjacentGridSquares(BasePlugin): """ Calculate the difference between adjacent grid squares within - a cube. The difference is calculated along the x and y axis + a cube. The difference is calculated along the x and y axes individually. """ @@ -182,29 +509,6 @@ def _axis_wraps_around_meridian(axis: Coord, cube: Cube) -> bool: """ return axis.circular and axis == cube.coord(axis="x") - @staticmethod - def _get_wrap_around_mean_point(points: ndarray) -> float: - """ - Calculates the midpoint between the two x coordinate points nearest the meridian. - - args: - points: - The x coordinate points of the cube. - - returns: - The x value of the midpoint between the two x coordinate points nearest the meridian. - """ - # The values of max and min azimuth doesn't matter as long as there is 360 degrees - # between them. - min_azimuth = -180 - max_azimuth = 180 - extra_mean_point = circmean([points[-1], points[0]], max_azimuth, min_azimuth) - extra_mean_point = np.round(extra_mean_point, 4) - if extra_mean_point < points[-1]: - # Ensures that the longitudinal coordinate is monotonically increasing - extra_mean_point += 360 - return extra_mean_point - @staticmethod def _update_metadata(diff_cube: Cube, coord_name: str, cube_name: str) -> None: """Rename cube, add attribute and cell method to describe difference. @@ -245,6 +549,13 @@ def create_difference_cube( """ axis = cube.coord(coord_name) points = axis.points + if self._axis_wraps_around_meridian(axis, cube): + points = np.hstack((points, 360 + points[0])) + if type(axis.coord_system) != GeogCS: + raise NotImplementedError( + "DifferenceBetweenAdjacentGridSquares does not support cubes with " + "circular x-axis that do not use a geographic (i.e. latlon) coordinate system." + ) mean_points = (points[1:] + points[:-1]) / 2 # Copy cube metadata and coordinates into a new cube. @@ -324,7 +635,7 @@ def process(self, cube: Cube) -> Tuple[Cube, Cube]: return tuple(diffs) -class GradientBetweenAdjacentGridSquares(BasePlugin): +class GradientBetweenAdjacentGridSquares(PostProcessingPlugin): """Calculate the gradient between adjacent grid squares within a cube. The gradient is calculated along the x and y axis @@ -336,67 +647,47 @@ def __init__(self, regrid: bool = False) -> None: Args: regrid: If True, the gradient cube is regridded to match the spatial - dimensions of the input cube. If False, the length of the - spatial dimensions of the gradient cube are one less than for - the input cube. + dimensions of the input cube. If False, the two output gradient cubes will have + different spatial coords such that the coord matching the gradient axis will + represent the midpoint of the input cube and will have one fewer points. + If the x-axis is marked as circular, the gradient between the last and first points + is also included. """ self.regrid = regrid @staticmethod - def _create_output_cube( - gradient: ndarray, diff: Cube, cube: Cube, axis: str - ) -> Cube: + def _create_output_cube(gradient: Cube, name: str) -> Cube: """ - Create the output gradient cube. + Create the output gradient cube, inheriting all metadata from source, but discarding + the "form_of_difference" attribute. Args: gradient: Gradient values used in the data array of the resulting cube. - diff: - Cube containing differences along the x or y axis - cube: - Cube with correct output dimensions - axis: - Short-hand reference for the x or y coordinate, as allowed by - iris.util.guess_coord_axis. + name: + Name to apply to the output cube. Returns: A cube of the gradients in the coordinate direction specified. """ + attributes = gradient.attributes + attributes.pop("form_of_difference") grad_cube = create_new_diagnostic_cube( - "gradient_of_" + cube.name(), - cube.units / diff.coord(axis=axis).units, - diff, - MANDATORY_ATTRIBUTE_DEFAULTS, - data=gradient, + name, + gradient.units, + gradient, + generate_mandatory_attributes([gradient]), + optional_attributes=attributes, + data=gradient.data, ) return grad_cube - @staticmethod - def _gradient_from_diff(diff: Cube, axis: str) -> ndarray: - """ - Calculate the gradient along the x or y axis from differences between - adjacent grid squares. - - Args: - diff: - Cube containing differences along the x or y axis - axis: - Short-hand reference for the x or y coordinate, as allowed by - iris.util.guess_coord_axis. - - Returns: - Array of the gradients in the coordinate direction specified. - """ - grid_spacing = np.diff(diff.coord(axis=axis).points)[0] - gradient = diff.data / grid_spacing - return gradient - def process(self, cube: Cube) -> Tuple[Cube, Cube]: """ Calculate the gradient along the x and y axes and return the result in separate cubes. The difference along each axis is - calculated using numpy.diff. + calculated using numpy.diff. This is then divided by the distance + between grid points along the same axis to get the gradient. Args: cube: @@ -404,15 +695,16 @@ def process(self, cube: Cube) -> Tuple[Cube, Cube]: Returns: - Cube after the gradients have been calculated along the - x axis. + x-axis. - Cube after the gradients have been calculated along the - y axis. + y-axis. """ gradients = [] diffs = DifferenceBetweenAdjacentGridSquares()(cube) - for axis, diff in zip(["x", "y"], diffs): - gradient = self._gradient_from_diff(diff, axis) - grad_cube = self._create_output_cube(gradient, diff, cube, axis) + distances = DistanceBetweenGridSquares()(cube) + for diff, distance in zip(diffs, distances): + gradient = diff / distance + grad_cube = self._create_output_cube(gradient, "gradient_of_" + cube.name()) if self.regrid: grad_cube = grad_cube.regrid(cube, iris.analysis.Linear()) gradients.append(grad_cube) diff --git a/improver_tests/utilities/spatial/test_DifferenceBetweenAdjacentGridSquares.py b/improver_tests/utilities/spatial/test_DifferenceBetweenAdjacentGridSquares.py index 4ff663261e..6db261d072 100644 --- a/improver_tests/utilities/spatial/test_DifferenceBetweenAdjacentGridSquares.py +++ b/improver_tests/utilities/spatial/test_DifferenceBetweenAdjacentGridSquares.py @@ -8,6 +8,7 @@ import iris import numpy as np +import pytest from iris.coords import CellMethod from iris.cube import Cube from iris.tests import IrisTest @@ -30,8 +31,8 @@ def setUp(self): ) self.plugin = DifferenceBetweenAdjacentGridSquares() - def test_y_dimension(self): - """Test differences calculated along the y dimension.""" + def test_y_dimension_equalarea(self): + """Test differences calculated along the y dimension, equalarea grid.""" points = self.cube.coord(axis="y").points expected_y_coords = (points[1:] + points[:-1]) / 2 result = self.plugin.create_difference_cube( @@ -41,8 +42,8 @@ def test_y_dimension(self): self.assertArrayAlmostEqual(result.coord(axis="y").points, expected_y_coords) self.assertArrayEqual(result.data, self.diff_in_y_array) - def test_x_dimension(self): - """Test differences calculated along the x dimension.""" + def test_x_dimension_equalarea(self): + """Test differences calculated along the x dimension, equalarea grid.""" diff_array = np.array([[1, 1], [2, 2], [5, 5]]) points = self.cube.coord(axis="x").points expected_x_coords = (points[1:] + points[:-1]) / 2 @@ -53,6 +54,19 @@ def test_x_dimension(self): self.assertArrayAlmostEqual(result.coord(axis="x").points, expected_x_coords) self.assertArrayEqual(result.data, diff_array) + def test_x_dimension_equalarea_circular(self): + """Test differences calculated along the x dimension when x is circular, equalarea grid.""" + diff_array = np.array([[1, 1], [2, 2], [5, 5]]) + self.cube.coord(axis="x").circular = True + with pytest.raises( + NotImplementedError, + match="DifferenceBetweenAdjacentGridSquares does not support cubes with circular " + "x-axis that do not use a geographic", + ): + self.plugin.create_difference_cube( + self.cube, "projection_x_coordinate", diff_array + ) + def test_x_dimension_for_circular_latlon_cube(self): """Test differences calculated along the x dimension for a cube which is circular in x.""" test_cube_data = np.array([[1, 2, 3], [2, 4, 6], [5, 10, 15]]) @@ -258,8 +272,13 @@ def test_3d_cube(self): self.assertArrayEqual(result[1].data, expected_y) def test_circular_non_geographic_cube_raises_approprate_exception(self): + """Check for error and message with projection coord and circular x axis""" self.cube.coord(axis="x").circular = True - with self.assertRaises(ValueError): + with self.assertRaisesRegex( + NotImplementedError, + "DifferenceBetweenAdjacentGridSquares does not support cubes with " + r"circular x-axis that do not use a geographic \(i.e. latlon\) coordinate system.", + ): self.plugin.process(self.cube)