collapsed preserving masked aux coords (SciTools#6719)

* collapsed preserving masked aux coords

* test coverage

* add whatsnew entry

Author: bjlittle

docs/src/whatsnew/latest.rst

@@ -48,6 +48,10 @@ This document explains the changes made to Iris for this release
    on the correct dimensions, and merge no longer fails when trying to add
    them to a dimension of the wrong length. (:issue:`2076`, :pull:`6688`)
 
+#. `@bjlittle`_ added support for preserving masked auxiliary coordinates when
+   using :meth:`~iris.cube.Cube.aggregated_by` or :meth:`~iris.cube.Cube.collapsed`.
+   (:issue:`6473`, :pull:`6706`, :pull:`6719`)
+
 
 💣 Incompatible Changes
 =======================

lib/iris/coords.py

@@ -2219,7 +2219,7 @@ def serialize(x, axis):
                 item = self.core_points()
 
             # Determine the array library for stacking
-            al = da if _lazy.is_lazy_data(item) else np
+            al = da if _lazy.is_lazy_data(item) else ma
 
             # Calculate the bounds and points along the right dims
             bounds = al.stack(
@@ -2231,6 +2231,12 @@ def serialize(x, axis):
             )
             points = al.array(bounds.sum(axis=-1) * 0.5, dtype=self.dtype)
 
+            if ma.isMaskedArray(points) and not np.any(points.mask):
+                points = points.data
+
+            if ma.isMaskedArray(bounds) and not np.any(bounds.mask):
+                bounds = bounds.data
+
             # Create the new collapsed coordinate.
             coord = self.copy(points=points, bounds=bounds)
         return coord

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

@@ -56,6 +56,11 @@ def is_real_data(array):
     return isinstance(array, np.ndarray)
 
 
+def is_masked_data(array):
+    # Check the array is a masked array.
+    return ma.isMaskedArray(array)
+
+
 def arrays_share_data(a1, a2):
     # Check whether 2 real arrays with the same content view the same data.
     # For an ndarray x, x.base will either be None (if x owns its data) or a
@@ -134,6 +139,12 @@ def assert_is_lazy_array(self, array, msg=None):
             msg = f"Array {array} is not a lazy array"
         assert is_lazy_data(array), msg
 
+    def assert_is_masked_array(self, array, msg=None):
+        # Check that the arg is a masked array.
+        if not msg:
+            msg = f"Array {array} is not a masked array"
+        assert is_masked_data(array), msg
+
     def assert_equal_real_arrays_and_dtypes(self, a1, a2):
         # Check that two arrays are real, equal, and have same dtype.
         self.assert_is_real_array(a1)

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

@@ -506,6 +506,46 @@ def test_lazy_points(self):
         assert collapsed_coord.has_lazy_bounds()
         assert collapsed_coord.has_lazy_points()
 
+    def test_numeric_masked_all(self):
+        coord = AuxCoord(
+            points=ma.array(
+                [[1, 2, 4, 5], [4, 5, 7, 8], [7, 8, 10, 11]],
+                mask=True,
+            ),
+        )
+
+        collapsed_coord = coord.collapsed()
+        self.assert_is_masked_array(collapsed_coord.points)
+        self.assert_is_masked_array(collapsed_coord.bounds)
+        expected = AuxCoord(
+            ma.array([-1], mask=True), bounds=ma.array([[-1, -1]], mask=[[True, True]])
+        )
+        assert collapsed_coord == expected
+
+        collapsed_coord = coord.collapsed(dims_to_collapse=1)
+        self.assert_is_masked_array(collapsed_coord.points)
+        self.assert_is_masked_array(collapsed_coord.bounds)
+        expected = AuxCoord(
+            points=ma.array([-1, -1, -1], mask=[True, True, True]),
+            bounds=ma.array(
+                [[-1, -1], [-1, -1], [-1, -1]],
+                mask=[[True, True], [True, True], [True, True]],
+            ),
+        )
+        assert collapsed_coord == expected
+
+        collapsed_coord = coord.collapsed(dims_to_collapse=0)
+        self.assert_is_masked_array(collapsed_coord.points)
+        self.assert_is_masked_array(collapsed_coord.bounds)
+        expected = AuxCoord(
+            points=ma.array([-1, -1, -1, -1], mask=[True, True, True, True]),
+            bounds=ma.array(
+                [[-1, -1], [-1, -1], [-1, -1], [-1, -1]],
+                mask=[[True, True], [True, True], [True, True], [True, True]],
+            ),
+        )
+        assert collapsed_coord == expected
+
     def test_numeric_nd(self):
         coord = AuxCoord(points=np.array([[1, 2, 4, 5], [4, 5, 7, 8], [7, 8, 10, 11]]))
 
@@ -530,6 +570,42 @@ def test_numeric_nd(self):
             np.array([[1, 7], [2, 8], [4, 10], [5, 11]]),
         )
 
