Make extract_levels preprocessor function lazy

environment.yml

@@ -35,7 +35,7 @@ dependencies:
   - py-cordex
   - pybtex
   - python>=3.8
-  - python-stratify
+  - python-stratify>=0.3
   - pyyaml
   - requests
   - scipy>=1.6

esmvalcore/preprocessor/_other.py

@@ -4,6 +4,7 @@
 from collections import defaultdict
 
 import dask.array as da
+import numpy as np
 
 logger = logging.getLogger(__name__)
 
@@ -71,3 +72,16 @@ def grouper(product):
 
     grouped = _groupby(products, keyfunc=grouper)
     return grouped.items()
+
+
+def get_array_module(*args):
+    """Return the best matching array module.
+
+    If at least one of the arguments is a :class:`dask.array.Array` object,
+    the :mod:`dask.array` module is returned. In all other cases the
+    :mod:`numpy` module is returned.
+    """
+    for arg in args:
+        if isinstance(arg, da.Array):
+            return da
+    return np

esmvalcore/preprocessor/_regrid.py

@@ -14,13 +14,13 @@
 import iris
 import numpy as np
 import stratify
-from dask import array as da
 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
+from ._other import get_array_module
 from ._regrid_esmpy import ESMF_REGRID_METHODS
 from ._regrid_esmpy import regrid as esmpy_regrid
 from ._supplementary_vars import add_ancillary_variable, add_cell_measure
@@ -797,22 +797,19 @@ def _vertical_interpolate(cube, src_levels, levels, interpolation,
                                               cube.coord_dims(src_levels))
 
     # force mask onto data as nan's
-    cube.data = da.ma.filled(cube.core_data(), np.nan)
+    npx = get_array_module(cube.core_data())
+    data = npx.ma.filled(cube.core_data(), np.nan)
 
     # Now perform the actual vertical interpolation.
     new_data = stratify.interpolate(levels,
                                     src_levels_broadcast,
-                                    cube.core_data(),
+                                    data,
                                     axis=z_axis,
                                     interpolation=interpolation,
                                     extrapolation=extrapolation)
 
     # Calculate the mask based on the any NaN values in the interpolated data.
-    mask = np.isnan(new_data)
-
-    if np.any(mask):
-        # Ensure that the data is masked appropriately.
-        new_data = np.ma.array(new_data, mask=mask, fill_value=_MDI)
+    new_data = npx.ma.masked_where(npx.isnan(new_data), new_data)
 
     # Construct the resulting cube with the interpolated data.
     return _create_cube(cube, new_data, src_levels, levels.astype(float))

setup.py

@@ -58,7 +58,7 @@
         'scipy>=1.6',
         'scitools-iris>=3.4.0',
         'shapely[vectorized]',
-        'stratify',
+        'stratify>=0.3',
         'yamale',
     ],
     # Test dependencies

tests/integration/preprocessor/_regrid/test_extract_levels.py

@@ -15,6 +15,7 @@
 
 
 class Test(tests.Test):
+
     def setUp(self):
         """Prepare tests."""
         shape = (3, 2, 2)
@@ -57,20 +58,52 @@ def test_interpolation__linear(self):
         levels = [0.5, 1.5]
         scheme = 'linear'
         result = extract_levels(self.cube, levels, scheme)
-        expected = np.array([[[[2., 3.], [4., 5.]], [[6., 7.], [8., 9.]]],
-                             [[[14., 15.], [16., 17.]], [[18., 19.],
-                                                         [20., 21.]]]])
+        expected = np.ma.array([
+            [
+                [[2., 3.], [4., 5.]],
+                [[6., 7.], [8., 9.]],
+            ],
+            [
+                [[14., 15.], [16., 17.]],
+                [[18., 19.], [20., 21.]],
+            ],
+        ])
         self.assert_array_equal(result.data, expected)
         self.shape[self.z_dim] = len(levels)
         self.assertEqual(result.shape, tuple(self.shape))
 
+    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)
+        self.assertTrue(result.has_lazy_data())
+        expected = np.ma.array([
+            [
+                [[2., 3.], [4., 5.]],
+                [[6., 7.], [8., 9.]],
+            ],
+            [
+                [[14., 15.], [16., 17.]],
+                [[18., 19.], [20., 21.]],
+            ],
+        ])
+        self.assert_array_equal(result.data, expected)
+
     def test_interpolation__nearest(self):
         levels = [0.49, 1.51]
         scheme = 'nearest'
         result = extract_levels(self.cube, levels, scheme)
