Compute volume from cell_area if available

esmvalcore/preprocessor/_volume.py

@@ -14,9 +14,8 @@
 import numpy as np
 from iris.coords import AuxCoord, CellMeasure
 from iris.cube import Cube
-from iris.exceptions import CoordinateMultiDimError
 
-from ._area import compute_area_weights
+from ._area import _try_adding_calculated_cell_area
 from ._shared import (
     get_iris_aggregator,
     get_normalized_cube,
@@ -108,9 +107,9 @@ def calculate_volume(cube: Cube) -> da.core.Array:
 
     Note
     ----
-    This only works if the grid cell areas can be calculated (i.e., latitude
-    and longitude are 1D) and if the depth coordinate is 1D or 4D with first
-    dimension 1.
+    It gets the cell_area from the cube if it is available. If not, it
+    calculates it from the grid. This only works if the grid cell areas can
+    be calculated (i.e., latitude and longitude are 1D).
 
     Parameters
     ----------
@@ -125,33 +124,46 @@ def calculate_volume(cube: Cube) -> da.core.Array:
     """
     # Load depth field and figure out which dim is which
     depth = cube.coord(axis='z')
-    z_dim = cube.coord_dims(depth)[0]
+    z_dim = cube.coord_dims(depth)
+    # Assert z has length > 0
+    if not z_dim:
+        raise ValueError("Cannot compute volume with length 0 Z-axis")
+
+    # Guess bounds if missing
+    if not depth.has_bounds():
+        depth.guess_bounds()
+    if depth.core_bounds().shape[-1] != 2:
+        raise ValueError(
+            f"Z axis bounds shape found {depth.core_bounds().shape}. "
+            "Bounds should be 2 in the last dimension to compute the "
+            "thickness.")
 
     # Calculate Z-direction thickness
     thickness = depth.core_bounds()[..., 1] - depth.core_bounds()[..., 0]
 
-    # Try to calculate grid cell area
-    try:
-        area = da.array(compute_area_weights(cube))
-    except CoordinateMultiDimError:
-        logger.error(
-            "Supplementary variables are needed to calculate grid cell "
-            "areas for irregular grid of cube %s",
-            cube.summary(shorten=True),
-        )
-        raise
+    # Get or calculate the horizontal areas of the cube
+    has_cell_measure = bool(cube.cell_measures('cell_area'))
+    _try_adding_calculated_cell_area(cube)
+    area = cube.cell_measure('cell_area')
+    area_dim = cube.cell_measure_dims(area)
 
-    # Try to calculate grid cell volume as area * thickness
-    if thickness.ndim == 1 and z_dim == 1:
-        grid_volume = area * thickness[None, :, None, None]
-    elif thickness.ndim == 4 and z_dim == 1:
-        grid_volume = area * thickness[:, :]
-    else:
-        raise ValueError(
-            f"Supplementary variables are needed to calculate grid cell "
-            f"volumes for cubes with {thickness.ndim:d}D depth coordinate, "
-            f"got cube {cube.summary(shorten=True)}"
-        )
+    # Make sure input cube has not been modified
+    if not has_cell_measure:
+        cube.remove_cell_measure('cell_area')
+
+    # Get slices to make area and thickness match with the same number
+    # of dimensions as the cube
+    z_slice = tuple(
+        slice(None) if i in z_dim else None
+        for i in range(cube.ndim)
+    )
+    area_slice = tuple(
+        slice(None) if i in area_dim else None
+        for i in range(cube.ndim)
+    )
+
+    # Calculate grid cell volume as area * thickness
+    grid_volume = area.lazy_data()[area_slice] * thickness[z_slice]
 
     return grid_volume
 
@@ -170,13 +182,17 @@ def _try_adding_calculated_ocean_volume(cube: Cube) -> None:
 
     grid_volume = calculate_volume(cube)
 
+    # add measure only in the dimensions that have lenght > 1
+    data_dims = [i for i, n in enumerate(grid_volume.shape) if n > 1]
+    grid_volume = grid_volume.squeeze()
+
     cell_measure = CellMeasure(
         grid_volume,
         standard_name='ocean_volume',
         units='m3',
         measure='volume',
     )
-    cube.add_cell_measure(cell_measure, np.arange(cube.ndim))
+    cube.add_cell_measure(cell_measure, data_dims)
 
 
 @register_supplementaries(
@@ -198,10 +214,11 @@ def volume_statistics(
     cube:
         Input cube. The input cube should have a
         :class:`iris.coords.CellMeasure` named ``'ocean_volume'``, unless it
-        has regular 1D latitude and longitude coordinates so the cell volumes
-        can be computed by using :func:`iris.analysis.cartography.area_weights`
-        to compute the cell areas and multiplying those by the cell thickness,
-        computed from the bounds of the vertical coordinate.
+        has :class:`iris.coords.CellMeasure` named ``'cell_area'`` or
+        has regular 1D latitude and longitude coordinates so the cell areas
+        can be computed using :func:`iris.analysis.cartography.area_weights`.
+        The volume will be computed from the area multiplied by the
+        thickness, computed from the bounds of the vertical coordinate.
     operator:
         The operation. Used to determine the :class:`iris.analysis.Aggregator`
         object used to calculate the statistics. Currently, only `mean` is
@@ -227,6 +244,20 @@ def volume_statistics(
     # TODO: Add other operations.
     if operator != 'mean':
         raise ValueError(f"Volume operator {operator} not recognised.")
+    # get z, y, x coords
+    z_axis = cube.coord(axis='Z')
+    y_axis = cube.coord(axis='Y')
+    x_axis = cube.coord(axis='X')
+
+    # assert z axis only uses 1 dimension more than x, y axis
+    xy_dims = tuple({*cube.coord_dims(y_axis), *cube.coord_dims(x_axis)})
+    xyz_dims = tuple({*cube.coord_dims(z_axis), *xy_dims})
+    if len(xyz_dims) > len(xy_dims) + 1:
+        raise ValueError(
+            f"X and Y axis coordinates depend on {xy_dims} dimensions, "
+            f"while X, Y, and Z axis depends on {xyz_dims} dimensions. "
+            "This may indicate Z axis depending on other dimension than"
+            "space that could provoke invalid aggregation...")
 
     (agg, agg_kwargs) = get_iris_aggregator(operator, **operator_kwargs)
     agg_kwargs = update_weights_kwargs(
@@ -237,11 +268,7 @@ def volume_statistics(
         _try_adding_calculated_ocean_volume,
     )
 
-    result = cube.collapsed(
-        [cube.coord(axis='Z'), cube.coord(axis='Y'), cube.coord(axis='X')],
-        agg,
-        **agg_kwargs,
-    )
+    result = cube.collapsed([z_axis, y_axis, x_axis], agg, **agg_kwargs)
     if normalize is not None:
         result = get_normalized_cube(cube, result, normalize)
 

tests/unit/preprocessor/_volume/test_volume.py

@@ -55,6 +55,10 @@ def setUp(self):
                                               [25., 250.]],
                                       units='m',
                                       attributes={'positive': 'down'})
+        zcoord_nobounds = iris.coords.DimCoord([0.5, 5., 50.],
+                                               long_name='zcoord',
+                                               units='m',
+                                               attributes={'positive': 'down'})
         zcoord_4d = iris.coords.AuxCoord(
             np.broadcast_to([[[[0.5]], [[5.]], [[50.]]]], (2, 3, 2, 2)),
             long_name='zcoord',
@@ -65,6 +69,18 @@ def setUp(self):
             units='m',
             attributes={'positive': 'down'},
         )
+        zcoord_3d_invalid_bounds = iris.coords.AuxCoord(
+            np.broadcast_to([[[0.5]], [[5.]], [[50.]]], (3, 2, 2)),
+            long_name='zcoord',
+            bounds=np.broadcast_to(
+                [[[[0., 2.5, 2.5, 3.]]],
+                 [[[2.5, 25., 25., 30.]]],
+                 [[[25., 250., 250., 300.]]]],
+                (3, 2, 2, 4),
+            ),
+            units='m',
+            attributes={'positive': 'down'},
+        )
         lons2 = iris.coords.DimCoord([1.5, 2.5],
                                      standard_name='longitude',
                                      bounds=[[1., 2.], [2., 3.]],
@@ -76,6 +92,15 @@ def setUp(self):
                                      units='degrees_north',
                                      coord_system=coord_sys)
 
+        lons2d = iris.coords.AuxCoord([[1.5, 2.5], [1.2, 2.7]],
+                                      standard_name='longitude',
+                                      units='degrees_east',
+                                      coord_system=coord_sys)
+        lats2d = iris.coords.AuxCoord([[1.5, 2.5], [1.2, 2.7]],
+                                      standard_name='latitude',
+                                      units='degrees_north',
+                                      coord_system=coord_sys)
+
         coords_spec3 = [(zcoord, 0), (lats2, 1), (lons2, 2)]
         self.grid_3d = iris.cube.Cube(data1, dim_coords_and_dims=coords_spec3)
 
@@ -100,6 +125,28 @@ def setUp(self):
             units='kg m-3',
         )
 
+        self.grid_4d_znobounds = iris.cube.Cube(
+            data2,
+            dim_coords_and_dims=[
+                (time, 0), (zcoord_nobounds, 1), (lats2, 2), (lons2, 3)
+            ],
+            units='kg m-3',
+        )
+
+        self.grid_4d_irregular = iris.cube.Cube(
+            data2,
+            dim_coords_and_dims=[(time, 0), (zcoord, 1)],
+            aux_coords_and_dims=[(lats2d, (2, 3)), (lons2d, (2, 3))],
+            units='kg m-3',
+        )
+
+        self.grid_invalid_z_bounds = iris.cube.Cube(
+            data2,
+            dim_coords_and_dims=[(time, 0), (lats2, 2), (lons2, 3)],
+            aux_coords_and_dims=[(zcoord_3d_invalid_bounds, (1, 2, 3))],
+            units='kg m-3',
+        )
+
         # allow iris to figure out the axis='z' coordinate
         iris.util.guess_coord_axis(self.grid_3d.coord('zcoord'))
         iris.util.guess_coord_axis(self.grid_4d.coord('zcoord'))
@@ -344,6 +391,20 @@ def test_volume_statistics(self):
         self.assertFalse(self.grid_4d.cell_measures('ocean_volume'))
         self.assertFalse(result.cell_measures('ocean_volume'))
 
+    def test_volume_nolevbounds(self):
+        """Test to take the volume weighted average of a cube with no bounds
+        in the z axis.
+        """
+
+        self.assertFalse(self.grid_4d_znobounds.coord(axis='z').has_bounds())
+        result = volume_statistics(self.grid_4d_znobounds, 'mean')
+
+        expected = np.ma.array([1., 1.], mask=False)
+        self.assert_array_equal(result.data, expected)
+        self.assertEqual(result.units, 'kg m-3')
+        self.assertFalse(self.grid_4d.cell_measures('ocean_volume'))
+        self.assertFalse(result.cell_measures('ocean_volume'))
+
     def test_volume_statistics_cell_measure(self):
         """Test to take the volume weighted average of a (2,3,2,2) cube.
 
@@ -443,49 +504,100 @@ def test_volume_statistics_wrong_operator_fail(self):
                          str(err.exception))
 
     def test_volume_statistics_2d_lat_fail(self):
