Lazy iris.cube.Cube.rolling_window (#5795)

* Lazy iris.cube.Cube.rolling_window

* Fix test

* Add whatsnew

* Move test to tests.unit

---------

Co-authored-by: Martin Yeo <[email protected]>

Author: bouweandela

docs/src/whatsnew/latest.rst

@@ -55,14 +55,14 @@ This document explains the changes made to Iris for this release
 🚀 Performance Enhancements
 ===========================
 
-#. N/A
-
 #. `@bouweandela`_ added the option to specify the Dask chunks of the target
    array in :func:`iris.util.broadcast_to_shape`. (:pull:`5620`)
 
 #. `@schlunma`_ allowed :func:`iris.analysis.cartography.area_weights` to
    return dask arrays with arbitrary chunks. (:pull:`5658`)
 
+#. `@bouweandela`_ made :meth:`iris.cube.Cube.rolling_window` work with lazy
+   data. (:pull:`5795`)
 
 🔥 Deprecations
 ===============

lib/iris/cube.py

@@ -4552,12 +4552,6 @@ def rolling_window(self, coord, aggregator, window, **kwargs):
         -------
         :class:`iris.cube.Cube`.
 
-        Notes
-        -----
-        .. note::
-
-            This operation does not yet have support for lazy evaluation.
-
         Examples
         --------
             >>> import iris, iris.analysis
@@ -4661,7 +4655,7 @@ def rolling_window(self, coord, aggregator, window, **kwargs):
         # this will add an extra dimension to the data at dimension + 1 which
         # represents the rolled window (i.e. will have a length of window)
         rolling_window_data = iris.util.rolling_window(
-            self.data, window=window, axis=dimension
+            self.core_data(), window=window, axis=dimension
         )
 
         # now update all of the coordinates to reflect the aggregation
@@ -4680,7 +4674,7 @@ def rolling_window(self, coord, aggregator, window, **kwargs):
                     "coordinate." % coord_.name()
                 )
 
-            new_bounds = iris.util.rolling_window(coord_.points, window)
+            new_bounds = iris.util.rolling_window(coord_.core_points(), window)
 
             if np.issubdtype(new_bounds.dtype, np.str_):
                 # Handle case where the AuxCoord contains string. The points
@@ -4726,9 +4720,12 @@ def rolling_window(self, coord, aggregator, window, **kwargs):
                 kwargs["weights"] = iris.util.broadcast_to_shape(
                     weights, rolling_window_data.shape, (dimension + 1,)
                 )
-        data_result = aggregator.aggregate(
-            rolling_window_data, axis=dimension + 1, **kwargs
-        )
+
+        if aggregator.lazy_func is not None and self.has_lazy_data():
+            agg_method = aggregator.lazy_aggregate
+        else:
+            agg_method = aggregator.aggregate
+        data_result = agg_method(rolling_window_data, axis=dimension + 1, **kwargs)
         result = aggregator.post_process(new_cube, data_result, [coord], **kwargs)
         return result
 

lib/iris/tests/unit/cube/test_Cube.py

@@ -877,7 +877,7 @@ def setUp(self):
         self.cell_measure = CellMeasure([0, 1, 2, 0, 1, 2], long_name="bar")
         self.multi_dim_cube.add_cell_measure(self.cell_measure, 1)
 
-        self.mock_agg = mock.Mock(spec=Aggregator)
+        self.mock_agg = mock.Mock(spec=Aggregator, lazy_func=None)
         self.mock_agg.aggregate = mock.Mock(return_value=np.empty([4]))
         self.mock_agg.post_process = mock.Mock(side_effect=lambda x, y, z: x)
 
@@ -919,6 +919,21 @@ def test_kwargs(self):
         )
         self.assertMaskedArrayEqual(expected_result, res_cube.data)
 
+    def test_lazy(self):
+        window = 2
+        self.cube.data = da.ma.masked_array(
+            self.cube.data, mask=([True, False, False, False, True, False])
+        )
+        res_cube = self.cube.rolling_window("val", iris.analysis.MEAN, window, mdtol=0)
+        self.assertTrue(self.cube.has_lazy_data())
+        self.assertTrue(res_cube.has_lazy_data())
+        expected_result = ma.array(
+            [-99.0, 1.5, 2.5, -99.0, -99.0],
+            mask=[True, False, False, True, True],
+            dtype=np.float64,
+        )
+        self.assertMaskedArrayEqual(expected_result, res_cube.data)
+
     def test_ancillary_variables_and_cell_measures_kept(self):
         res_cube = self.multi_dim_cube.rolling_window("val", self.mock_agg, 3)
         self.assertEqual(res_cube.ancillary_variables(), [self.ancillary_variable])