-        expected = np.array([[[[0., 1.], [2., 3.]], [[8., 9.], [10., 11.]]],
-                             [[[12., 13.], [14., 15.]], [[20., 21.],
-                                                         [22., 23.]]]])
+        expected = np.ma.array([
+            [
+                [[0., 1.], [2., 3.]],
+                [[8., 9.], [10., 11.]],
+            ],
+            [
+                [[12., 13.], [14., 15.]],
+                [[20., 21.], [22., 23.]],
+            ],
+        ])
         self.assert_array_equal(result.data, expected)
         self.shape[self.z_dim] = len(levels)
         self.assertEqual(result.shape, tuple(self.shape))

tests/unit/preprocessor/_other/test_other.py

@@ -2,6 +2,7 @@
 
 import unittest
 
+import dask.array as da
 import iris.coord_categorisation
 import iris.coords
 import numpy as np
@@ -10,11 +11,16 @@
 from numpy.testing import assert_array_equal
 
 from esmvalcore.preprocessor import PreprocessorFile
-from esmvalcore.preprocessor._other import _group_products, clip
+from esmvalcore.preprocessor._other import (
+    _group_products,
+    clip,
+    get_array_module,
+)
 
 
 class TestOther(unittest.TestCase):
     """Test class for _other."""
+
     def test_clip(self):
         """Test clip function."""
         cube = Cube(np.array([-10, 0, 10]))
@@ -67,5 +73,23 @@ def test_group_products_string_list():
     assert grouped_by_list == grouped_by_string
 
 
+def test_get_array_module_da():
+
+    npx = get_array_module(da.array([1, 2]))
+    assert npx is da
+
+
+def test_get_array_module_np():
+
+    npx = get_array_module(np.array([1, 2]))
+    assert npx is np
+
+
+def test_get_array_module_mixed():
+
+    npx = get_array_module(da.array([1]), np.array([1]))
+    assert npx is da
+
+
 if __name__ == '__main__':
     unittest.main()

tests/unit/preprocessor/_regrid/test_extract_levels.py

@@ -1,6 +1,4 @@
 """Unit tests for :func:`esmvalcore.preprocessor.regrid.extract_levels`."""
-import os
-import tempfile
 import unittest
 from unittest import mock
 
@@ -35,8 +33,6 @@ def setUp(self):
         self.schemes = [
             'linear', 'nearest', 'linear_extrapolate', 'nearest_extrapolate',
         ]
-        descriptor, self.filename = tempfile.mkstemp('.nc')
-        os.close(descriptor)
 
     def test_invalid_scheme__unknown(self):
         levels = mock.sentinel.levels
@@ -160,7 +156,7 @@ def test_interpolation(self):
         # Check the _create_cube args ...
         self.assertEqual(len(args), 4)
         self.assertEqual(args[0], self.cube)
-        self.assert_array_equal(args[1], new_data)
+        self.assert_array_equal(args[1], np.ma.array(new_data))
         self.assert_array_equal(args[2],
                                 self.cube.coord(axis='z', dim_coords=True))
         self.assert_array_equal(args[3], levels)
@@ -281,15 +277,8 @@ def test_interpolation__masked(self):
         masked = ma.empty(self.shape)
         masked.mask = mask
         cube = _make_cube(masked, dtype=self.dtype)
-        # save cube to test the lazy data interpolation too
-        iris.save(cube, self.filename)
         with mock.patch('stratify.interpolate',
                         return_value=new_data) as mocker:
-            # first test lazy
-            loaded_cube = iris.load_cube(self.filename)
-            result_from_lazy = extract_levels(loaded_cube, levels, scheme)
-            self.assertEqual(result_from_lazy, self.created_cube)
-            # then test realized
             result = extract_levels(cube, levels, scheme)
             self.assertEqual(result, self.created_cube)
             args, kwargs = mocker.call_args
@@ -300,7 +289,7 @@ def test_interpolation__masked(self):
             src_levels_broadcast = np.broadcast_to(pts.reshape(self.z, 1, 1),
                                                    cube.shape)
             self.assert_array_equal(args[1], src_levels_broadcast)
-            self.assert_array_equal(args[2], cube.data)
+            self.assert_array_equal(args[2], np.ma.filled(masked, np.nan))
             # Check the stratify.interpolate kwargs ...
             self.assertEqual(
                 kwargs, dict(axis=0, interpolation=scheme,