Add plev/altitude conversion to extract_levels

doc/recipe/preprocessor.rst

@@ -243,6 +243,13 @@ the name of the desired coordinate:
           scheme: linear_horizontal_extrapolate_vertical
           coordinate: air_pressure
 
+If ``coordinate`` is specified, pressure levels (if present) can be converted
+to height levels and vice versa using the US standard atmosphere. E.g.
+``coordinate = altitude`` will convert existing pressure levels
+(air_pressure) to height levels (altitude);
+``coordinate = air_pressure`` will convert existing height levels
+(altitude) to pressure levels (air_pressure).
+
 * See also :func:`esmvalcore.preprocessor.extract_levels`.
 * See also :func:`esmvalcore.preprocessor.get_cmor_levels`.
 

esmvalcore/cmor/_fixes/shared.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import warnings
+from functools import lru_cache
 
 import dask.array as da
 import iris
@@ -14,21 +15,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _get_altitude_to_pressure_func():
-    """Get function converting altitude [m] to air pressure [Pa]."""
-    base_dir = os.path.dirname(os.path.abspath(__file__))
-    source_file = os.path.join(base_dir, 'us_standard_atmosphere.csv')
-    data_frame = pd.read_csv(source_file, comment='#')
-    func = interp1d(data_frame['Altitude [m]'],
-                    data_frame['Pressure [Pa]'],
-                    kind='cubic',
-                    fill_value='extrapolate')
-    return func
-
-
-altitude_to_pressure = _get_altitude_to_pressure_func()  # noqa
-
-
 class AtmosphereSigmaFactory(iris.aux_factory.AuxCoordFactory):
     """Defines an atmosphere sigma coordinate factory."""
 
@@ -257,8 +243,12 @@ def add_plev_from_altitude(cube):
         height_coord = cube.coord('altitude')
         if height_coord.units != 'm':
             height_coord.convert_units('m')
+        altitude_to_pressure = get_altitude_to_pressure_func()
         pressure_points = altitude_to_pressure(height_coord.core_points())
-        pressure_bounds = altitude_to_pressure(height_coord.core_bounds())
+        if height_coord.core_bounds() is None:
+            pressure_bounds = None
+        else:
+            pressure_bounds = altitude_to_pressure(height_coord.core_bounds())
         pressure_coord = iris.coords.AuxCoord(pressure_points,
                                               bounds=pressure_bounds,
                                               var_name='plev',
@@ -272,6 +262,43 @@ def add_plev_from_altitude(cube):
         "available")
 
 
+def add_altitude_from_plev(cube):
+    """Add altitude coordinate from pressure level coordinate.
+
+    Parameters
+    ----------
+    cube : iris.cube.Cube
+        Input cube.
+
+    Raises
+    ------
+    ValueError
+        ``cube`` does not contain coordinate ``air_pressure``.
+
+    """
+    if cube.coords('air_pressure'):
+        plev_coord = cube.coord('air_pressure')
+        if plev_coord.units != 'Pa':
+            plev_coord.convert_units('Pa')
+        pressure_to_altitude = get_pressure_to_altitude_func()
+        altitude_points = pressure_to_altitude(plev_coord.core_points())
+        if plev_coord.core_bounds() is None:
+            altitude_bounds = None
+        else:
+            altitude_bounds = pressure_to_altitude(plev_coord.core_bounds())
+        altitude_coord = iris.coords.AuxCoord(altitude_points,
+                                              bounds=altitude_bounds,
+                                              var_name='alt',
+                                              standard_name='altitude',
+                                              long_name='altitude',
+                                              units='m')
+        cube.add_aux_coord(altitude_coord, cube.coord_dims(plev_coord))
+        return
+    raise ValueError(
+        "Cannot add 'altitude' coordinate, 'air_pressure' coordinate not "
+        "available")
+
+
 def add_scalar_depth_coord(cube, depth=0.0):
     """Add scalar coordinate 'depth' with value of `depth`m."""
     logger.debug("Adding depth coordinate (%sm)", depth)
