Add preprocessor to extract surface values from 3D atmospheric variables (#2641)

Co-authored-by: Manuel Schlund <32543114+schlunma@users.noreply.github.com>
Co-authored-by: Manuel Schlund <manuel.schlund@dlr.de>

Author: jlenh

.zenodo.json

@@ -219,6 +219,11 @@
       "affiliation": "ECCC, Canada",
       "name": "Garcia Perdomo, Karen",
       "orcid": "0009-0004-2333-3358"
+    },
+    {
+      "affiliation": "SMHI, Sweden",
+      "name": "Lenhardt, Julien",
+      "orcid": "https://orcid.org/0000-0002-9949-3989"
     }
   ],
   "description": "ESMValCore: A community tool for pre-processing data from Earth system models in CMIP and running analysis scripts.",

CITATION.cff

@@ -223,6 +223,11 @@ authors:
     family-names: Garcia Perdomo
     given-names: Karen
     orcid: "https://orcid.org/0009-0004-2333-3358"
+  -
+    affiliation: "SMHI, Sweden"
+    family-names: Lenhardt
+    given-names: Julien
+    orcid: "https://orcid.org/0000-0002-9949-3989"
 
 cff-version: 1.2.0
 date-released: 2025-02-27

doc/recipe/preprocessor.rst

