From 65d60cef3de833ccc7573929bcca1f0b41bf6456 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 15 Mar 2022 15:43:28 +0100 Subject: [PATCH 01/12] Use dask and iris --- esmvalcore/preprocessor/_volume.py | 166 ++++------------------------- 1 file changed, 19 insertions(+), 147 deletions(-) diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py index 367cc55f3c..0a8157048d 100644 --- a/esmvalcore/preprocessor/_volume.py +++ b/esmvalcore/preprocessor/_volume.py @@ -4,7 +4,6 @@ depth or height regions; constructing volumetric averages; """ import logging -from copy import deepcopy import dask.array as da import iris @@ -52,90 +51,6 @@ def extract_volume(cube, z_min, z_max): return cube.extract(z_constraint) -def _create_cube_time(src_cube, data, times): - """Generate a new cube with the volume averaged data. - - The resultant cube is seeded with `src_cube` metadata and coordinates, - excluding any source coordinates that span the associated vertical - dimension. The `times` of interpolation are used along with the - associated source cube time coordinate metadata to add a new - time coordinate to the resultant cube. - - Based on the _create_cube method from _regrid.py. - - Parameters - ---------- - src_cube : cube - The source cube that was vertically interpolated. - data : array - The payload resulting from interpolating the source cube - over the specified times. - times : array - The array of times. - - Returns - ------- - cube - .. note:: - If there is only one level of interpolation, the resultant cube - will be collapsed over the associated vertical dimension, and a - scalar vertical coordinate will be added. - """ - # Get the source cube vertical coordinate and associated dimension. - src_times = src_cube.coord('time') - t_dim, = src_cube.coord_dims(src_times) - - if data.shape[t_dim] != len(times): - emsg = ('Mismatch between data and times for data dimension {!r}, ' - 'got data shape {!r} with times shape {!r}.') - raise ValueError(emsg.format(t_dim, data.shape, times.shape)) - - # Construct the resultant cube with the interpolated data - # and the source cube metadata. - kwargs = deepcopy(src_cube.metadata)._asdict() - result = iris.cube.Cube(data, **kwargs) - - # Add the appropriate coordinates to the cube, excluding - # any coordinates that span the z-dimension of interpolation. - for coord in src_cube.dim_coords: - [dim] = src_cube.coord_dims(coord) - if dim != t_dim: - result.add_dim_coord(coord.copy(), dim) - - for coord in src_cube.aux_coords: - dims = src_cube.coord_dims(coord) - if t_dim not in dims: - result.add_aux_coord(coord.copy(), dims) - - for coord in src_cube.derived_coords: - dims = src_cube.coord_dims(coord) - if t_dim not in dims: - result.add_aux_coord(coord.copy(), dims) - - # Construct the new vertical coordinate for the interpolated - # z-dimension, using the associated source coordinate metadata. - metadata = src_times.metadata - - kwargs = { - 'standard_name': metadata.standard_name, - 'long_name': metadata.long_name, - 'var_name': metadata.var_name, - 'units': metadata.units, - 'attributes': metadata.attributes, - 'coord_system': metadata.coord_system, - 'climatological': metadata.climatological, - } - - try: - coord = iris.coords.DimCoord(times, **kwargs) - result.add_dim_coord(coord, t_dim) - except ValueError: - coord = iris.coords.AuxCoord(times, **kwargs) - result.add_aux_coord(coord, t_dim) - - return result - - def calculate_volume(cube): """Calculate volume from a cube. @@ -162,7 +77,7 @@ def calculate_volume(cube): # #### # Calculate grid volume: - area = iris.analysis.cartography.area_weights(cube) + area = da.array(iris.analysis.cartography.area_weights(cube)) if thickness.ndim == 1 and z_dim == 1: grid_volume = area * thickness[None, :, None, None] if thickness.ndim == 4 and z_dim == 1: @@ -198,9 +113,8 @@ def volume_statistics(cube, operator): # TODO: Test sigma coordinates. # TODO: Add other operations. - # #### - # Load z coordinate field and figure out which dim is which. - t_dim = cube.coord_dims('time')[0] + if operator != 'mean': + raise ValueError(f'Volume operator {operator} not recognised') try: grid_volume = cube.cell_measure('ocean_volume').core_data() @@ -214,65 +128,23 @@ def volume_statistics(cube, operator): if cube.data.shape != grid_volume.shape: raise ValueError('Cube shape ({}) doesn`t match grid volume shape ' - '({})'.format(cube.data.shape, grid_volume.shape)) - - # ##### - # Calculate global volume weighted average - result = [] - # ##### - # iterate over time and z-coordinate dimensions. - for time_itr in range(cube.shape[t_dim]): - # #### - # create empty output arrays - column = [] - depth_volume = [] + f'({cube.shape, grid_volume.shape})') + + + yx_dims = cube.coord_dims(cube.coord(axis='Y')) + if len(yx_dims) == 1: + yx_dims += cube.coord_dims(cube.coord(axis='X')) + column = cube.collapsed( + [cube.coord(axis='Y'), cube.coord(axis='X')], + iris.analysis.MEAN, + weights=grid_volume) + layer_volume = da.sum(grid_volume, axis=yx_dims) + result = column.collapsed( + cube.coord(axis='Z'), + iris.analysis.MEAN, + weights=layer_volume) - # #### - # iterate over time and z-coordinate dimensions. - for z_itr in range(cube.shape[1]): - # #### - # Calculate weighted mean for this time and layer - if operator == 'mean': - total = cube[time_itr, z_itr].collapsed( - [cube.coord(axis='z'), 'longitude', 'latitude'], - iris.analysis.MEAN, - weights=grid_volume[time_itr, z_itr]).data - else: - raise ValueError('Volume operator ({}) not ' - 'recognised.'.format(operator)) - column.append(total) - - try: - layer_vol = np.ma.masked_where(cube[time_itr, z_itr].data.mask, - grid_volume[time_itr, - z_itr]).sum() - - except AttributeError: - # #### - # No mask in the cube data. - layer_vol = grid_volume.sum() - depth_volume.append(layer_vol) - # #### - # Calculate weighted mean over the water volumn - column = np.ma.array(column) - depth_volume = np.ma.array(depth_volume) - result.append(np.ma.average(column, weights=depth_volume)) - - # #### - # Send time series and dummy cube to cube creating tool. - times = np.array(cube.coord('time').points.astype(float)) - result = np.ma.array(result) - - # ##### - # Create a small dummy output array for the output cube - if operator == 'mean': - src_cube = cube[:2, :2].collapsed( - [cube.coord(axis='z'), 'longitude', 'latitude'], - iris.analysis.MEAN, - weights=grid_volume[:2, :2], - ) - - return _create_cube_time(src_cube, result, times) + return result def depth_integration(cube): From e5006aad78eb307a7c5d352072f6457a5137ffbb Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 15 Mar 2022 16:01:23 +0100 Subject: [PATCH 02/12] Fix flake --- esmvalcore/preprocessor/_volume.py | 1 - 1 file changed, 1 deletion(-) diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py index 0a8157048d..c49fe68917 100644 --- a/esmvalcore/preprocessor/_volume.py +++ b/esmvalcore/preprocessor/_volume.py @@ -130,7 +130,6 @@ def volume_statistics(cube, operator): raise ValueError('Cube shape ({}) doesn`t match grid volume shape ' f'({cube.shape, grid_volume.shape})') - yx_dims = cube.coord_dims(cube.coord(axis='Y')) if len(yx_dims) == 1: yx_dims += cube.coord_dims(cube.coord(axis='X')) From 8cc49660b3bf78847193486664b96d3f0e1c93e5 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 16 Mar 2022 18:11:32 +0100 Subject: [PATCH 03/12] Add mask --- esmvalcore/preprocessor/_volume.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py index c49fe68917..a93b4acc42 100644 --- a/esmvalcore/preprocessor/_volume.py +++ b/esmvalcore/preprocessor/_volume.py @@ -137,11 +137,14 @@ def volume_statistics(cube, operator): [cube.coord(axis='Y'), cube.coord(axis='X')], iris.analysis.MEAN, weights=grid_volume) - layer_volume = da.sum(grid_volume, axis=yx_dims) + masked_volume = da.ma.masked_where( + da.ma.getmaskarray(cube.lazy_data), + grid_volume) + depth_volume = da.sum(masked_volume, axis=yx_dims) result = column.collapsed( cube.coord(axis='Z'), iris.analysis.MEAN, - weights=layer_volume) + weights=depth_volume) return result From 3162d55e4b2b47a3779330e08f5a6c368abbf54e Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 16 Mar 2022 18:15:38 +0100 Subject: [PATCH 04/12] Fix --- esmvalcore/preprocessor/_volume.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py index a93b4acc42..134c68bd12 100644 --- a/esmvalcore/preprocessor/_volume.py +++ b/esmvalcore/preprocessor/_volume.py @@ -138,7 +138,7 @@ def volume_statistics(cube, operator): iris.analysis.MEAN, weights=grid_volume) masked_volume = da.ma.masked_where( - da.ma.getmaskarray(cube.lazy_data), + da.ma.getmaskarray(cube.lazy_data()), grid_volume) depth_volume = da.sum(masked_volume, axis=yx_dims) result = column.collapsed( From 7b31b439f0c1407fd17ed20b5e347cbc64f1fb35 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 23 Mar 2022 10:07:45 +0100 Subject: [PATCH 05/12] Improve test coverage --- esmvalcore/preprocessor/_volume.py | 2 +- tests/unit/preprocessor/_volume/test_volume.py | 10 +++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py index 134c68bd12..a21b41602d 100644 --- a/esmvalcore/preprocessor/_volume.py +++ b/esmvalcore/preprocessor/_volume.py @@ -114,7 +114,7 @@ def volume_statistics(cube, operator): # TODO: Add other operations. if operator != 'mean': - raise ValueError(f'Volume operator {operator} not recognised') + raise ValueError(f'Volume operator {operator} not recognised.') try: grid_volume = cube.cell_measure('ocean_volume').core_data() diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py index 82a755660d..c9393dd8f5 100644 --- a/tests/unit/preprocessor/_volume/test_volume.py +++ b/tests/unit/preprocessor/_volume/test_volume.py @@ -20,7 +20,7 @@ class Test(tests.Test): def setUp(self): """Prepare tests""" - coord_sys = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS) + coord_sys = iris.coord_systems.GeogCS(62563658) data1 = np.ones((3, 2, 2)) data2 = np.ma.ones((2, 3, 2, 2)) data3 = np.ma.ones((4, 3, 2, 2)) @@ -156,6 +156,14 @@ def test_volume_statistics_masked_timestep(self): expected = np.ma.array([1., 1], mask=[True, False]) self.assert_array_equal(result.data, expected) + def test_volume_statistics_wrong_operator(self): + with self.assertRaises(ValueError) as err: + volume_statistics(self.grid_4d, 'wrong') + self.assertEqual( + 'Volume operator wrong not recognised.', + str(err.exception)) + + def test_depth_integration_1d(self): """Test to take the depth integration of a 3 layer cube.""" result = depth_integration(self.grid_3d[:, 0, 0]) From 4db4cff0c2b9fb67812c14cf5cdfee43196b725a Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 23 Mar 2022 10:16:28 +0100 Subject: [PATCH 06/12] Fix flake --- tests/unit/preprocessor/_volume/test_volume.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py index c9393dd8f5..2b516ea688 100644 --- a/tests/unit/preprocessor/_volume/test_volume.py +++ b/tests/unit/preprocessor/_volume/test_volume.py @@ -163,7 +163,6 @@ def test_volume_statistics_wrong_operator(self): 'Volume operator wrong not recognised.', str(err.exception)) - def test_depth_integration_1d(self): """Test to take the depth integration of a 3 layer cube.""" result = depth_integration(self.grid_3d[:, 0, 0]) From 83264efa3d56cc2e57eb2cf7700d5080d48586a7 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Wed, 23 Mar 2022 10:34:06 +0100 Subject: [PATCH 07/12] Fix typo --- tests/unit/preprocessor/_volume/test_volume.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py index 2b516ea688..81f6281b7f 100644 --- a/tests/unit/preprocessor/_volume/test_volume.py +++ b/tests/unit/preprocessor/_volume/test_volume.py @@ -20,7 +20,7 @@ class Test(tests.Test): def setUp(self): """Prepare tests""" - coord_sys = iris.coord_systems.GeogCS(62563658) + coord_sys = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS) data1 = np.ones((3, 2, 2)) data2 = np.ma.ones((2, 3, 2, 2)) data3 = np.ma.ones((4, 3, 2, 2)) From 8a1f6d1b80269d6e977394907393cd5c8ebb760f Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 31 May 2022 14:26:24 +0200 Subject: [PATCH 08/12] Add test --- .../unit/preprocessor/_volume/test_volume.py | 41 +++++++++++++++---- 1 file changed, 33 insertions(+), 8 deletions(-) diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py index 81f6281b7f..b219f6661b 100644 --- a/tests/unit/preprocessor/_volume/test_volume.py +++ b/tests/unit/preprocessor/_volume/test_volume.py @@ -7,12 +7,12 @@ from cf_units import Unit import tests -from esmvalcore.preprocessor._volume import (volume_statistics, +from esmvalcore.preprocessor._volume import (calculate_volume, depth_integration, extract_trajectory, extract_transect, extract_volume, - calculate_volume) + volume_statistics) class Test(tests.Test): @@ -109,6 +109,22 @@ def test_volume_statistics(self): expected = np.ma.array([1., 1.], mask=False) self.assert_array_equal(result.data, expected) + def test_volume_statistics_weights(self): + data = np.arange(1, 25).reshape(2, 3, 2, 2) + self.grid_4d.data = data + measure = iris.coords.CellMeasure( + data, + standard_name='ocean_volume', + units='m3', + measure='volume' + ) + self.grid_4d.add_cell_measure(measure, range(0, measure.ndim)) + result = volume_statistics(self.grid_4d, 'mean') + expected = np.ma.array( + [8.333333333333334, 19.144144144144143], + mask=[False, False]) + self.assert_array_equal(result.data, expected) + def test_volume_statistics_cell_measure(self): """ Test to take the volume weighted average of a (2,3,2,2) cube. @@ -156,12 +172,21 @@ def test_volume_statistics_masked_timestep(self): expected = np.ma.array([1., 1], mask=[True, False]) self.assert_array_equal(result.data, expected) - def test_volume_statistics_wrong_operator(self): - with self.assertRaises(ValueError) as err: - volume_statistics(self.grid_4d, 'wrong') - self.assertEqual( - 'Volume operator wrong not recognised.', - str(err.exception)) + def test_volume_statistics_weights(self): + data = np.ma.arange(1, 25).reshape(2, 3, 2, 2) + self.grid_4d.data = data + measure = iris.coords.CellMeasure( + data, + standard_name='ocean_volume', + units='m3', + measure='volume' + ) + self.grid_4d.add_cell_measure(measure, range(0, measure.ndim)) + result = volume_statistics(self.grid_4d, 'mean') + expected = np.ma.array( + [8.333333333333334, 19.144144144144143], + mask=[False, False]) + self.assert_array_equal(result.data, expected) def test_depth_integration_1d(self): """Test to take the depth integration of a 3 layer cube.""" From 0f157474d3099b0220adc95b5ca22a739da2ea4e Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 31 May 2022 14:37:27 +0200 Subject: [PATCH 09/12] Remove duplicated test --- .../unit/preprocessor/_volume/test_volume.py | 62 +++++++------------ 1 file changed, 21 insertions(+), 41 deletions(-) diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py index b219f6661b..cea010cbcf 100644 --- a/tests/unit/preprocessor/_volume/test_volume.py +++ b/tests/unit/preprocessor/_volume/test_volume.py @@ -7,19 +7,21 @@ from cf_units import Unit import tests -from esmvalcore.preprocessor._volume import (calculate_volume, - depth_integration, - extract_trajectory, - extract_transect, - extract_volume, - volume_statistics) +from esmvalcore.preprocessor._volume import ( + calculate_volume, + depth_integration, + extract_trajectory, + extract_transect, + extract_volume, + volume_statistics, +) class Test(tests.Test): - """Test class for _volume_pp""" + """Test class for _volume_pp.""" def setUp(self): - """Prepare tests""" + """Prepare tests.""" coord_sys = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS) data1 = np.ones((3, 2, 2)) data2 = np.ma.ones((2, 3, 2, 2)) @@ -86,9 +88,8 @@ def test_extract_volume(self): self.assert_array_equal(result.data, expected) def test_extract_volume_mean(self): - """ - Test to extract the top two layers and compute the - weighted average of a cube.""" + """Test to extract the top two layers and compute the weighted average + of a cube.""" grid_volume = calculate_volume(self.grid_4d) measure = iris.coords.CellMeasure( grid_volume, @@ -109,25 +110,9 @@ def test_volume_statistics(self): expected = np.ma.array([1., 1.], mask=False) self.assert_array_equal(result.data, expected) - def test_volume_statistics_weights(self): - data = np.arange(1, 25).reshape(2, 3, 2, 2) - self.grid_4d.data = data - measure = iris.coords.CellMeasure( - data, - standard_name='ocean_volume', - units='m3', - measure='volume' - ) - self.grid_4d.add_cell_measure(measure, range(0, measure.ndim)) - result = volume_statistics(self.grid_4d, 'mean') - expected = np.ma.array( - [8.333333333333334, 19.144144144144143], - mask=[False, False]) - self.assert_array_equal(result.data, expected) - def test_volume_statistics_cell_measure(self): - """ - Test to take the volume weighted average of a (2,3,2,2) cube. + """Test to take the volume weighted average of a (2,3,2,2) cube. + The volume measure is pre-loaded in the cube. """ grid_volume = calculate_volume(self.grid_4d) @@ -142,31 +127,26 @@ def test_volume_statistics_cell_measure(self): self.assert_array_equal(result.data, expected) def test_volume_statistics_long(self): - """ - Test to take the volume weighted average of a (4,3,2,2) cube. + """Test to take the volume weighted average of a (4,3,2,2) cube. - This extra time is needed, as the volume average calculation uses - different methods for small and large cubes. + This extra time is needed, as the volume average calculation + uses different methods for small and large cubes. """ result = volume_statistics(self.grid_4d_2, 'mean') expected = np.ma.array([1., 1., 1., 1.], mask=False) self.assert_array_equal(result.data, expected) def test_volume_statistics_masked_level(self): - """ - Test to take the volume weighted average of a (2,3,2,2) cube - where the last depth level is fully masked. - """ + """Test to take the volume weighted average of a (2,3,2,2) cube where + the last depth level is fully masked.""" self.grid_4d.data[:, -1, :, :] = np.ma.masked_all((2, 2, 2)) result = volume_statistics(self.grid_4d, 'mean') expected = np.ma.array([1., 1.], mask=False) self.assert_array_equal(result.data, expected) def test_volume_statistics_masked_timestep(self): - """ - Test to take the volume weighted average of a (2,3,2,2) cube - where the first timestep is fully masked. - """ + """Test to take the volume weighted average of a (2,3,2,2) cube where + the first timestep is fully masked.""" self.grid_4d.data[0, :, :, :] = np.ma.masked_all((3, 2, 2)) result = volume_statistics(self.grid_4d, 'mean') expected = np.ma.array([1., 1], mask=[True, False]) From bf9eda7eb8884fda441274f392d19242c9cd0772 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 31 May 2022 14:47:05 +0200 Subject: [PATCH 10/12] Add docstring --- tests/unit/preprocessor/_volume/test_volume.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py index cea010cbcf..5aed3a9c6b 100644 --- a/tests/unit/preprocessor/_volume/test_volume.py +++ b/tests/unit/preprocessor/_volume/test_volume.py @@ -153,6 +153,10 @@ def test_volume_statistics_masked_timestep(self): self.assert_array_equal(result.data, expected) def test_volume_statistics_weights(self): + """Test to take the volume weighted average of a (2,3,2,2) cube. + + The data and weights are not defined as arrays of ones. + """ data = np.ma.arange(1, 25).reshape(2, 3, 2, 2) self.grid_4d.data = data measure = iris.coords.CellMeasure( From 3a2ee991c22498797540c3e2c55fbff9b9b7c483 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Tue, 31 May 2022 15:02:14 +0200 Subject: [PATCH 11/12] Add back accidentally removed test --- tests/unit/preprocessor/_volume/test_volume.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/unit/preprocessor/_volume/test_volume.py b/tests/unit/preprocessor/_volume/test_volume.py index 5aed3a9c6b..4c59fa82f8 100644 --- a/tests/unit/preprocessor/_volume/test_volume.py +++ b/tests/unit/preprocessor/_volume/test_volume.py @@ -172,6 +172,13 @@ def test_volume_statistics_weights(self): mask=[False, False]) self.assert_array_equal(result.data, expected) + def test_volume_statistics_wrong_operator(self): + with self.assertRaises(ValueError) as err: + volume_statistics(self.grid_4d, 'wrong') + self.assertEqual( + 'Volume operator wrong not recognised.', + str(err.exception)) + def test_depth_integration_1d(self): """Test to take the depth integration of a 3 layer cube.""" result = depth_integration(self.grid_3d[:, 0, 0]) From 8d8963377a8988263a82e1f2ba1e7661b64b3b67 Mon Sep 17 00:00:00 2001 From: sloosvel Date: Fri, 17 Jun 2022 16:12:20 +0200 Subject: [PATCH 12/12] Collapse all axis at once --- esmvalcore/preprocessor/_volume.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py index a21b41602d..1b6f140598 100644 --- a/esmvalcore/preprocessor/_volume.py +++ b/esmvalcore/preprocessor/_volume.py @@ -130,21 +130,13 @@ def volume_statistics(cube, operator): raise ValueError('Cube shape ({}) doesn`t match grid volume shape ' f'({cube.shape, grid_volume.shape})') - yx_dims = cube.coord_dims(cube.coord(axis='Y')) - if len(yx_dims) == 1: - yx_dims += cube.coord_dims(cube.coord(axis='X')) - column = cube.collapsed( - [cube.coord(axis='Y'), cube.coord(axis='X')], - iris.analysis.MEAN, - weights=grid_volume) masked_volume = da.ma.masked_where( da.ma.getmaskarray(cube.lazy_data()), grid_volume) - depth_volume = da.sum(masked_volume, axis=yx_dims) - result = column.collapsed( - cube.coord(axis='Z'), + result = cube.collapsed( + [cube.coord(axis='Z'), cube.coord(axis='Y'), cube.coord(axis='X')], iris.analysis.MEAN, - weights=depth_volume) + weights=masked_volume) return result