From 0f63288682717a7d8ac7f38a5718a1f7df81a986 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Wed, 28 Jun 2023 16:40:20 +0200
Subject: [PATCH 01/11] More lazy extract_levels

---
 esmvalcore/preprocessor/_regrid.py | 48 ++++++++++++++++++------------
 1 file changed, 29 insertions(+), 19 deletions(-)

diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py
index daa6b6acfd..f3c29715cf 100644
--- a/esmvalcore/preprocessor/_regrid.py
+++ b/esmvalcore/preprocessor/_regrid.py
@@ -9,7 +9,7 @@
 from copy import deepcopy
 from decimal import Decimal
 from pathlib import Path
-from typing import Dict
+from typing import Dict, Optional
 
 import dask.array as da
 import iris
@@ -859,20 +859,27 @@ def _vertical_interpolate(cube, src_levels, levels, interpolation,
 
     # Broadcast the 1d source cube vertical coordinate to fully
     # describe the spatial extent that will be interpolated.
-    src_levels_broadcast = broadcast_to_shape(src_levels.points, cube.shape,
-                                              cube.coord_dims(src_levels))
+    src_levels_broadcast = broadcast_to_shape(
+        src_levels.core_points(),
+        cube.shape,
+        cube.coord_dims(src_levels),
+    )
 
     # force mask onto data as nan's
     npx = get_array_module(cube.core_data())
     data = npx.ma.filled(cube.core_data(), np.nan)
 
+    if isinstance(src_levels_broadcast, da.Array):
+        levels = da.asarray(levels)
     # Now perform the actual vertical interpolation.
-    new_data = stratify.interpolate(levels,
-                                    src_levels_broadcast,
-                                    data,
-                                    axis=z_axis,
-                                    interpolation=interpolation,
-                                    extrapolation=extrapolation)
+    new_data = stratify.interpolate(
+        levels,
+        src_levels_broadcast,
+        data,
+        axis=z_axis,
+        interpolation=interpolation,
+        extrapolation=extrapolation,
+    )
 
     # Calculate the mask based on the any NaN values in the interpolated data.
     new_data = npx.ma.masked_where(npx.isnan(new_data), new_data)
@@ -942,19 +949,21 @@ def parse_vertical_scheme(scheme):
     return scheme, extrap_scheme
 
 
-def extract_levels(cube,
-                   levels,
-                   scheme,
-                   coordinate=None,
-                   rtol=1e-7,
-                   atol=None):
+def extract_levels(
+    cube: iris.cube.Cube,
+    levels: np.typing.ArrayLike,
+    scheme: str,
+    coordinate: Optional[str] = None,
+    rtol: Optional[float] = 1e-7,
+    atol: Optional[float] = None,
+):
     """Perform vertical interpolation.
 
     Parameters
     ----------
     cube : iris.cube.Cube
         The source cube to be vertically interpolated.
-    levels : ArrayLike
+    levels : np.typing.ArrayLike
         One or more target levels for the vertical interpolation. Assumed
         to be in the same S.I. units of the source cube vertical dimension
         coordinate. If the requested levels are sufficiently close to the
@@ -997,7 +1006,8 @@ def extract_levels(cube,
     interpolation, extrapolation = parse_vertical_scheme(scheme)
 
     # Ensure we have a non-scalar array of levels.
-    levels = np.array(levels, ndmin=1)
+    if not isinstance(levels, da.Array):
+        levels = np.array(levels, ndmin=1)
 
     # Get the source cube vertical coordinate, if available.
     if coordinate:
@@ -1016,10 +1026,10 @@ def extract_levels(cube,
         src_levels = cube.coord(axis='z', dim_coords=True)
 
     if (src_levels.shape == levels.shape and np.allclose(
-            src_levels.points,
+            src_levels.core_points(),
             levels,
             rtol=rtol,
-            atol=1e-7 * np.mean(src_levels.points) if atol is None else atol,
+            atol=1e-7 * np.mean(src_levels.core_points()) if atol is None else atol,
     )):
         # Only perform vertical extraction/interpolation if the source
         # and target levels are not "similar" enough.

From 8d06564851ebdd543359e08e3bef697ca7eb19f2 Mon Sep 17 00:00:00 2001
From: Valeriu Predoi <valeriu.predoi@gmail.com>
Date: Thu, 29 Jun 2023 12:39:32 +0100
Subject: [PATCH 02/11] line too long flake

---
 esmvalcore/preprocessor/_regrid.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py
index f3c29715cf..2076f4917c 100644
--- a/esmvalcore/preprocessor/_regrid.py
+++ b/esmvalcore/preprocessor/_regrid.py
@@ -1029,7 +1029,8 @@ def extract_levels(
             src_levels.core_points(),
             levels,
             rtol=rtol,
-            atol=1e-7 * np.mean(src_levels.core_points()) if atol is None else atol,
+            atol=1e-7 * np.mean(src_levels.core_points())
+            if atol is None else atol,
     )):
         # Only perform vertical extraction/interpolation if the source
         # and target levels are not "similar" enough.

From adad828dc603a0b5571aa79b91263ff3fde12ad5 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Fri, 30 Jun 2023 20:10:24 +0200
Subject: [PATCH 03/11] Some improvements

---
 esmvalcore/preprocessor/_regrid.py | 92 +++++++++++++++++++++++-------
 1 file changed, 72 insertions(+), 20 deletions(-)

diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py
index 2076f4917c..d221dc6182 100644
--- a/esmvalcore/preprocessor/_regrid.py
+++ b/esmvalcore/preprocessor/_regrid.py
@@ -8,8 +8,9 @@
 import ssl
 from copy import deepcopy
 from decimal import Decimal
+from functools import partial
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 
 import dask.array as da
 import iris
@@ -17,7 +18,6 @@
 import stratify
 from geopy.geocoders import Nominatim
 from iris.analysis import AreaWeighted, Linear, Nearest, UnstructuredNearest
-from iris.util import broadcast_to_shape
 
 from ..cmor._fixes.shared import add_altitude_from_plev, add_plev_from_altitude
 from ..cmor.table import CMOR_TABLES
@@ -851,27 +851,79 @@ def _create_cube(src_cube, data, src_levels, levels):
     return result
 
 
+def is_lazy_masked_data(array):
+    """Similar to `iris._lazy_data.is_lazy_masked_data`."""
+    return isinstance(array, da.Array) and isinstance(
+        da.utils.meta_from_array(array), np.ma.MaskedArray)
+
+
+def broadcast_to_shape(array, shape, chunks, dim_map):
+    """Copy of `iris.util.broadcast_to_shape` that allows specifying chunks."""
+    if isinstance(array, da.Array):
+        chunks = list(chunks)
+        for src_idx, tgt_idx in enumerate(dim_map):
+            # Only use the specified chunks along new dimensions.
+            chunks[tgt_idx] = array.chunks[src_idx]
+        broadcast_to = partial(da.broadcast_to, chunks=chunks)
+    else:
+        broadcast_to = np.broadcast_to
+
+    n_orig_dims = len(array.shape)
+    n_new_dims = len(shape) - n_orig_dims
+    array = array.reshape(array.shape + (1,) * n_new_dims)
+
+    # Get dims in required order.
+    array = np.moveaxis(array, range(n_orig_dims), dim_map)
+    new_array = broadcast_to(array, shape)
+
+    if np.ma.isMA(array):
+        # broadcast_to strips masks so we need to handle them explicitly.
+        mask = np.ma.getmask(array)
+        if mask is np.ma.nomask:
+            new_mask = np.ma.nomask
+        else:
+            new_mask = np.broadcast_to(mask, shape)
+        new_array = np.ma.array(new_array, mask=new_mask)
+
+    elif is_lazy_masked_data(array):
+        # broadcast_to strips masks so we need to handle them explicitly.
+        mask = da.ma.getmaskarray(array)
+        new_mask = broadcast_to(mask, shape)
+        new_array = da.ma.masked_array(new_array, new_mask)
+
+    return new_array
+
+
 def _vertical_interpolate(cube, src_levels, levels, interpolation,
                           extrapolation):
     """Perform vertical interpolation."""
     # Determine the source levels and axis for vertical interpolation.
     z_axis, = cube.coord_dims(cube.coord(axis='z', dim_coords=True))
 
-    # Broadcast the 1d source cube vertical coordinate to fully
-    # describe the spatial extent that will be interpolated.
+    if cube.has_lazy_data():
+        # Make source levels lazy if cube has lazy data.
+        src_points = src_levels.lazy_points()
+    else:
+        src_points = src_levels.core_points()
+
+    # Make the target levels lazy if the input data is lazy.
+    if cube.has_lazy_data() and isinstance(src_points, da.Array):
+        levels = da.asarray(levels)
+
+    # Broadcast the source cube vertical coordinate to fully describe the
+    # spatial extent that will be interpolated.
     src_levels_broadcast = broadcast_to_shape(
-        src_levels.core_points(),
-        cube.shape,
-        cube.coord_dims(src_levels),
+        src_points,
+        shape=cube.shape,
+        chunks=cube.lazy_data().chunks if cube.has_lazy_data() else None,
+        dim_map=cube.coord_dims(src_levels),
     )
 
     # force mask onto data as nan's
     npx = get_array_module(cube.core_data())
     data = npx.ma.filled(cube.core_data(), np.nan)
 
-    if isinstance(src_levels_broadcast, da.Array):
-        levels = da.asarray(levels)
-    # Now perform the actual vertical interpolation.
+    # Perform vertical interpolation.
     new_data = stratify.interpolate(
         levels,
         src_levels_broadcast,
@@ -951,42 +1003,42 @@ def parse_vertical_scheme(scheme):
 
 def extract_levels(
     cube: iris.cube.Cube,
-    levels: np.typing.ArrayLike,
+    levels: Union[np.typing.ArrayLike, da.Array],
     scheme: str,
     coordinate: Optional[str] = None,
-    rtol: Optional[float] = 1e-7,
+    rtol: float = 1e-7,
     atol: Optional[float] = None,
 ):
     """Perform vertical interpolation.
 
     Parameters
     ----------
-    cube : iris.cube.Cube
+    cube:
         The source cube to be vertically interpolated.
-    levels : np.typing.ArrayLike
+    levels:
         One or more target levels for the vertical interpolation. Assumed
         to be in the same S.I. units of the source cube vertical dimension
         coordinate. If the requested levels are sufficiently close to the
         levels of the cube, cube slicing will take place instead of
         interpolation.
-    scheme : str
+    scheme:
         The vertical interpolation scheme to use. Choose from
         'linear',
         'nearest',
         'linear_extrapolate',
         'nearest_extrapolate'.
-    coordinate :  optional str
+    coordinate:
         The coordinate to interpolate. If specified, pressure levels
         (if present) can be converted to height levels and vice versa using
         the US standard atmosphere. E.g. 'coordinate = altitude' will convert
         existing pressure levels (air_pressure) to height levels (altitude);
         'coordinate = air_pressure' will convert existing height levels
         (altitude) to pressure levels (air_pressure).
-    rtol : float
+    rtol:
         Relative tolerance for comparing the levels in `cube` to the requested
         levels. If the levels are sufficiently close, the requested levels
         will be assigned to the cube and no interpolation will take place.
-    atol : float
+    atol:
         Absolute tolerance for comparing the levels in `cube` to the requested
         levels. If the levels are sufficiently close, the requested levels
         will be assigned to the cube and no interpolation will take place.
@@ -1029,8 +1081,8 @@ def extract_levels(
             src_levels.core_points(),
             levels,
             rtol=rtol,
-            atol=1e-7 * np.mean(src_levels.core_points())
-            if atol is None else atol,
+            atol=1e-7 *
+            np.mean(src_levels.core_points()) if atol is None else atol,
     )):
         # Only perform vertical extraction/interpolation if the source
         # and target levels are not "similar" enough.

