From ba3ac6d8ca03a4b3d3e23fe42df40bfbaefdbc4a Mon Sep 17 00:00:00 2001 From: Wolfgang Hayek Date: Sat, 1 Apr 2023 01:31:12 +1300 Subject: [PATCH 1/3] Enable lazy computation of wind vector rotation (#4972) * Lazy rotate_winds and associated tests * Data ownership bug fix and extra tests * Fix input data test * Code formatting * Add whats new entry * Add missing author name entry --------- Co-authored-by: Bill Little --- docs/src/whatsnew/latest.rst | 8 ++- lib/iris/analysis/cartography.py | 45 ++++++++++----- .../analysis/cartography/test_rotate_winds.py | 57 +++++++++++++++++++ 3 files changed, 96 insertions(+), 14 deletions(-) diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index 904cb106ed..16069ddf6d 100644 --- a/docs/src/whatsnew/latest.rst +++ b/docs/src/whatsnew/latest.rst @@ -49,6 +49,10 @@ This document explains the changes made to Iris for this release #. `@lbdreyer`_ and `@trexfeathers`_ (reviewer) added :func:`iris.plot.hist` and :func:`iris.quickplot.hist`. (:pull:`5189`) +#. `@tinyendian`_ edited :func:`~iris.analysis.cartography.rotate_winds` to + enable lazy computation of rotated wind vector components (:issue:`4934`, + :pull:`4972`) + 🐛 Bugs Fixed ============= @@ -156,9 +160,11 @@ This document explains the changes made to Iris for this release .. _@ed-hawkins: https://github.com/ed-hawkins .. _@scottrobinson02: https://github.com/scottrobinson02 .. _@agriyakhetarpal: https://github.com/agriyakhetarpal +.. _@tinyendian: https://github.com/tinyendian + .. comment Whatsnew resources in alphabetical order: .. _#ShowYourStripes: https://showyourstripes.info/s/globe/ -.. _README.md: https://github.com/SciTools/iris#----- +.. _README.md: https://github.com/SciTools/iris#----- \ No newline at end of file diff --git a/lib/iris/analysis/cartography.py b/lib/iris/analysis/cartography.py index a8e90a63ad..fffb11f295 100644 --- a/lib/iris/analysis/cartography.py +++ b/lib/iris/analysis/cartography.py @@ -15,6 +15,7 @@ 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 @@ -1206,9 +1207,15 @@ 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() + # 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())) @@ -1236,8 +1243,16 @@ def rotate_winds(u_cube, v_cube, target_cs): apply_mask = mask.any() 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: + 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) # Project vectors with u, v components one horiz slice at a time and # insert into the resulting cubes. @@ -1250,16 +1265,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..e522e4f077 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,62 @@ 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 lazy data 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() + + # 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()) + # 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 + 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, 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, 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) + + # 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) + + 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 87869626e6ff955c7188eb049751b5f18254452b Mon Sep 17 00:00:00 2001 From: Bill Little Date: Fri, 31 Mar 2023 16:33:20 +0100 Subject: [PATCH 2/3] raise dask min pin (#5225) * raise dask min pin * replace workaround * add whatsnew entry --- docs/src/whatsnew/latest.rst | 7 ++++++- lib/iris/analysis/cartography.py | 8 ++------ requirements/py310.yml | 2 +- requirements/py38.yml | 2 +- requirements/py39.yml | 2 +- setup.cfg | 2 +- 6 files changed, 12 insertions(+), 11 deletions(-) diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index 16069ddf6d..51adfb9d8d 100644 --- a/docs/src/whatsnew/latest.rst +++ b/docs/src/whatsnew/latest.rst @@ -151,6 +151,10 @@ This document explains the changes made to Iris for this release #. `@trexfeathers`_ moved the benchmark runner conveniences from ``noxfile.py`` to a dedicated ``benchmarks/bm_runner.py``. (:pull:`5215`) +#. `@bjlittle`_ follow-up to :pull:`4972`, enforced ``dask>=2022.09.0`` minimum + pin for first use of `dask.array.ma.empty_like`_ and replaced `@tinyendian`_ + workaround. (:pull:`5225`) + .. comment Whatsnew author names (@github name) in alphabetical order. Note that, @@ -167,4 +171,5 @@ This document explains the changes made to Iris for this release Whatsnew resources in alphabetical order: .. _#ShowYourStripes: https://showyourstripes.info/s/globe/ -.. _README.md: https://github.com/SciTools/iris#----- \ No newline at end of file +.. _README.md: https://github.com/SciTools/iris#----- +.. _dask.array.ma.empty_like: https://docs.dask.org/en/stable/generated/dask.array.ma.empty_like.html diff --git a/lib/iris/analysis/cartography.py b/lib/iris/analysis/cartography.py index fffb11f295..5b11495d5a 100644 --- a/lib/iris/analysis/cartography.py +++ b/lib/iris/analysis/cartography.py @@ -1244,12 +1244,8 @@ 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.empty_like(ut_cube.core_data())) + vt_cube = vt_cube.copy(data=da.ma.empty_like(vt_cube.core_data())) else: ut_cube.data = ma.asanyarray(ut_cube.data) vt_cube.data = ma.asanyarray(vt_cube.data) diff --git a/requirements/py310.yml b/requirements/py310.yml index e3bada6596..49b8ae78ab 100644 --- a/requirements/py310.yml +++ b/requirements/py310.yml @@ -14,7 +14,7 @@ dependencies: - cartopy >=0.21 - cf-units >=3.1 - cftime >=1.5 - - dask-core >=2.26 + - dask-core >=2022.9.0 - matplotlib >=3.5 - netcdf4 - numpy >=1.19 diff --git a/requirements/py38.yml b/requirements/py38.yml index 9393060113..1e3c61f369 100644 --- a/requirements/py38.yml +++ b/requirements/py38.yml @@ -14,7 +14,7 @@ dependencies: - cartopy >=0.21 - cf-units >=3.1 - cftime >=1.5 - - dask-core >=2.26 + - dask-core >=2022.9.0 - matplotlib >=3.5 - netcdf4 - numpy >=1.19 diff --git a/requirements/py39.yml b/requirements/py39.yml index 349784ec46..6551eee495 100644 --- a/requirements/py39.yml +++ b/requirements/py39.yml @@ -14,7 +14,7 @@ dependencies: - cartopy >=0.21 - cf-units >=3.1 - cftime >=1.5 - - dask-core >=2.26 + - dask-core >=2022.9.0 - matplotlib >=3.5 - netcdf4 - numpy >=1.19 diff --git a/setup.cfg b/setup.cfg index ba9844f5d8..aac726776c 100644 --- a/setup.cfg +++ b/setup.cfg @@ -50,7 +50,7 @@ install_requires = cartopy>=0.21 cf-units>=3.1 cftime>=1.5.0 - dask[array]>=2.26 + dask[array]>=2022.9.0 matplotlib>=3.5 netcdf4 numpy>=1.19 From 6b25c3a21af383da173c304f09b38352e359322f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 3 Apr 2023 10:31:27 +0100 Subject: [PATCH 3/3] Bump scitools/workflows from 2023.03.2 to 2023.03.3 (#5227) Bumps [scitools/workflows](https://github.com/scitools/workflows) from 2023.03.2 to 2023.03.3. - [Release notes](https://github.com/scitools/workflows/releases) - [Commits](https://github.com/scitools/workflows/compare/2023.03.2...2023.03.3) --- updated-dependencies: - dependency-name: scitools/workflows dependency-type: direct:production ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/refresh-lockfiles.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/refresh-lockfiles.yml b/.github/workflows/refresh-lockfiles.yml index 66e316397a..cb875d6c6f 100644 --- a/.github/workflows/refresh-lockfiles.yml +++ b/.github/workflows/refresh-lockfiles.yml @@ -14,5 +14,5 @@ on: jobs: refresh_lockfiles: - uses: scitools/workflows/.github/workflows/refresh-lockfiles.yml@2023.03.2 + uses: scitools/workflows/.github/workflows/refresh-lockfiles.yml@2023.03.3 secrets: inherit