Make difference module handle circular cubes

improver/utilities/spatial.py

@@ -5,6 +5,7 @@
 """ Provides support utilities."""
 
 import copy
+import warnings
 from typing import List, Optional, Tuple, Union
 
 import cartopy.crs as ccrs
@@ -13,11 +14,13 @@
 import numpy as np
 from cartopy.crs import CRS
 from cf_units import Unit
-from iris.coords import AuxCoord, CellMethod
+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
@@ -164,6 +167,45 @@ class DifferenceBetweenAdjacentGridSquares(BasePlugin):
     individually.
     """
 
+    @staticmethod
+    def _axis_wraps_around_meridian(axis: Coord, cube: Cube) -> bool:
+        """Returns true if the cube is 'circular' with the given axis wrapping around, i.e. if there
+        is a smooth transition between 180 degrees and -180 degrees on the axis.
+
+        Args:
+            axis:
+                Axis to check for circularity.
+            cube:
+                The cube to which the axis belongs.
+
+        Returns:
+            True if the axis wraps around the meridian; false otherwise.
+        """
+        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.
@@ -183,13 +225,11 @@ 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)
 
-    @staticmethod
     def create_difference_cube(
-        cube: Cube, coord_name: str, diff_along_axis: ndarray
+        self, cube: Cube, coord_name: str, diff_along_axis: ndarray
     ) -> Cube:
         """
-        Put the difference array into a cube with the appropriate
-        metadata.
+        Put the difference array into a cube with the appropriate metadata.
 
         Args:
             cube:
@@ -204,8 +244,21 @@ def create_difference_cube(
             Cube containing the differences calculated along the
             specified axis.
         """
-        points = cube.coord(coord_name).points
+        axis = cube.coord(coord_name)
+        points = axis.points
         mean_points = (points[1:] + points[:-1]) / 2
+        if self._axis_wraps_around_meridian(axis, cube):
+            if type(axis.coord_system) != GeogCS:
+                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."
+                )
+            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
@@ -226,8 +279,7 @@ def create_difference_cube(
             diff_cube.add_aux_coord(coord.copy(), dims)
         return diff_cube
 
-    @staticmethod
-    def calculate_difference(cube: Cube, coord_name: str) -> ndarray:
+    def calculate_difference(self, cube: Cube, coord_name: str) -> ndarray:
         """
         Calculate the difference along the axis specified by the
         coordinate.
@@ -242,8 +294,21 @@ def calculate_difference(cube: Cube, coord_name: str) -> ndarray:
             Array after the differences have been calculated along the
             specified axis.
         """
-        diff_axis = cube.coord_dims(coord_name)[0]
-        diff_along_axis = np.diff(cube.data, axis=diff_axis)
+        diff_axis = cube.coord(name_or_coord=coord_name)
+        diff_axis_number = cube.coord_dims(coord_name)[0]
+        diff_along_axis = np.diff(cube.data, axis=diff_axis_number)
+        if self._axis_wraps_around_meridian(diff_axis, cube):
+            # Get wrap-around difference:
+            first_column = np.take(cube.data, indices=0, axis=diff_axis_number)
+            last_column = np.take(cube.data, indices=-1, axis=diff_axis_number)
+            wrap_around_diff = first_column - last_column
+            # Apply wrap-around difference vector to diff array:
+            if diff_axis_number == 0:
+                diff_along_axis = np.vstack([diff_along_axis, wrap_around_diff])
+            elif diff_axis_number == 1:
+                diff_along_axis = np.hstack(
+                    [diff_along_axis, wrap_around_diff.reshape([-1, 1])]
+                )
         return diff_along_axis
 
     def process(self, cube: Cube) -> Tuple[Cube, Cube]:
@@ -265,9 +330,8 @@ def process(self, cube: Cube) -> Tuple[Cube, Cube]:
         diffs = []
         for axis in ["x", "y"]:
             coord_name = cube.coord(axis=axis).name()
-            diff_cube = self.create_difference_cube(
-                cube, coord_name, self.calculate_difference(cube, coord_name)
-            )
+            difference = self.calculate_difference(cube, coord_name)
+            diff_cube = self.create_difference_cube(cube, coord_name, difference)
             self._update_metadata(diff_cube, coord_name, cube.name())
             diffs.append(diff_cube)
         return tuple(diffs)

improver_tests/utilities/test_DifferenceBetweenAdjacentGridSquares.py

@@ -33,26 +33,74 @@ def setUp(self):
     def test_y_dimension(self):
         """Test differences calculated along the y dimension."""
         points = self.cube.coord(axis="y").points
-        expected_y = (points[1:] + points[:-1]) / 2
+        expected_y_coords = (points[1:] + points[:-1]) / 2
         result = self.plugin.create_difference_cube(
             self.cube, "projection_y_coordinate", self.diff_in_y_array
         )
         self.assertIsInstance(result, Cube)
-        self.assertArrayAlmostEqual(result.coord(axis="y").points, expected_y)
+        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."""
         diff_array = np.array([[1, 1], [2, 2], [5, 5]])
         points = self.cube.coord(axis="x").points
-        expected_x = (points[1:] + points[:-1]) / 2
+        expected_x_coords = (points[1:] + points[:-1]) / 2
         result = self.plugin.create_difference_cube(
             self.cube, "projection_x_coordinate", diff_array
         )
         self.assertIsInstance(result, Cube)
-        self.assertArrayAlmostEqual(result.coord(axis="x").points, expected_x)
+        self.assertArrayAlmostEqual(result.coord(axis="x").points, expected_x_coords)
         self.assertArrayEqual(result.data, 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]])