@@ -345,6 +372,26 @@ def cube_to_aux_coord(cube):
     )
 
 
+@lru_cache()
+def get_altitude_to_pressure_func():
+    """Get function converting altitude [m] to air pressure [Pa].
+
+    Returns
+    -------
+    callable
+        Function that converts altitude to air pressure.
+
+    """
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    source_file = os.path.join(base_dir, 'us_standard_atmosphere.csv')
+    data_frame = pd.read_csv(source_file, comment='#')
+    func = interp1d(data_frame['Altitude [m]'],
+                    data_frame['Pressure [Pa]'],
+                    kind='cubic',
+                    fill_value='extrapolate')
+    return func
+
+
 def get_bounds_cube(cubes, coord_var_name):
     """Find bound cube for a given variable in a :class:`iris.cube.CubeList`.
 
@@ -382,6 +429,26 @@ def get_bounds_cube(cubes, coord_var_name):
         f"cubes\n{cubes}")
 
 
+@lru_cache()
+def get_pressure_to_altitude_func():
+    """Get function converting air pressure [Pa] to altitude [m].
+
+    Returns
+    -------
+    callable
+        Function that converts air pressure to altitude.
+
+    """
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    source_file = os.path.join(base_dir, 'us_standard_atmosphere.csv')
+    data_frame = pd.read_csv(source_file, comment='#')
+    func = interp1d(data_frame['Pressure [Pa]'],
+                    data_frame['Altitude [m]'],
+                    kind='cubic',
+                    fill_value='extrapolate')
+    return func
+
+
 def fix_bounds(cube, cubes, coord_var_names):
     """Fix bounds for cube that could not be read correctly by :mod:`iris`.
 

esmvalcore/cmor/fixes.py

@@ -1,8 +1,13 @@
 """Functions for fixing specific issues with datasets."""
 
-from ._fixes.shared import add_plev_from_altitude, add_sigma_factory
+from ._fixes.shared import (
+    add_altitude_from_plev,
+    add_plev_from_altitude,
+    add_sigma_factory,
+)
 
 __all__ = [
+    'add_altitude_from_plev',
     'add_plev_from_altitude',
     'add_sigma_factory',
 ]

esmvalcore/preprocessor/_regrid.py

@@ -4,13 +4,14 @@
 import re
 from copy import deepcopy
 
+import iris
 import numpy as np
 import stratify
-import iris
 from dask import array as da
 from iris.analysis import AreaWeighted, Linear, Nearest, UnstructuredNearest
 from iris.util import broadcast_to_shape
 
+from ..cmor._fixes.shared import add_altitude_from_plev, add_plev_from_altitude
 from ..cmor.fix import fix_file, fix_metadata
 from ..cmor.table import CMOR_TABLES
 from ._io import concatenate_callback, load
@@ -475,7 +476,12 @@ def extract_levels(cube, levels, scheme, coordinate=None):
         'nearest_horizontal_extrapolate_vertical',
         'linear_horizontal_extrapolate_vertical'.
     coordinate :  optional str
-        The coordinate to interpolate
+        The coordinate to interpolate. If specified, pressure levels
+        (if present) can be converted to height levels and vice versa using
+        the US standard atmosphere. E.g. 'coordinate = altitude' will convert
+        existing pressure levels (air_pressure) to height levels (altitude);
+        'coordinate = air_pressure' will convert existing height levels
+        (altitude) to pressure levels (air_pressure).
 
     Returns
     -------
@@ -506,6 +512,16 @@ def extract_levels(cube, levels, scheme, coordinate=None):
 
     # Get the source cube vertical coordinate, if available.
     if coordinate:
+        coord_names = [coord.name() for coord in cube.coords()]
+        if coordinate not in coord_names:
+            # Try to calculate air_pressure from altitude coordinate or
+            # vice versa using US standard atmosphere for conversion.
+            if coordinate == 'air_pressure' and 'altitude' in coord_names:
+                # Calculate pressure level coordinate from altitude.
+                add_plev_from_altitude(cube)
+            if coordinate == 'altitude' and 'air_pressure' in coord_names:
+                # Calculate altitude coordinate from pressure levels.
+                add_altitude_from_plev(cube)
         src_levels = cube.coord(coordinate)
     else:
         src_levels = cube.coord(axis='z', dim_coords=True)

tests/integration/cmor/_fixes/test_shared.py

@@ -4,24 +4,28 @@
 import pytest
 from cf_units import Unit
 
-from esmvalcore.cmor._fixes.shared import (_get_altitude_to_pressure_func,
-                                           add_aux_coords_from_cubes,
-                                           add_plev_from_altitude,
-                                           add_scalar_depth_coord,
-                                           add_scalar_height_coord,
-                                           add_scalar_typeland_coord,
-                                           add_scalar_typesea_coord,
-                                           add_sigma_factory,
-                                           altitude_to_pressure,
-                                           cube_to_aux_coord, fix_bounds,
-                                           get_bounds_cube, round_coordinates)
+from esmvalcore.cmor._fixes.shared import (
+    add_altitude_from_plev,
+    add_aux_coords_from_cubes,
+    add_plev_from_altitude,
+    add_scalar_depth_coord,
+    add_scalar_height_coord,
+    add_scalar_typeland_coord,
+    add_scalar_typesea_coord,
+    add_sigma_factory,
+    cube_to_aux_coord,
+    fix_bounds,
+    get_altitude_to_pressure_func,
+    get_bounds_cube,
+    get_pressure_to_altitude_func,
+    round_coordinates,
+)
 from esmvalcore.iris_helpers import var_name_constraint
 
 
-@pytest.mark.parametrize('func', [altitude_to_pressure,
-                                  _get_altitude_to_pressure_func()])
-def test_altitude_to_pressure_func(func):
+def test_altitude_to_pressure_func():
     """Test altitude to pressure function."""
+    func = get_altitude_to_pressure_func()
     assert callable(func)
     np.testing.assert_allclose(func(-6000.0), 196968.01058487315)
     np.testing.assert_allclose(func(-5000.0), 177687.0)
@@ -33,6 +37,20 @@ def test_altitude_to_pressure_func(func):
                                [101325.0, 100129.0])
 
 
+def test_pressure_to_altitude_func():
+    """Test pressure to altitude function."""
+    func = get_pressure_to_altitude_func()
+    assert callable(func)
+    np.testing.assert_allclose(func(200000.0), -6166.332306480035)
+    np.testing.assert_allclose(func(177687.0), -5000.0)
+    np.testing.assert_allclose(func(101325.0), 0.0, atol=1.0e-7)
+    np.testing.assert_allclose(func(1000.0), 31054.63120206961)
+    np.testing.assert_allclose(func(75.9448), 50000)
+    np.testing.assert_allclose(func(0.1), 91607.36011892557)
+    np.testing.assert_allclose(func(np.array([101325.0, 177687.0])),
+                               [0.0, -5000.0], atol=1.0e-7)
+
+
 TEST_ADD_AUX_COORDS_FROM_CUBES = [
     ({}, 1),
     ({'x': ()}, 0),
@@ -82,20 +100,31 @@ def test_add_aux_coords_from_cubes(coord_dict, output):
 
 
 ALT_COORD = iris.coords.AuxCoord([0.0], bounds=[[-100.0, 500.0]],
+                                 var_name='alt', long_name='altitude',
                                  standard_name='altitude', units='m')
+ALT_COORD_NB = iris.coords.AuxCoord([0.0], var_name='alt',
+                                    long_name='altitude',
+                                    standard_name='altitude', units='m')
 ALT_COORD_KM = iris.coords.AuxCoord([0.0], bounds=[[-0.1, 0.5]],
+                                    var_name='alt', long_name='altitude',
                                     standard_name='altitude', units='km')
 P_COORD = iris.coords.AuxCoord([101325.0], bounds=[[102532.0, 95460.8]],
                                var_name='plev', standard_name='air_pressure',
                                long_name='pressure', units='Pa')
+P_COORD_NB = iris.coords.AuxCoord([101325.0], var_name='plev',
+                                  standard_name='air_pressure',
+                                  long_name='pressure', units='Pa')
 CUBE_ALT = iris.cube.Cube([1.0], var_name='x',
                           aux_coords_and_dims=[(ALT_COORD, 0)])
+CUBE_ALT_NB = iris.cube.Cube([1.0], var_name='x',
+                             aux_coords_and_dims=[(ALT_COORD_NB, 0)])
 CUBE_ALT_KM = iris.cube.Cube([1.0], var_name='x',
                              aux_coords_and_dims=[(ALT_COORD_KM, 0)])
 
 
 TEST_ADD_PLEV_FROM_ALTITUDE = [
     (CUBE_ALT.copy(), P_COORD.copy()),
+    (CUBE_ALT_NB.copy(), P_COORD_NB.copy()),
     (CUBE_ALT_KM.copy(), P_COORD.copy()),
     (iris.cube.Cube(0.0), None),
 ]
@@ -115,6 +144,58 @@ def test_add_plev_from_altitude(cube, output):
     add_plev_from_altitude(cube)
     air_pressure_coord = cube.coord('air_pressure')
     assert air_pressure_coord == output
+    assert cube.coords('altitude')
+
+
+P_COORD_HPA = iris.coords.AuxCoord([1013.25], bounds=[[1025.32, 954.60]],
+                                   var_name='plev',
+                                   standard_name='air_pressure',
+                                   long_name='pressure', units='hPa')
+CUBE_PLEV = iris.cube.Cube([1.0], var_name='x',
+                           aux_coords_and_dims=[(P_COORD, 0)])
+CUBE_PLEV_NB = iris.cube.Cube([1.0], var_name='x',
+                              aux_coords_and_dims=[(P_COORD_NB, 0)])
+CUBE_PLEV_HPA = iris.cube.Cube([1.0], var_name='x',
+                               aux_coords_and_dims=[(P_COORD_HPA, 0)])
+
+
+TEST_ADD_ALTITUDE_FROM_PLEV = [
+    (CUBE_PLEV.copy(), ALT_COORD.copy()),
+    (CUBE_PLEV_NB.copy(), ALT_COORD_NB.copy()),
+    (CUBE_PLEV_HPA.copy(), ALT_COORD.copy()),
+    (iris.cube.Cube(0.0), None),
+]
+
+
+@pytest.mark.parametrize('cube,output', TEST_ADD_ALTITUDE_FROM_PLEV)
+def test_add_altitude_from_plev(cube, output):
+    """Test adding of altitude coordinate."""
+    if output is None:
+        with pytest.raises(ValueError) as err:
+            add_altitude_from_plev(cube)
+        msg = ("Cannot add 'altitude' coordinate, 'air_pressure' coordinate "
+               "not available")
+        assert str(err.value) == msg
+        return
+    assert not cube.coords('altitude')
+    add_altitude_from_plev(cube)
+    altitude_coord = cube.coord('altitude')
+    metadata_list = [
+        'var_name',
+        'standard_name',
+        'long_name',
+        'units',
+        'attributes',
+    ]
+    for attr in metadata_list:
+        assert getattr(altitude_coord, attr) == getattr(output, attr)
+    np.testing.assert_allclose(altitude_coord.points, output.points, atol=1e-7)
+    if output.bounds is None:
+        assert altitude_coord.bounds is None
+    else:
+        np.testing.assert_allclose(altitude_coord.bounds, output.bounds,
+                                   rtol=1e-3)
+    assert cube.coords('air_pressure')
 
 
 DIM_COORD = iris.coords.DimCoord([3.141592],