From 7e2ed684c21b8f497c8dea3026bebdc3210a9e0a Mon Sep 17 00:00:00 2001 From: mspelman07 <99179165+mspelman07@users.noreply.github.com> Date: Tue, 10 Oct 2023 14:47:13 +0100 Subject: [PATCH] Wet bulb freezing level (#1949) * Updates ExtractPressureLevel to also be able to extract a height level and updates unit tests * Update plugin to use stratify's extrapolation * Adds acceptance tests and cli for wet bulb freezing level. ALlso rename unit test file name. * Updates to hail size plugin * Formatting * Update checksums and uncommenting * update docstring formatting * doc string updated * review response * formatting review changes --------- Co-authored-by: Marcus Spelman --- improver/cli/wet_bulb_freezing_level.py | 66 +++++ .../psychrometric_calculations/hail_size.py | 10 +- improver/utilities/cube_extraction.py | 259 ++++++++---------- improver_tests/acceptance/SHA256SUMS | 2 + .../test_wet_bulb_freezing_level.py | 54 ++++ ...tPressureLevel.py => test_ExtractLevel.py} | 155 +++++++---- 6 files changed, 342 insertions(+), 204 deletions(-) create mode 100644 improver/cli/wet_bulb_freezing_level.py create mode 100644 improver_tests/acceptance/test_wet_bulb_freezing_level.py rename improver_tests/utilities/cube_extraction/{test_ExtractPressureLevel.py => test_ExtractLevel.py} (58%) diff --git a/improver/cli/wet_bulb_freezing_level.py b/improver/cli/wet_bulb_freezing_level.py new file mode 100644 index 0000000000..09883a772f --- /dev/null +++ b/improver/cli/wet_bulb_freezing_level.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------- +# (C) British Crown copyright. The Met Office. +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. +"""CLI to extract wet-bulb freezing level from wet-bulb temperature on height levels""" + +from improver import cli + + +@cli.clizefy +@cli.with_output +def process(wet_bulb_temperature: cli.inputcube): + """Module to generate wet-bulb freezing level. + + The height level at which the wet-bulb temperature first drops below 273.15K + (0 degrees Celsius) is extracted from the wet-bulb temperature cube starting from + the ground and ascending through height levels. + + In grid squares where the temperature never goes below 273.15K the highest + height level on the cube is returned. In grid squares where the temperature + starts below 273.15K the lowest height on the cube is returned. + + Args: + wet_bulb_temperature (iris.cube.Cube): + Cube of wet-bulb air temperatures over multiple height levels. + + Returns: + iris.cube.Cube: + Cube of wet-bulb freezing level. + + """ + from improver.utilities.cube_extraction import ExtractLevel + + wet_bulb_freezing_level = ExtractLevel( + positive_correlation=False, value_of_level=273.15 + )(wet_bulb_temperature) + wet_bulb_freezing_level.rename("wet_bulb_freezing_level") + + return wet_bulb_freezing_level diff --git a/improver/psychrometric_calculations/hail_size.py b/improver/psychrometric_calculations/hail_size.py index 2df2b76363..7d9974187e 100644 --- a/improver/psychrometric_calculations/hail_size.py +++ b/improver/psychrometric_calculations/hail_size.py @@ -48,7 +48,7 @@ saturated_humidity, ) from improver.utilities.cube_checker import assert_spatial_coords_match -from improver.utilities.cube_extraction import ExtractPressureLevel +from improver.utilities.cube_extraction import ExtractLevel from improver.utilities.cube_manipulation import enforce_coordinate_ordering @@ -525,8 +525,8 @@ def process( wet_bulb_zero_height_asl, orography, ) - extract_pressure = ExtractPressureLevel( - value_of_pressure_level=268.15, positive_correlation=True + extract_pressure = ExtractLevel( + value_of_level=268.15, positive_correlation=True ) pressure_at_268 = extract_pressure(temperature_on_pressure) @@ -534,9 +534,7 @@ def process( temperature_at_268.rename("temperature_of_atmosphere_at_268.15K") temperature_at_268.remove_coord("pressure") temperature = np.full_like( - temperature_at_268.data, - extract_pressure.value_of_pressure_level, - dtype=np.float32, + temperature_at_268.data, extract_pressure.value_of_level, dtype=np.float32, ) temperature = np.ma.masked_where(np.ma.getmask(pressure_at_268), temperature) temperature_at_268.data = temperature diff --git a/improver/utilities/cube_extraction.py b/improver/utilities/cube_extraction.py index 75f1ec4088..f0cf3072c3 100644 --- a/improver/utilities/cube_extraction.py +++ b/improver/utilities/cube_extraction.py @@ -30,13 +30,13 @@ # POSSIBILITY OF SUCH DAMAGE. """ Utilities to parse a list of constraints and extract matching subcube """ -import operator from ast import literal_eval from typing import Callable, Dict, List, Optional, Tuple, Union import numpy as np from iris import Constraint from iris.cube import Cube +from iris.exceptions import CoordinateNotFoundError from improver import BasePlugin from improver.metadata.constants import FLOAT_DTYPE @@ -312,49 +312,53 @@ def thin_cube(cube: Cube, thinning_dict: Dict[str, int]) -> Cube: return cube[tuple(slices)] -class ExtractPressureLevel(BasePlugin): +class ExtractLevel(BasePlugin): """ - Class for extracting a pressure surface from data-on-pressure-levels. + Class for extracting a pressure or height surface from data-on-levels. """ def __init__( - self, positive_correlation: bool, value_of_pressure_level: float, + self, positive_correlation: bool, value_of_level: float, ): """Sets up Class - Args: - positive_correlation: - Set to True when the variable generally increases as pressure increases - value_of_pressure_level: - The value of the input cube for which the pressure level is required + + Args: + positive_correlation: + Set to True when the variable generally increases as pressure increase + or when the variable generally increases as height decreases. + value_of_level: + The value of the input cube for which the pressure or height level + is required """ self.positive_correlation = positive_correlation - self.value_of_pressure_level = value_of_pressure_level + self.value_of_level = value_of_level + self.coordinate = None - @staticmethod - def pressure_grid(variable_on_pressure: Cube) -> np.ndarray: - """Creates a pressure grid of the same shape as variable_on_pressure cube. + def coordinate_grid(self, variable_on_levels: Cube) -> np.ndarray: + """Creates a pressure or height grid of the same shape as variable_on_levels cube. It is populated at every grid square and for every realization with - a column of all pressure levels taken from variable_on_pressure's pressure coordinate + a column of all pressure or height levels taken from variable_on_levels's pressure or height + coordinate Args: - variable_on_pressure: - Cube of some variable with pressure levels + variable_on_levels: + Cube of some variable with pressure or height levels Returns: - An n dimensional array with the same dimensions as variable_on_pressure containing, - at every grid square and for every realization, a column of all pressure levels - taken from variable_on_pressure's pressure coordinate + An n dimensional array with the same dimensions as variable_on_levels containing, + at every grid square and for every realization, a column of all pressure or + height levels taken from variable_on_levels's pressure or height coordinate """ - required_shape = variable_on_pressure.shape - pressure_points = variable_on_pressure.coord("pressure").points - (pressure_axis,) = variable_on_pressure.coord_dims("pressure") - pressure_shape = np.ones_like(required_shape) - pressure_shape[pressure_axis] = required_shape[pressure_axis] - pressure_array = np.broadcast_to( - pressure_points.reshape(pressure_shape), required_shape + required_shape = variable_on_levels.shape + coordinate_points = variable_on_levels.coord(self.coordinate).points + (coordinate_axis,) = variable_on_levels.coord_dims(self.coordinate) + coordinate_shape = np.ones_like(required_shape) + coordinate_shape[coordinate_axis] = required_shape[coordinate_axis] + coordinate_array = np.broadcast_to( + coordinate_points.reshape(coordinate_shape), required_shape ) - return pressure_array + return coordinate_array def fill_invalid(self, cube: Cube): """Populate any invalid values in the source data with the neighbouring value in that @@ -364,177 +368,148 @@ def fill_invalid(self, cube: Cube): Args: cube: - Cube of variable on pressure levels (3D) (modified in-place). + Cube of variable on levels (3D) (modified in-place). """ if np.isfinite(cube.data).all() and not np.ma.is_masked(cube.data): return data = np.ma.masked_invalid(cube.data) - (pressure_axis,) = cube.coord_dims("pressure") - pressure_points = cube.coord("pressure").points + (coordinate_axis,) = cube.coord_dims(self.coordinate) + coordinate_points = cube.coord(self.coordinate).points # Find the least significant increment to use as an offset for filling missing values. # We don't really care so long as it is non-zero and has the same sign. - increasing_p_order = np.all(np.diff(cube.coord("pressure").points) > 0) - sign = 1 if self.positive_correlation == increasing_p_order else -1 + increasing_order = np.all(np.diff(cube.coord(self.coordinate).points) > 0) + sign = 1 if self.positive_correlation == increasing_order else -1 v_increment = sign * 10 ** (-cube.attributes.get("least_significant_digit", 2)) self._one_way_fill( - data, pressure_axis, pressure_points, v_increment, reverse=True + data, coordinate_axis, coordinate_points, v_increment, reverse=True ) cube.data = data if np.isfinite(cube.data).all() and not np.ma.is_masked(cube.data): # We've filled everything in. No need to try again. return self._one_way_fill( - data, pressure_axis, pressure_points, v_increment, reverse=False + data, coordinate_axis, coordinate_points, v_increment, reverse=False ) cube.data = data @staticmethod def _one_way_fill( data: np.ma.MaskedArray, - pressure_axis: int, - pressure_points: np.ndarray, + coordinate_axis: int, + coordinate_points: np.ndarray, v_increment: np.ndarray, reverse: bool = False, ): """ - Scans through the pressure axis forwards or backwards, filling any missing data with the - previous value plus (or minus in reverse) the specified increment. Running this in both - directions will therefore populate all columns so long as there is at least one valid - data point to start with. + Scans through the pressure or height axis forwards or backwards, filling any missing + data with the previous value plus (or minus in reverse) the specified increment. + Running this in both directions will therefore populate all columns so long as there + is at least one valid data point to start with. """ - last_p_slice = [slice(None)] * data.ndim + last_slice = [slice(None)] * data.ndim if reverse: local_increment = -v_increment - first = len(pressure_points) - 2 + first = len(coordinate_points) - 2 last = -1 step = -1 - last_p_slice[pressure_axis] = slice( - len(pressure_points) - 1, len(pressure_points) + last_slice[coordinate_axis] = slice( + len(coordinate_points) - 1, len(coordinate_points) ) else: local_increment = v_increment first = 1 - last = len(pressure_points) + last = len(coordinate_points) step = 1 - last_p_slice[pressure_axis] = slice(0, 1) + last_slice[coordinate_axis] = slice(0, 1) - for p in range(first, last, step): - p_slice = [slice(None)] * data.ndim - p_slice[pressure_axis] = slice(p, p + 1) - data[p_slice] = np.ma.where( - data.mask[p_slice], data[last_p_slice] + local_increment, data[p_slice], + for c in range(first, last, step): + c_slice = [slice(None)] * data.ndim + c_slice[coordinate_axis] = slice(c, c + 1) + data[c_slice] = np.ma.where( + data.mask[c_slice], data[last_slice] + local_increment, data[c_slice], ) - last_p_slice = p_slice - - def fill_in_bounds( - self, pressure_of_variable: np.ma.MaskedArray, source_cube: Cube - ) -> np.ndarray: - """Update any undefined pressure_of_variable values with the maximum or minimum - pressure for that column. This occurs when the requested variable value falls above - or below the entire pressure column. - Maximum pressure is chosen if the maximum data value in the column is lower than - the value of self.value_of_pressure_level. - - Args: - pressure_of_variable: - 2D array of the pressure at the required variable value, masked True - where this method needs to fill it in. (modified in-place) - source_cube: - Cube of variable on pressure levels (3D) which shows where the lowest - and highest valid values are. - Returns: - Updated pressure_of_variable array - """ - if not np.ma.is_masked(pressure_of_variable): - return pressure_of_variable.data - pressure_coord = source_cube.coord("pressure") - max_pressure = pressure_coord.points.max() - min_pressure = pressure_coord.points.min() - (pressure_axis,) = source_cube.coord_dims("pressure") - max_pressure_index = np.argmax(pressure_coord.points) - # The values at the maximum pressure will be compared with the requested variable value - # using an appropriate operator based on whether value and pressure are increasing. - comparator = operator.lt if self.positive_correlation else operator.gt - values_at_max_pressure = source_cube.data.take( - axis=pressure_axis, indices=max_pressure_index - ) - # First, fill in missing values at the maximum end of the pressure coordinate: - pressure_of_variable = np.ma.where( - np.logical_and( - comparator(values_at_max_pressure, self.value_of_pressure_level), - pressure_of_variable.mask, - ), - max_pressure, - pressure_of_variable, - ) - # Now fill in remaining missing values which should be at the minimum end: - pressure_of_variable = np.where( - pressure_of_variable.mask, min_pressure, pressure_of_variable, - ) - return pressure_of_variable + last_slice = c_slice - def _make_pressure_cube( - self, result_data: np.array, variable_on_pressure_levels: Cube + def _make_template_cube( + self, result_data: np.array, variable_on_levels: Cube ) -> Cube: - """Creates a cube of the variable on a pressure level based on the input cube""" - pressure_cube = next( - variable_on_pressure_levels.slices_over(["pressure"]) - ).copy(result_data) - pressure_cube.rename( - "pressure_of_atmosphere_at_" - f"{self.value_of_pressure_level}{variable_on_pressure_levels.units}" + """Creates a cube of the variable on a pressure or height level based on the input cube""" + template_cube = next(variable_on_levels.slices_over([self.coordinate])).copy( + result_data ) - pressure_cube.units = variable_on_pressure_levels.coord("pressure").units - pressure_cube.remove_coord("pressure") - return pressure_cube + if self.coordinate == "pressure": + template_cube.rename( + "pressure_of_atmosphere_at_" + f"{self.value_of_level}{variable_on_levels.units}" + ) + else: + template_cube.rename( + "height_at_" f"{self.value_of_level}{variable_on_levels.units}" + ) + + template_cube.units = variable_on_levels.coord(self.coordinate).units + template_cube.remove_coord(self.coordinate) + return template_cube - def process(self, variable_on_pressure_levels: Cube) -> Cube: - """Extracts the pressure level where the environment - variable first intersects self.value_of_pressure_level starting at a pressure value - near the surface and ascending in altitude from there. - Where the pressure surface falls outside the available data, the maximum or minimum - pressure will be returned, even if the source data has no value at that point. + def process(self, variable_on_levels: Cube) -> Cube: + """Extracts the pressure or height level (depending on which is present on the cube) + where the environment variable first intersects self.value_of_level starting at a + pressure or height value near the surface and ascending in altitude from there. + Where the surface falls outside the available data, the maximum or minimum + of the surface will be returned, even if the source data has no value at that point. Args: - variable_on_pressure_levels: - A cube of data on pressure levels + variable_on_levels: + A cube of data on pressure or height levels Returns: - A cube of the environment pressure at self.value_of_pressure_level + A cube of the environment pressure or height at self.value_of_level + + Raises: + NotImplementError: + If variable_on_levels has both a height and pressure coordinate """ - from stratify import interpolate + from stratify import EXTRAPOLATE_NEAREST, interpolate + + coord_list = [coord.name() for coord in variable_on_levels.coords()] + if "height" in coord_list and "pressure" in coord_list: + raise NotImplementedError( + """Input Cube has both a pressure and height coordinate. + Only one of these should be present on the cube.""" + ) - self.fill_invalid(variable_on_pressure_levels) + try: + self.coordinate = variable_on_levels.coord("pressure").name() + except CoordinateNotFoundError: + self.coordinate = variable_on_levels.coord("height").name() - if "realization" in [c.name() for c in variable_on_pressure_levels.dim_coords]: - slicer = variable_on_pressure_levels.slices_over((["realization"])) - one_slice = variable_on_pressure_levels.slices_over( - (["realization"]) - ).next() + self.fill_invalid(variable_on_levels) + + if "realization" in [c.name() for c in variable_on_levels.dim_coords]: + slicer = variable_on_levels.slices_over((["realization"])) + one_slice = variable_on_levels.slices_over((["realization"])).next() has_r_coord = True else: - slicer = [variable_on_pressure_levels] - one_slice = variable_on_pressure_levels + slicer = [variable_on_levels] + one_slice = variable_on_levels has_r_coord = False - p_grid = self.pressure_grid(one_slice).astype(np.float32) - (pressure_axis,) = one_slice.coord_dims("pressure") + grid = self.coordinate_grid(one_slice).astype(np.float32) + (coordinate_axis,) = one_slice.coord_dims(self.coordinate) result_data = np.empty_like( - variable_on_pressure_levels.slices_over((["pressure"])).next().data + variable_on_levels.slices_over(([self.coordinate])).next().data ) for i, zyx_slice in enumerate(slicer): interp_data = interpolate( - np.array([self.value_of_pressure_level], dtype=np.float32), + np.array([self.value_of_level], dtype=np.float32), zyx_slice.data.data, - p_grid, - axis=pressure_axis, - ).squeeze(axis=pressure_axis) + grid, + axis=coordinate_axis, + rising=False, + extrapolation=EXTRAPOLATE_NEAREST, + ).squeeze(axis=coordinate_axis) if has_r_coord: result_data[i] = interp_data else: result_data = interp_data - result_data = np.ma.masked_invalid(result_data) - result_data = self.fill_in_bounds(result_data, variable_on_pressure_levels) - pressure_cube = self._make_pressure_cube( - result_data, variable_on_pressure_levels - ) - return pressure_cube + output_cube = self._make_template_cube(result_data, variable_on_levels) + return output_cube diff --git a/improver_tests/acceptance/SHA256SUMS b/improver_tests/acceptance/SHA256SUMS index 38b6099454..8516326944 100644 --- a/improver_tests/acceptance/SHA256SUMS +++ b/improver_tests/acceptance/SHA256SUMS @@ -788,6 +788,8 @@ dc56c3290bae734db3bf479628e5da0b93ccada06a68606b22e25a07a4fc24b5 ./weighted_ble dc84cdaafb5a0b8c48e648a6734be019cc4b7b9593c00607667059a71a4c571d ./weighted_blending/three_models/nc/20190101T0400Z-PT0001H00M-precip_rate.nc e535f9b75ca96622b96916bc249a4aa919e5216fcf548a7f6c90748538bfe22a ./weighted_blending/three_models/ukvx/20190101T0400Z-PT0002H00M-precip_rate.nc 32042d75b04735b7f4b8db22f793cacb8a71c455f9442dc14c029d78fe59a80c ./weighted_blending/weights_from_dict/kgo.nc +2e422d6952a2cef43937ad37137500efefe3ca1e16adfaa3608a0b49f5d482db ./wet-bulb-freezing-level/kgo.nc +3d6b8720c087c4ef8daf32ff761695e0f6169104fb134d03c26dd17a5e4d4a02 ./wet-bulb-freezing-level/wet_bulb_temperature.nc ef62d78071045163858387292ec56c11aa92e1fa6f5580c1e4ac485cd01979fc ./wet-bulb-temperature-integral/basic/input.nc 4143e7583d66c6e776e43e619622eab00fdee3293becadfaaa00ddd72147ac6a ./wet-bulb-temperature-integral/basic/with_id_attr/kgo.nc f854f85ef58e2a0dc41c3bcdf940a23daa2782a50d3eaf798335fefa8bea21b8 ./wet-bulb-temperature-integral/basic/without_id_attr/kgo.nc diff --git a/improver_tests/acceptance/test_wet_bulb_freezing_level.py b/improver_tests/acceptance/test_wet_bulb_freezing_level.py new file mode 100644 index 0000000000..598d9e6b0c --- /dev/null +++ b/improver_tests/acceptance/test_wet_bulb_freezing_level.py @@ -0,0 +1,54 @@ +# -*- coding: utf-8 -*- +# ----------------------------------------------------------------------------- +# (C) British Crown copyright. The Met Office. +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. +"""Tests for the wet-bulb-freezing-level CLI""" + +import pytest + +from . import acceptance as acc + +pytestmark = [pytest.mark.acc, acc.skip_if_kgo_missing] +CLI = acc.cli_name_with_dashes(__file__) +run_cli = acc.run_cli(CLI) + + +def test_basic(tmp_path): + """Test basic wet bulb freezing level calculation""" + kgo_dir = acc.kgo_root() / "wet-bulb-freezing-level" + kgo_path = kgo_dir / "kgo.nc" + output_path = tmp_path / "output.nc" + + args = [ + kgo_dir / "wet_bulb_temperature.nc", + "--output", + output_path, + ] + run_cli(args) + acc.compare(output_path, kgo_path) diff --git a/improver_tests/utilities/cube_extraction/test_ExtractPressureLevel.py b/improver_tests/utilities/cube_extraction/test_ExtractLevel.py similarity index 58% rename from improver_tests/utilities/cube_extraction/test_ExtractPressureLevel.py rename to improver_tests/utilities/cube_extraction/test_ExtractLevel.py index e341b4edd8..b1eaecf9fd 100644 --- a/improver_tests/utilities/cube_extraction/test_ExtractPressureLevel.py +++ b/improver_tests/utilities/cube_extraction/test_ExtractLevel.py @@ -28,7 +28,7 @@ # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE # POSSIBILITY OF SUCH DAMAGE. -"""Unit tests for the ExtractPressureLevel plugin""" +"""Unit tests for the ExtractLevel plugin""" from collections.abc import Iterable import numpy as np @@ -36,7 +36,7 @@ from iris.cube import Cube from improver.synthetic_data.set_up_test_cubes import set_up_variable_cube -from improver.utilities.cube_extraction import ExtractPressureLevel +from improver.utilities.cube_extraction import ExtractLevel LOCAL_MANDATORY_ATTRIBUTES = { "title": "unit test data", @@ -64,20 +64,41 @@ def temperature_on_pressure_levels() -> Cube: return t_cube -def metadata_check(pressure_slice_cube: Cube, value: float, units: str): +@pytest.fixture +def temperature_on_height_levels() -> Cube: + """Set up a r, h, y, x cube of temperature of atmosphere on height levels""" + temperatures = np.array([300, 286, 280, 274, 267, 262, 257, 245], dtype=np.float32) + data = np.broadcast_to( + temperatures.reshape((1, len(temperatures), 1, 1)), (2, len(temperatures), 3, 2) + ) + t_cube = set_up_variable_cube( + data, + height_levels=np.arange(0, 8000, 1000), + name="temperature_on_height_levels", + units="K", + attributes=LOCAL_MANDATORY_ATTRIBUTES, + ) + return t_cube + + +def metadata_check(cube_slice: Cube, value: float, units: str, coordinate: str): """Checks the cube produced by plugin has the expected metadata.""" - assert pressure_slice_cube.long_name == f"pressure_of_atmosphere_at_{value}{units}" - assert pressure_slice_cube.units == "Pa" - assert pressure_slice_cube.attributes == { + if coordinate == "pressure": + assert cube_slice.long_name == f"pressure_of_atmosphere_at_{value}{units}" + assert cube_slice.units == "Pa" + else: + assert cube_slice.long_name == f"height_at_{value}{units}" + assert cube_slice.units == "m" + assert cube_slice.attributes == { "title": "unit test data", "source": "unit test", "institution": "somewhere", } -def cube_shape_check_with_realizations(pressure_slice_cube): +def cube_shape_check_with_realizations(cube_slice): """Checks cube coordinates and dimensions when two realizations are present""" - coord_names = [coord.name() for coord in pressure_slice_cube.coords()] + coord_names = [coord.name() for coord in cube_slice.coords()] assert coord_names == [ "realization", "latitude", @@ -86,12 +107,12 @@ def cube_shape_check_with_realizations(pressure_slice_cube): "forecast_reference_time", "time", ] - assert pressure_slice_cube.shape == (2, 3, 2) + assert cube_slice.shape == (2, 3, 2) -def cube_shape_check_without_realizations(pressure_slice_cube): +def cube_shape_check_without_realizations(cube_slice): """Checks cube coordinates and dimensions when realization is a scalar coord""" - coord_names = [coord.name() for coord in pressure_slice_cube.coords()] + coord_names = [coord.name() for coord in cube_slice.coords()] assert coord_names == [ "latitude", "longitude", @@ -100,97 +121,106 @@ def cube_shape_check_without_realizations(pressure_slice_cube): "realization", "time", ] - assert pressure_slice_cube.shape == (3, 2) + assert cube_slice.shape == (3, 2) +@pytest.mark.parametrize("coordinate", ("pressure", "height")) @pytest.mark.parametrize("least_significant_digit", (0, None)) -@pytest.mark.parametrize("reverse_pressure", (False, True)) +@pytest.mark.parametrize("reverse_coordinate", (False, True)) @pytest.mark.parametrize( "special_value", (None, np.nan, True, np.inf, (np.nan, np.nan)) ) @pytest.mark.parametrize("with_realization", (True, False)) @pytest.mark.parametrize( - "temperature,expected_p_index", + "temperature,expected_index", ( - (280, 2), # Exactly matches a pressure value - (277, 2.5), # Half way between pressure values - (301, 0), # Temperature above max snaps to pressure at max - (244, 7), # Temperature below min snaps to pressure at min + (280, 2), # Exactly matches a value + (277, 2.5), # Half way between values + (301, 0), # Temperature above max snaps to at max + (244, 7), # Temperature below min snaps to at min ), ) def test_basic( temperature, - temperature_on_pressure_levels, - expected_p_index, + request, + expected_index, with_realization, special_value, - reverse_pressure, + reverse_coordinate, least_significant_digit, + coordinate, ): - """Tests the ExtractPressureLevel plugin with values for temperature and - temperature on pressure levels to check for expected result. - Tests behaviour when temperature and/or pressure increase or decrease along - the pressure axis. + """Tests the ExtractLevel plugin with values for temperature and + temperature on levels to check for expected result. + Tests behaviour when temperature and/or pressure/height increase or decrease along + the pressure/height axis. Tests behaviour with different special values in the temperature data. Tests behaviour with and without a realization coordinate. + Tests behaviour when extracting a pressure level or height level. Also checks the metadata of the output cube""" + if coordinate == "height": + temperature_on_levels = request.getfixturevalue("temperature_on_height_levels") + else: + temperature_on_levels = request.getfixturevalue( + "temperature_on_pressure_levels" + ) special_value_index = 0 positive_correlation = True - if reverse_pressure: + if reverse_coordinate: # Flip the pressure coordinate for this test. We also swap which end the # special value goes, so we can test _one_way_fill in both modes. - temperature_on_pressure_levels.coord( - "pressure" - ).points = temperature_on_pressure_levels.coord("pressure").points[::-1] + temperature_on_levels.coord(coordinate).points = temperature_on_levels.coord( + coordinate + ).points[::-1] special_value_index = -1 positive_correlation = False + + if coordinate == "height": + # height has the opposite correlation to pressure + positive_correlation = not positive_correlation expected = np.interp( - expected_p_index, - range(len(temperature_on_pressure_levels.coord("pressure").points)), - temperature_on_pressure_levels.coord("pressure").points, - ) - expected_data = np.full_like( - temperature_on_pressure_levels.data[:, 0, ...], expected + expected_index, + range(len(temperature_on_levels.coord(coordinate).points)), + temperature_on_levels.coord(coordinate).points, ) + expected_data = np.full_like(temperature_on_levels.data[:, 0, ...], expected) + if special_value is True: # This is a proxy for setting a mask=True entry - temperature_on_pressure_levels.data = np.ma.MaskedArray( - temperature_on_pressure_levels.data, mask=False + temperature_on_levels.data = np.ma.MaskedArray( + temperature_on_levels.data, mask=False ) - temperature_on_pressure_levels.data.mask[ - 0, special_value_index, 0, 0 - ] = special_value + temperature_on_levels.data.mask[0, special_value_index, 0, 0] = special_value elif special_value is None: pass else: - temperature_on_pressure_levels.data = temperature_on_pressure_levels.data.copy() + temperature_on_levels.data = temperature_on_levels.data.copy() if isinstance(special_value, Iterable): # This catches the test case where two consecutive special values are to be used if special_value_index < 0: - temperature_on_pressure_levels.data[0, -2:, 0, 0] = special_value + temperature_on_levels.data[0, -2:, 0, 0] = special_value else: - temperature_on_pressure_levels.data[0, 0:2, 0, 0] = special_value + temperature_on_levels.data[0, 0:2, 0, 0] = special_value else: - temperature_on_pressure_levels.data[ - 0, special_value_index, 0, 0 - ] = special_value + temperature_on_levels.data[0, special_value_index, 0, 0] = special_value if not with_realization: - temperature_on_pressure_levels = temperature_on_pressure_levels[0] + temperature_on_levels = temperature_on_levels[0] expected_data = expected_data[0] if least_significant_digit: - temperature_on_pressure_levels.attributes[ + temperature_on_levels.attributes[ "least_significant_digit" ] = least_significant_digit - result = ExtractPressureLevel( - value_of_pressure_level=temperature, positive_correlation=positive_correlation - )(temperature_on_pressure_levels) + result = ExtractLevel( + value_of_level=temperature, positive_correlation=positive_correlation + )(temperature_on_levels) + assert not np.ma.is_masked(result.data) np.testing.assert_array_almost_equal(result.data, expected_data) - metadata_check(result, temperature, temperature_on_pressure_levels.units) + metadata_check(result, temperature, temperature_on_levels.units, coordinate) if with_realization: cube_shape_check_with_realizations(result) else: @@ -213,7 +243,7 @@ def test_basic( def test_only_one_point( temperature_on_pressure_levels, index, expected, special_value, ): - """Tests the ExtractPressureLevel plugin with the unusual case that only one layer has + """Tests the ExtractLevel plugin with the unusual case that only one layer has a valid value. """ temperature_on_pressure_levels = temperature_on_pressure_levels[0] @@ -233,8 +263,21 @@ def test_only_one_point( expected_data = np.full_like(temperature_on_pressure_levels.data[0, ...], 80000) expected_data[0, 0] = expected - result = ExtractPressureLevel( - value_of_pressure_level=280, positive_correlation=True - )(temperature_on_pressure_levels) + result = ExtractLevel(value_of_level=280, positive_correlation=True)( + temperature_on_pressure_levels + ) assert not np.ma.is_masked(result.data) np.testing.assert_array_almost_equal(result.data, expected_data) + + +def test_both_pressure_and_height_error(temperature_on_height_levels): + """Tests an error is raised if both pressure and height coordinates are present + on the input cube""" + temperature_on_height_levels.coord("realization").rename("pressure") + with pytest.raises( + NotImplementedError, + match="Input Cube has both a pressure and height coordinate.", + ): + ExtractLevel(value_of_level=277, positive_correlation=True)( + temperature_on_height_levels + )