Lazy iris.cube.Cube.rolling_window

bouweandela · Mar 4, 2024 · 3427c19 · 3427c19
1 parent c6151e8
commit 3427c19
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 12 deletions.
diff --git a/lib/iris/cube.py b/lib/iris/cube.py
@@ -4502,12 +4502,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
@@ -4611,7 +4605,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
@@ -4630,7 +4624,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
@@ -4676,9 +4670,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
 

diff --git a/lib/iris/tests/test_analysis.py b/lib/iris/tests/test_analysis.py
@@ -1452,6 +1452,32 @@ def test_longitude_masked(self):
 
         self.assertMaskedArrayEqual(expected_result, res_cube.data)
 
+    def test_longitude_masked_lazy(self):
+        self.cube.data = ma.array(
+            self.cube.data,
+            mask=[
+                [True, True, True, True],
+                [True, False, True, True],
+                [False, False, False, False],
+            ],
+        )
+        self.cube.data = self.cube.lazy_data()
+        res_cube = self.cube.rolling_window("longitude", iris.analysis.MEAN, window=2)
+
+        expected_result = np.ma.array(
+            [[-99.0, -99.0, -99.0], [12.0, 12.0, -99.0], [15.0, 11.0, 8.0]],
+            mask=[
+                [True, True, True],
+                [False, False, True],
+                [False, False, False],
+            ],
+            dtype=np.float64,
+        )
+
+        self.assertTrue(self.cube.has_lazy_data())
+        self.assertTrue(res_cube.has_lazy_data())
+        self.assertMaskedArrayEqual(expected_result, res_cube.data)
+
     def test_longitude_circular(self):
         cube = self.cube
         cube.coord("longitude").circular = True

diff --git a/lib/iris/tests/unit/cube/test_Cube.py b/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)