Update humidity mixing ratio to work with pressure and height level data

improver/psychrometric_calculations/psychrometric_calculations.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 from iris.cube import Cube, CubeList
+from iris.exceptions import ConstraintMismatchError
 from numpy import ndarray
 from scipy.optimize import newton
 
@@ -22,7 +23,10 @@
     generate_mandatory_attributes,
 )
 from improver.utilities.common_input_handle import as_cubelist
-from improver.utilities.cube_manipulation import sort_coord_in_cube
+from improver.utilities.cube_manipulation import (
+    enforce_coordinate_ordering,
+    sort_coord_in_cube,
+)
 from improver.utilities.interpolation import interpolate_missing_data
 from improver.utilities.mathematical_operations import fast_linear_fit
 from improver.utilities.spatial import OccurrenceWithinVicinity
@@ -315,23 +319,62 @@ def _make_humidity_cube(self, data: np.ndarray) -> Cube:
         )
         return cube
 
+    def generate_pressure_cube(self) -> None:
+        """Generate a pressure cube from the pressure coordinate on the temperature cube"""
+        coord_list = [coord.name() for coord in self.temperature.coords()]
+        pressure_list = CubeList()
+        for temp_slice in self.temperature.slices_over("pressure"):
+            pressure_value = temp_slice.coord("pressure").points
+            temp_slice.data = np.broadcast_to(pressure_value, temp_slice.shape)
+            pressure_list.append(temp_slice)
+        self.pressure = pressure_list.merge_cube()
+        enforce_coordinate_ordering(self.pressure, coord_list)
+        self.pressure.rename("surface_air_pressure")
+        self.pressure.units = "Pa"
+
     def process(self, *cubes: Union[Cube, CubeList]) -> Cube:
         """
-        Calculates the humidity mixing ratio from the inputs.
+        Calculates the humidity mixing ratio from the inputs. The inputs can be on height levels
+        or pressure levels, and the output will be on the same levels. If the input cubes are on
+        pressure levels then a pressure cube doesn't need to be provided and the pressure levels
+        will be inferred from the pressure coordinate on the temperature cube.
 
         Args:
             cubes:
-                Cubes of temperature (K), pressure (Pa) and relative humidity (1)
+                Cubes of temperature (K) and relative humidity (1). A cube of pressure (Pa) must also
+                br provided unless there is a pressure coordinate in the temperature and relative humidity cubes.
 
         Returns:
-            Cube of humidity mixing ratio
+            Cube of humidity mixing ratio on same levels as input cubes
 
         """
         cubes = as_cubelist(*cubes)
-        (self.temperature, self.pressure, self.rel_humidity) = cubes.extract_cubes(
-            ["air_temperature", "surface_air_pressure", "relative_humidity"]
+
+        (self.temperature, self.rel_humidity) = cubes.extract_cubes(
+            ["air_temperature", "relative_humidity"]
         )
 
+        try:
+            self.pressure = cubes.extract_cube("surface_air_pressure")
+        except ConstraintMismatchError:
+            # If no pressure cube is provided, check if pressure is a coordinate in the temperature and relative humidity cubes
+            temp_coord = [
+                True
+                for coord in self.temperature.coords()
+                if coord.name() == "pressure"
+            ]
+            rh_coord = [
+                True
+                for coord in self.rel_humidity.coords()
+                if coord.name() == "pressure"
+            ]
+            if any(temp_coord) and any(rh_coord):
+                self.generate_pressure_cube()
+            else:
+                raise ValueError(
+                    "No pressure cube called 'surface_air_pressure' found and no pressure coordinate found in temperature or relative humidity cubes"
+                )
+
         self.mandatory_attributes = generate_mandatory_attributes(
             [self.temperature, self.pressure, self.rel_humidity]
         )

improver_tests/psychrometric_calculations/test_HumidityMixingRatio.py

@@ -136,6 +136,69 @@ def test_basic(
     assert np.isclose(result.data, expected, atol=1e-7).all()
 
 
+def test_height_levels():
+    """Check that the plugin works with height level data"""
+    temperature = set_up_variable_cube(
+        np.full((1, 2, 2, 2), fill_value=293, dtype=np.float32),
+        name="air_temperature",
+        units="K",
+        attributes=LOCAL_MANDATORY_ATTRIBUTES,
+        height_levels=[100, 400],
+    )
+    pressure_cube = set_up_variable_cube(
+        np.full((1, 2, 2, 2), fill_value=100000, dtype=np.float32),
+        name="surface_air_pressure",
+        units="Pa",
+        attributes=LOCAL_MANDATORY_ATTRIBUTES,
+        height_levels=[100, 400],
+    )
+    rel_humidity = set_up_variable_cube(
+        np.full((1, 2, 2, 2), fill_value=1.0, dtype=np.float32),
+        name="relative_humidity",
+        units="1",
+        attributes=LOCAL_MANDATORY_ATTRIBUTES,
+        height_levels=[100, 400],
+    )
+    result = HumidityMixingRatio()([temperature, pressure_cube, rel_humidity])
+    metadata_ok(result, temperature)
+    assert np.isclose(result.data, 1.459832e-2, atol=1e-7).all()
+
+
+def test_pressure_levels():
+    """Check that the plugin works with pressure level data when pressure cube is not provided"""
+    temperature = set_up_variable_cube(
+        np.full((1, 2, 2, 2), fill_value=293, dtype=np.float32),
+        name="air_temperature",
+        units="K",
+        attributes=LOCAL_MANDATORY_ATTRIBUTES,
+        height_levels=[95000, 100000],
+        pressure=True,
+    )
+    rel_humidity = set_up_variable_cube(
+        np.full((1, 2, 2, 2), fill_value=1.0, dtype=np.float32),
+        name="relative_humidity",
+        units="1",
+        attributes=LOCAL_MANDATORY_ATTRIBUTES,
+        height_levels=[95000, 100000],
+        pressure=True,
+    )
+    result = HumidityMixingRatio()([temperature, rel_humidity])
+    metadata_ok(result, temperature)
+    assert np.isclose(result.data[:, 0], 1.537017e-2, atol=1e-7).all()
+    assert np.isclose(result.data[:, 1], 1.459832e-2, atol=1e-7).all()
+
+
+def test_error_raised_no_pressure_coordinate_or_pressure_cube(
+    temperature, rel_humidity
+):
+    """Check that the plugin raises an error if there is no pressure coordinate and no pressure cube"""
+    with pytest.raises(
+        ValueError,
+        match="No pressure cube called 'surface_air_pressure' found and no pressure coordinate",
+    ):
+        HumidityMixingRatio()([temperature, rel_humidity])
+
+
 @pytest.mark.parametrize("model_id_attr", ("mosg__model_configuration", None))
 def test_model_id_attr(temperature, pressure, rel_humidity, model_id_attr):
     """Check that tests pass if model_id_attr is set on inputs and is applied or not"""