diff --git a/improver/psychrometric_calculations/psychrometric_calculations.py b/improver/psychrometric_calculations/psychrometric_calculations.py index 43d7143f44..6708c37edd 100644 --- a/improver/psychrometric_calculations/psychrometric_calculations.py +++ b/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,58 @@ 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 + be 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_flag = any( + coord.name() == "pressure" for coord in self.temperature.coords() + ) + rh_coord_flag = any( + coord.name() == "pressure" for coord in self.rel_humidity.coords() + ) + if temp_coord_flag & rh_coord_flag: + 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] ) diff --git a/improver_tests/psychrometric_calculations/test_HumidityMixingRatio.py b/improver_tests/psychrometric_calculations/test_HumidityMixingRatio.py index 7c818e7801..840aa965c1 100644 --- a/improver_tests/psychrometric_calculations/test_HumidityMixingRatio.py +++ b/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"""