diff --git a/environment.yml b/environment.yml index 26809991e1..66c4782c88 100644 --- a/environment.yml +++ b/environment.yml @@ -35,7 +35,7 @@ dependencies: - py-cordex - pybtex - python>=3.9 - - python-stratify + - python-stratify>=0.3 - pyyaml - requests - scipy>=1.6 diff --git a/esmvalcore/preprocessor/_other.py b/esmvalcore/preprocessor/_other.py index 36bbccafb6..14594ba7cf 100644 --- a/esmvalcore/preprocessor/_other.py +++ b/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 diff --git a/esmvalcore/preprocessor/_regrid.py b/esmvalcore/preprocessor/_regrid.py index a14abfbe63..4972a7ce8f 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/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)) diff --git a/setup.py b/setup.py index 05d5a9ac55..5743b3f603 100755 --- a/setup.py +++ b/setup.py @@ -58,7 +58,7 @@ 'scipy>=1.6', 'scitools-iris>=3.4.0', 'shapely[vectorized]', - 'stratify', + 'stratify>=0.3', 'yamale', ], # Test dependencies diff --git a/tests/integration/preprocessor/_regrid/test_extract_levels.py b/tests/integration/preprocessor/_regrid/test_extract_levels.py index 8687f19228..b0869244a6 100644 --- a/tests/integration/preprocessor/_regrid/test_extract_levels.py +++ b/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)) diff --git a/tests/unit/preprocessor/_other/test_other.py b/tests/unit/preprocessor/_other/test_other.py index ab98895c02..803b66d997 100644 --- a/tests/unit/preprocessor/_other/test_other.py +++ b/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() diff --git a/tests/unit/preprocessor/_regrid/test_extract_levels.py b/tests/unit/preprocessor/_regrid/test_extract_levels.py index 638d61fbcd..f53fbac708 100644 --- a/tests/unit/preprocessor/_regrid/test_extract_levels.py +++ b/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,