From 3f79e1f8f195a8649ec2babd5c19611d2a11bb96 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Mon, 15 Jan 2024 15:45:26 +0100
Subject: [PATCH 04/11] Ignore type error

---
 esmvalcore/preprocessor/_regrid.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py
index 76ce1e5fce..0efe28b01a 100644
--- a/esmvalcore/preprocessor/_regrid.py
+++ b/esmvalcore/preprocessor/_regrid.py
@@ -1070,7 +1070,9 @@ def extract_levels(
             set(levels).issubset(set(src_levels.points)):
         # If all target levels exist in the source cube, simply extract them.
         name = src_levels.name()
-        coord_values = {name: lambda cell: cell.point in set(levels)}
+        coord_values = {
+            name: lambda cell: cell.point in set(levels)  # type: ignore
+        }
         constraint = iris.Constraint(coord_values=coord_values)
         result = cube.extract(constraint)
         # Ensure the constraint did not fail.

From 9bd168597bd886f28812423f6a398be71f484cf1 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Mon, 15 Jan 2024 16:26:40 +0100
Subject: [PATCH 05/11] Update Codecov orb to see if that fixes failure to
 upload coverage

---
 .circleci/config.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 8e5f38b84f..ec9fbf38ba 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -3,7 +3,7 @@ version: 2.1
 
 orbs:
   coverage-reporter: codacy/coverage-reporter@13.13.7
-  codecov: codecov/codecov@3.2.5
+  codecov: codecov/codecov@3.3.0
 
 commands:
   check_changes:

From ec0a4abe209032b312e0edce27aee40da604da09 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Tue, 16 Jan 2024 18:09:16 +0100
Subject: [PATCH 06/11] Copy tests and implementation from
 https://github.com/SciTools/iris/pull/5620

---
 esmvalcore/preprocessor/_regrid.py            |  23 ++--
 .../_regrid/test_broadcast_to_shape.py        | 112 ++++++++++++++++++
 2 files changed, 125 insertions(+), 10 deletions(-)
 create mode 100644 tests/unit/preprocessor/_regrid/test_broadcast_to_shape.py

diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py
index 0efe28b01a..f157653cf7 100644
--- a/esmvalcore/preprocessor/_regrid.py
+++ b/esmvalcore/preprocessor/_regrid.py
@@ -834,16 +834,19 @@ def is_lazy_masked_data(array):
         da.utils.meta_from_array(array), np.ma.MaskedArray)
 
 