-        # Create dummy 2D latitude from depth
-        new_lat_coord = self.grid_4d_z.coord('zcoord')[0, 0, :, :]
-        new_lat_coord.rename('latitude')
-        self.grid_4d_z.remove_coord('latitude')
-        self.grid_4d_z.add_aux_coord(new_lat_coord, (2, 3))
         with self.assertRaises(CoordinateMultiDimError):
-            volume_statistics(self.grid_4d_z, 'mean')
+            volume_statistics(self.grid_4d_irregular, 'mean')
+
+    def test_volume_statistics_2d_lat_cellarea(self):
+        measure = iris.coords.CellMeasure(np.arange(1, 5).reshape(2, 2),
+                                          standard_name='cell_area',
+                                          units='m2',
+                                          measure='area')
+        self.grid_4d_irregular.add_cell_measure(measure, (2, 3))
 
-    def test_volume_statistics_4d_depth_fail(self):
-        # Fails because depth coord dims are (0, ...), but must be (1, ...)
+        result = volume_statistics(self.grid_4d_irregular, 'mean')
+        expected = np.ma.array([1., 1.], mask=False)
+        self.assert_array_equal(result.data, expected)
+        self.assertEqual(result.units, 'kg m-3')
+
+        data = np.ma.arange(1, 25).reshape(2, 3, 2, 2)
+        self.grid_4d_irregular.data = data
+
+        result = volume_statistics(self.grid_4d_irregular, 'mean')
+        expected = np.ma.array([10.56, 22.56], mask=False)
+        self.assert_array_equal(result.data, expected)
+        self.assertEqual(result.units, 'kg m-3')
+
+    def test_volume_statistics_invalid_bounds(self):
+        """Test z-axis bounds is not 2 in last dimension"""
+
+        with self.assertRaises(ValueError) as err:
+            volume_statistics(self.grid_invalid_z_bounds, 'mean')
+        self.assertIn(
+            "Z axis bounds shape found (3, 2, 2, 4). Bounds should be "
+            "2 in the last dimension to compute the thickness.",
+            str(err.exception)
+        )
+
+    def test_volume_statistics_z_axis_time_error(self):
+        # Fails because depth z-axis coord depends on time dimensions
+        # which would aggregate also along that dimension
         with self.assertRaises(ValueError) as err:
             volume_statistics(self.grid_4d_z, 'mean')
         self.assertIn(
-            "Supplementary variables are needed to calculate grid cell "
-            "volumes for cubes with 4D depth coordinate, got cube ",
-            str(err.exception),
+            "X and Y axis coordinates depend on (2, 3) dimensions, "
+            "while X, Y, and Z axis depends on (0, 1, 2, 3) dimensions. "
+            "This may indicate Z axis depending on other dimension than"
+            "space that could provoke invalid aggregation...",
+            str(err.exception)
         )
 
