diff --git a/esmvalcore/cmor/_fixes/shared.py b/esmvalcore/cmor/_fixes/shared.py index b50deb422b..b658fa5cad 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 8b3c483154..4c84cedfce 100644 --- a/esmvalcore/preprocessor/_regrid.py +++ b/esmvalcore/preprocessor/_regrid.py @@ -11,8 +11,9 @@ import warnings from copy import deepcopy from decimal import Decimal +from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Optional import dask.array as da import iris @@ -21,7 +22,6 @@ from geopy.geocoders import Nominatim from iris.analysis import AreaWeighted, Linear, Nearest from iris.cube import Cube -from iris.util import broadcast_to_shape from esmvalcore.cmor._fixes.shared import ( add_altitude_from_plev, @@ -935,28 +935,90 @@ 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, dim_map, chunks=None): + """Copy of `iris.util.broadcast_to_shape` that allows specifying chunks.""" + if isinstance(array, da.Array): + 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 = partial(np.broadcast_to, shape=shape) + + 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(array) + + 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 = 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(mask) + 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. - src_levels_broadcast = broadcast_to_shape(src_levels.points, cube.shape, - cube.coord_dims(src_levels)) + 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() + + # Broadcast the source cube vertical coordinate to fully describe the + # spatial extent that will be interpolated. + src_levels_broadcast = broadcast_to_shape( + src_points, + shape=cube.shape, + chunks=cube.lazy_data().chunks if cube.has_lazy_data() else None, + 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) - # Now perform the actual vertical interpolation. - new_data = stratify.interpolate(levels, - src_levels_broadcast, - data, - axis=z_axis, - interpolation=interpolation, - extrapolation=extrapolation) + # Perform vertical interpolation. + 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) @@ -1026,42 +1088,74 @@ def parse_vertical_scheme(scheme): return scheme, extrap_scheme -def extract_levels(cube, - levels, - scheme, - coordinate=None, - rtol=1e-7, - atol=None): +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: np.typing.ArrayLike | da.Array, + scheme: str, + coordinate: Optional[str] = None, + 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 : 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. @@ -1081,29 +1175,37 @@ 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) - - # 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) + if not isinstance(levels, da.Array): + levels = np.array(levels, ndmin=1) + + # 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.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. @@ -1114,7 +1216,9 @@ def extract_levels(cube, 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. 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_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) diff --git a/tests/unit/preprocessor/_regrid/test_extract_levels.py b/tests/unit/preprocessor/_regrid/test_extract_levels.py index f53fbac708..4825350f17 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,53 @@ def test_interpolation__masked(self): self.assertEqual(kwargs, dict()) -if __name__ == '__main__': - unittest.main() +def test_rechunk_aux_factory_dependencies(): + + delta = iris.coords.AuxCoord( + 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", + ) + 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