-def broadcast_to_shape(array, shape, chunks, dim_map):
+def broadcast_to_shape(array, shape, dim_map, chunks=None):
     """Copy of `iris.util.broadcast_to_shape` that allows specifying chunks."""
     if isinstance(array, da.Array):
-        chunks = list(chunks)
-        for src_idx, tgt_idx in enumerate(dim_map):
-            # Only use the specified chunks along new dimensions.
-            chunks[tgt_idx] = array.chunks[src_idx]
-        broadcast_to = partial(da.broadcast_to, chunks=chunks)
+        if chunks is not None:
+            chunks = list(chunks)
+            for src_idx, tgt_idx in enumerate(dim_map):
+                # Only use the specified chunks along new dimensions or on
+                # dimensions that have size 1 in the source array.
+                if array.shape[src_idx] != 1:
+                    chunks[tgt_idx] = array.chunks[src_idx]
+        broadcast = partial(da.broadcast_to, shape=shape, chunks=chunks)
     else:
-        broadcast_to = np.broadcast_to
+        broadcast = partial(np.broadcast_to, shape=shape)
 
     n_orig_dims = len(array.shape)
     n_new_dims = len(shape) - n_orig_dims
@@ -851,7 +854,7 @@ def broadcast_to_shape(array, shape, chunks, dim_map):
 
     # Get dims in required order.
     array = np.moveaxis(array, range(n_orig_dims), dim_map)
