Reinstating support for (partial) collapse of lazy coordinate points / bounds. Added some tests to catch this in the future.

duncanwp · duncanwp · commit 2e31ba30ccd9 · 2018-05-24T20:58:55.000+01:00
(cherry picked from commit 386edc7)
diff --git a/lib/iris/coords.py b/lib/iris/coords.py
@@ -1172,6 +1172,7 @@ def collapsed(self, dims_to_collapse=None):
 
         Replaces the points & bounds with a simple bounded region.
         """
+        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)):
@@ -1210,13 +1211,17 @@ 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 = np.concatenate(self.core_bounds()) if self.has_bounds() \
+            # Determine the right array method for stacking
+            stack_method = da.stack if self.has_bounds() \
+                                       and is_lazy_data(self.core_bounds()) \
+                else np.stack
+
+            item = stack_method(self.core_bounds()) if self.has_bounds() \
                 else self.core_points()
 
             # 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
+            bounds = stack_method([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.
diff --git a/lib/iris/tests/test_analysis.py b/lib/iris/tests/test_analysis.py
@@ -357,6 +357,7 @@ def test_max(self):
                                                 [118, 123],
                                                 [124, 129]]))
 
+
 class TestAggregator_mdtol_keyword(tests.IrisTest):
     def setUp(self):
         data = ma.array([[1, 2], [4, 5]], dtype=np.float32,
diff --git a/lib/iris/tests/unit/coords/test_Coord.py b/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.
@@ -300,6 +301,49 @@ def test_numeric_nd(self):
                                                                 [4, 10],
                                                                 [5, 11]]))
 
+    def test_lazy_nd_points(self):
+        import dask.array as da
+        self.setupTestArrays((3, 4))
+        coord = AuxCoord(self.pts_lazy)
+
+        collapsed_coord = coord.collapsed()
+
+        self.assertTrue(collapsed_coord.has_lazy_points())
+        self.assertFalse(collapsed_coord.has_lazy_bounds())
+
+        self.assertArrayEqual(collapsed_coord.points, da.array([55]))
+        self.assertArrayEqual(collapsed_coord.bounds, da.array([[0, 110]]))
+
+    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):
     def setUp(self):
