From 8376b1db69649a10ac5c98c0583c460c3684c6b2 Mon Sep 17 00:00:00 2001 From: Wolfgang Hayek Date: Fri, 16 Sep 2022 13:08:56 +1200 Subject: [PATCH 1/6] Lazy rotate_winds and associated tests --- lib/iris/analysis/cartography.py | 49 +++++++++++++------ .../analysis/cartography/test_rotate_winds.py | 47 ++++++++++++++++++ 2 files changed, 80 insertions(+), 16 deletions(-) diff --git a/lib/iris/analysis/cartography.py b/lib/iris/analysis/cartography.py index f38e48354d..952ba85927 100644 --- a/lib/iris/analysis/cartography.py +++ b/lib/iris/analysis/cartography.py @@ -17,6 +17,7 @@ import cf_units import numpy as np import numpy.ma as ma +import dask.array as da import iris.coord_systems import iris.coords @@ -1169,12 +1170,6 @@ def rotate_winds(u_cube, v_cube, target_cs): x = x.transpose() y = y.transpose() - # Create resulting cubes. - ut_cube = u_cube.copy() - vt_cube = v_cube.copy() - ut_cube.rename("transformed_{}".format(u_cube.name())) - vt_cube.rename("transformed_{}".format(v_cube.name())) - # Get distance scalings for source crs. ds_dx1, ds_dy1 = _crs_distance_differentials(src_crs, x, y) @@ -1197,10 +1192,28 @@ def rotate_winds(u_cube, v_cube, target_cs): src_crs, x, y, target_crs, ds, dx2, dy2 ) apply_mask = mask.any() + + # Create resulting cubes - both will be lazy if at least one input cube is + lazy_output = u_cube.has_lazy_data() or v_cube.has_lazy_data() if apply_mask: # Make masked arrays to accept masking. - ut_cube.data = ma.asanyarray(ut_cube.data) - vt_cube.data = ma.asanyarray(vt_cube.data) + if lazy_output: + # Use da.ma.empty_like to preserve Dask array attributes like chunking + ut_cube = u_cube.copy(data=da.ma.empty_like(u_cube.lazy_data())) + vt_cube = v_cube.copy(data=da.ma.empty_like(v_cube.lazy_data())) + else: + ut_cube = u_cube.copy(data=ma.asanyarray(u_cube.data)) + vt_cube = v_cube.copy(data=ma.asanyarray(v_cube.data)) + else: + if lazy_output: + ut_cube = u_cube.copy(data=u_cube.lazy_data()) + vt_cube = v_cube.copy(data=v_cube.lazy_data()) + else: + ut_cube = u_cube.copy() + vt_cube = v_cube.copy() + + ut_cube.rename("transformed_{}".format(u_cube.name())) + vt_cube.rename("transformed_{}".format(v_cube.name())) # Project vectors with u, v components one horiz slice at a time and # insert into the resulting cubes. @@ -1213,16 +1226,20 @@ def rotate_winds(u_cube, v_cube, target_cs): for dim in dims: index[dim] = slice(None, None) index = tuple(index) - u = u_cube.data[index] - v = v_cube.data[index] + u = u_cube.core_data()[index] + v = v_cube.core_data()[index] ut, vt = _transform_distance_vectors(u, v, ds, dx2, dy2) if apply_mask: - ut = ma.asanyarray(ut) - ut[mask] = ma.masked - vt = ma.asanyarray(vt) - vt[mask] = ma.masked - ut_cube.data[index] = ut - vt_cube.data[index] = vt + if lazy_output: + ut = da.ma.masked_array(ut, mask=mask) + vt = da.ma.masked_array(vt, mask=mask) + else: + ut = ma.asanyarray(ut) + ut[mask] = ma.masked + vt = ma.asanyarray(vt) + vt[mask] = ma.masked + ut_cube.core_data()[index] = ut + vt_cube.core_data()[index] = vt # Calculate new coords of locations in target coordinate system. xyz_tran = target_crs.transform_points(src_crs, x, y) diff --git a/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py b/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py index 7952b3bb46..5540287298 100644 --- a/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py +++ b/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py @@ -511,5 +511,52 @@ def test_non_earth_semimajor_axis(self): rotate_winds(u, v, other_cs) +class TestLazyRotateWinds(tests.IrisTest): + def _compare_lazy_rotate_winds(self, masked): + # Compute wind rotation with NumPy data first, then with Dask arrays, + # and compare results + + # Choose target coord system that will (not) lead to masked results + if masked: + coord_sys = iris.coord_systems.OSGB() + else: + coord_sys = iris.coord_systems.GeogCS(6371229) + + u, v = uv_cubes() + u_lazy = u.copy(data=u.lazy_data()) + v_lazy = v.copy(data=v.lazy_data()) + + ut_ref, vt_ref = rotate_winds(u.copy(), v.copy(), coord_sys) + self.assertFalse(ut_ref.has_lazy_data()) + self.assertFalse(vt_ref.has_lazy_data()) + # Check if choice of target coordinates leads to (no) masking + self.assertTrue(ma.isMaskedArray(ut_ref.data) == masked) + + # Results are lazy if at least one component is lazy + ut, vt = rotate_winds(u_lazy.copy(), v.copy(), coord_sys) + self.assertTrue(ut.has_lazy_data()) + self.assertTrue(vt.has_lazy_data()) + self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) + self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) + + ut, vt = rotate_winds(u.copy(), v_lazy.copy(),coord_sys) + self.assertTrue(ut.has_lazy_data()) + self.assertTrue(vt.has_lazy_data()) + self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) + self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) + + ut, vt = rotate_winds(u_lazy.copy(), v_lazy.copy(), coord_sys) + self.assertTrue(ut.has_lazy_data()) + self.assertTrue(vt.has_lazy_data()) + self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) + self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) + + def test_lazy_rotate_winds_masked(self): + self._compare_lazy_rotate_winds(True) + + def test_lazy_rotate_winds_notmasked(self): + self._compare_lazy_rotate_winds(False) + + if __name__ == "__main__": tests.main() From e767ac328fe74d9a0fac90101dc383ee5dd3b83d Mon Sep 17 00:00:00 2001 From: Wolfgang Hayek Date: Sat, 17 Sep 2022 11:58:34 +1200 Subject: [PATCH 2/6] Data ownership bug fix and extra tests --- lib/iris/analysis/cartography.py | 34 +++++++++---------- .../analysis/cartography/test_rotate_winds.py | 23 +++++++++---- 2 files changed, 33 insertions(+), 24 deletions(-) diff --git a/lib/iris/analysis/cartography.py b/lib/iris/analysis/cartography.py index 952ba85927..7fac792b1b 100644 --- a/lib/iris/analysis/cartography.py +++ b/lib/iris/analysis/cartography.py @@ -1170,6 +1170,18 @@ def rotate_winds(u_cube, v_cube, target_cs): x = x.transpose() y = y.transpose() + # Create resulting cubes - produce lazy output data if at least + # one input cube has lazy data + lazy_output = u_cube.has_lazy_data() or v_cube.has_lazy_data() + if lazy_output: + ut_cube = u_cube.copy(data=da.empty_like(u_cube.lazy_data())) + vt_cube = v_cube.copy(data=da.empty_like(v_cube.lazy_data())) + else: + ut_cube = u_cube.copy() + vt_cube = v_cube.copy() + ut_cube.rename("transformed_{}".format(u_cube.name())) + vt_cube.rename("transformed_{}".format(v_cube.name())) + # Get distance scalings for source crs. ds_dx1, ds_dy1 = _crs_distance_differentials(src_crs, x, y) @@ -1192,28 +1204,14 @@ def rotate_winds(u_cube, v_cube, target_cs): src_crs, x, y, target_crs, ds, dx2, dy2 ) apply_mask = mask.any() - - # Create resulting cubes - both will be lazy if at least one input cube is - lazy_output = u_cube.has_lazy_data() or v_cube.has_lazy_data() if apply_mask: # Make masked arrays to accept masking. if lazy_output: - # Use da.ma.empty_like to preserve Dask array attributes like chunking - ut_cube = u_cube.copy(data=da.ma.empty_like(u_cube.lazy_data())) - vt_cube = v_cube.copy(data=da.ma.empty_like(v_cube.lazy_data())) + ut_cube = ut_cube.copy(data=da.ma.masked_array(ut_cube.core_data())) + vt_cube = vt_cube.copy(data=da.ma.masked_array(vt_cube.core_data())) else: - ut_cube = u_cube.copy(data=ma.asanyarray(u_cube.data)) - vt_cube = v_cube.copy(data=ma.asanyarray(v_cube.data)) - else: - if lazy_output: - ut_cube = u_cube.copy(data=u_cube.lazy_data()) - vt_cube = v_cube.copy(data=v_cube.lazy_data()) - else: - ut_cube = u_cube.copy() - vt_cube = v_cube.copy() - - ut_cube.rename("transformed_{}".format(u_cube.name())) - vt_cube.rename("transformed_{}".format(v_cube.name())) + ut_cube.data = ma.asanyarray(ut_cube.data) + vt_cube.data = ma.asanyarray(vt_cube.data) # Project vectors with u, v components one horiz slice at a time and # insert into the resulting cubes. diff --git a/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py b/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py index 5540287298..2ec52cd31b 100644 --- a/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py +++ b/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py @@ -523,34 +523,45 @@ def _compare_lazy_rotate_winds(self, masked): coord_sys = iris.coord_systems.GeogCS(6371229) u, v = uv_cubes() - u_lazy = u.copy(data=u.lazy_data()) - v_lazy = v.copy(data=v.lazy_data()) - ut_ref, vt_ref = rotate_winds(u.copy(), v.copy(), coord_sys) + # Create lazy cubes, setting explicit chunksizes to test if + # Dask array metadata is preserved + u_lazy = u.copy(data=u.lazy_data().rechunk([2,1])) + v_lazy = v.copy(data=v.lazy_data().rechunk([1,2])) + + ut_ref, vt_ref = rotate_winds(u, v, coord_sys) self.assertFalse(ut_ref.has_lazy_data()) self.assertFalse(vt_ref.has_lazy_data()) # Check if choice of target coordinates leads to (no) masking self.assertTrue(ma.isMaskedArray(ut_ref.data) == masked) # Results are lazy if at least one component is lazy - ut, vt = rotate_winds(u_lazy.copy(), v.copy(), coord_sys) + ut, vt = rotate_winds(u_lazy, v, coord_sys) self.assertTrue(ut.has_lazy_data()) self.assertTrue(vt.has_lazy_data()) + self.assertTrue(ut.core_data().chunksize == (2,1)) self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) - ut, vt = rotate_winds(u.copy(), v_lazy.copy(),coord_sys) + ut, vt = rotate_winds(u, v_lazy, coord_sys) self.assertTrue(ut.has_lazy_data()) self.assertTrue(vt.has_lazy_data()) + self.assertTrue(vt.core_data().chunksize == (1,2)) self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) - ut, vt = rotate_winds(u_lazy.copy(), v_lazy.copy(), coord_sys) + ut, vt = rotate_winds(u_lazy, v_lazy, coord_sys) self.assertTrue(ut.has_lazy_data()) self.assertTrue(vt.has_lazy_data()) + self.assertTrue(ut.core_data().chunksize == (2,1)) + self.assertTrue(vt.core_data().chunksize == (1,2)) self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) + # Check that input data has not been modified + self.assertArrayAllClose(u.data, u_lazy.data, rtol=1e-5) + self.assertArrayAllClose(v.data, v_lazy.data, rtol=1e-5) + def test_lazy_rotate_winds_masked(self): self._compare_lazy_rotate_winds(True) From e0ae4a24846f5d049cdeae9e396f12faedff9bc8 Mon Sep 17 00:00:00 2001 From: Wolfgang Hayek Date: Mon, 19 Sep 2022 10:19:19 +1200 Subject: [PATCH 3/6] Fix input data test --- .../analysis/cartography/test_rotate_winds.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py b/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py index 2ec52cd31b..a87db46f8d 100644 --- a/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py +++ b/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py @@ -513,8 +513,7 @@ def test_non_earth_semimajor_axis(self): class TestLazyRotateWinds(tests.IrisTest): def _compare_lazy_rotate_winds(self, masked): - # Compute wind rotation with NumPy data first, then with Dask arrays, - # and compare results + # Compute wind rotation with lazy data and compare results # Choose target coord system that will (not) lead to masked results if masked: @@ -524,15 +523,15 @@ def _compare_lazy_rotate_winds(self, masked): u, v = uv_cubes() - # Create lazy cubes, setting explicit chunksizes to test if - # Dask array metadata is preserved - u_lazy = u.copy(data=u.lazy_data().rechunk([2,1])) - v_lazy = v.copy(data=v.lazy_data().rechunk([1,2])) + # Create deep copy of the cubes with rechunked lazy data to check if + # input data is modified, and if Dask metadata is preserved + u_lazy = u.copy(data=u.copy().lazy_data().rechunk([2,1])) + v_lazy = v.copy(data=v.copy().lazy_data().rechunk([1,2])) ut_ref, vt_ref = rotate_winds(u, v, coord_sys) self.assertFalse(ut_ref.has_lazy_data()) self.assertFalse(vt_ref.has_lazy_data()) - # Check if choice of target coordinates leads to (no) masking + # Ensure that choice of target coordinates leads to (no) masking self.assertTrue(ma.isMaskedArray(ut_ref.data) == masked) # Results are lazy if at least one component is lazy @@ -558,7 +557,7 @@ def _compare_lazy_rotate_winds(self, masked): self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) - # Check that input data has not been modified + # Ensure that input data has not been modified self.assertArrayAllClose(u.data, u_lazy.data, rtol=1e-5) self.assertArrayAllClose(v.data, v_lazy.data, rtol=1e-5) From 1bc362811893f980f309416d5ccd9ee005571ce3 Mon Sep 17 00:00:00 2001 From: Wolfgang Hayek Date: Mon, 19 Sep 2022 10:47:10 +1200 Subject: [PATCH 4/6] Code formatting --- lib/iris/analysis/cartography.py | 10 +++++++--- .../unit/analysis/cartography/test_rotate_winds.py | 12 ++++++------ 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/lib/iris/analysis/cartography.py b/lib/iris/analysis/cartography.py index 7fac792b1b..f493dc7fa0 100644 --- a/lib/iris/analysis/cartography.py +++ b/lib/iris/analysis/cartography.py @@ -15,9 +15,9 @@ import cartopy.crs as ccrs import cartopy.img_transform import cf_units +import dask.array as da import numpy as np import numpy.ma as ma -import dask.array as da import iris.coord_systems import iris.coords @@ -1207,8 +1207,12 @@ def rotate_winds(u_cube, v_cube, target_cs): if apply_mask: # Make masked arrays to accept masking. if lazy_output: - ut_cube = ut_cube.copy(data=da.ma.masked_array(ut_cube.core_data())) - vt_cube = vt_cube.copy(data=da.ma.masked_array(vt_cube.core_data())) + ut_cube = ut_cube.copy( + data=da.ma.masked_array(ut_cube.core_data()) + ) + vt_cube = vt_cube.copy( + data=da.ma.masked_array(vt_cube.core_data()) + ) else: ut_cube.data = ma.asanyarray(ut_cube.data) vt_cube.data = ma.asanyarray(vt_cube.data) diff --git a/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py b/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py index a87db46f8d..e522e4f077 100644 --- a/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py +++ b/lib/iris/tests/unit/analysis/cartography/test_rotate_winds.py @@ -525,8 +525,8 @@ def _compare_lazy_rotate_winds(self, masked): # Create deep copy of the cubes with rechunked lazy data to check if # input data is modified, and if Dask metadata is preserved - u_lazy = u.copy(data=u.copy().lazy_data().rechunk([2,1])) - v_lazy = v.copy(data=v.copy().lazy_data().rechunk([1,2])) + u_lazy = u.copy(data=u.copy().lazy_data().rechunk([2, 1])) + v_lazy = v.copy(data=v.copy().lazy_data().rechunk([1, 2])) ut_ref, vt_ref = rotate_winds(u, v, coord_sys) self.assertFalse(ut_ref.has_lazy_data()) @@ -538,22 +538,22 @@ def _compare_lazy_rotate_winds(self, masked): ut, vt = rotate_winds(u_lazy, v, coord_sys) self.assertTrue(ut.has_lazy_data()) self.assertTrue(vt.has_lazy_data()) - self.assertTrue(ut.core_data().chunksize == (2,1)) + self.assertTrue(ut.core_data().chunksize == (2, 1)) self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) ut, vt = rotate_winds(u, v_lazy, coord_sys) self.assertTrue(ut.has_lazy_data()) self.assertTrue(vt.has_lazy_data()) - self.assertTrue(vt.core_data().chunksize == (1,2)) + self.assertTrue(vt.core_data().chunksize == (1, 2)) self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) ut, vt = rotate_winds(u_lazy, v_lazy, coord_sys) self.assertTrue(ut.has_lazy_data()) self.assertTrue(vt.has_lazy_data()) - self.assertTrue(ut.core_data().chunksize == (2,1)) - self.assertTrue(vt.core_data().chunksize == (1,2)) + self.assertTrue(ut.core_data().chunksize == (2, 1)) + self.assertTrue(vt.core_data().chunksize == (1, 2)) self.assertArrayAllClose(ut.data, ut_ref.data, rtol=1e-5) self.assertArrayAllClose(vt.data, vt_ref.data, rtol=1e-5) From 915eafe1dd048e2504629b791af39b5aa7b6ac3c Mon Sep 17 00:00:00 2001 From: Wolfgang Hayek Date: Mon, 19 Sep 2022 12:03:06 +1200 Subject: [PATCH 5/6] Add whats new entry --- docs/src/whatsnew/latest.rst | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index 73f7e60351..7aaa41eaf8 100644 --- a/docs/src/whatsnew/latest.rst +++ b/docs/src/whatsnew/latest.rst @@ -35,6 +35,9 @@ This document explains the changes made to Iris for this release non-existing paths, and added expansion functionality to :func:`~iris.io.save`. (:issue:`4772`, :pull:`4913`) +#. `@tinyendian`_ edited :func:`~iris.analysis.cartography.rotate_winds` to + enable lazy computation of rotated wind vector components (:issue:`4934`, + :pull:`4972`) 🐛 Bugs Fixed ============= @@ -126,4 +129,4 @@ This document explains the changes made to Iris for this release .. _NEP18: https://numpy.org/neps/nep-0018-array-function-protocol.html .. _pypa/setuptools#1684: https://github.com/pypa/setuptools/issues/1684 .. _SciTools/cartopy@fcb784d: https://github.com/SciTools/cartopy/commit/fcb784daa65d95ed9a74b02ca292801c02bc4108 -.. _SciTools/cartopy@8860a81: https://github.com/SciTools/cartopy/commit/8860a8186d4dc62478e74c83f3b2b3e8f791372e \ No newline at end of file +.. _SciTools/cartopy@8860a81: https://github.com/SciTools/cartopy/commit/8860a8186d4dc62478e74c83f3b2b3e8f791372e From 03a791259ab3dbbb15b8e7f62f01970c25291860 Mon Sep 17 00:00:00 2001 From: Wolfgang Hayek Date: Mon, 19 Sep 2022 12:50:20 +1200 Subject: [PATCH 6/6] Add missing author name entry --- docs/src/whatsnew/latest.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index 7aaa41eaf8..2d941f2a1c 100644 --- a/docs/src/whatsnew/latest.rst +++ b/docs/src/whatsnew/latest.rst @@ -118,6 +118,7 @@ This document explains the changes made to Iris for this release Whatsnew author names (@github name) in alphabetical order. Note that, core dev names are automatically included by the common_links.inc: +.. _@tinyendian: https://github.com/tinyendian .. _@TTV-Intrepid: https://github.com/TTV-Intrepid