Merge pull request #618 from ESMValGroup/change_volume_stats_to_masked_result

Change volume stats to handle and output masked array result

Author: mattiarighi

esmvalcore/preprocessor/_volume.py

@@ -78,13 +78,10 @@ def _create_cube_time(src_cube, data, 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')
@@ -268,12 +265,14 @@ def volume_statistics(
             depth_volume.append(layer_vol)
         # ####
         # Calculate weighted mean over the water volumn
-        result.append(np.average(column, weights=depth_volume))
+        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.array(result)
+    result = np.ma.array(result)
 
     # #####
     # Create a small dummy output array for the output cube

tests/unit/preprocessor/_volume/test_volume.py

@@ -86,7 +86,7 @@ def test_extract_volume(self):
     def test_volume_statistics(self):
         """Test to take the volume weighted average of a (2,3,2,2) cube."""
         result = volume_statistics(self.grid_4d, 'mean')
-        expected = np.array([1., 1.])
+        expected = np.ma.array([1., 1.], mask=False)
         self.assert_array_equal(result.data, expected)
 
     def test_volume_statistics_long(self):
@@ -97,7 +97,27 @@ def test_volume_statistics_long(self):
         different methods for small and large cubes.
         """
         result = volume_statistics(self.grid_4d_2, 'mean')
-        expected = np.array([1., 1., 1., 1.])
+        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.
+        """
+        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.
+        """
+        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])
         self.assert_array_equal(result.data, expected)
 
     def test_depth_integration_1d(self):