Repository: esmvalgroup/esmvalcore
Issue: Derivation of co2s is not correct

Right now, the calculation of `co2s` (CO2 mole fraction at surface level) is scientifically wrong. Instead of calculating the mole fraction at the surface level, it calculates the mole fraction at the lowest unmasked level. For the correct calculation, it is necessary to interpolate `co2` on the surface air pressure (`ps`). Thanks @bettina-gier for raising this and @axel-lauer for the advice!


Changed files:
tests/unit/preprocessor/_derive/test_co2s.py, esmvalcore/preprocessor/_derive/co2s.py

diff --git a/esmvalcore/preprocessor/_derive/co2s.py b/esmvalcore/preprocessor/_derive/co2s.py
index 8a1dd0949..42623206f 100644
--- a/esmvalcore/preprocessor/_derive/co2s.py
+++ b/esmvalcore/preprocessor/_derive/co2s.py
@@ -1,39 +1,96 @@
 """Derivation of variable ``co2s``."""
 import dask.array as da
 import iris
+import numpy as np
+import stratify
 
 from ._baseclass import DerivedVariableBase
 
 
+def _get_first_unmasked_data(array, axis):
+    """Get first unmasked value of an array along an axis."""
+    mask = da.ma.getmaskarray(array)
+    numerical_mask = da.where(mask, -1.0, 1.0)
+    indices_first_positive = da.argmax(numerical_mask, axis=axis)
+    indices = da.meshgrid(
+        *[da.arange(array.shape[i]) for i in range(array.ndim) if i != axis],
+        indexing='ij')
+    indices.insert(axis, indices_first_positive)
+    first_unmasked_data = np.array(array)[tuple(indices)]
+    return first_unmasked_data
+
+
 class DerivedVariable(DerivedVariableBase):
-    """Derivation of variable ``co2s``."""
+    """Derivation of variable ``co2s``.
+
+    Use linear interpolation/extrapolation and surface air pressure to
+    calculate CO2 mole fraction at surface.
+
+    Note
+    ----
+    In some cases, ``co2`` data is masked. In these cases, the masked values
+    correspond to values where the pressure level is higher than the surface
+    air pressure (e.g. the 1000 hPa level for grid cells with high elevation).
+    To obtain an unmasked ``co2s`` field, it is necessary to fill these masked
+    values accordingly, i.e. with the lowest unmasked value for each grid cell.
+
+    """
 
     @staticmethod
     def required(project):
         """Declare the variables needed for derivation."""
-        required = [{'short_name': 'co2'}]
+        required = [{'short_name': 'co2'}, {'short_name': 'ps'}]
         return required
 
     @staticmethod
     def calculate(cubes):
         """Compute mole fraction of CO2 at surface."""