@@ -311,6 +311,7 @@ Preprocessor                                                          Variable s
 :ref:`weighting_landsea_fraction<land/sea fraction weighting>` [#f3]_ ``sftlf``, ``sftof``           land_area_fraction, sea_area_fraction
 :ref:`distance_metric<distance_metric>` [#f5]_                        ``areacella``, ``areacello``   cell_area
 :ref:`histogram<histogram>` [#f5]_                                    ``areacella``, ``areacello``   cell_area
+:ref:`extract_surface_from_atm<extract_surface_from_atm>` [#f3]_      ``ps``                         surface_air_pressure
 ===================================================================== ============================== =====================================
 
 .. [#f3] This preprocessor requires at least one of the mentioned supplementary
@@ -2118,7 +2119,6 @@ See also :func:`esmvalcore.preprocessor.extract_location`.
 
 ``zonal_statistics``
 --------------------
-
 The function calculates the zonal statistics by applying an operator
 along the longitude coordinate.
 
@@ -2223,6 +2223,7 @@ The ``_volume.py`` module contains the following preprocessor functions:
 * ``extract_transect``: Extract data along a line of constant latitude or
   longitude.
 * ``extract_trajectory``: Extract data along a specified trajectory.
+* ``extract_surface_from_atm``: Extract atmospheric data at the surface.
 
 
 ``extract_volume``
@@ -2406,6 +2407,22 @@ Note that this function uses the expensive ``interpolate`` method from
 See also :func:`esmvalcore.preprocessor.extract_trajectory`.
 
 
+.. _extract_surface_from_atm:
+
+``extract_surface_from_atm``
+----------------------------
+
+This function extracts data at the surface for an atmospheric variable.
+
+The function returns the interpolated value of an input filed at the corresponding surface pressure
+given by the surface air pressure ``ps``.
+
+The required supplementary surface air pressure ``ps`` is attached to
+the main dataset as described in :ref:`supplementary_variables`.
+
+See also :func:`esmvalcore.preprocessor.extract_surface_from_atm`.
+
+
 .. _cycles:
 
 Cycles

esmvalcore/preprocessor/__init__.py

@@ -86,6 +86,7 @@
 from ._volume import (
     axis_statistics,
     depth_integration,
+    extract_surface_from_atm,
     extract_trajectory,
     extract_transect,
     extract_volume,
@@ -122,6 +123,8 @@
     "resample_time",
     # Level extraction
     "extract_levels",
+    # Extract surface
+    "extract_surface_from_atm",
     # Weighting
     "weighting_landsea_fraction",
     # Mask landsea (fx or Natural Earth)

esmvalcore/preprocessor/_derive/co2s.py

@@ -1,27 +1,11 @@
 """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",
-    )
+from esmvalcore.preprocessor._supplementary_vars import add_ancillary_variable
+from esmvalcore.preprocessor._volume import extract_surface_from_atm
 
-    indices = list(indices)
-    indices.insert(axis, indices_first_positive)
-    first_unmasked_data = np.array(array)[tuple(indices)]
-    return first_unmasked_data
+from ._baseclass import DerivedVariableBase
 
 
 class DerivedVariable(DerivedVariableBase):
@@ -55,56 +39,9 @@ def calculate(cubes):
         ps_cube = cubes.extract_cube(
             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))
+        # Add ps as AncillaryVariable in the CO2 cube
+        add_ancillary_variable(cube=co2_cube, ancillary_cube=ps_cube)
+        # Extract surface from 3D atmospheric cube
+        co2s_cube = extract_surface_from_atm(co2_cube)
         co2s_cube.convert_units("1e-6")
         return co2s_cube

esmvalcore/preprocessor/_regrid.py

@@ -18,6 +18,7 @@
 
 import dask.array as da
 import iris
+import iris.coords
 import numpy as np
 import stratify
 from geopy.geocoders import Nominatim
@@ -1020,7 +1021,7 @@ def _create_cube(src_cube, data, src_levels, levels):
     z_coord = src_cube.coord(axis="z", dim_coords=True)
     (z_dim,) = src_cube.coord_dims(z_coord)
 
-    if data.shape[z_dim] != levels.size:
+    if (len(levels.shape) == 1) and (data.shape[z_dim] != levels.size):
         emsg = (
             "Mismatch between data and levels for data dimension {!r}, "
             "got data shape {!r} with levels shape {!r}."
@@ -1068,7 +1069,10 @@ def _create_cube(src_cube, data, src_levels, levels):
         result.add_dim_coord(coord, z_dim)
     except ValueError:
         coord = iris.coords.AuxCoord(levels, **kwargs)
-        result.add_aux_coord(coord, z_dim)
+        result.add_aux_coord(
+            coord,
+            z_dim if len(levels.shape) == 1 else np.arange(len(coord.shape)),
+        )
 
     # Collapse the z-dimension for the scalar case.
     if levels.size == 1:
@@ -1151,7 +1155,24 @@ def _preserve_fx_vars(cube, result):
                     result.var_name,
                 )
             else:
-                add_ancillary_variable(result, ancillary_var)
+                # Create cube and add coordinates to ancillary variable
+                ancillary_coords: list[
+                    tuple[iris.coords.AncillaryVariable, int]
+                ] = []
+                for i, coord in enumerate(cube.coords()):
+                    if i in ancillary_dims:
+                        coord_idx = len(ancillary_coords)
+                        ancillary_coords.append((coord.copy(), coord_idx))
+                ancillary_cube = iris.cube.Cube(
+                    ancillary_var.core_data(),
+                    standard_name=ancillary_var.standard_name,
+                    long_name=ancillary_var.long_name,
+                    units=ancillary_var.units,
+                    var_name=ancillary_var.var_name,
+                    attributes=ancillary_var.attributes,
+                    dim_coords_and_dims=ancillary_coords,
+                )
+                add_ancillary_variable(result, ancillary_cube)
 
 
 def parse_vertical_scheme(scheme):
@@ -1288,8 +1309,10 @@ def extract_levels(
         result = cube
         # Set the levels to the requested values
         src_levels.points = levels
-    elif len(src_levels.shape) == 1 and set(levels).issubset(
-        set(src_levels.points)
+    elif (
+        len(src_levels.shape) == 1
+        and len(levels.shape) == 1
+        and set(levels.flatten()).issubset(set(src_levels.points))
     ):
         # If all target levels exist in the source cube, simply extract them.
         name = src_levels.name()

esmvalcore/preprocessor/_supplementary_vars.py

@@ -45,7 +45,7 @@ def wrapper(func):
 
 def add_cell_measure(
     cube: Cube,
-    cell_measure_cube: Cube,
+    cell_measure_cube: Cube | iris.coords.CellMeasure,
     measure: Literal["area", "volume"],
 ) -> None:
     """Add cell measure to cube (in-place).
@@ -60,7 +60,7 @@ def add_cell_measure(
     cube:
         Iris cube with input data.
     cell_measure_cube:
-        Iris cube with cell measure data.
+        Iris cube or cell measure coord with cell measure data.
     measure:
         Name of the measure, can be 'area' or 'volume'.
 
@@ -101,39 +101,73 @@ def add_cell_measure(
     )
 
 
-def add_ancillary_variable(cube: Cube, ancillary_cube: Cube) -> None:
+def add_ancillary_variable(
+    cube: Cube, ancillary_cube: Cube | iris.coords.AncillaryVariable
+) -> None:
     """Add ancillary variable to cube (in-place).
 
-    Note
-    ----
-    This assumes that the ancillary variable spans the rightmost dimensions of
-    the cube.
-
     Parameters
     ----------
     cube:
         Iris cube with input data.
     ancillary_cube:
-        Iris cube with ancillary data.
+        Iris cube or AncillaryVariable with ancillary data.
 
     Returns
     -------
     iris.cube.Cube
         Cube with added ancillary variables
     """
-    coord_dims = tuple(range(cube.ndim - len(ancillary_cube.shape), cube.ndim))
-    ancillary_data = ancillary_cube.core_data()
+    try:
+        ancillary_var = iris.coords.AncillaryVariable(
+            ancillary_cube.core_data(),
+            standard_name=ancillary_cube.standard_name,
+            units=ancillary_cube.units,
+            var_name=ancillary_cube.var_name,
+            attributes=ancillary_cube.attributes,
+            long_name=ancillary_cube.long_name,
+        )
+    except AttributeError as err:
+        msg = (
+            f"Failed to add {ancillary_cube} to {cube} as ancillary var."
+            "ancillary_cube should be either an iris.cube.Cube or an "
+            "iris.coords.AncillaryVariable object."
+        )
+        raise ValueError(msg) from err
+    # Match the coordinates of the ancillary cube to coordinates and
+    # dimensions in the input cube before adding the ancillary variable.
+    data_dims: list[int | None] = []
+    if isinstance(ancillary_cube, iris.coords.AncillaryVariable):
+        start_dim = cube.ndim - len(ancillary_var.shape)
+        data_dims = list(range(start_dim, cube.ndim))
+    else:
+        data_dims = [None] * ancillary_cube.ndim
+        for coord in ancillary_cube.coords():
+            try:
+                for ancillary_dim, cube_dim in zip(
+                    ancillary_cube.coord_dims(coord),
+                    cube.coord_dims(coord),
+                    strict=False,
+                ):
+                    data_dims[ancillary_dim] = cube_dim
+            except iris.exceptions.CoordinateNotFoundError:
+                logger.debug(
+                    "%s from ancillary cube not found in cube coords.", coord
+                )
+        if None in data_dims:
+            none_dims = ", ".join(
+                str(i) for i, d in enumerate(data_dims) if d is None
+            )
+            msg = (
+                f"Failed to add {ancillary_cube} to {cube} as ancillary var."
+                f"No coordinate associated with ancillary cube dimensions"
+                f"{none_dims}"
+            )
+            raise ValueError(msg)
     if ancillary_cube.has_lazy_data():
-        cube_chunks = tuple(cube.lazy_data().chunks[d] for d in coord_dims)
-        ancillary_data = ancillary_data.rechunk(cube_chunks)
-    ancillary_var = iris.coords.AncillaryVariable(
-        ancillary_data,
-        standard_name=ancillary_cube.standard_name,
-        units=ancillary_cube.units,
-        var_name=ancillary_cube.var_name,
-        attributes=ancillary_cube.attributes,
-    )
-    cube.add_ancillary_variable(ancillary_var, coord_dims)
+        cube_chunks = tuple(cube.lazy_data().chunks[d] for d in data_dims)
+        ancillary_var.data = ancillary_cube.lazy_data().rechunk(cube_chunks)
+    cube.add_ancillary_variable(ancillary_var, data_dims)
     logger.debug(
         "Added %s as ancillary variable in cube of %s.",
         ancillary_cube.var_name,