+        test_cube_x_grid_spacing = 120
+        test_cube = set_up_variable_cube(
+            test_cube_data,
+            "latlon",
+            x_grid_spacing=test_cube_x_grid_spacing,
+            name="wind_speed",
+            units="m s-1",
+        )
+        test_cube.coord(axis="x").circular = True
+        expected_diff_array = np.array([[1, 1, -2], [2, 2, -4], [5, 5, -10]])
+        expected_x_coords = np.array(
+            [-60, 60, 180]
+        )  # Original data are at [-120, 0, 120], therefore differences are at [-60, 60, 180].
+        result = self.plugin.create_difference_cube(
+            test_cube, "longitude", expected_diff_array
+        )
+        self.assertIsInstance(result, Cube)
+        self.assertArrayAlmostEqual(result.coord(axis="x").points, expected_x_coords)
+        self.assertArrayEqual(result.data, expected_diff_array)
+
+    def test_x_dimension_for_circular_latlon_cube_360_degree_coord(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]])
+        test_cube_x_grid_spacing = 120
+        test_cube = set_up_variable_cube(
+            test_cube_data,
+            "latlon",
+            x_grid_spacing=test_cube_x_grid_spacing,
+            name="wind_speed",
+            units="m s-1",
+        )
+        test_cube.coord(axis="x").bounds = [[0, 120], [120, 240], [240, 360]]
+        test_cube.coord(axis="x").points = [60, 120, 300]
+        test_cube.coord(axis="x").circular = True
+        expected_diff_array = np.array([[1, 1, -2], [2, 2, -4], [5, 5, -10]])
+        expected_x_coords = np.array(
+            [90, 210, 360]
+        )  # Original data are at [60, 120, 300], therefore differences are at [90, 210, 360].
+        result = self.plugin.create_difference_cube(
+            test_cube, "longitude", expected_diff_array
+        )
+        self.assertIsInstance(result, Cube)
+        self.assertArrayAlmostEqual(result.coord(axis="x").points, expected_x_coords)
+        self.assertArrayEqual(result.data, expected_diff_array)
+
     def test_othercoords(self):
         """Test that other coords are transferred properly"""
         time_coord = self.cube.coord("time")
@@ -65,20 +113,40 @@ def test_othercoords(self):
 
 
 class Test_calculate_difference(IrisTest):
-
     """Test the calculate_difference method."""
 
     def setUp(self):
         """Set up cube."""
-        data = np.array([[1, 2, 3], [2, 4, 6], [5, 10, 15]])
+        data = np.array([[1, 2, 3, 4], [2, 4, 6, 8], [5, 10, 15, 20]])
         self.cube = set_up_variable_cube(
-            data, name="wind_speed", units="m s-1", spatial_grid="equalarea",
+            data, "equalarea", name="wind_speed", units="m s-1",
         )
         self.plugin = DifferenceBetweenAdjacentGridSquares()
 
     def test_x_dimension(self):
         """Test differences calculated along the x dimension."""
