Implement lazy area weights (#2354)

Co-authored-by: Bouwe Andela <[email protected]>

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,14 +298,50 @@ 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",
                 warning.message, 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'):

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()