Merge pull request #94 from ACCESS-NRI/92/refactor-plevs

Merge 92/Refactor pressure level modifications.

.coveragerc

@@ -1,6 +1,9 @@
 [run]
 branch = True
 
+# skip coverage testing for external source files
+omit = umpost/_version.py
+
 [html]
 title = Coverage report: um2nc-standalone
 directory = coverage_html

test/test_um2netcdf.py

@@ -1049,21 +1049,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
@@ -1079,7 +1089,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)
 
@@ -1108,6 +1117,74 @@ 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)
+
+    # ensure no cube is returned if Cube not modified in fix_pressure_levels()
+    assert um2nc.fix_pressure_levels(cube) is None
+
+    c_pressure = cube.coord('pressure')
+    assert c_pressure.attributes["positive"] == "down"
+    assert all(c_pressure.points == [1.0, 0.0])
+
+
+@pytest.mark.skip
+def test_fix_pressure_levels_reverse_pressure(get_fake_cube_coords):
+    # TODO: test is broken due to fiddly mocking problems (see below)
+
+    m_pressure = mock.Mock()
+    # m_pressure.ndim = 1
+    _add_attrs_points(m_pressure, [0.000001, 1.000001])
+    extra = {"pressure": m_pressure}
+    cube = get_fake_cube_coords(extra)
+    # cube.ndim = 3
+
+    # TODO: testing gets odd here at the um2nc & iris "boundary":
+    #   * A mock reverse() needs to flip pressure.points & return a modified cube.
+    #     Creating a mock to verifying these attributes is unproductive.
+    #   * Using the real reverse() requires several additional cube attributes
+    #     (see commented out ndim etc above). It requires __getitem__() for
+    #     https://github.com/SciTools/iris/blob/main/lib/iris/util.py#L612
+    #
+    # TODO: this leaves a few options:
+    #    * ignore unit testing this branch (rely on integration testing?)
+    #    * replace iris with an adapter?
+    #    * fix/refactor the function later?
+    #
+    # The test is disabled awaiting a solution...
+
+    with mock.patch("iris.util.reverse"):
+        mod_cube = um2nc.fix_pressure_levels(cube)
+
+    assert mod_cube is not None
+    assert mod_cube != cube
+    c_pressure = mod_cube.coord('pressure')
+    assert c_pressure.attributes["positive"] == "down"
+    assert all(c_pressure.points == [1.0, 0.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

@@ -321,17 +321,8 @@ def fix_latlon_coords(cube, grid_type, dlat, dlon):
 
 # 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:
@@ -902,6 +893,42 @@ 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!
+        # TODO: add an iris.util.monotonic() check here?
+        return iris.util.reverse(cube, 'pressure')
+
+
 MAX_NP_INT32 = np.iinfo(np.int32).max
 MIN_NP_INT32 = np.iinfo(np.int32).min