diff --git a/esmvalcore/preprocessor/_mask.py b/esmvalcore/preprocessor/_mask.py index cdedcc6391..f7db814009 100644 --- a/esmvalcore/preprocessor/_mask.py +++ b/esmvalcore/preprocessor/_mask.py @@ -8,23 +8,33 @@ import logging import os +from collections.abc import Iterable +from typing import Literal, Optional import cartopy.io.shapereader as shpreader import dask.array as da import iris +import iris.util import numpy as np import shapely.vectorized as shp_vect from iris.analysis import Aggregator +from iris.cube import Cube from iris.util import rolling_window +from esmvalcore.preprocessor._shared import get_array_module + from ._supplementary_vars import register_supplementaries logger = logging.getLogger(__name__) -def _get_fx_mask(fx_data, fx_option, mask_type): +def _get_fx_mask( + fx_data: np.ndarray | da.Array, + fx_option: Literal['land', 'sea', 'landsea', 'ice'], + mask_type: Literal['sftlf', 'sftof', 'sftgif'], +) -> np.ndarray | da.Array: """Build a percentage-thresholded mask from an fx file.""" - inmask = da.zeros_like(fx_data, bool) + inmask = np.zeros_like(fx_data, bool) # respects dask through dispatch if mask_type == 'sftlf': if fx_option == 'land': # Mask land out @@ -50,22 +60,29 @@ def _get_fx_mask(fx_data, fx_option, mask_type): return inmask -def _apply_fx_mask(fx_mask, var_data): - """Apply the fx data extracted mask on the actual processed data.""" - # Apply mask across - old_mask = da.ma.getmaskarray(var_data) - mask = old_mask | fx_mask - var_data = da.ma.masked_array(var_data, mask=mask) - # maybe fill_value=1e+20 - - return var_data +def _apply_mask( + mask: np.ndarray | da.Array, + array: np.ndarray | da.Array, + dim_map: Optional[Iterable[int]] = None, +) -> np.ndarray | da.Array: + """Apply a (broadcasted) mask on an array.""" + npx = get_array_module(mask, array) + if dim_map is not None: + if isinstance(array, da.Array): + chunks = array.chunks + else: + chunks = None + mask = iris.util.broadcast_to_shape( + mask, array.shape, dim_map, chunks=chunks + ) + return npx.ma.masked_where(mask, array) @register_supplementaries( variables=['sftlf', 'sftof'], required='prefer_at_least_one', ) -def mask_landsea(cube, mask_out): +def mask_landsea(cube: Cube, mask_out: Literal['land', 'sea']) -> Cube: """Mask out either land mass or sea (oceans, seas and lakes). It uses dedicated ancillary variables (sftlf or sftof) or, @@ -78,16 +95,15 @@ def mask_landsea(cube, mask_out): Parameters ---------- - cube: iris.cube.Cube - data cube to be masked. If the cube has an + cube: + Data cube to be masked. If the cube has an :class:`iris.coords.AncillaryVariable` with standard name ``'land_area_fraction'`` or ``'sea_area_fraction'`` that will be used. If both are present, only the 'land_area_fraction' will be used. If the ancillary variable is not available, the mask will be calculated from Natural Earth shapefiles. - - mask_out: str - either "land" to mask out land mass or "sea" to mask out seas. + mask_out: + Either ``'land'`` to mask out land mass or ``'sea'`` to mask out seas. Returns ------- @@ -112,35 +128,40 @@ def mask_landsea(cube, mask_out): } # preserve importance order: try stflf first then sftof - fx_cube = None + ancillary_var = None try: - fx_cube = cube.ancillary_variable('land_area_fraction') + ancillary_var = cube.ancillary_variable('land_area_fraction') except iris.exceptions.AncillaryVariableNotFoundError: try: - fx_cube = cube.ancillary_variable('sea_area_fraction') + ancillary_var = cube.ancillary_variable('sea_area_fraction') except iris.exceptions.AncillaryVariableNotFoundError: - logger.debug('Ancillary variables land/sea area fraction not ' - 'found in cube. Check fx_file availability.') - - if fx_cube: - fx_cube_data = da.broadcast_to(fx_cube.core_data(), cube.shape) - landsea_mask = _get_fx_mask(fx_cube_data, mask_out, - fx_cube.var_name) - cube.data = _apply_fx_mask(landsea_mask, cube.core_data()) - logger.debug("Applying land-sea mask: %s", fx_cube.var_name) + logger.debug( + "Ancillary variables land/sea area fraction not found in " + "cube. Check fx_file availability." + ) + + if ancillary_var: + landsea_mask = _get_fx_mask( + ancillary_var.core_data(), mask_out, ancillary_var.var_name + ) + cube.data = _apply_mask( + landsea_mask, + cube.core_data(), + cube.ancillary_variable_dims(ancillary_var), + ) + logger.debug("Applying land-sea mask: %s", ancillary_var.var_name) else: if cube.coord('longitude').points.ndim < 2: - cube = _mask_with_shp(cube, shapefiles[mask_out], [ - 0, - ]) + cube = _mask_with_shp(cube, shapefiles[mask_out], [0]) logger.debug( "Applying land-sea mask from Natural Earth shapefile: \n%s", shapefiles[mask_out], ) else: - msg = ("Use of shapefiles with irregular grids not yet " - "implemented, land-sea mask not applied.") - raise ValueError(msg) + raise ValueError( + "Use of shapefiles with irregular grids not yet implemented, " + "land-sea mask not applied." + ) return cube @@ -149,7 +170,7 @@ def mask_landsea(cube, mask_out): variables=['sftgif'], required='require_at_least_one', ) -def mask_landseaice(cube, mask_out): +def mask_landseaice(cube: Cube, mask_out: Literal['landsea', 'ice']) -> Cube: """Mask out either landsea (combined) or ice. Function that masks out either landsea (land and seas) or ice (Antarctica, @@ -159,13 +180,13 @@ def mask_landseaice(cube, mask_out): Parameters ---------- - cube: iris.cube.Cube - data cube to be masked. It should have an + cube: + Data cube to be masked. It should have an :class:`iris.coords.AncillaryVariable` with standard name ``'land_ice_area_fraction'``. - mask_out: str - either "landsea" to mask out landsea or "ice" to mask out ice. + Either ``'landsea'`` to mask out land and oceans or ``'ice'`` to mask + out ice. Returns ------- @@ -178,20 +199,26 @@ def mask_landseaice(cube, mask_out): Error raised if landsea-ice mask not found as an ancillary variable. """ # sftgif is the only one so far but users can set others - fx_cube = None + ancillary_var = None try: - fx_cube = cube.ancillary_variable('land_ice_area_fraction') + ancillary_var = cube.ancillary_variable('land_ice_area_fraction') except iris.exceptions.AncillaryVariableNotFoundError: - logger.debug('Ancillary variable land ice area fraction ' - 'not found in cube. Check fx_file availability.') - if fx_cube: - fx_cube_data = da.broadcast_to(fx_cube.core_data(), cube.shape) - landice_mask = _get_fx_mask(fx_cube_data, mask_out, fx_cube.var_name) - cube.data = _apply_fx_mask(landice_mask, cube.core_data()) + logger.debug( + "Ancillary variable land ice area fraction not found in cube. " + "Check fx_file availability." + ) + if ancillary_var: + landseaice_mask = _get_fx_mask( + ancillary_var.core_data(), mask_out, ancillary_var.var_name + ) + cube.data = _apply_mask( + landseaice_mask, + cube.core_data(), + cube.ancillary_variable_dims(ancillary_var), + ) logger.debug("Applying landsea-ice mask: sftgif") else: - msg = "Landsea-ice mask could not be found. Stopping. " - raise ValueError(msg) + raise ValueError("Landsea-ice mask could not be found. Stopping.") return cube @@ -285,9 +312,10 @@ def _mask_with_shp(cube, shapefilename, region_indices=None): # Create a set of x,y points from the cube # 1D regular grids if cube.coord('longitude').points.ndim < 2: - x_p, y_p = da.meshgrid( + x_p, y_p = np.meshgrid( cube.coord(axis='X').points, - cube.coord(axis='Y').points) + cube.coord(axis='Y').points, + ) # 2D irregular grids; spit an error for now else: msg = ("No fx-files found (sftlf or sftof)!" @@ -296,14 +324,14 @@ def _mask_with_shp(cube, shapefilename, region_indices=None): raise ValueError(msg) # Wrap around longitude coordinate to match data - x_p_180 = da.where(x_p >= 180., x_p - 360., x_p) + x_p_180 = np.where(x_p >= 180., x_p - 360., x_p) # the NE mask has no points at x = -180 and y = +/-90 # so we will fool it and apply the mask at (-179, -89, 89) instead - x_p_180 = da.where(x_p_180 == -180., x_p_180 + 1., x_p_180) + x_p_180 = np.where(x_p_180 == -180., x_p_180 + 1., x_p_180) - y_p_0 = da.where(y_p == -90., y_p + 1., y_p) - y_p_90 = da.where(y_p_0 == 90., y_p_0 - 1., y_p_0) + y_p_0 = np.where(y_p == -90., y_p + 1., y_p) + y_p_90 = np.where(y_p_0 == 90., y_p_0 - 1., y_p_0) mask = None for region in regions: @@ -313,13 +341,14 @@ def _mask_with_shp(cube, shapefilename, region_indices=None): else: mask |= shp_vect.contains(region, x_p_180, y_p_90) - mask = da.array(mask) - iris.util.broadcast_to_shape(mask, cube.shape, cube.coord_dims('latitude') - + cube.coord_dims('longitude')) + if cube.has_lazy_data(): + mask = da.array(mask) - old_mask = da.ma.getmaskarray(cube.core_data()) - mask = old_mask | mask - cube.data = da.ma.masked_array(cube.core_data(), mask=mask) + cube.data = _apply_mask( + mask, + cube.core_data(), + cube.coord_dims('latitude') + cube.coord_dims('longitude'), + ) return cube diff --git a/esmvalcore/preprocessor/_shared.py b/esmvalcore/preprocessor/_shared.py index bf6f1fd088..67c3b89f5d 100644 --- a/esmvalcore/preprocessor/_shared.py +++ b/esmvalcore/preprocessor/_shared.py @@ -329,10 +329,11 @@ def get_weights( # Time weights: lengths of time interval if 'time' in coords: - weights *= broadcast_to_shape( + weights = weights * broadcast_to_shape( npx.array(get_time_weights(cube)), cube.shape, cube.coord_dims('time'), + chunks=cube.lazy_data().chunks if cube.has_lazy_data() else None, ) # Latitude weights: cell areas @@ -350,10 +351,17 @@ def get_weights( f"variable)" ) try_adding_calculated_cell_area(cube) - weights *= broadcast_to_shape( - cube.cell_measure('cell_area').core_data(), + area_weights = cube.cell_measure('cell_area').core_data() + if cube.has_lazy_data(): + area_weights = da.array(area_weights) + chunks = cube.lazy_data().chunks + else: + chunks = None + weights = weights * broadcast_to_shape( + area_weights, cube.shape, cube.cell_measure_dims('cell_area'), + chunks=chunks, ) return weights diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py index 473abb0190..840af3214e 100644 --- a/esmvalcore/preprocessor/_volume.py +++ b/esmvalcore/preprocessor/_volume.py @@ -161,9 +161,10 @@ def calculate_volume(cube: Cube) -> da.core.Array: try_adding_calculated_cell_area(cube) area = cube.cell_measure('cell_area').copy() area_dim = cube.cell_measure_dims(area) - - # Ensure cell area is in square meters as the units area.convert_units('m2') + area_array = area.core_data() + if cube.has_lazy_data(): + area_array = da.array(area_array) # Make sure input cube has not been modified if not has_cell_measure: @@ -171,9 +172,11 @@ def calculate_volume(cube: Cube) -> da.core.Array: chunks = cube.core_data().chunks if cube.has_lazy_data() else None area_arr = broadcast_to_shape( - area.core_data(), cube.shape, area_dim, chunks=chunks) + area_array, cube.shape, area_dim, chunks=chunks + ) thickness_arr = broadcast_to_shape( - thickness, cube.shape, z_dim, chunks=chunks) + thickness, cube.shape, z_dim, chunks=chunks + ) grid_volume = area_arr * thickness_arr return grid_volume diff --git a/tests/integration/preprocessor/_mask/test_mask.py b/tests/integration/preprocessor/_mask/test_mask.py index 4e5e51167b..1b085519eb 100644 --- a/tests/integration/preprocessor/_mask/test_mask.py +++ b/tests/integration/preprocessor/_mask/test_mask.py @@ -4,10 +4,12 @@ """ from pathlib import Path +import dask.array as da import iris import iris.fileformats import numpy as np import pytest +from iris.coords import AuxCoord from esmvalcore.preprocessor import ( PreprocessorFile, @@ -64,55 +66,59 @@ def setUp(self): self.mock_data = np.ma.empty((4, 3, 3)) self.mock_data[:] = 10. - def test_components_fx_var(self): + @pytest.mark.parametrize('lazy_fx', [True, False]) + @pytest.mark.parametrize('lazy', [True, False]) + def test_components_fx_var(self, lazy, lazy_fx): """Test compatibility of ancillary variables.""" - self.fx_mask.var_name = 'sftlf' - self.fx_mask.standard_name = 'land_area_fraction' + if lazy: + cube_data = da.array(self.new_cube_data) + else: + cube_data = self.new_cube_data + fx_cube = self.fx_mask.copy() + if lazy_fx: + fx_cube.data = fx_cube.lazy_data() + + # mask_landsea + fx_cube.var_name = 'sftlf' + fx_cube.standard_name = 'land_area_fraction' new_cube_land = iris.cube.Cube( - self.new_cube_data, - dim_coords_and_dims=self.cube_coords_spec - ) - new_cube_land = add_supplementary_variables( - new_cube_land, - [self.fx_mask], - ) - result_land = mask_landsea( - new_cube_land, - 'land', + cube_data, dim_coords_and_dims=self.cube_coords_spec ) + new_cube_land = add_supplementary_variables(new_cube_land, [fx_cube]) + result_land = mask_landsea(new_cube_land, 'land') assert isinstance(result_land, iris.cube.Cube) + assert result_land.has_lazy_data() is (lazy or lazy_fx) - self.fx_mask.var_name = 'sftgif' - self.fx_mask.standard_name = 'land_ice_area_fraction' + # mask_landseaice + fx_cube.var_name = 'sftgif' + fx_cube.standard_name = 'land_ice_area_fraction' new_cube_ice = iris.cube.Cube( - self.new_cube_data, - dim_coords_and_dims=self.cube_coords_spec - ) - new_cube_ice = add_supplementary_variables( - new_cube_ice, - [self.fx_mask], - ) - result_ice = mask_landseaice( - new_cube_ice, - 'ice', + cube_data, dim_coords_and_dims=self.cube_coords_spec ) + new_cube_ice = add_supplementary_variables(new_cube_ice, [fx_cube]) + result_ice = mask_landseaice(new_cube_ice, 'ice') assert isinstance(result_ice, iris.cube.Cube) + assert result_ice.has_lazy_data() is (lazy or lazy_fx) - def test_mask_landsea(self): + @pytest.mark.parametrize('lazy', [True, False]) + def test_mask_landsea(self, lazy): """Test mask_landsea func.""" + if lazy: + cube_data = da.array(self.new_cube_data) + else: + cube_data = self.new_cube_data + self.fx_mask.var_name = 'sftlf' self.fx_mask.standard_name = 'land_area_fraction' new_cube_land = iris.cube.Cube( - self.new_cube_data, - dim_coords_and_dims=self.cube_coords_spec + cube_data, dim_coords_and_dims=self.cube_coords_spec ) new_cube_land = add_supplementary_variables( new_cube_land, [self.fx_mask], ) new_cube_sea = iris.cube.Cube( - self.new_cube_data, - dim_coords_and_dims=self.cube_coords_spec + cube_data, dim_coords_and_dims=self.cube_coords_spec ) new_cube_sea = add_supplementary_variables( new_cube_sea, @@ -128,6 +134,8 @@ def test_mask_landsea(self): new_cube_sea, 'sea', ) + assert result_land.has_lazy_data() is lazy + assert result_sea.has_lazy_data() is lazy expected = np.ma.empty((2, 3, 3)) expected.data[:] = 200. expected.mask = np.ones((2, 3, 3), bool) @@ -143,17 +151,17 @@ def test_mask_landsea(self): # mask with shp files new_cube_land = iris.cube.Cube( - self.new_cube_data, - dim_coords_and_dims=self.cube_coords_spec + cube_data, dim_coords_and_dims=self.cube_coords_spec ) new_cube_sea = iris.cube.Cube( - self.new_cube_data, - dim_coords_and_dims=self.cube_coords_spec + cube_data, dim_coords_and_dims=self.cube_coords_spec ) result_land = mask_landsea(new_cube_land, 'land') result_sea = mask_landsea(new_cube_sea, 'sea') # bear in mind all points are in the ocean + assert result_land.has_lazy_data() is lazy + assert result_sea.has_lazy_data() is lazy np.ma.set_fill_value(result_land.data, 1e+20) np.ma.set_fill_value(result_sea.data, 1e+20) expected.mask = np.zeros((3, 3), bool) @@ -161,19 +169,87 @@ def test_mask_landsea(self): expected.mask = np.ones((3, 3), bool) assert_array_equal(result_sea.data, expected) - def test_mask_landseaice(self): + @pytest.mark.parametrize('lazy', [True, False]) + def test_mask_landsea_transposed_fx(self, lazy): + """Test mask_landsea func.""" + if lazy: + cube_data = da.array(self.new_cube_data) + else: + cube_data = self.new_cube_data + cube = iris.cube.Cube( + cube_data, dim_coords_and_dims=self.cube_coords_spec + ) + self.fx_mask.var_name = 'sftlf' + self.fx_mask.standard_name = 'land_area_fraction' + cube = add_supplementary_variables(cube, [self.fx_mask]) + cube.transpose([2, 1, 0]) + + result = mask_landsea(cube, 'land') + + assert result.has_lazy_data() is lazy + expected = np.ma.array( + np.full((3, 3, 2), 200.0), mask=np.ones((3, 3, 2), bool) + ) + expected.mask[2, 1, :] = False + assert_array_equal(result.data, expected) + + @pytest.mark.parametrize('lazy', [True, False]) + def test_mask_landsea_transposed_shp(self, lazy): + """Test mask_landsea func.""" + if lazy: + cube_data = da.array(self.new_cube_data) + else: + cube_data = self.new_cube_data + cube = iris.cube.Cube( + cube_data, dim_coords_and_dims=self.cube_coords_spec + ) + cube.transpose([2, 1, 0]) + + result = mask_landsea(cube, 'land') + + assert result.has_lazy_data() is lazy + expected = np.ma.array( + np.full((3, 3, 2), 200.0), mask=np.zeros((3, 3, 2), bool) + ) + assert_array_equal(result.data, expected) + + def test_mask_landsea_multidim_fail(self): + """Test mask_landsea func.""" + lon_coord = AuxCoord(np.ones((3, 3)), standard_name='longitude') + cube = iris.cube.Cube( + self.new_cube_data, + dim_coords_and_dims=[(self.zcoord, 0), (self.lats, 1)], + aux_coords_and_dims=[(lon_coord, (1, 2))], + ) + + msg = ( + "Use of shapefiles with irregular grids not yet implemented, " + "land-sea mask not applied." + ) + with pytest.raises(ValueError, match=msg): + mask_landsea(cube, 'land') + + @pytest.mark.parametrize('lazy', [True, False]) + def test_mask_landseaice(self, lazy): """Test mask_landseaice func.""" + if lazy: + cube_data = da.array(self.new_cube_data).rechunk((1, 3, 3)) + else: + cube_data = self.new_cube_data + self.fx_mask.var_name = 'sftgif' self.fx_mask.standard_name = 'land_ice_area_fraction' new_cube_ice = iris.cube.Cube( - self.new_cube_data, - dim_coords_and_dims=self.cube_coords_spec + cube_data, dim_coords_and_dims=self.cube_coords_spec ) new_cube_ice = add_supplementary_variables( new_cube_ice, [self.fx_mask], ) result_ice = mask_landseaice(new_cube_ice, 'ice') + assert result_ice.has_lazy_data() is lazy + if lazy: + assert result_ice.lazy_data().chunksize == (1, 3, 3) expected = np.ma.empty((2, 3, 3)) expected.data[:] = 200. expected.mask = np.ones((2, 3, 3), bool) @@ -182,6 +258,19 @@ def test_mask_landseaice(self): np.ma.set_fill_value(expected, 1e+20) assert_array_equal(result_ice.data, expected) + def test_mask_landseaice_multidim_fail(self): + """Test mask_landseaice func.""" + lon_coord = AuxCoord(np.ones((3, 3)), standard_name='longitude') + cube = iris.cube.Cube( + self.new_cube_data, + dim_coords_and_dims=[(self.zcoord, 0), (self.lats, 1)], + aux_coords_and_dims=[(lon_coord, (1, 2))], + ) + + msg = "Landsea-ice mask could not be found. Stopping." + with pytest.raises(ValueError, match=msg): + mask_landseaice(cube, 'ice') + @pytest.mark.parametrize('lazy', [True, False]) def test_mask_fillvalues(self, mocker, lazy): """Test the fillvalues mask: func mask_fillvalues.""" diff --git a/tests/unit/preprocessor/_compare_with_refs/test_compare_with_refs.py b/tests/unit/preprocessor/_compare_with_refs/test_compare_with_refs.py index 7def9afe1b..e30a4634fd 100644 --- a/tests/unit/preprocessor/_compare_with_refs/test_compare_with_refs.py +++ b/tests/unit/preprocessor/_compare_with_refs/test_compare_with_refs.py @@ -364,6 +364,7 @@ def test_reference_none_cubes(regular_cubes): ) +@pytest.mark.parametrize('lazy_weights', [True, False]) @pytest.mark.parametrize( 'metric,data,ref_data,long_name,var_name,units', TEST_DISTANCE_METRICS ) @@ -376,9 +377,14 @@ def test_distance_metric( long_name, var_name, units, + lazy_weights, ): """Test `distance_metric`.""" - regular_cubes[0].add_cell_measure(AREA_WEIGHTS, (1, 2)) + regular_cubes[0].add_cell_measure(AREA_WEIGHTS.copy(), (1, 2)) + if lazy_weights: + regular_cubes[0].cell_measure('cell_area').data = ( + regular_cubes[0].cell_measure('cell_area').lazy_data() + ) ref_product = PreprocessorFile( ref_cubes, 'REF', {'reference_for_metric': True} ) @@ -390,6 +396,10 @@ def test_distance_metric( out_products = distance_metric(products, metric) + assert ( + regular_cubes[0].cell_measure('cell_area').has_lazy_data() is + lazy_weights + ) assert isinstance(out_products, set) out_dict = products_set_to_dict(out_products) assert len(out_dict) == 3 @@ -407,6 +417,7 @@ def test_distance_metric( out_cube = product_a.cubes[0] assert out_cube.shape == () assert out_cube.dtype == np.float32 + assert not out_cube.has_lazy_data() assert_allclose(out_cube.data, np.array(data, dtype=np.float32)) assert out_cube.var_name == var_name assert out_cube.long_name == long_name @@ -425,6 +436,7 @@ def test_distance_metric( out_cube = product_b.cubes[0] assert out_cube.shape == () assert out_cube.dtype == np.float32 + assert not out_cube.has_lazy_data() assert_allclose(out_cube.data, np.array(data, dtype=np.float32)) assert out_cube.var_name == var_name assert out_cube.long_name == long_name @@ -445,6 +457,7 @@ def test_distance_metric( out_cube = product_ref.cubes[0] assert out_cube.shape == () assert out_cube.dtype == np.float32 + assert not out_cube.has_lazy_data() assert_allclose(out_cube.data, ref_data) assert out_cube.var_name == var_name assert out_cube.long_name == long_name @@ -531,22 +544,40 @@ def test_distance_metric_lazy( assert product_a.mock_ancestors == {ref_product} +@pytest.mark.parametrize('lazy_weights', [True, False]) @pytest.mark.parametrize( 'metric,data,_,long_name,var_name,units', TEST_DISTANCE_METRICS ) def test_distance_metric_cubes( - regular_cubes, ref_cubes, metric, data, _, long_name, var_name, units + regular_cubes, + ref_cubes, + metric, + data, + _, + long_name, + var_name, + units, + lazy_weights, ): """Test `distance_metric` with cubes.""" - regular_cubes[0].add_cell_measure(AREA_WEIGHTS, (1, 2)) + regular_cubes[0].add_cell_measure(AREA_WEIGHTS.copy(), (1, 2)) + if lazy_weights: + regular_cubes[0].cell_measure('cell_area').data = ( + regular_cubes[0].cell_measure('cell_area').lazy_data() + ) out_cubes = distance_metric(regular_cubes, metric, reference=ref_cubes[0]) + assert ( + regular_cubes[0].cell_measure('cell_area').has_lazy_data() is + lazy_weights + ) assert isinstance(out_cubes, CubeList) assert len(out_cubes) == 1 out_cube = out_cubes[0] assert out_cube.shape == () assert out_cube.dtype == np.float32 + assert not out_cube.has_lazy_data() assert_allclose(out_cube.data, np.array(data, dtype=np.float32)) assert out_cube.var_name == var_name assert out_cube.long_name == long_name @@ -557,12 +588,22 @@ def test_distance_metric_cubes( ) +@pytest.mark.parametrize('lazy_weights', [True, False]) @pytest.mark.parametrize('lazy', [True, False]) @pytest.mark.parametrize( 'metric,data,_,long_name,var_name,units', TEST_DISTANCE_METRICS ) def test_distance_metric_masked_data( - regular_cubes, ref_cubes, metric, data, _, long_name, var_name, units, lazy + regular_cubes, + ref_cubes, + metric, + data, + _, + long_name, + var_name, + units, + lazy, + lazy_weights, ): """Test `distance_metric` with masked data.""" # Test cube @@ -585,7 +626,11 @@ def test_distance_metric_masked_data( np.ma.masked_invalid(cube_data), dim_coords_and_dims=coord_specs ) cube.metadata = regular_cubes[0].metadata - cube.add_cell_measure(AREA_WEIGHTS, (1, 2)) + cube.add_cell_measure(AREA_WEIGHTS.copy(), (1, 2)) + if lazy_weights: + cube.cell_measure('cell_area').data = ( + cube.cell_measure('cell_area').lazy_data() + ) # Ref cube ref_cube = cube.copy() @@ -604,6 +649,7 @@ def test_distance_metric_masked_data( out_cubes = distance_metric([cube], metric, reference=ref_cube) + assert cube.cell_measure('cell_area').has_lazy_data() is lazy_weights assert isinstance(out_cubes, CubeList) assert len(out_cubes) == 1 out_cube = out_cubes[0] @@ -630,17 +676,31 @@ def test_distance_metric_masked_data( ) +@pytest.mark.parametrize('lazy_weights', [True, False]) @pytest.mark.parametrize('lazy', [True, False]) @pytest.mark.parametrize( 'metric,_,__,long_name,var_name,units', TEST_DISTANCE_METRICS ) def test_distance_metric_fully_masked_data( - regular_cubes, ref_cubes, metric, _, __, long_name, var_name, units, lazy + regular_cubes, + ref_cubes, + metric, + _, + __, + long_name, + var_name, + units, + lazy, + lazy_weights, ): """Test `distance_metric` with fully_masked data.""" cube = regular_cubes[0] cube.data = np.ma.masked_invalid(np.full(cube.shape, np.nan)) - cube.add_cell_measure(AREA_WEIGHTS, (1, 2)) + cube.add_cell_measure(AREA_WEIGHTS.copy(), (1, 2)) + if lazy_weights: + cube.cell_measure('cell_area').data = ( + cube.cell_measure('cell_area').lazy_data() + ) ref_cube = ref_cubes[0] if lazy: @@ -649,6 +709,7 @@ def test_distance_metric_fully_masked_data( out_cubes = distance_metric([cube], metric, reference=ref_cube) + assert cube.cell_measure('cell_area').has_lazy_data() is lazy_weights assert isinstance(out_cubes, CubeList) assert len(out_cubes) == 1 out_cube = out_cubes[0] diff --git a/tests/unit/preprocessor/_mask/test_mask.py b/tests/unit/preprocessor/_mask/test_mask.py index 44cb0246f9..229657b309 100644 --- a/tests/unit/preprocessor/_mask/test_mask.py +++ b/tests/unit/preprocessor/_mask/test_mask.py @@ -2,18 +2,22 @@ import unittest -import numpy as np - import iris import iris.fileformats -import tests +import numpy as np from cf_units import Unit -from esmvalcore.preprocessor._mask import (_apply_fx_mask, - count_spells, _get_fx_mask, - mask_above_threshold, - mask_below_threshold, - mask_glaciated, mask_inside_range, - mask_outside_range) + +import tests +from esmvalcore.preprocessor._mask import ( + _apply_mask, + _get_fx_mask, + count_spells, + mask_above_threshold, + mask_below_threshold, + mask_glaciated, + mask_inside_range, + mask_outside_range, +) class Test(tests.Test): @@ -48,11 +52,12 @@ def setUp(self): def test_apply_fx_mask_on_nonmasked_data(self): """Test _apply_fx_mask func.""" dummy_fx_mask = np.ma.array((True, False, True)) - app_mask = _apply_fx_mask(dummy_fx_mask, - self.time_cube.data[0:3].astype('float64')) - app_mask = app_mask.compute() - fixed_mask = np.ma.array(self.time_cube.data[0:3].astype('float64'), - mask=dummy_fx_mask) + app_mask = _apply_mask( + dummy_fx_mask, self.time_cube.data[0:3].astype('float64') + ) + fixed_mask = np.ma.array( + self.time_cube.data[0:3].astype('float64'), mask=dummy_fx_mask + ) self.assert_array_equal(fixed_mask, app_mask) def test_apply_fx_mask_on_masked_data(self): @@ -60,8 +65,7 @@ def test_apply_fx_mask_on_masked_data(self): dummy_fx_mask = np.ma.array((True, True, True)) masked_data = np.ma.array(self.time_cube.data[0:3].astype('float64'), mask=np.ma.array((False, True, False))) - app_mask = _apply_fx_mask(dummy_fx_mask, masked_data) - app_mask = app_mask.compute() + app_mask = _apply_mask(dummy_fx_mask, masked_data) fixed_mask = np.ma.array(self.time_cube.data[0:3].astype('float64'), mask=dummy_fx_mask) self.assert_array_equal(fixed_mask, app_mask) diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py index f055d2f7b2..d45b9bd744 100644 --- a/tests/unit/preprocessor/_volume/test_volume.py +++ b/tests/unit/preprocessor/_volume/test_volume.py @@ -402,6 +402,7 @@ def test_extract_volume_mean(self): """Test to extract the top two layers and compute the weighted average of a cube.""" grid_volume = calculate_volume(self.grid_4d) + assert isinstance(grid_volume, np.ndarray) measure = iris.coords.CellMeasure(grid_volume, standard_name='ocean_volume', units='m3',