Repository: ESMValGroup/ESMValCore

Issue:
Weights for `climate_statistics` and `axis_statistics` are not lazy
**Describe the bug**
`iris.coords.DimCoords` are [never lazy](https://scitools-iris.readthedocs.io/en/stable/userguide/real_and_lazy_data.html#coordinates). Thus, the results of `core_points()` and `core_bounds()` will always be numpy arrays. Since we use these during the calculation of the time and axis weights (e.g., [here](https://github.com/ESMValGroup/ESMValCore/blob/bb7866e6a136f6e011bc6419c4e3ed9eec25721e/esmvalcore/preprocessor/_time.py#L411)), the resulting weights array will always be a numpy array.

When collapsing over a coordinate with weights, iris internally broadcasts this array to the cube shape. Thus, for very large cubes, this leads to very large numpy arrays. I found this issue while trying to convert the time mean for 20 years of ERA5 data (~18 GiB of data). With the current code, this just hangs indefinitely, with a simple fix it runs in less than 2 mins.

I'll open a PR with a fix in a moment.

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/ESMValGroup/ESMValCore/pull/2346?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 94.11%. Comparing base [(`bb7866e`)](https://app.codecov.io/gh/ESMValGroup/ESMValCore/commit/bb7866e6a136f6e011bc6419c4e3ed9eec25721e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup) to head [(`7623c49`)](https://app.codecov.io/gh/ESMValGroup/ESMValCore/pull/2346?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup).


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #2346   +/-   ##
=======================================
  Coverage   94.11%   94.11%           
=======================================
  Files         241      241           
  Lines       13376    13381    +5     
=======================================
+ Hits        12589    12594    +5     
  Misses        787      787           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ESMValGroup/ESMValCore/pull/2346?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ESMValGroup).


Patch:
diff --git a/esmvalcore/preprocessor/_time.py b/esmvalcore/preprocessor/_time.py
index ca4e95ce5..ee11e1146 100644
--- a/esmvalcore/preprocessor/_time.py
+++ b/esmvalcore/preprocessor/_time.py
@@ -392,8 +392,10 @@ def get_time_weights(cube: Cube) -> np.ndarray | da.core.Array:
 
     Returns
     -------
-    np.ndarray or da.core.Array
-        Array of time weights for averaging.
+    np.ndarray or da.Array
+        Array of time weights for averaging. Returns a
+        :class:`dask.array.Array` if the input cube has lazy data; a
+        :class:`numpy.ndarray` otherwise.
 
     """
     time = cube.coord('time')
@@ -408,7 +410,9 @@ def get_time_weights(cube: Cube) -> np.ndarray | da.core.Array:
         )
 
     # Extract 1D time weights (= lengths of time intervals)
-    time_weights = time.core_bounds()[:, 1] - time.core_bounds()[:, 0]
+    time_weights = time.lazy_bounds()[:, 1] - time.lazy_bounds()[:, 0]
+    if not cube.has_lazy_data():
+        time_weights = time_weights.compute()
     return time_weights
 
 
diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py
index 5f6460919..878bc03c6 100644
--- a/esmvalcore/preprocessor/_volume.py
+++ b/esmvalcore/preprocessor/_volume.py
@@ -352,8 +352,11 @@ def axis_statistics(
 
 def _add_axis_stats_weights_coord(cube, coord, coord_dims):
     """Add weights for axis_statistics to cube (in-place)."""
+    weights = np.abs(coord.lazy_bounds()[:, 1] - coord.lazy_bounds()[:, 0])
+    if not cube.has_lazy_data():
+        weights = weights.compute()
     weights_coord = AuxCoord(
-        np.abs(coord.core_bounds()[..., 1] - coord.core_bounds()[..., 0]),
+        weights,
         long_name='_axis_statistics_weights_',
         units=coord.units,
     )


Test patch:
diff --git a/tests/unit/preprocessor/_time/test_time.py b/tests/unit/preprocessor/_time/test_time.py
index a89c0ba76..a1ae135c2 100644
--- a/tests/unit/preprocessor/_time/test_time.py
+++ b/tests/unit/preprocessor/_time/test_time.py
@@ -1881,6 +1881,17 @@ def test_get_time_weights():
     """Test ``get_time_weights`` for complex cube."""
     cube = _make_cube()
     weights = get_time_weights(cube)
+    assert isinstance(weights, np.ndarray)
+    assert weights.shape == (2, )
+    np.testing.assert_allclose(weights, [15.0, 30.0])
+
+
+def test_get_time_weights_lazy():
+    """Test ``get_time_weights`` for complex cube with lazy data."""
+    cube = _make_cube()
+    cube.data = cube.lazy_data()
+    weights = get_time_weights(cube)
+    assert isinstance(weights, da.Array)
     assert weights.shape == (2, )
     np.testing.assert_allclose(weights, [15.0, 30.0])
 
diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py
index f0b730561..fccd065a3 100644
--- a/tests/unit/preprocessor/_volume/test_volume.py
+++ b/tests/unit/preprocessor/_volume/test_volume.py
@@ -10,6 +10,7 @@ from iris.exceptions import CoordinateMultiDimError
 
 import tests
 from esmvalcore.preprocessor._volume import (
+    _add_axis_stats_weights_coord,
     axis_statistics,
     calculate_volume,
     depth_integration,
@@ -106,6 +107,29 @@ class Test(tests.Test):
         iris.util.guess_coord_axis(self.grid_4d_2.coord('zcoord'))
         iris.util.guess_coord_axis(self.grid_4d_z.coord('zcoord'))
 
+    def test_add_axis_stats_weights_coord(self):
+        """Test _add_axis_stats_weights_coord."""
+        assert not self.grid_4d.coords('_axis_statistics_weights_')
+        coord = self.grid_4d.coord('zcoord')
+        coord_dims = self.grid_4d.coord_dims('zcoord')
+        _add_axis_stats_weights_coord(self.grid_4d, coord, coord_dims)
+        weights_coord = self.grid_4d.coord('_axis_statistics_weights_')
+        assert not weights_coord.has_lazy_points()
+        assert weights_coord.units == 'm'
+        np.testing.assert_allclose(weights_coord.points, [2.5, 22.5, 225.0])
+
+    def test_add_axis_stats_weights_coord_lazy(self):
+        """Test _add_axis_stats_weights_coord."""
+        self.grid_4d.data = self.grid_4d.lazy_data()
+        assert not self.grid_4d.coords('_axis_statistics_weights_')
+        coord = self.grid_4d.coord('zcoord')
+        coord_dims = self.grid_4d.coord_dims('zcoord')
+        _add_axis_stats_weights_coord(self.grid_4d, coord, coord_dims)
+        weights_coord = self.grid_4d.coord('_axis_statistics_weights_')
+        assert weights_coord.has_lazy_points()
+        assert weights_coord.units == 'm'
+        np.testing.assert_allclose(weights_coord.points, [2.5, 22.5, 225.0])
+
     def test_axis_statistics_mean(self):
         """Test axis statistics with operator mean."""
         data = np.ma.arange(1, 25).reshape(2, 3, 2, 2)