-        cube = cubes.extract_strict(
+        co2_cube = cubes.extract_strict(
             iris.Constraint(name='mole_fraction_of_carbon_dioxide_in_air'))
-        mask = da.ma.getmaskarray(cube.core_data())
-        if not mask.any():
-            cube = cube[:, 0, :, :]
-        else:
-            numerical_mask = da.where(mask, -1.0, 1.0)
-            indices_first_positive = da.argmax(numerical_mask, axis=1)
-            indices = da.meshgrid(
-                da.arange(cube.shape[0]),
-                da.arange(cube.shape[2]),
-                da.arange(cube.shape[3]),
-                indexing='ij',
-            )
-            indices.insert(1, indices_first_positive)
-            surface_data = cube.data[tuple(indices)]
-            cube = cube[:, 0, :, :]
-            cube.data = surface_data
-        cube.convert_units('1e-6')
-        return cube
+        ps_cube = cubes.extract_strict(
+            iris.Constraint(name='surface_air_pressure'))
+
+        # Fill masked data if necessary (interpolation fails with masked data)
+        (z_axis,) = co2_cube.coord_dims(co2_cube.coord(axis='Z',
+                                                       dim_coords=True))
+        mask = da.ma.getmaskarray(co2_cube.core_data())
+        if mask.any():
+            first_unmasked_data = _get_first_unmasked_data(
+                co2_cube.core_data(), axis=z_axis)
+            dim_map = [dim for dim in range(co2_cube.ndim) if dim != z_axis]
+            first_unmasked_data = iris.util.broadcast_to_shape(
+                first_unmasked_data, co2_cube.shape, dim_map)
+            co2_cube.data = da.where(mask, first_unmasked_data,
+                                     co2_cube.core_data())
+
+        # Interpolation (not supported for dask arrays)
+        air_pressure_coord = co2_cube.coord('air_pressure')
+        original_levels = iris.util.broadcast_to_shape(
+            air_pressure_coord.points, co2_cube.shape,
+            co2_cube.coord_dims(air_pressure_coord))
+        target_levels = np.expand_dims(ps_cube.data, axis=z_axis)
+        co2s_data = stratify.interpolate(
+            target_levels,
+            original_levels,
+            co2_cube.data,
+            axis=z_axis,
+            interpolation='linear',
+            extrapolation='linear',
+        )
+        co2s_data = np.squeeze(co2s_data, axis=z_axis)
+
+        # Construct co2s cube
+        indices = [slice(None)] * co2_cube.ndim
+        indices[z_axis] = 0
+        co2s_cube = co2_cube[tuple(indices)]
+        co2s_cube.data = co2s_data
+        if co2s_cube.coords('air_pressure'):
+            co2s_cube.remove_coord('air_pressure')
+        ps_coord = iris.coords.AuxCoord(ps_cube.data,
+                                        var_name='plev',
+                                        standard_name='air_pressure',
+                                        long_name='pressure',
+                                        units=ps_cube.units)
+        co2s_cube.add_aux_coord(ps_coord, np.arange(co2s_cube.ndim))
+        co2s_cube.convert_units('1e-6')
+        return co2s_cube
diff --git a/tests/unit/preprocessor/_derive/test_co2s.py b/tests/unit/preprocessor/_derive/test_co2s.py
index ad05aebd1..3fd364edd 100644
--- a/tests/unit/preprocessor/_derive/test_co2s.py
+++ b/tests/unit/preprocessor/_derive/test_co2s.py
@@ -7,28 +7,51 @@ import pytest
 import esmvalcore.preprocessor._derive.co2s as co2s
 
 
-def get_coord_spec():
+def get_coord_spec(include_plev=True):
     """Coordinate specs for cubes."""
     time_coord = iris.coords.DimCoord([0], var_name='time',
                                       standard_name='time',
                                       units='days since 0000-01-01 00:00:00')
-    plev_coord = iris.coords.DimCoord([123456.0, 50000.0, 1000.0],
-                                      var_name='plev',
-                                      standard_name='air_pressure', units='Pa')
     lat_coord = iris.coords.DimCoord([0.0, 1.0], var_name='latitude',
                                      standard_name='latitude', units='degrees')
     lon_coord = iris.coords.DimCoord([0.0, 1.0], var_name='longitude',
                                      standard_name='longitude',
                                      units='degrees')
-    coord_spec = [
-        (time_coord, 0),
-        (plev_coord, 1),
-        (lat_coord, 2),
-        (lon_coord, 3),
-    ]
+    if include_plev:
+        plev_coord = iris.coords.DimCoord([100000.0, 90000.0, 50000.0],
+                                          var_name='plev',
+                                          standard_name='air_pressure',
+                                          units='Pa')
+        coord_spec = [
+            (time_coord, 0),
+            (plev_coord, 1),
+            (lat_coord, 2),
+            (lon_coord, 3),
+        ]
+    else:
+        coord_spec = [
+            (time_coord, 0),
+            (lat_coord, 1),
+            (lon_coord, 2),
+        ]
     return coord_spec
 
 
+def get_ps_cube():
+    """Surface air pressure cube."""
+    ps_data = [[[105000.0, 50000.0],
+                [95000.0, 60000.0]]]
+    coord_spec = get_coord_spec(include_plev=False)
+    cube = iris.cube.Cube(
+        ps_data,
+        var_name='ps',
+        standard_name='surface_air_pressure',
+        units='Pa',
+        dim_coords_and_dims=coord_spec,
+    )
+    return cube
+
+
 @pytest.fixture
 def masked_cubes():
     """Masked CO2 cube."""
@@ -39,14 +62,15 @@ def masked_cubes():
                                     [80.0, -1.0]],
                                    [[100.0, 50.0],
                                     [30.0, 10.0]]]], 0.0)