+    def test_numeric_masked(self):
+        coord = AuxCoord(
+            points=ma.array(
+                [[1, 2, 4, 5], [4, 5, 7, 8], [7, 8, 10, 11]],
+                mask=[
+                    [True, False, False, True],
+                    [True, True, True, True],
+                    [False, True, True, False],
+                ],
+            ),
+        )
+
+        collapsed_coord = coord.collapsed()
+        _shared_utils.assert_array_equal(collapsed_coord.points, np.array([6]))
+        _shared_utils.assert_array_equal(collapsed_coord.bounds, np.array([[2, 11]]))
+
+        collapsed_coord = coord.collapsed(dims_to_collapse=1)
+        self.assert_is_masked_array(collapsed_coord.points)
+        self.assert_is_masked_array(collapsed_coord.bounds)
+        expected = AuxCoord(
+            points=ma.array([3, -1, 9], mask=[False, True, False]),
+            bounds=ma.array(
+                [[2, 4], [-1, -1], [7, 11]],
+                mask=[[False, False], [True, True], [False, False]],
+            ),
+        )
+        assert collapsed_coord == expected
+
+        collapsed_coord = coord.collapsed(dims_to_collapse=0)
+        _shared_utils.assert_array_equal(
+            collapsed_coord.points, np.array([7, 2, 4, 11])
+        )
+        _shared_utils.assert_array_equal(
+            collapsed_coord.bounds, np.array([[7, 7], [2, 2], [4, 4], [11, 11]])
+        )
+
     def test_numeric_nd_bounds_all(self):
         self.setup_test_arrays((3, 4))
         coord = AuxCoord(self.pts_real, bounds=self.bds_real)
@@ -538,6 +614,22 @@ def test_numeric_nd_bounds_all(self):
         _shared_utils.assert_array_equal(collapsed_coord.points, np.array([55]))
         _shared_utils.assert_array_equal(collapsed_coord.bounds, np.array([[-2, 112]]))
 
+    def test_numeric_no_masked_bounds_all(self):
+        self.setup_test_arrays((3, 4), masked=True)
+        coord = AuxCoord(self.no_masked_pts_real, bounds=self.no_masked_bds_real)
+
+        collapsed_coord = coord.collapsed()
+        _shared_utils.assert_array_equal(collapsed_coord.points, np.array([55]))
+        _shared_utils.assert_array_equal(collapsed_coord.bounds, np.array([[-2, 112]]))
+
+    def test_numeric_masked_bounds_all(self):
+        self.setup_test_arrays((3, 4), masked=True)
+        coord = AuxCoord(self.masked_pts_real, bounds=self.masked_bds_real)
+
+        collapsed_coord = coord.collapsed()
+        _shared_utils.assert_array_equal(collapsed_coord.points, np.array([75]))
+        _shared_utils.assert_array_equal(collapsed_coord.bounds, np.array([[38, 112]]))
+
     def test_numeric_nd_bounds_second(self):
         self.setup_test_arrays((3, 4))
         coord = AuxCoord(self.pts_real, bounds=self.bds_real)
@@ -547,6 +639,30 @@ def test_numeric_nd_bounds_second(self):
             collapsed_coord.bounds, np.array([[-2, 32], [38, 72], [78, 112]])
         )
 