-    def test_volume_statistics_2d_depth_fail(self):
+        # Fails because x axis is missing
+        grid_3d_no_x = self.grid_4d_z[..., 0]
+        with self.assertRaises(ValueError) as err:
+            volume_statistics(grid_3d_no_x, 'mean')
+        self.assertIn(
+            "X and Y axis coordinates depend on (2,) dimensions, "
+            "while X, Y, and Z axis depends on (0, 1, 2) dimensions. "
+            "This may indicate Z axis depending on other dimension than"
+            "space that could provoke invalid aggregation...",
+            str(err.exception)
+        )
+
+    def test_volume_statistics_missing_axis(self):
+        # x axis is missing
+        grid_no_x = self.grid_4d[..., 0]
+        volume_statistics(grid_no_x, 'mean')
+
+        # y axis is missing
+        grid_no_y = self.grid_4d[..., 0, :]
+        volume_statistics(grid_no_y, 'mean')
+
+        # z axis is missing
+        grid_no_z = self.grid_4d[:, 0]
+        with self.assertRaises(ValueError) as err:
+            volume_statistics(grid_no_z, 'mean')
+        self.assertIn("Cannot compute volume with length 0 Z-axis",
+                      str(err.exception))
+
+    def test_volume_statistics_2d_depth(self):
         # Create new 2D depth coord
         new_z_coord = self.grid_4d_z.coord('zcoord')[0, :, :, 0]
         self.grid_4d_z.remove_coord('zcoord')
         self.grid_4d_z.add_aux_coord(new_z_coord, (1, 2))
-        with self.assertRaises(ValueError) as err:
-            volume_statistics(self.grid_4d_z, 'mean')
-        self.assertIn(
-            "Supplementary variables are needed to calculate grid cell "
-            "volumes for cubes with 2D depth coordinate, got cube ",
-            str(err.exception),
-        )
+        result = volume_statistics(self.grid_4d, 'mean')
+        expected = np.ma.array([1., 1.], mask=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."""
         result = depth_integration(self.grid_3d[:, 0, 0])
         expected = np.ones((1, 1)) * 250.
-        print(result.data, expected.data)
         self.assert_array_equal(result.data, expected)
 
     def test_depth_integration_3d(self):
         """Test to take the depth integration of a 3 layer cube."""
         result = depth_integration(self.grid_3d)
         expected = np.ones((2, 2)) * 250.
-        print(result.data, expected.data)
         self.assert_array_equal(result.data, expected)
 
     def test_extract_transect_latitude(self):