Broadcast properly cell_measures when using extract_shape with `d…

…ecomposed: True` (#2348) Co-authored-by: Manuel Schlund <32543114+schlunma@users.noreply.github.com>
ESMValGroup · Apr 3, 2024 · fcd5455 · fcd5455
1 parent c8555d4
commit fcd5455
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 0 deletions.
diff --git a/esmvalcore/preprocessor/_area.py b/esmvalcore/preprocessor/_area.py
@@ -938,6 +938,23 @@ def _mask_cube(cube: Cube, masks: dict[str, np.ndarray]) -> Cube:
     result = fix_coordinate_ordering(cubelist.merge_cube())
     if cube.cell_measures():
         for measure in cube.cell_measures():
+            # Cell measures that are time-dependent, with 4 dimension and
+            # an original shape of (time, depth, lat, lon), need to be
+            # broadcasted to the cube with 5 dimensions and shape
+            # (time, shape_id, depth, lat, lon)
+            if measure.ndim > 3 and result.ndim > 4:
+                data = measure.core_data()
+                data = da.expand_dims(data, axis=(1,))
+                data = da.broadcast_to(data, result.shape)
+                measure = iris.coords.CellMeasure(
+                    data,
+                    standard_name=measure.standard_name,
+                    long_name=measure.long_name,
+                    units=measure.units,
+                    measure=measure.measure,
+                    var_name=measure.var_name,
+                    attributes=measure.attributes,
+                )
             add_cell_measure(result, measure, measure.measure)
     if cube.ancillary_variables():
         for ancillary_variable in cube.ancillary_variables():

diff --git a/tests/unit/preprocessor/_area/test_area.py b/tests/unit/preprocessor/_area/test_area.py
@@ -1399,5 +1399,51 @@ def test_meridional_statistics_invalid_norm_fail(make_testcube):
         meridional_statistics(make_testcube, 'sum', normalize='x')
 
 
+def test_time_dependent_volcello():
+    coord_sys = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS)
+    data = np.ma.ones((2, 3, 2, 2))
+
+    time = iris.coords.DimCoord([15, 45],
+                                standard_name='time',
+                                bounds=[[1., 30.], [30., 60.]],
+                                units=Unit('days since 1950-01-01',
+                                calendar='gregorian'))
+
+    zcoord = iris.coords.DimCoord([0.5, 5., 50.],
+                                  long_name='zcoord',
+                                  bounds=[[0., 2.5], [2.5, 25.],
+                                          [25., 250.]],
+                                  units='m',
+                                  attributes={'positive': 'down'})
+    lons = iris.coords.DimCoord([1.5, 2.5],
+                                standard_name='longitude',
+                                bounds=[[1., 2.], [2., 3.]],
+                                units='degrees_east',
+                                coord_system=coord_sys)
+    lats = iris.coords.DimCoord([1.5, 2.5],
+                                standard_name='latitude',
+                                bounds=[[1., 2.], [2., 3.]],
+                                units='degrees_north',
+                                coord_system=coord_sys)
+    coords_spec4 = [(time, 0), (zcoord, 1), (lats, 2), (lons, 3)]
+    cube = iris.cube.Cube(data, dim_coords_and_dims=coords_spec4)
+    volcello = iris.coords.CellMeasure(
+        data,
+        standard_name='ocean_volume',
+        units='m3',
+        measure='volume')
+    cube.add_cell_measure(volcello, range(0, volcello.ndim))
+    cube = extract_shape(
+        cube,
+        'AR6',
+        method='contains',
+        crop=False,
+        decomposed=True,
+        ids={'Acronym': ['EAO', 'WAF']},
+    )
+
+    assert cube.shape == cube.cell_measure('ocean_volume').shape
+
+
 if __name__ == '__main__':
     unittest.main()