-        expected = np.array([[1, 1], [2, 2], [5, 5]])
+        expected = np.array([[1, 1, 1], [2, 2, 2], [5, 5, 5]])
+        result = self.plugin.calculate_difference(
+            self.cube, self.cube.coord(axis="x").name()
+        )
+        self.assertIsInstance(result, np.ndarray)
+        self.assertArrayEqual(result, expected)
+
+    def test_x_dimension_wraps_around_meridian(self):
+        """Test differences calculated along the x dimension for a cube which is circular in x."""
+        self.cube.coord(axis="x").circular = True
+        expected = np.array([[1, 1, 1, -3], [2, 2, 2, -6], [5, 5, 5, -15]])
+        result = self.plugin.calculate_difference(
+            self.cube, self.cube.coord(axis="x").name()
+        )
+        self.assertIsInstance(result, np.ndarray)
+        self.assertArrayEqual(result, expected)
+
+    def test_x_dimension_wraps_around_meridian_cube_axes_flipped(self):
+        """Test differences calculated along the x dimension for a cube which is circular in x."""
+        self.cube.coord(axis="x").circular = True
+        self.cube.transpose()
+        expected = np.array([[1, 1, 1, -3], [2, 2, 2, -6], [5, 5, 5, -15]]).transpose()
         result = self.plugin.calculate_difference(
             self.cube, self.cube.coord(axis="x").name()
         )
@@ -87,7 +155,7 @@ def test_x_dimension(self):
 
     def test_y_dimension(self):
         """Test differences calculated along the y dimension."""
-        expected = np.array([[1, 2, 3], [3, 6, 9]])
+        expected = np.array([[1, 2, 3, 4], [3, 6, 9, 12]])
         result = self.plugin.calculate_difference(
             self.cube, self.cube.coord(axis="y").name()
         )
@@ -96,9 +164,11 @@ def test_y_dimension(self):
 
     def test_missing_data(self):
         """Test that the result is as expected when data is missing."""
-        data = np.array([[1, 2, 3], [np.nan, 4, 6], [5, 10, 15]], dtype=np.float32)
+        data = np.array(
+            [[1, 2, 3, 4], [np.nan, 4, 6, 8], [5, 10, 15, 20]], dtype=np.float32
+        )
         self.cube.data = data
-        expected = np.array([[np.nan, 2, 3], [np.nan, 6, 9]])
+        expected = np.array([[np.nan, 2, 3, 4], [np.nan, 6, 9, 12]])
         result = self.plugin.calculate_difference(
             self.cube, self.cube.coord(axis="y").name()
         )
@@ -108,10 +178,13 @@ def test_missing_data(self):
     def test_masked_data(self):
         """Test that the result is as expected when data is masked."""
         data = ma.array(
-            [[1, 2, 3], [2, 4, 6], [5, 10, 15]], mask=[[0, 0, 0], [1, 0, 0], [0, 0, 0]]
+            [[1, 2, 3, 4], [2, 4, 6, 8], [5, 10, 15, 20]],
+            mask=[[0, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 0]],
         )
         self.cube.data = data
-        expected = ma.array([[1, 2, 3], [3, 6, 9]], mask=[[1, 0, 0], [1, 0, 0]])
+        expected = ma.array(
+            [[1, 2, 3, 4], [3, 6, 9, 12]], mask=[[1, 0, 0, 0], [1, 0, 0, 0]]
+        )
         result = self.plugin.calculate_difference(
             self.cube, self.cube.coord(axis="y").name()
         )
@@ -121,7 +194,6 @@ def test_masked_data(self):
 
 
 class Test_process(IrisTest):
-
     """Test the process method."""
 
     def setUp(self):
@@ -185,6 +257,11 @@ def test_3d_cube(self):
         self.assertIsInstance(result[1], iris.cube.Cube)
         self.assertArrayEqual(result[1].data, expected_y)
 
+    def test_circular_non_geographic_cube_raises_approprate_exception(self):
+        self.cube.coord(axis="x").circular = True
+        with self.assertRaises(ValueError):
+            self.plugin.process(self.cube)
+
 
 if __name__ == "__main__":
     unittest.main()