From 6ae291f7aaaaa520060ef4cdef4ca25dcbe55b3a Mon Sep 17 00:00:00 2001 From: mspelman07 <99179165+mspelman07@users.noreply.github.com> Date: Tue, 16 Jul 2024 15:17:54 +0100 Subject: [PATCH] Make difference module handle circular cubes (#2016) * Resolved merge conflict * Added another failing test for difference module. Differences calculated correctly but need to add extra dimension for the wrap-around distance. * Resolved merge conflicts * Auto-formatting * Import Bugfix * Ran isort against spatial.py * Added test for appropriate error for unhandled cubes and made code match it. * iSort on spatial.py again * Another attempt at autoformatting * Tidied up * Removed confusing variable name * Autoformatting * Tried manually sorting import order as isort still failing in github actions despite running locally and pushing * And again * And again * Linting suggestion * Review Actions * Ran isort to fix gha * Added test for cubes with flipped axes (lonitude increasing along columns not rows and vice versa for latitude). Fixed code to handle this case. * black, isort, flake8 * Ran isort in isolation * Fixed test failures resulting from breaking changes made by upstream PR * Black * updates as per review comments and adds test for 0-360 degree case * update comment * formatting --------- Co-authored-by: Marcus Spelman --- improver/utilities/spatial.py | 457 ++++-------------- ...st_DifferenceBetweenAdjacentGridSquares.py | 29 +- 2 files changed, 94 insertions(+), 392 deletions(-) diff --git a/improver/utilities/spatial.py b/improver/utilities/spatial.py index b4d4df7e9b..3c98dd660b 100644 --- a/improver/utilities/spatial.py +++ b/improver/utilities/spatial.py @@ -5,7 +5,7 @@ """ Provides support utilities.""" import copy -from abc import ABC, abstractmethod +import warnings from typing import List, Optional, Tuple, Union import cartopy.crs as ccrs @@ -14,21 +14,20 @@ import numpy as np from cartopy.crs import CRS from cf_units import Unit -from iris.coord_systems import CoordSystem, GeogCS -from iris.coords import AuxCoord, CellMethod, Coord, DimCoord +from iris.coord_systems import GeogCS +from iris.coords import AuxCoord, CellMethod, Coord from iris.cube import Cube, CubeList from numpy import ndarray from numpy.ma import MaskedArray from scipy.ndimage.filters import maximum_filter +from scipy.stats import circmean 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, - generate_mandatory_attributes, -) +from improver.metadata.utilities import create_new_diagnostic_cube from improver.utilities.cube_checker import check_cube_coordinates, spatial_coords_match from improver.utilities.cube_manipulation import enforce_coordinate_ordering @@ -127,8 +126,9 @@ 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, "m", axis=axis) + grid_spacing_metres = calculate_grid_spacing(cube, "metres", axis=axis) grid_cells = distance / abs(grid_spacing_metres) if return_int: @@ -160,336 +160,10 @@ 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 axes + a cube. The difference is calculated along the x and y axis individually. """ @@ -509,6 +183,29 @@ 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. @@ -528,7 +225,7 @@ def _update_metadata(diff_cube: Cube, coord_name: str, cube_name: str) -> None: diff_cube.attributes["form_of_difference"] = "forward_difference" diff_cube.rename("difference_of_" + cube_name) - def create_difference_cube(self, + def create_difference_cube( self, cube: Cube, coord_name: str, diff_along_axis: ndarray ) -> Cube: """ @@ -549,14 +246,19 @@ def create_difference_cube(self, """ axis = cube.coord(coord_name) points = axis.points + mean_points = (points[1:] + points[:-1]) / 2 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." + warnings.warn( + "DifferenceBetweenAdjacentGridSquares does not fully support cubes with " + "circular x-axis that do not use a geographic (i.e. latlon) coordinate system. " + "Such cubes will be handled as if they were not circular, meaning that the " + "differences cube returned will have one fewer points along the specified axis" + "than the input cube." ) - mean_points = (points[1:] + points[:-1]) / 2 + else: + extra_mean_point = self._get_wrap_around_mean_point(points) + mean_points = np.hstack([mean_points, extra_mean_point]) # Copy cube metadata and coordinates into a new cube. # Create a new coordinate for the coordinate along which the @@ -635,7 +337,7 @@ def process(self, cube: Cube) -> Tuple[Cube, Cube]: return tuple(diffs) -class GradientBetweenAdjacentGridSquares(PostProcessingPlugin): +class GradientBetweenAdjacentGridSquares(BasePlugin): """Calculate the gradient between adjacent grid squares within a cube. The gradient is calculated along the x and y axis @@ -647,47 +349,67 @@ 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 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. + 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. """ self.regrid = regrid @staticmethod - def _create_output_cube(gradient: Cube, name: str) -> Cube: + def _create_output_cube( + gradient: ndarray, diff: Cube, cube: Cube, axis: str + ) -> Cube: """ - Create the output gradient cube, inheriting all metadata from source, but discarding - the "form_of_difference" attribute. + Create the output gradient cube. Args: gradient: Gradient values used in the data array of the resulting cube. - name: - Name to apply to the output 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. 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( - name, - gradient.units, - gradient, - generate_mandatory_attributes([gradient]), - optional_attributes=attributes, - data=gradient.data, + "gradient_of_" + cube.name(), + cube.units / diff.coord(axis=axis).units, + diff, + MANDATORY_ATTRIBUTE_DEFAULTS, + data=gradient, ) 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. This is then divided by the distance - between grid points along the same axis to get the gradient. + calculated using numpy.diff. Args: cube: @@ -695,16 +417,15 @@ 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) - distances = DistanceBetweenGridSquares()(cube) - for diff, distance in zip(diffs, distances): - gradient = diff / distance - grad_cube = self._create_output_cube(gradient, "gradient_of_" + cube.name()) + 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) 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 6db261d072..4ff663261e 100644 --- a/improver_tests/utilities/spatial/test_DifferenceBetweenAdjacentGridSquares.py +++ b/improver_tests/utilities/spatial/test_DifferenceBetweenAdjacentGridSquares.py @@ -8,7 +8,6 @@ import iris import numpy as np -import pytest from iris.coords import CellMethod from iris.cube import Cube from iris.tests import IrisTest @@ -31,8 +30,8 @@ def setUp(self): ) self.plugin = DifferenceBetweenAdjacentGridSquares() - def test_y_dimension_equalarea(self): - """Test differences calculated along the y dimension, equalarea grid.""" + def test_y_dimension(self): + """Test differences calculated along the y dimension.""" points = self.cube.coord(axis="y").points expected_y_coords = (points[1:] + points[:-1]) / 2 result = self.plugin.create_difference_cube( @@ -42,8 +41,8 @@ def test_y_dimension_equalarea(self): self.assertArrayAlmostEqual(result.coord(axis="y").points, expected_y_coords) self.assertArrayEqual(result.data, self.diff_in_y_array) - def test_x_dimension_equalarea(self): - """Test differences calculated along the x dimension, equalarea grid.""" + def test_x_dimension(self): + """Test differences calculated along the x dimension.""" 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 @@ -54,19 +53,6 @@ def test_x_dimension_equalarea(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]]) @@ -272,13 +258,8 @@ 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.assertRaisesRegex( - NotImplementedError, - "DifferenceBetweenAdjacentGridSquares does not support cubes with " - r"circular x-axis that do not use a geographic \(i.e. latlon\) coordinate system.", - ): + with self.assertRaises(ValueError): self.plugin.process(self.cube)