-    cube = iris.cube.Cube(
+    co2_cube = iris.cube.Cube(
         co2_data,
         var_name='co2',
         standard_name='mole_fraction_of_carbon_dioxide_in_air',
         units='1e-6',
         dim_coords_and_dims=coord_spec,
     )
-    return iris.cube.CubeList([cube])
+    ps_cube = get_ps_cube()
+    return iris.cube.CubeList([co2_cube, ps_cube])
 
 
 @pytest.fixture
@@ -59,14 +83,15 @@ def unmasked_cubes():
                            [70.0, 5.0]],
                           [[100.0, 50.0],
                            [30.0, 1.0]]]])
-    cube = iris.cube.Cube(
+    co2_cube = iris.cube.Cube(
         co2_data,
         var_name='co2',
         standard_name='mole_fraction_of_carbon_dioxide_in_air',
         units='1e-8',
         dim_coords_and_dims=coord_spec,
     )
-    return iris.cube.CubeList([cube])
+    ps_cube = get_ps_cube()
+    return iris.cube.CubeList([co2_cube, ps_cube])
 
 
 def test_co2_calculate_masked_cubes(masked_cubes):
@@ -75,11 +100,16 @@ def test_co2_calculate_masked_cubes(masked_cubes):
     out_cube = derived_var.calculate(masked_cubes)
     assert not np.ma.is_masked(out_cube.data)
     np.testing.assert_allclose(out_cube.data,
-                               [[[170.0, 100.0],
+                               [[[180.0, 50.0],
                                  [80.0, 10.0]]])
     assert out_cube.units == '1e-6'
-    np.testing.assert_allclose(out_cube.coord('air_pressure').points,
-                               123456.0)
+    plev_coord = out_cube.coord('air_pressure')
+    assert plev_coord.var_name == 'plev'
+    assert plev_coord.standard_name == 'air_pressure'
+    assert plev_coord.long_name == 'pressure'
+    assert plev_coord.units == 'Pa'
+    np.testing.assert_allclose(plev_coord.points,
+                               [[[105000.0, 50000.0], [95000.0, 60000.0]]])
 
 
 def test_co2_calculate_unmasked_cubes(unmasked_cubes):
@@ -88,8 +118,13 @@ def test_co2_calculate_unmasked_cubes(unmasked_cubes):
     out_cube = derived_var.calculate(unmasked_cubes)
     assert not np.ma.is_masked(out_cube.data)
     np.testing.assert_allclose(out_cube.data,
-                               [[[2.0, 1.0],
-                                 [0.8, 0.09]]])
+                               [[[2.25, 0.50],
+                                 [0.75, 0.02]]])
     assert out_cube.units == '1e-6'
-    np.testing.assert_allclose(out_cube.coord('air_pressure').points,
-                               123456.0)
+    plev_coord = out_cube.coord('air_pressure')
+    assert plev_coord.var_name == 'plev'
+    assert plev_coord.standard_name == 'air_pressure'
+    assert plev_coord.long_name == 'pressure'
+    assert plev_coord.units == 'Pa'
+    np.testing.assert_allclose(plev_coord.points,
+                               [[[105000.0, 50000.0], [95000.0, 60000.0]]])