Partial collapse of multi-dimensional coordinates (#3028)

Adding support for partial collapse of multi-dimensional coordinates

docs/iris/src/whatsnew/contributions_2.1/newfeature_2018-May-03_multidim_collapse.txt

@@ -0,0 +1,2 @@
+* The partial collapse of multi-dimensional auxiliary coordinates is now
+  supported. Collapsed bounds span the range of the collapsed dimension(s).

lib/iris/coords.py

@@ -1169,20 +1169,14 @@ def collapsed(self, dims_to_collapse=None):
         the specified dimensions.
 
         Replaces the points & bounds with a simple bounded region.
-
-        .. note::
-            You cannot partially collapse a multi-dimensional coordinate. See
-            :ref:`cube.collapsed <partially_collapse_multi-dim_coord>` for more
-            information.
-
         """
+        import dask.array as da
+        # Ensure dims_to_collapse is a tuple to be able to pass
+        # through to numpy
         if isinstance(dims_to_collapse, (int, np.integer)):
-            dims_to_collapse = [dims_to_collapse]
-
-        if dims_to_collapse is not None and \
-                set(range(self.ndim)) != set(dims_to_collapse):
-            raise ValueError('Cannot partially collapse a coordinate (%s).'
-                             % self.name())
+            dims_to_collapse = (dims_to_collapse, )
+        if isinstance(dims_to_collapse, list):
+            dims_to_collapse = tuple(dims_to_collapse)
 
         if np.issubdtype(self.dtype, np.str_):
             # Collapse the coordinate by serializing the points and
@@ -1215,28 +1209,20 @@ def serialize(x):
                     'Metadata may not be fully descriptive for {!r}.'
                 warnings.warn(msg.format(self.name()))
 
-            # Create bounds for the new collapsed coordinate.
-            item = self.core_bounds() if self.has_bounds() \
+            # Determine the array library for stacking
+            al = da if self.has_bounds() \
+                and _lazy.is_lazy_data(self.core_bounds()) else np
+
+            item = al.concatenate(self.core_bounds()) if self.has_bounds() \
                 else self.core_points()
-            lower, upper = item.min(), item.max()
-            bounds_dtype = item.dtype
-            # Ensure 2D shape of new bounds.
-            bounds = np.empty((1, 2), 'object')
-            bounds[0, 0] = lower
-            bounds[0, 1] = upper
-            # Create points for the new collapsed coordinate.
-            points_dtype = self.dtype
-            points = (float(lower) + float(upper)) * 0.5
+
+            # Calculate the bounds and points along the right dims
+            bounds = al.stack([item.min(axis=dims_to_collapse),
+                               item.max(axis=dims_to_collapse)]).T
+            points = al.array(bounds.sum(axis=-1) * 0.5, dtype=self.dtype)
 
             # Create the new collapsed coordinate.
-            if _lazy.is_lazy_data(item):
-                bounds = _lazy.multidim_lazy_stack(bounds)
-                coord = self.copy(points=points, bounds=bounds)
-            else:
-                bounds = np.concatenate(bounds)
-                bounds = np.array(bounds, dtype=bounds_dtype)
-                coord = self.copy(points=np.array(points, dtype=points_dtype),
-                                  bounds=bounds)
+            coord = self.copy(points=points, bounds=bounds)
         return coord
 
     def _guess_bounds(self, bound_position=0.5):

lib/iris/cube.py

@@ -3160,22 +3160,6 @@ def collapsed(self, coords, aggregator, **kwargs):
 
                 cube.collapsed(['latitude', 'longitude'],
                                iris.analysis.VARIANCE)
-
-        .. _partially_collapse_multi-dim_coord:
-
-        .. note::
-            You cannot partially collapse a multi-dimensional coordinate. Doing
-            so would result in a partial collapse of the multi-dimensional
-            coordinate. Instead you must either:
-                 * collapse in a single operation all cube axes that the
-                   multi-dimensional coordinate spans,
-                 * remove the multi-dimensional coordinate from the cube before
-                   performing the collapse operation, or
-                 * not collapse the coordinate at all.
-
-            Multi-dimensional derived coordinates will not prevent a successful
-            collapse operation.
-
         """
         # Convert any coordinate names to coordinates
         coords = self._as_list_of_coords(coords)

lib/iris/tests/test_analysis.py

@@ -1,4 +1,4 @@
-# (C) British Crown Copyright 2010 - 2017, Met Office
+# (C) British Crown Copyright 2010 - 2018, Met Office
 #
 # This file is part of Iris.
 #
@@ -310,6 +310,54 @@ def test_sum(self):
         np.testing.assert_array_equal(cube.data, np.array([6, 18, 17]))
 
 
+class TestAuxCoordCollapse(tests.IrisTest):
+
+    def setUp(self):
+        self.cube_with_aux_coord = tests.stock.simple_4d_with_hybrid_height()
+
+        # Guess bounds to get the weights
+        self.cube_with_aux_coord.coord('grid_latitude').guess_bounds()
+        self.cube_with_aux_coord.coord('grid_longitude').guess_bounds()
+
+    def test_max(self):
+        cube = self.cube_with_aux_coord.collapsed('grid_latitude',
+                                                  iris.analysis.MAX)
+        np.testing.assert_array_equal(cube.coord('surface_altitude').points,
+                                      np.array([112, 113, 114,
+                                                115, 116, 117]))
+
+        np.testing.assert_array_equal(cube.coord('surface_altitude').bounds,
+                                      np.array([[100, 124],
+                                                [101, 125],
+                                                [102, 126],
+                                                [103, 127],
+                                                [104, 128],
+                                                [105, 129]]))
+
+        # Check collapsing over the whole coord still works
+        cube = self.cube_with_aux_coord.collapsed('altitude',
+                                                  iris.analysis.MAX)
+
+        np.testing.assert_array_equal(cube.coord('surface_altitude').points,
+                                      np.array([114]))
+
+        np.testing.assert_array_equal(cube.coord('surface_altitude').bounds,
+                                      np.array([[100, 129]]))
+
+        cube = self.cube_with_aux_coord.collapsed('grid_longitude',
+                                                  iris.analysis.MAX)
+
+        np.testing.assert_array_equal(cube.coord('surface_altitude').points,
+                                      np.array([102, 108, 114, 120, 126]))
+
+        np.testing.assert_array_equal(cube.coord('surface_altitude').bounds,
+                                      np.array([[100, 105],
+                                                [106, 111],
+                                                [112, 117],
+                                                [118, 123],
+                                                [124, 129]]))
+
+
 class TestAggregator_mdtol_keyword(tests.IrisTest):
     def setUp(self):
         data = ma.array([[1, 2], [4, 5]], dtype=np.float32,

lib/iris/tests/unit/coords/test_Coord.py

@@ -31,6 +31,7 @@
 from iris.coords import DimCoord, AuxCoord, Coord
 from iris.tests import mock
 from iris.exceptions import UnitConversionError
+from iris.tests.unit.coords import CoordTestMixin
 
 
 Pair = collections.namedtuple('Pair', 'points bounds')
@@ -232,7 +233,7 @@ def test_time_as_object(self):
                                      mock.sentinel.upper))])
 
 
