diff --git a/esmvalcore/preprocessor/_area.py b/esmvalcore/preprocessor/_area.py index 61f0402b39..39b01ea9a3 100644 --- a/esmvalcore/preprocessor/_area.py +++ b/esmvalcore/preprocessor/_area.py @@ -20,6 +20,7 @@ from iris.cube import Cube, CubeList from iris.exceptions import CoordinateMultiDimError, CoordinateNotFoundError +from esmvalcore.preprocessor._regrid import broadcast_to_shape from esmvalcore.preprocessor._shared import ( get_iris_aggregator, get_normalized_cube, @@ -297,7 +298,13 @@ def compute_area_weights(cube): category=UserWarning, module='iris.analysis.cartography', ) - weights = iris.analysis.cartography.area_weights(cube) + # TODO: replace the following line with + # weights = iris.analysis.cartography.area_weights( + # cube, compute=not cube.has_lazy_data() + # ) + # once https://github.com/SciTools/iris/pull/5658 is available + weights = _get_area_weights(cube) + for warning in caught_warnings: logger.debug( "%s while computing area weights of the following cube:\n%s", @@ -305,6 +312,36 @@ def compute_area_weights(cube): return weights +def _get_area_weights(cube: Cube) -> np.ndarray | da.Array: + """Get area weights. + + For non-lazy data, simply use the according iris function. For lazy data, + calculate area weights for a single lat-lon slice and broadcast it to the + correct shape. + + Note + ---- + This is a temporary workaround to get lazy area weights. Can be removed + once https://github.com/SciTools/iris/pull/5658 is available. + + """ + if not cube.has_lazy_data(): + return iris.analysis.cartography.area_weights(cube) + + lat_lon_dims = sorted( + tuple(set(cube.coord_dims('latitude') + cube.coord_dims('longitude'))) + ) + lat_lon_slice = next(cube.slices(['latitude', 'longitude'], ordered=False)) + weights_2d = iris.analysis.cartography.area_weights(lat_lon_slice) + weights = broadcast_to_shape( + da.array(weights_2d), + cube.shape, + lat_lon_dims, + chunks=cube.lazy_data().chunks, + ) + return weights + + def _try_adding_calculated_cell_area(cube: Cube) -> None: """Try to add calculated cell measure 'cell_area' to cube (in-place).""" if cube.cell_measures('cell_area'): diff --git a/tests/unit/preprocessor/_area/test_area.py b/tests/unit/preprocessor/_area/test_area.py index f7f507e496..7eed4938be 100644 --- a/tests/unit/preprocessor/_area/test_area.py +++ b/tests/unit/preprocessor/_area/test_area.py @@ -2,6 +2,7 @@ import unittest from pathlib import Path +import dask.array as da import fiona import iris import numpy as np @@ -17,6 +18,7 @@ import tests from esmvalcore.preprocessor._area import ( _crop_cube, + _get_area_weights, _get_requested_geometries, _update_shapefile_path, area_statistics, @@ -1445,5 +1447,22 @@ def test_time_dependent_volcello(): assert cube.shape == cube.cell_measure('ocean_volume').shape +@pytest.mark.parametrize('lazy', [True, False]) +def test_get_area_weights(lazy): + """Test _get_area_weights.""" + cube = _create_sample_full_cube() + if lazy: + cube.data = cube.lazy_data() + weights = _get_area_weights(cube) + if lazy: + assert isinstance(weights, da.Array) + assert weights.chunks == cube.lazy_data().chunks + else: + assert isinstance(weights, np.ndarray) + np.testing.assert_allclose( + weights, iris.analysis.cartography.area_weights(cube) + ) + + if __name__ == '__main__': unittest.main()