+    def test_numeric_no_masked_bounds_second(self):
+        self.setup_test_arrays((3, 4), masked=True)
+        coord = AuxCoord(self.no_masked_pts_real, bounds=self.no_masked_bds_real)
+        collapsed_coord = coord.collapsed(dims_to_collapse=1)
+        _shared_utils.assert_array_equal(collapsed_coord.points, np.array([15, 55, 95]))
+        _shared_utils.assert_array_equal(
+            collapsed_coord.bounds, np.array([[-2, 32], [38, 72], [78, 112]])
+        )
+
+    def test_numeric_masked_bounds_second(self):
+        self.setup_test_arrays((3, 4), masked=True)
+        coord = AuxCoord(self.masked_pts_real, bounds=self.masked_bds_real)
+        collapsed_coord = coord.collapsed(dims_to_collapse=1)
+        self.assert_is_masked_array(collapsed_coord.points)
+        self.assert_is_masked_array(collapsed_coord.bounds)
+        expected = AuxCoord(
+            ma.array([-1, 55, 95], mask=[True, False, False]),
+            bounds=ma.array(
+                [[-1, -1], [38, 72], [78, 112]],
+                mask=[[True, True], [False, False], [False, False]],
+            ),
+        )
+        assert collapsed_coord == expected
+
     def test_numeric_nd_bounds_first(self):
         self.setup_test_arrays((3, 4))
         coord = AuxCoord(self.pts_real, bounds=self.bds_real)
@@ -560,6 +676,28 @@ def test_numeric_nd_bounds_first(self):
             np.array([[-2, 82], [8, 92], [18, 102], [28, 112]]),
         )
 
+    def test_numeric_no_masked_bounds_first(self):
+        self.setup_test_arrays((3, 4), masked=True)
+        coord = AuxCoord(self.no_masked_pts_real, bounds=self.no_masked_bds_real)
+        collapsed_coord = coord.collapsed(dims_to_collapse=0)
+        _shared_utils.assert_array_equal(
+            collapsed_coord.points, np.array([40, 50, 60, 70])
+        )
+        _shared_utils.assert_array_equal(
+            collapsed_coord.bounds, np.array([[-2, 82], [8, 92], [18, 102], [28, 112]])
+        )
+
+    def test_numeric_masked_bounds_first(self):
+        self.setup_test_arrays((3, 4), masked=True)
+        coord = AuxCoord(self.masked_pts_real, bounds=self.masked_bds_real)
+        collapsed_coord = coord.collapsed(dims_to_collapse=0)
+        _shared_utils.assert_array_equal(
+            collapsed_coord.points, np.array([60, 70, 80, 90])
+        )
+        _shared_utils.assert_array_equal(
+            collapsed_coord.bounds, np.array([[38, 82], [48, 92], [58, 102], [68, 112]])
+        )
+
     def test_numeric_nd_bounds_last(self):
         self.setup_test_arrays((3, 4))
         coord = AuxCoord(self.pts_real, bounds=self.bds_real)
@@ -570,6 +708,30 @@ def test_numeric_nd_bounds_last(self):
             collapsed_coord.bounds, np.array([[-2, 32], [38, 72], [78, 112]])
         )
 
+    def test_numeric_no_masked_bounds_last(self):
+        self.setup_test_arrays((3, 4), masked=True)
+        coord = AuxCoord(self.no_masked_pts_real, bounds=self.no_masked_bds_real)
+        collapsed_coord = coord.collapsed(dims_to_collapse=-1)
+        _shared_utils.assert_array_equal(collapsed_coord.points, np.array([15, 55, 95]))
+        _shared_utils.assert_array_equal(
+            collapsed_coord.bounds, np.array([[-2, 32], [38, 72], [78, 112]])
+        )
+
+    def test_numeric_masked_bounds_last(self):
+        self.setup_test_arrays((3, 4), masked=True)
+        coord = AuxCoord(self.masked_pts_real, bounds=self.masked_bds_real)
+        collapsed_coord = coord.collapsed(dims_to_collapse=-1)
+        self.assert_is_masked_array(collapsed_coord.points)
+        self.assert_is_masked_array(collapsed_coord.bounds)
+        expected = AuxCoord(
+            ma.array([-1, 55, 95], mask=[True, False, False]),
+            bounds=ma.array(
+                [[-1, -1], [38, 72], [78, 112]],
+                mask=[[True, True], [False, False], [False, False]],
+            ),
+        )
+        assert collapsed_coord == expected
+
     def test_lazy_nd_bounds_all(self):
         self.setup_test_arrays((3, 4))
         coord = AuxCoord(self.pts_real, bounds=self.bds_lazy)