-    new_array = broadcast_to(array, shape)
+    new_array = broadcast(array)
 
     if np.ma.isMA(array):
         # broadcast_to strips masks so we need to handle them explicitly.
@@ -859,13 +862,13 @@ def broadcast_to_shape(array, shape, chunks, dim_map):
         if mask is np.ma.nomask:
             new_mask = np.ma.nomask
         else:
-            new_mask = np.broadcast_to(mask, shape)
+            new_mask = broadcast(mask)
         new_array = np.ma.array(new_array, mask=new_mask)
 
     elif is_lazy_masked_data(array):
         # broadcast_to strips masks so we need to handle them explicitly.
         mask = da.ma.getmaskarray(array)
-        new_mask = broadcast_to(mask, shape)
+        new_mask = broadcast(mask)
         new_array = da.ma.masked_array(new_array, new_mask)
 
     return new_array
diff --git a/tests/unit/preprocessor/_regrid/test_broadcast_to_shape.py b/tests/unit/preprocessor/_regrid/test_broadcast_to_shape.py
new file mode 100644
index 0000000000..a9b8f586bf
--- /dev/null
+++ b/tests/unit/preprocessor/_regrid/test_broadcast_to_shape.py
@@ -0,0 +1,112 @@
+# Copyright Iris contributors
+#
+# This file is part of Iris and is released under the BSD license.
+# See LICENSE in the root of the repository for full licensing details.
+"""Test function :func:`iris.util.broadcast_to_shape`."""
+
+from unittest import mock
+
+import dask
+import dask.array as da
+import numpy as np
+import numpy.ma as ma
+
+from esmvalcore.preprocessor._regrid import broadcast_to_shape
+from tests import assert_array_equal
+
+
+def test_same_shape():
+    # broadcast to current shape should result in no change
+    a = np.random.random([2, 3])
+    b = broadcast_to_shape(a, a.shape, (0, 1))
+    assert_array_equal(b, a)
+
+
+def test_added_dimensions():
+    # adding two dimensions, on at the front and one in the middle of
+    # the existing dimensions
+    a = np.random.random([2, 3])
+    b = broadcast_to_shape(a, (5, 2, 4, 3), (1, 3))
+    for i in range(5):
+        for j in range(4):
+            assert_array_equal(b[i, :, j, :], a)
+
+
+def test_added_dimensions_transpose():
+    # adding dimensions and having the dimensions of the input
+    # transposed
+    a = np.random.random([2, 3])
+    b = broadcast_to_shape(a, (5, 3, 4, 2), (3, 1))
+    for i in range(5):
+        for j in range(4):
+            assert_array_equal(b[i, :, j, :].T, a)
+
+
+@mock.patch.object(dask.base, "compute", wraps=dask.base.compute)
+def test_lazy_added_dimensions_transpose(mocked_compute):
+    # adding dimensions and having the dimensions of the input
+    # transposed
+    a = da.random.random([2, 3])
+    b = broadcast_to_shape(a, (5, 3, 4, 2), (3, 1))
+    mocked_compute.assert_not_called()
+    for i in range(5):
+        for j in range(4):
+            assert_array_equal(b[i, :, j, :].T.compute(), a.compute())
+
+
+def test_masked():
+    # masked arrays are also accepted
+    a = np.random.random([2, 3])
+    m = ma.array(a, mask=[[0, 1, 0], [0, 1, 1]])
+    b = broadcast_to_shape(m, (5, 3, 4, 2), (3, 1))
+    for i in range(5):
+        for j in range(4):
+            assert_array_equal(b[i, :, j, :].T, m)
+
+
+@mock.patch.object(dask.base, "compute", wraps=dask.base.compute)
+def test_lazy_masked(mocked_compute):
+    # masked arrays are also accepted
+    a = np.random.random([2, 3])
+    m = da.ma.masked_array(a, mask=[[0, 1, 0], [0, 1, 1]])
+    b = broadcast_to_shape(m, (5, 3, 4, 2), (3, 1))
+    mocked_compute.assert_not_called()
+    for i in range(5):
+        for j in range(4):
+            assert_array_equal(b[i, :, j, :].compute().T, m.compute())
+
+
+@mock.patch.object(dask.base, "compute", wraps=dask.base.compute)
+def test_lazy_chunks(mocked_compute):
+    # chunks can be specified along with the target shape and are only used
+    # along new dimensions or on dimensions that have size 1 in the source
+    # array.
+    m = da.ma.masked_array(
+        data=[[1, 2, 3, 4, 5]],
+        mask=[[0, 1, 0, 0, 0]],
+    ).rechunk((1, 2))
+    b = broadcast_to_shape(
+        m,
+        dim_map=(1, 2),
+        shape=(3, 4, 5),
+        chunks=(
+            1,  # used because target is new dim
+            2,  # used because input size 1
+            3,  # not used because broadcast does not rechunk
+        ),
+    )
+    mocked_compute.assert_not_called()
+    for i in range(3):
+        for j in range(4):
+            assert_array_equal(b[i, j, :].compute(), m[0].compute())
+    assert b.chunks == ((1, 1, 1), (2, 2), (2, 2, 1))
+
+
+def test_masked_degenerate():
+    # masked arrays can have degenerate masks too
+    a = np.random.random([2, 3])
+    m = ma.array(a)
+    b = broadcast_to_shape(m, (5, 3, 4, 2), (3, 1))
+    for i in range(5):
+        for j in range(4):
+            assert_array_equal(b[i, :, j, :].T, m)

