From 273859e990f98fa84038967fe19287e6f40dde84 Mon Sep 17 00:00:00 2001
From: Manuel Schlund <manuel.schlund@dlr.de>
Date: Wed, 28 Feb 2024 13:36:07 +0100
Subject: [PATCH 1/3] Workaround for lazy area weights

---
 esmvalcore/preprocessor/_area.py           | 39 +++++++++++++++++++++-
 tests/unit/preprocessor/_area/test_area.py | 18 ++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/esmvalcore/preprocessor/_area.py b/esmvalcore/preprocessor/_area.py
index 7d7bbe0c91..f7a973ef51 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().chunksize,
+    )
+    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 8835014ff2..cf26a36dee 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,
@@ -1399,5 +1401,21 @@ def test_meridional_statistics_invalid_norm_fail(make_testcube):
         meridional_statistics(make_testcube, 'sum', normalize='x')
 
 
+@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)
+    else:
+        assert isinstance(weights, np.ndarray)
+    np.testing.assert_allclose(
+        weights, iris.analysis.cartography.area_weights(cube)
+    )
+
+
 if __name__ == '__main__':
     unittest.main()

From 7d1f7d42999876602ef45ba59dc2be72fbb6b62a Mon Sep 17 00:00:00 2001
From: Manuel Schlund <32543114+schlunma@users.noreply.github.com>
Date: Tue, 16 Apr 2024 14:36:40 +0200
Subject: [PATCH 2/3] Update esmvalcore/preprocessor/_area.py

Co-authored-by: Bouwe Andela <b.andela@esciencecenter.nl>
---
 esmvalcore/preprocessor/_area.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/esmvalcore/preprocessor/_area.py b/esmvalcore/preprocessor/_area.py
index 5e8dcd29f6..39b01ea9a3 100644
--- a/esmvalcore/preprocessor/_area.py
+++ b/esmvalcore/preprocessor/_area.py
@@ -337,7 +337,7 @@ def _get_area_weights(cube: Cube) -> np.ndarray | da.Array:
         da.array(weights_2d),
         cube.shape,
         lat_lon_dims,
-        chunks=cube.lazy_data().chunksize,
+        chunks=cube.lazy_data().chunks,
     )
     return weights
 

From e645a6244043d286957804e72ef134edcfdf1ac3 Mon Sep 17 00:00:00 2001
From: Manuel Schlund <manuel.schlund@dlr.de>
Date: Tue, 16 Apr 2024 14:44:19 +0200
Subject: [PATCH 3/3] Also test chunk sizes

---
 tests/unit/preprocessor/_area/test_area.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit/preprocessor/_area/test_area.py b/tests/unit/preprocessor/_area/test_area.py
index 6dd14ab21a..7eed4938be 100644
--- a/tests/unit/preprocessor/_area/test_area.py
+++ b/tests/unit/preprocessor/_area/test_area.py
@@ -1456,6 +1456,7 @@ def test_get_area_weights(lazy):
     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(