-class Test_collapsed(tests.IrisTest):
+class Test_collapsed(tests.IrisTest, CoordTestMixin):
 
     def test_serialize(self):
         # Collapse a string AuxCoord, causing it to be serialised.
@@ -277,19 +278,58 @@ def test_dim_1d(self):
                                   [[coord.bounds.min(), coord.bounds.max()]])
 
     def test_numeric_nd(self):
-        # Contiguous only defined for 2d bounds.
-        coord = AuxCoord(points=np.array([3, 6, 9]),
-                         bounds=np.array([[1, 2, 4, 5],
+        coord = AuxCoord(points=np.array([[1, 2, 4, 5],
                                           [4, 5, 7, 8],
                                           [7, 8, 10, 11]]))
-        with self.assertRaises(ValueError):
-            coord.collapsed()
-
-    def test_collapsed_overflow(self):
-        coord = DimCoord(points=np.array([1493892000, 1493895600, 1493899200],
-                                         dtype=np.int32))
-        result = coord.collapsed()
-        self.assertEqual(result.points, 1493895600)
+
+        collapsed_coord = coord.collapsed()
+        self.assertArrayEqual(collapsed_coord.points, np.array([6]))
+        self.assertArrayEqual(collapsed_coord.bounds, np.array([[1, 11]]))
+
+        # Test partially collapsing one dimension...
+        collapsed_coord = coord.collapsed(1)
+        self.assertArrayEqual(collapsed_coord.points, np.array([3.,  6.,  9.]))
+        self.assertArrayEqual(collapsed_coord.bounds, np.array([[1,  5],
+                                                                [4,  8],
+                                                                [7, 11]]))
+
+        # ... and the other
+        collapsed_coord = coord.collapsed(0)
+        self.assertArrayEqual(collapsed_coord.points, np.array([4,  5,  7, 8]))
+        self.assertArrayEqual(collapsed_coord.bounds, np.array([[1,  7],
+                                                                [2,  8],
+                                                                [4, 10],
+                                                                [5, 11]]))
+
+    def test_lazy_nd_bounds(self):
+        import dask.array as da
+
+        self.setupTestArrays((3, 4))
+        coord = AuxCoord(self.pts_real, bounds=self.bds_lazy)
+
+        collapsed_coord = coord.collapsed()
+
+        # Note that the new points get recalculated from the lazy bounds
+        #  and so end up as lazy
+        self.assertTrue(collapsed_coord.has_lazy_points())
+        self.assertTrue(collapsed_coord.has_lazy_bounds())
+
+        self.assertArrayEqual(collapsed_coord.points, np.array([55]))
+        self.assertArrayEqual(collapsed_coord.bounds, da.array([[-2, 112]]))
+
+    def test_lazy_nd_points_and_bounds(self):
+        import dask.array as da
+
+        self.setupTestArrays((3, 4))
+        coord = AuxCoord(self.pts_lazy, bounds=self.bds_lazy)
+
+        collapsed_coord = coord.collapsed()
+
+        self.assertTrue(collapsed_coord.has_lazy_points())
+        self.assertTrue(collapsed_coord.has_lazy_bounds())
+
+        self.assertArrayEqual(collapsed_coord.points, da.array([55]))
+        self.assertArrayEqual(collapsed_coord.bounds, da.array([[-2, 112]]))
 
 
 class Test_is_compatible(tests.IrisTest):