Adding support for partial collapse of multi-dimensional coordinates

lib/iris/coords.py

@@ -1144,20 +1144,12 @@ 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.
-
         """
+        # 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
@@ -1191,27 +1183,16 @@ def serialize(x):
                 warnings.warn(msg.format(self.name()))
 
             # Create bounds for the new collapsed coordinate.
-            item = self.core_bounds() if self.has_bounds() \
+            item = np.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 = np.stack([item.min(axis=dims_to_collapse),
+                               item.max(axis=dims_to_collapse)]).T
+            points = item.mean(axis=dims_to_collapse, dtype=self.dtype)
 
             # Create the new collapsed coordinate.
-            if is_lazy_data(item):
-                bounds = 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

@@ -3111,22 +3111,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

@@ -310,6 +310,60 @@ def test_sum(self):
         np.testing.assert_array_equal(cube.data, np.array([6, 18, 17]))
 
 
+class TestAuxCoordCollapse(tests.IrisTest):
+
+    def setUp(self):
+        from iris.analysis.cartography import area_weights
+        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()
+
+        self.weights = area_weights(self.cube_with_aux_coord, normalize=False)
+        self.normalized_weights = area_weights(self.cube_with_aux_coord, normalize=True)
+
+        self.original_alt = self.cube_with_aux_coord.coord('altitude')
+        # [[100, 101, 102, 103, 104, 105],
+        #  [106, 107, 108, 109, 110, 111],
+        #  [112, 113, 114, 115, 116, 117],
+        #  [118, 119, 120, 121, 122, 123],
+        #  [124, 125, 126, 127, 128, 129]]
+
+    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

@@ -277,19 +277,28 @@ 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]]))
 
 
 class Test_is_compatible(tests.IrisTest):