From 3e1e80ecdb3784f25f4f2bc03bc8e0dbb0a1afa0 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Wed, 22 May 2024 10:53:17 +0200 Subject: [PATCH] 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 <40734014+trexfeathers@users.noreply.github.com> --- docs/src/whatsnew/latest.rst | 4 ++-- lib/iris/cube.py | 19 ++++++++----------- lib/iris/tests/unit/cube/test_Cube.py | 17 ++++++++++++++++- 3 files changed, 26 insertions(+), 14 deletions(-) diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index 31af80faa6..aca8511148 100644 --- a/docs/src/whatsnew/latest.rst +++ b/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 =============== diff --git a/lib/iris/cube.py b/lib/iris/cube.py index 8418a630b5..dc8a08e1bd 100644 --- a/lib/iris/cube.py +++ b/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 diff --git a/lib/iris/tests/unit/cube/test_Cube.py b/lib/iris/tests/unit/cube/test_Cube.py index ec94e346b2..408f8a9002 100644 --- 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) @@ -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])