Add plev/altitude conversion to extract_levels

esmvalcore/cmor/_fixes/shared.py

@@ -29,6 +29,21 @@ def _get_altitude_to_pressure_func():
 altitude_to_pressure = _get_altitude_to_pressure_func()  # noqa
 
 
+def _get_pressure_to_altitude_func():
+    """Get function converting air pressure [Pa] to altitude [m]."""
+    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
+
+
+pressure_to_altitude = _get_pressure_to_altitude_func()  # noqa
+
+
 class AtmosphereSigmaFactory(iris.aux_factory.AuxCoordFactory):
     """Defines an atmosphere sigma coordinate factory."""
 
@@ -258,7 +273,10 @@ def add_plev_from_altitude(cube):
         if height_coord.units != 'm':
             height_coord.convert_units('m')
         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 +290,42 @@ 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')
+        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)

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
@@ -506,6 +507,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,17 +4,24 @@
 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 (
+    _get_altitude_to_pressure_func,
+    _get_pressure_to_altitude_func,
+    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,
+    altitude_to_pressure,
+    cube_to_aux_coord,
+    fix_bounds,
+    get_bounds_cube,
+    pressure_to_altitude,
+    round_coordinates,
+)
 from esmvalcore.iris_helpers import var_name_constraint
 
 
@@ -33,6 +40,21 @@ def test_altitude_to_pressure_func(func):
                                [101325.0, 100129.0])
 
 
+@pytest.mark.parametrize('func', [pressure_to_altitude,
+                                  _get_pressure_to_altitude_func()])
+def test_pressure_to_altitude_func(func):
+    """Test pressure to altitude function."""
+    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 +104,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 +148,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],

tests/integration/preprocessor/_regrid/test_extract_levels.py

@@ -101,6 +101,34 @@ def test_interpolation__scalar_collapse(self):
         del self.shape[self.z_dim]
         self.assertEqual(result.shape, tuple(self.shape))
 
+    def test_add_alt_coord(self):
+        assert self.cube.coords('air_pressure')
+        assert not self.cube.coords('altitude')
+        result = extract_levels(self.cube, [1, 2],
+                                'linear_horizontal_extrapolate_vertical',
+                                coordinate='altitude')
+        assert not result.coords('air_pressure')
+        assert result.coords('altitude')
+        assert result.shape == (2, 2, 2, 2)
+        np.testing.assert_allclose(result.coord('altitude').points,
+                                   [1.0, 2.0])
+
+    def test_add_plev_coord(self):
+        self.cube.coord('air_pressure').standard_name = 'altitude'
+        self.cube.coord('altitude').var_name = 'alt'
+        self.cube.coord('altitude').long_name = 'altitude'
+        self.cube.coord('altitude').units = 'm'
+        assert not self.cube.coords('air_pressure')
+        assert self.cube.coords('altitude')
+        result = extract_levels(self.cube, [1, 2],
+                                'linear_horizontal_extrapolate_vertical',
+                                coordinate='air_pressure')
+        assert result.coords('air_pressure')
+        assert not result.coords('altitude')
+        assert result.shape == (2, 2, 2, 2)
+        np.testing.assert_allclose(result.coord('air_pressure').points,
+                                   [1.0, 2.0])
+
 
 if __name__ == '__main__':
     unittest.main()

tests/unit/cmor/test_fixes.py

@@ -1,12 +1,17 @@
 """Test individual fix functions."""
 import pytest
 
-import esmvalcore.cmor.fixes
+import esmvalcore.cmor._fixes.shared as shared
+import esmvalcore.cmor.fixes as fixes
 
 
 @pytest.mark.parametrize('func', [
+    'add_altitude_from_plev',
     'add_plev_from_altitude',
     'add_sigma_factory',
 ])
 def test_imports(func):
-    assert func in esmvalcore.cmor.fixes.__all__
+    assert func in fixes.__all__
+    fn_in_shared = getattr(shared, func)
+    fn_in_fixes = getattr(fixes, func)
+    assert fn_in_shared is fn_in_fixes