From 867e0d7a35646f161213f2296711bc3856dc0df8 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Wed, 24 Jan 2024 15:14:12 +0100
Subject: [PATCH 07/11] Improve chunking for derived coordinates and add lazy
 altitude-pressure level conversion

---
 esmvalcore/cmor/_fixes/shared.py   | 11 ++++-
 esmvalcore/preprocessor/_regrid.py | 72 ++++++++++++++++++++++--------
 2 files changed, 64 insertions(+), 19 deletions(-)

diff --git a/esmvalcore/cmor/_fixes/shared.py b/esmvalcore/cmor/_fixes/shared.py
index e9ae61499f..68a735a42c 100644
--- a/esmvalcore/cmor/_fixes/shared.py
+++ b/esmvalcore/cmor/_fixes/shared.py
@@ -74,7 +74,16 @@ def _map_on_filled(function, array):
     array = num_module.ma.filled(array, fill_value)
 
     # Apply function and return masked array
-    array = function(array)
+    if isinstance(array, da.Array):
+        array = da.map_blocks(
+            function,
+            array,
+            dtype=array.dtype,
+            enforce_ndim=True,
+            meta=da.utils.meta_from_array(array),
+        )
+    else:
+        array = function(array)
     return num_module.ma.masked_array(array, mask=mask)
 
 
diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py
index f157653cf7..4dd3fdd07d 100644
--- a/esmvalcore/preprocessor/_regrid.py
+++ b/esmvalcore/preprocessor/_regrid.py
@@ -886,10 +886,6 @@ def _vertical_interpolate(cube, src_levels, levels, interpolation,
     else:
         src_points = src_levels.core_points()
 
-    # Make the target levels lazy if the input data is lazy.
-    if cube.has_lazy_data() and isinstance(src_points, da.Array):
-        levels = da.asarray(levels)
-
     # Broadcast the source cube vertical coordinate to fully describe the
     # spatial extent that will be interpolated.
     src_levels_broadcast = broadcast_to_shape(
@@ -899,6 +895,10 @@ def _vertical_interpolate(cube, src_levels, levels, interpolation,
         dim_map=cube.coord_dims(src_levels),
     )
 
+    # Make the target levels lazy if the input data is lazy.
+    if cube.has_lazy_data() and isinstance(src_points, da.Array):
+        levels = da.asarray(levels)
+
     # force mask onto data as nan's
     npx = get_array_module(cube.core_data())
     data = npx.ma.filled(cube.core_data(), np.nan)
@@ -981,6 +981,36 @@ def parse_vertical_scheme(scheme):
     return scheme, extrap_scheme
 
 
+def _rechunk_aux_factory_dependencies(
+    cube: iris.cube.Cube,
+    coord_name: str,
+) -> iris.cube.Cube:
+    """Rechunk coordinate aux factory dependencies.
+
+    This ensures that the resulting coordinate has reasonably sized
+    chunks that are aligned with the cube data for optimal computational
+    performance.
+    """
+    # Workaround for https://github.com/SciTools/iris/issues/5457
+    try:
+        factory = cube.aux_factory(coord_name)
+    except iris.exceptions.CoordinateNotFoundError:
+        return cube
+
+    cube = cube.copy()
+    cube_chunks = cube.lazy_data().chunks
+    for coord in factory.dependencies.values():
+        coord_dims = cube.coord_dims(coord)
+        if coord_dims is not None:
+            coord = coord.copy()
+            chunks = tuple(cube_chunks[i] for i in coord_dims)
+            coord.points = coord.lazy_points().rechunk(chunks)
+            if coord.has_bounds():
+                coord.bounds = coord.lazy_bounds().rechunk(chunks + (None, ))
+            cube.replace_coord(coord)
+    return cube
+
+
 def extract_levels(
     cube: iris.cube.Cube,
     levels: Union[np.typing.ArrayLike, da.Array],
@@ -1041,21 +1071,27 @@ def extract_levels(
     if not isinstance(levels, da.Array):
         levels = np.array(levels, ndmin=1)
 
-    # Get the source cube vertical coordinate, if available.
-    if coordinate:
-        coord_names = [coord.name() for coord in cube.coords()]
-        if coordinate not in coord_names:
-            # Try to calculate air_pressure from altitude coordinate or
-            # vice versa using US standard atmosphere for conversion.
-            if coordinate == 'air_pressure' and 'altitude' in coord_names:
-                # Calculate pressure level coordinate from altitude.
-                add_plev_from_altitude(cube)
-            if coordinate == 'altitude' and 'air_pressure' in coord_names:
-                # Calculate altitude coordinate from pressure levels.
-                add_altitude_from_plev(cube)
-        src_levels = cube.coord(coordinate)
+    # Try to determine the name of the vertical coordinate automatically
+    if coordinate is None:
+        coordinate = cube.coord(axis='z', dim_coords=True).name()
+
+    # Add extra coordinates
+    coord_names = [coord.name() for coord in cube.coords()]
+    if coordinate in coord_names:
+        cube = _rechunk_aux_factory_dependencies(cube, coordinate)
     else:
-        src_levels = cube.coord(axis='z', dim_coords=True)
+        # Try to calculate air_pressure from altitude coordinate or
+        # vice versa using US standard atmosphere for conversion.
+        if coordinate == 'air_pressure' and 'altitude' in coord_names:
+            # Calculate pressure level coordinate from altitude.
+            cube = _rechunk_aux_factory_dependencies(cube, 'altitude')
+            add_plev_from_altitude(cube)
+        if coordinate == 'altitude' and 'air_pressure' in coord_names:
+            # Calculate altitude coordinate from pressure levels.
+            cube = _rechunk_aux_factory_dependencies(cube, 'air_pressure')
+            add_altitude_from_plev(cube)
+
+    src_levels = cube.coord(coordinate)
 
     if (src_levels.shape == levels.shape and np.allclose(
             src_levels.core_points(),

From db3acbbe4526fe95497bbb9714062a7417bdf75b Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Thu, 25 Jan 2024 09:46:14 +0100
Subject: [PATCH 08/11] Add test

---
 .../_regrid/test_extract_levels.py            | 54 +++++++++++++++++--
 1 file changed, 51 insertions(+), 3 deletions(-)

diff --git a/tests/unit/preprocessor/_regrid/test_extract_levels.py b/tests/unit/preprocessor/_regrid/test_extract_levels.py
index f53fbac708..b372f9c885 100644
--- a/tests/unit/preprocessor/_regrid/test_extract_levels.py
+++ b/tests/unit/preprocessor/_regrid/test_extract_levels.py
@@ -1,9 +1,10 @@
 """Unit tests for :func:`esmvalcore.preprocessor.regrid.extract_levels`."""
-import unittest
 from unittest import mock
 
+import dask.array as da
 import iris
 import numpy as np
+from iris.aux_factory import HybridPressureFactory
 from numpy import ma
 
 import tests
@@ -11,6 +12,7 @@
     _MDI,
     VERTICAL_SCHEMES,
     _preserve_fx_vars,
+    _rechunk_aux_factory_dependencies,
     extract_levels,
     parse_vertical_scheme,
 )
@@ -317,5 +319,51 @@ def test_interpolation__masked(self):
         self.assertEqual(kwargs, dict())
 
 
-if __name__ == '__main__':
-    unittest.main()
+def test_rechunk_aux_factory_dependencies():
+
+    delta = iris.coords.AuxCoord(
+        np.array([0.0, 1.0, 2.0], dtype=np.float64),
+        long_name="level_pressure",
+        units="Pa",
+    )
+    sigma = iris.coords.AuxCoord(
+        np.array([1.0, 0.9, 0.8], dtype=np.float64),
+        long_name="sigma",
+        units="1",
+    )
+    surface_air_pressure = iris.coords.AuxCoord(
+        np.arange(4).astype(np.float64).reshape(2, 2),
+        long_name="surface_air_pressure",
+        units="Pa",
+    )
+    factory = HybridPressureFactory(
+        delta=delta,
+        sigma=sigma,
+        surface_air_pressure=surface_air_pressure,
+    )
+
+    cube = iris.cube.Cube(
+        da.asarray(
+            np.arange(3 * 2 * 2).astype(np.float32).reshape(3, 2, 2),
+            chunks=(1, 2, 2),
+        ), )
+    cube.add_aux_coord(delta, 0)
+    cube.add_aux_coord(sigma, 0)
+    cube.add_aux_coord(surface_air_pressure, [1, 2])
+    cube.add_aux_factory(factory)
+
+    result = _rechunk_aux_factory_dependencies(cube, 'air_pressure')
+
+    # Check that the 'air_pressure' coordinate of the resulting cube has been
+    # rechunked:
+    assert (
+        (1, 1, 1),
+        (2, ),
+        (2, ),
+    ) == result.coord('air_pressure').core_points().chunks
+    # Check that the original cube has not been modified:
+    assert (
+        (3, ),
+        (2, ),
+        (2, ),
+    ) == cube.coord('air_pressure').core_points().chunks

From ee1eddf7f0058113cdcb667e54a3b6beda63a611 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Thu, 25 Jan 2024 10:24:27 +0100
Subject: [PATCH 09/11] Add more tests

---
 .../_regrid/test_extract_levels.py            | 22 ++++++++++++++-----
 .../_regrid/test_extract_levels.py            |  4 +++-
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/tests/integration/preprocessor/_regrid/test_extract_levels.py b/tests/integration/preprocessor/_regrid/test_extract_levels.py
index b0869244a6..287f019146 100644
--- a/tests/integration/preprocessor/_regrid/test_extract_levels.py
+++ b/tests/integration/preprocessor/_regrid/test_extract_levels.py
@@ -1,11 +1,9 @@
-"""
-Integration tests for the :func:`esmvalcore.preprocessor.regrid.extract_levels`
-function.
-
-"""
+"""Integration tests for the
+:func:`esmvalcore.preprocessor.regrid.extract_levels` function."""
 
 import unittest
 
+import dask.array as da
 import iris
 import numpy as np
 
@@ -76,8 +74,20 @@ def test_interpolation__linear_lazy(self):
         levels = [0.5, 1.5]
         scheme = 'linear'
         cube = self.cube.copy(self.cube.lazy_data())
-        result = extract_levels(cube, levels, scheme)
+        coord_name = 'multidimensional_vertical_coord'
+        coord_points = (
+            cube.coord('air_pressure').core_points().reshape(3, 1, 1) *
+            np.ones((3, 2, 2)))
+        cube.add_aux_coord(
+            iris.coords.AuxCoord(
+                da.asarray(coord_points),
+                long_name=coord_name,
+            ),
+            [1, 2, 3],
+        )
+        result = extract_levels(cube, levels, scheme, coordinate=coord_name)
         self.assertTrue(result.has_lazy_data())
+        self.assertTrue(cube.coord(coord_name).has_lazy_points())
         expected = np.ma.array([
             [
                 [[2., 3.], [4., 5.]],
diff --git a/tests/unit/preprocessor/_regrid/test_extract_levels.py b/tests/unit/preprocessor/_regrid/test_extract_levels.py
index b372f9c885..4825350f17 100644
--- a/tests/unit/preprocessor/_regrid/test_extract_levels.py
+++ b/tests/unit/preprocessor/_regrid/test_extract_levels.py
@@ -322,7 +322,9 @@ def test_interpolation__masked(self):
 def test_rechunk_aux_factory_dependencies():
 
     delta = iris.coords.AuxCoord(
-        np.array([0.0, 1.0, 2.0], dtype=np.float64),
+        points=np.array([0.0, 1.0, 2.0], dtype=np.float64),
+        bounds=np.array([[-0.5, 0.5], [0.5, 1.5], [1.5, 2.5]],
+                        dtype=np.float64),
         long_name="level_pressure",
         units="Pa",
     )

From 1f476daf9f946c4ad8a56c50414a5bbbb60ce637 Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Tue, 30 Jan 2024 11:34:40 +0100
Subject: [PATCH 10/11] Use modern type hints

---
 esmvalcore/preprocessor/_regrid.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py
index 4dd3fdd07d..b91adc1568 100644
--- a/esmvalcore/preprocessor/_regrid.py
+++ b/esmvalcore/preprocessor/_regrid.py
@@ -1,4 +1,5 @@
 """Horizontal and vertical regridding module."""
+from __future__ import annotations
 
 import importlib
 import inspect
@@ -10,7 +11,7 @@
 from decimal import Decimal
 from functools import partial
 from pathlib import Path
-from typing import Dict, Optional, Union
+from typing import Dict, Optional
 
 import dask.array as da
 import iris
@@ -1013,7 +1014,7 @@ def _rechunk_aux_factory_dependencies(
 
 def extract_levels(
     cube: iris.cube.Cube,
-    levels: Union[np.typing.ArrayLike, da.Array],
+    levels: np.typing.ArrayLike | da.Array,
     scheme: str,
     coordinate: Optional[str] = None,
     rtol: float = 1e-7,

From a2a7d0dd2b320db1ff5914713b9c4633b9a68b0d Mon Sep 17 00:00:00 2001
From: Bouwe Andela <b.andela@esciencecenter.nl>
Date: Tue, 30 Jan 2024 11:45:31 +0100
Subject: [PATCH 11/11] Undo unrelated change

---
 .circleci/config.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index ec9fbf38ba..8e5f38b84f 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -3,7 +3,7 @@ version: 2.1
 
 orbs:
   coverage-reporter: codacy/coverage-reporter@13.13.7
-  codecov: codecov/codecov@3.3.0
+  codecov: codecov/codecov@3.2.5
 
 commands:
   check_changes: