92/Refactor pressure level modifications

test/test_um2netcdf.py

@@ -685,21 +685,31 @@ def level_heights():
     #     fix_level_coords() only accesses height array[0]
     return [20.0003377]
 
+
 @pytest.fixture
 def level_coords(level_heights):
     return {um2nc.MODEL_LEVEL_NUM: iris.coords.DimCoord(range(1, 39)),
             um2nc.LEVEL_HEIGHT: iris.coords.DimCoord(level_heights),
             um2nc.SIGMA: iris.coords.AuxCoord(np.array([0.99771646]))}
 
+
 @pytest.fixture
 def get_fake_cube_coords(level_coords):
 
     @dataclass
     class FakeCubeCoords:
         """Test object to represent a cube with a coords() access function."""
 
+        def __init__(self, custom_coord: dict = None):
+            if custom_coord:
+                level_coords.update(custom_coord)
+
         def coord(self, key):
-            return level_coords[key]
+            if key in level_coords:
+                return level_coords[key]
+
+            msg = f"{self.__class__}: lacks coord for '{key}'"
+            raise iris.exceptions.CoordinateNotFoundError(msg)
 
     # return class for instantiation in tests
     return FakeCubeCoords
@@ -715,7 +725,6 @@ def test_fix_level_coord_modify_cube_with_rho(level_heights,
     assert cube.coord(um2nc.LEVEL_HEIGHT).var_name is None
     assert cube.coord(um2nc.SIGMA).var_name is None
 
-
     rho = np.ones(z_sea_theta_data.shape) * level_heights[0]
     um2nc.fix_level_coord(cube, rho, z_sea_theta_data)
 
@@ -744,6 +753,58 @@ def test_fix_level_coord_skipped_if_no_levels(z_sea_rho_data, z_sea_theta_data):
     um2nc.fix_level_coord(m_cube, z_sea_rho_data, z_sea_theta_data)
 
 
+# tests - fix pressure level data
+
+def test_fix_pressure_levels_no_pressure_coord(get_fake_cube_coords):
+    cube = get_fake_cube_coords()
+
+    with pytest.raises(iris.exceptions.CoordinateNotFoundError):
+        cube.coord("pressure")  # ensure missing 'pressure' coord
+
+    # fix function should return if there is no pressure coord to modify
+    assert um2nc.fix_pressure_levels(cube) is None  # should just exit
+
+
+def _add_attrs_points(m_plevs: mock.MagicMock, points):
+    # NB: iris attributes appear to be added via mixins, so it's easier but
+    #     less desirable to rely on mock attrs here
+    setattr(m_plevs, "attributes", {"positive": None})
+    setattr(m_plevs, "points", points)
+
+
+def test_fix_pressure_levels_do_rounding(get_fake_cube_coords):
+    m_pressure = mock.Mock()
+    _add_attrs_points(m_pressure, [1.000001, 0.000001])
+    extra = {"pressure": m_pressure}
+    cube = get_fake_cube_coords(extra)
+
+    um2nc.fix_pressure_levels(cube)
+
+    # TODO: test flaw, this verifies pressure coord but ignores fix_pressure_levels()
+    #       returning a new cube if the pressure is reversed. This is verified
+    #       in command line testing though
+    c_pressure = cube.coord('pressure')
+    assert c_pressure.attributes["positive"] == "down"
+    assert all(c_pressure.points == [1.0, 0.0])
+
+
+def test_fix_pressure_levels_reverse_pressure(get_fake_cube_coords):
+    m_pressure = mock.Mock()
+    _add_attrs_points(m_pressure, [0.000001, 1.000001])
+    extra = {"pressure": m_pressure}
+    cube = get_fake_cube_coords(extra)
+
+    with mock.patch("iris.util.reverse"):
+        um2nc.fix_pressure_levels(cube)
+
+    # TODO: test flaw, this verifies pressure coord but ignores fix_pressure_levels()
+    #       returning a new cube if the pressure is reversed. This is verified
+    #       in command line testing though
+    c_pressure = cube.coord('pressure')
+    assert c_pressure.attributes["positive"] == "down"
+    assert all(c_pressure.points == [0.0, 1.0])
+
+
 # int64 to int32 data conversion tests
 # NB: skip float64 to float32 overflow as float32 min/max is huge: -/+ 3.40e+38
 @pytest.mark.parametrize("array,_operator,bound",

umpost/um2netcdf.py

@@ -147,17 +147,8 @@ def _add_coord_bounds(coord):
 
 # TODO: split cube ops into functions, this will likely increase process() workflow steps
 def cubewrite(cube, sman, compression, use64bit, verbose):
-    try:
-        plevs = cube.coord('pressure')
-        plevs.attributes['positive'] = 'down'
-        plevs.convert_units('Pa')
-        # Otherwise they're off by 1e-10 which looks odd in ncdump
-        plevs.points = np.round(plevs.points, 5)
-        if plevs.points[0] < plevs.points[-1]:
-            # Flip to get pressure decreasing as in CMIP6 standard
-            cube = iris.util.reverse(cube, 'pressure')
-    except iris.exceptions.CoordinateNotFoundError:
-        pass
+    # TODO: move into process() AND if a new cube is returned, swap into filtered cube list
+    cube = fix_pressure_levels(cube) or cube  # NB: use new cube if pressure points are modified
 
     # TODO: flag warnings as an error for the driver script?
     if not use64bit:
@@ -734,6 +725,41 @@ def fix_level_coord(cube, z_rho, z_theta, tol=1e-6):
                 c_sigma.var_name = 'sigma_theta'
 
 
+def fix_pressure_levels(cube, decimals=5):
+    """
+    Reformat pressure level data for NetCDF output.
+
+    This converts units, rounds small fractional errors & ensures pressure is
+    decreasing (following the CMIP6 standard).
+
+    Parameters
+    ----------
+    cube : iris Cube (modifies in place)
+    decimals : number of decimals to round to
+
+    Returns
+    -------
+    None if cube lacks pressure coord or is modified in place, otherwise a new
+    cube if the pressure levels are reversed.
+    """
+    try:
+        pressure = cube.coord('pressure')
+    except iris.exceptions.CoordinateNotFoundError:
+        return
+
+    # update existing cube metadata in place
+    pressure.attributes['positive'] = 'down'
+    pressure.convert_units('Pa')
+
+    # Round small fractions otherwise coordinates are off by 1e-10 in ncdump output
+    pressure.points = np.round(pressure.points, decimals)
+
+    if pressure.points[0] < pressure.points[-1]:
+        # Flip to get pressure decreasing as per CMIP6 standard
+        # NOTE: returns a new cube!
+        return iris.util.reverse(cube, 'pressure')
+
+
 MAX_NP_INT32 = np.iinfo(np.int32).max
 MIN_NP_INT32 = np.iinfo(np.int32).min