From 89fcd30bf9e86241eb7542f837e3302f0d7dd358 Mon Sep 17 00:00:00 2001 From: mspelman07 <99179165+mspelman07@users.noreply.github.com> Date: Tue, 27 Jun 2023 11:56:14 +0100 Subject: [PATCH] Mobt515 cloud base height spot extraction (#1911) * Adds plugin to adjust height above ground level forecasts and relevant tests * corrects handling of NAN data * formatting and removes un-needed test * formatting * formatting * update docstring * removes comment * add neighbour_selection_method to cli * correct spelling * formatting * changes following review * formatting * Updates following review comments --------- Co-authored-by: Marcus Spelman --- improver/cli/apply_height_adjustment.py | 88 ++++++ improver/spotdata/apply_lapse_rate.py | 27 +- improver/spotdata/height_adjustment.py | 213 ++++++++++++++ improver_tests/acceptance/SHA256SUMS | 5 + .../test_apply_height_adjustment.py | 55 ++++ .../spotdata/test_HeightAdjustment.py | 260 ++++++++++++++++++ 6 files changed, 636 insertions(+), 12 deletions(-) create mode 100644 improver/cli/apply_height_adjustment.py create mode 100644 improver/spotdata/height_adjustment.py create mode 100644 improver_tests/acceptance/test_apply_height_adjustment.py create mode 100644 improver_tests/spotdata/test_HeightAdjustment.py diff --git a/improver/cli/apply_height_adjustment.py b/improver/cli/apply_height_adjustment.py new file mode 100644 index 0000000000..f281ebf922 --- /dev/null +++ b/improver/cli/apply_height_adjustment.py @@ -0,0 +1,88 @@ +#!/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. +"""Script to apply height adjustments to spot data.""" + +from improver import cli + + +@cli.clizefy +@cli.with_output +def process( + spot_cube: cli.inputcube, + neighbour: cli.inputcube, + *, + land_constraint: bool = False, + similar_altitude: bool = False, +): + """Apply height adjustment to account for the difference between site altitude and + grid square orography. The spot forecast contains information representative of the + associated grid point. This needs to be adjusted to reflect the true site altitude. + + Args: + spot_cube (iris.cube.Cube): + A cube of spot forecasts. If this is a cube of probabilities + then the units of the threshold coordinate must be convertible to + metres as this is expected to represent a vertical coordinate. + If this is a cube of percentiles or realizations then the + units of the cube must be convertible to metres as the cube is + expected to represent a vertical profile. + neighbour (iris.cube.Cube): + A cube containing information about spot-data neighbours and + the spot site information. + land_constraint (bool): + Use to select the nearest-with-land-constraint neighbour-selection + method from the neighbour_cube. This means that the grid points + should be land points except for sites where none were found within + the search radius when the neighbour cube was created. May be used + with similar_altitude. + similar_altitude (bool): + Use to select the nearest-with-height-constraint + neighbour-selection method from the neighbour_cube. These are grid + points that were found to be the closest in altitude to the spot + site within the search radius defined when the neighbour cube was + created. May be used with land_constraint. + + Returns: + iris.cube.Cube: + A cube of spot data values with the same metadata as spot_cube but with data + adjusted to be relative to site height rather than orography grid square + height + """ + from improver.spotdata.height_adjustment import SpotHeightAdjustment + from improver.spotdata.neighbour_finding import NeighbourSelection + + neighbour_selection_method = NeighbourSelection( + land_constraint=land_constraint, minimum_dz=similar_altitude + ).neighbour_finding_method_name() + + result = SpotHeightAdjustment(neighbour_selection_method)(spot_cube, neighbour) + return result diff --git a/improver/spotdata/apply_lapse_rate.py b/improver/spotdata/apply_lapse_rate.py index dcf9f4d32b..a9c298b901 100644 --- a/improver/spotdata/apply_lapse_rate.py +++ b/improver/spotdata/apply_lapse_rate.py @@ -200,17 +200,6 @@ def broadcast_fixed_lapse_rate(self, spot_data_cube: Cube) -> np.ndarray: """Create an array of fixed lapse rate values""" return np.full(spot_data_cube.shape, self.fixed_lapse_rate, dtype=np.float32) - def extract_vertical_displacements(self, neighbour_cube: Cube) -> Cube: - """Extract vertical displacements between the model orography and sites.""" - method_constraint = iris.Constraint( - neighbour_selection_method_name=self.neighbour_selection_method - ) - data_constraint = iris.Constraint(grid_attributes_key="vertical_displacement") - vertical_displacement = neighbour_cube.extract( - method_constraint & data_constraint - ) - return vertical_displacement - def process( self, spot_data_cube: Cube, @@ -259,7 +248,9 @@ def process( spot_data_cube, neighbour_cube, gridded_lapse_rate_cube ) - vertical_displacement = self.extract_vertical_displacements(neighbour_cube) + vertical_displacement = extract_vertical_displacements( + neighbour_cube, self.neighbour_selection_method + ) new_temperatures = ( spot_data_cube.data @@ -268,3 +259,15 @@ def process( ) ).astype(np.float32) return spot_data_cube.copy(data=new_temperatures) + + +def extract_vertical_displacements( + neighbour_cube: Cube, neighbour_selection_method_name: str +) -> Cube: + """Extract vertical displacements between the model orography and sites.""" + method_constraint = iris.Constraint( + neighbour_selection_method_name=neighbour_selection_method_name + ) + data_constraint = iris.Constraint(grid_attributes_key="vertical_displacement") + vertical_displacement = neighbour_cube.extract(method_constraint & data_constraint) + return vertical_displacement diff --git a/improver/spotdata/height_adjustment.py b/improver/spotdata/height_adjustment.py new file mode 100644 index 0000000000..8b26171666 --- /dev/null +++ b/improver/spotdata/height_adjustment.py @@ -0,0 +1,213 @@ +# -*- 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. + +"""Applies height adjustment for height above ground level spot forecasts.""" + +from itertools import product + +import iris +import numpy as np +from iris.coords import DimCoord +from iris.cube import Cube +from iris.exceptions import CoordinateNotFoundError +from scipy.interpolate import LinearNDInterpolator + +from improver import BasePlugin +from improver.metadata.probabilistic import find_threshold_coordinate, is_probability +from improver.spotdata.apply_lapse_rate import extract_vertical_displacements +from improver.utilities.cube_manipulation import enforce_coordinate_ordering + + +class SpotHeightAdjustment(BasePlugin): + """ + Class to adjust spot extracted "height above ground level" forecasts to account + for differences between site height and orography grid square height. The spot + forecast contains information representative of the associated grid point and as + such this needs to be adjusted to reflect the true site altitude. + + For realization or percentile data the vertical displacement is added on to + each realization or percentile. + + For probability data the data is interpolated between thresholds for each site and + the equivalent set of thresholds relative to the site altitude extracted. Any + new threshold that is above or below the original set of thresholds uses the highest + or lowest threshold's probability from the original cube respectively for that spot. + """ + + def __init__(self, neighbour_selection_method: str = "nearest",) -> None: + """ + Args: + neighbour_selection_method: + The neighbour cube may contain one or several sets of grid + coordinates that match a spot site. These are determined by + the neighbour finding method employed. This keyword is used to + extract the desired set of coordinates from the neighbour cube. + """ + self.neighbour_selection_method = neighbour_selection_method + self.threshold_coord = DimCoord + self.units = None + + def adjust_prob_cube(self, spot_cube: Cube, vertical_displacement: Cube) -> Cube: + """ + Adjust probability spot forecasts based on the vertical displacement of sites + in relation to orography. + + Args: + spot_cube: + A cube of spot forecasts. If this is a cube of probabilities + then the units of the threshold coordinate must be convertible to + metres as this is expected to represent a vertical coordinate. There + must be at least two thresholds on this cube. + If this is a cube of percentiles or realizations then the + units of the cube must be convertible to metres as the cube is + expected to represent a vertical profile. + vertical_displacement: + A cube containing information about the difference between spot + data site height and the orography grid square height. + + Returns: + A cube with the same metadata and shape as spot_cube but with probabilities + adjusted to be relative to the site altitude rather than grid square altitude. + """ + + enforce_coordinate_ordering( + spot_cube, ["spot_index", self.threshold_coord.name()] + ) + n_sites = len(spot_cube.coord("spot_index").points) + n_thresholds = len(self.threshold_coord.points) + + thresholds = self.threshold_coord.points + spot_index = spot_cube.coord("spot_index").points + shape = spot_cube.shape + + # Maximum probability over all thresholds for each site. + broadcast_max = np.transpose( + np.broadcast_to(np.amax(spot_cube.data, axis=1), (n_thresholds, n_sites)) + ) + # Minimum probability over all thresholds for each site. + broadcast_min = np.transpose( + np.broadcast_to(np.amin(spot_cube.data, axis=1), (n_thresholds, n_sites)) + ) + + broadcast_thresholds = np.broadcast_to(thresholds, (n_sites, n_thresholds)) + broadcast_vertical_displacement = np.transpose( + np.broadcast_to(vertical_displacement.data, (n_thresholds, n_sites),) + ) + desired_thresholds = broadcast_thresholds + broadcast_vertical_displacement.data + + # creates a list of pairs of values of spot index with the thresholds that need to + # be calculated for the spot index + coord = list(product(spot_index, thresholds)) + needed_pair = [] + for index, threshold in zip(spot_index, desired_thresholds): + needed_pair.extend(list(product([index], threshold))) + + # interpolate across the cube and request needed thresholds + interp = LinearNDInterpolator(coord, spot_cube.data.flatten()) + spot_data = np.reshape(interp(needed_pair), shape) + + # Points outside the range of the original data return NAN. These points are replaced + # with the highest or lowest along the axis depending on the whether the vertical + # displacement was positive or negative + indices = np.where(np.isnan(spot_data)) + spot_data[indices] = np.where( + broadcast_vertical_displacement[indices] > 0, + broadcast_max[indices], + broadcast_min[indices], + ) + spot_cube.data = spot_data + return spot_cube + + def process(self, spot_cube: Cube, neighbour: Cube,) -> Cube: + """ + Adjusts spot forecast data to be relative to site height rather than + grid square orography height. + + Args: + spot_cube: + A cube of spot forecasts. If this is a cube of probabilities + then the units of the threshold coordinate must be convertible to + metres as this is expected to represent a vertical coordinate. There + must be at least two thresholds on this cube. + If this is a cube of percentiles or realizations then the + units of the cube must be convertible to metres as the cube is + expected to represent a vertical profile. + neighbour: + A cube containing information about the spot data sites and + their grid point neighbours. + Returns: + A cube of the same shape as spot_data but with data adjusted to account for + the difference between site height and orography height. + + Raises: + ValueError: + If spot_cube is a probability cube and there are fewer than two thresholds. + """ + vertical_displacement = extract_vertical_displacements( + neighbour_cube=neighbour, + neighbour_selection_method_name=self.neighbour_selection_method, + ) + + if is_probability(spot_cube): + threshold_coord = find_threshold_coordinate(spot_cube) + self.threshold_coord = threshold_coord.copy() + + if len(self.threshold_coord.points) < 2: + raise ValueError( + f"There are fewer than 2 thresholds present in this cube, this is " + f"{spot_cube.coord(threshold_coord).points}. At least two thresholds " + "are needed for interpolation" + ) + + self.units = self.threshold_coord.units + self.threshold_coord.convert_units("m") + + try: + cube_slices = [x for x in spot_cube.slices_over("realization")] + except CoordinateNotFoundError: + cube_slices = [spot_cube] + coord_list = [c.name() for c in spot_cube.dim_coords] + spot_data = iris.cube.CubeList() + for cube_slice in cube_slices: + spot_data.append( + self.adjust_prob_cube(cube_slice, vertical_displacement) + ) + spot_cube = spot_data.merge_cube() + enforce_coordinate_ordering(spot_cube, coord_list) + + else: + self.units = spot_cube.units + spot_cube.convert_units("m") + + spot_cube.data = spot_cube.data + vertical_displacement.data + spot_cube.convert_units(self.units) + spot_cube.data = spot_cube.data.astype(np.float32) + return spot_cube diff --git a/improver_tests/acceptance/SHA256SUMS b/improver_tests/acceptance/SHA256SUMS index 005ba9a470..5b816ab3a8 100644 --- a/improver_tests/acceptance/SHA256SUMS +++ b/improver_tests/acceptance/SHA256SUMS @@ -52,6 +52,11 @@ a0112b5a48ba99a1a4a345d43b0c453caaf25181504fb9a13786b5722f84cc10 ./apply-emos-c 965a1f0565f9192d95eb01d0a61dc7bede6902f5e1a0481d92611e677841a139 ./apply-emos-coefficients/truncated_normal/input.nc feb2d93cf96b05889571bcb348dfbb4d60cfc1f3d9b7343dcf7d85cf34339746 ./apply-emos-coefficients/truncated_normal/kgo.nc d2cab1d3d8aa588be08a3d7d65e95e859fed37daa767e8d4a2bdaae25702b9a8 ./apply-emos-coefficients/truncated_normal/truncated_normal_coefficients.nc +8330e013c590d2cadf4f39d35a40c967d46f5695097ef36dfb96fb89069af18b ./apply-height-adjustment/input_prob.nc +0caffe65e2178cbec6e818a3a0c17122c3cb706f124afa6e33e328a66f215457 ./apply-height-adjustment/input_realization.nc +2fda31975f564f8f2a79cb1a7fe4b662e04017f4b63c70d15685e359edf07e8d ./apply-height-adjustment/kgo_prob.nc +df52902fd9cfab5c2f583f13af401e0c737ec7e73e9bb7907df82f230184f34f ./apply-height-adjustment/kgo_realization.nc +5495b2382fb2d33f34c726134496b0f198cc6378031b50c0e450f5ade5c17f83 ./apply-height-adjustment/neighbours.nc abc718f2469ecbe0c4cd84e26f0adde87a65b1ae5fc17da60df27d9a3fe4871c ./apply-lapse-rate/basic/kgo.nc c7eb9bab2ad43ac19ecc071730479a9f27a58992fcb22050743d34cdf2ad9639 ./apply-lapse-rate/basic/ukvx_lapse_rate.nc 34efbac81d20f8cbae8f7881d984ce7fc9fb59d3f54e478611a914bf94b80dfc ./apply-lapse-rate/basic/ukvx_orography.nc diff --git a/improver_tests/acceptance/test_apply_height_adjustment.py b/improver_tests/acceptance/test_apply_height_adjustment.py new file mode 100644 index 0000000000..adb6be4955 --- /dev/null +++ b/improver_tests/acceptance/test_apply_height_adjustment.py @@ -0,0 +1,55 @@ +# -*- 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 apply-height-adjustment 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) + + +@pytest.mark.parametrize("input_type", ("prob", "realization")) +def test_basic(tmp_path, input_type): + """Test apply-height-adjustment for a probability and realization input cube""" + kgo_dir = acc.kgo_root() / "apply-height-adjustment/" + kgo_path = kgo_dir / f"kgo_{input_type}.nc" + output_path = tmp_path / "output.nc" + args = [ + kgo_dir / f"input_{input_type}.nc", + kgo_dir / "neighbours.nc", + "--output", + output_path, + ] + run_cli(args) + acc.compare(output_path, kgo_path) diff --git a/improver_tests/spotdata/test_HeightAdjustment.py b/improver_tests/spotdata/test_HeightAdjustment.py new file mode 100644 index 0000000000..277fa85803 --- /dev/null +++ b/improver_tests/spotdata/test_HeightAdjustment.py @@ -0,0 +1,260 @@ +# -*- 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. +"""Unit tests for SpotHeightAdjustment plugin""" + +import iris +import numpy as np +import pytest +from iris.coords import DimCoord +from iris.cube import Cube +from iris.exceptions import CoordinateNotFoundError + +from improver.spotdata.build_spotdata_cube import build_spotdata_cube +from improver.spotdata.height_adjustment import SpotHeightAdjustment +from improver.utilities.cube_manipulation import enforce_coordinate_ordering + +name = "cloud_base_height_assuming_only_consider_cloud_area_fraction_greater_than_4p5_oktas" + + +@pytest.fixture() +def prob_cube() -> Cube: + """Set up a spot data cube of probabilities""" + altitude = np.array([256.5, 359.1, 301.8, 406.2]) + latitude = np.linspace(58.0, 59.5, 4) + longitude = np.linspace(-0.25, 0.5, 4) + wmo_id = ["03854", "03962", "03142", "03331"] + threshold_coord = DimCoord( + points=[50, 100, 1000], var_name="threshold", long_name=name, units="m", + ) + + data = np.asarray([[0.1, 0, 0.2, 0.1], [0.7, 0, 0.3, 0.2], [1, 0.4, 0.4, 0.9]]) + + prob_cube = build_spotdata_cube( + data, + name="probability_of_" + name + "_below_threshold", + units="1", + altitude=altitude, + latitude=latitude, + longitude=longitude, + wmo_id=wmo_id, + additional_dims=[threshold_coord], + ) + enforce_coordinate_ordering(prob_cube, ["spot_index", name]) + return prob_cube + + +@pytest.fixture() +def prob_cube_realizations(prob_cube) -> Cube: + """Set up a spot data cube of probabilities with a realization coordinate""" + realization0 = iris.coords.DimCoord([0], standard_name="realization", units=1) + realization1 = iris.coords.DimCoord([1], standard_name="realization", units=1) + + cube0 = prob_cube.copy() + cube1 = prob_cube.copy() + cube0.add_aux_coord(realization0) + cube1.add_aux_coord(realization1) + cubes = iris.cube.CubeList([cube0, cube1]) + cube = cubes.merge_cube() + return cube + + +@pytest.fixture() +def realization_cube() -> Cube: + """Set up a spot data cube with a realization coordinate. The units of this cube + are set to feet so unit conversion can be tested within the plugin""" + altitude = np.array([256.5, 359.1, 301.8, 406.2]) + latitude = np.linspace(58.0, 59.5, 4) + longitude = np.linspace(-0.25, 0.5, 4) + wmo_id = ["03854", "03962", "03142", "03331"] + realization_coord = DimCoord(points=[0, 1, 2], var_name="realization", units="1") + + data = np.asarray( + [[1000, 4000, -200, 100], [3000, 10000, 0, 200], [4000, 11000, 30, 150]] + ) + + rea_cube = build_spotdata_cube( + data, + name=name, + units="ft", + altitude=altitude, + latitude=latitude, + longitude=longitude, + wmo_id=wmo_id, + additional_dims=[realization_coord], + ) + return rea_cube + + +@pytest.fixture() +def percentile_cube() -> Cube: + """Set up a spot data cube with a percentile coordinate""" + altitude = np.array([256.5, 359.1, 301.8, 406.2]) + latitude = np.linspace(58.0, 59.5, 4) + longitude = np.linspace(-0.25, 0.5, 4) + wmo_id = ["03854", "03962", "03142", "03331"] + percentile_coord = DimCoord(points=[25, 50, 75], var_name="percentile", units="%") + + data = np.asarray( + [[1400, 4000, -500, 100], [2000, 10000, 0, 100], [3000, 11000, 30, 150]] + ) + + perc_cube = build_spotdata_cube( + data, + name=name, + units="m", + altitude=altitude, + latitude=latitude, + longitude=longitude, + wmo_id=wmo_id, + additional_dims=[percentile_coord], + ) + return perc_cube + + +@pytest.fixture() +def neighbour_cube() -> Cube: + """Set up a neighbour cube with vertical displacement coordinate""" + neighbours = np.array([[[0.0, -100.0, 0.0, 100.0]]]) + + altitudes = np.array([0, 1, 3, 2]) + latitudes = np.array([10, 10, 20, 20]) + longitudes = np.array([10, 10, 20, 20]) + wmo_ids = np.arange(4) + neighbour_cube = build_spotdata_cube( + neighbours, + "grid_neighbours", + 1, + altitudes, + latitudes, + longitudes, + wmo_ids, + neighbour_methods=["nearest"], + grid_attributes=["vertical_displacement"], + ) + return neighbour_cube + + +@pytest.mark.parametrize("order", (True, False)) +@pytest.mark.parametrize( + "cube_name, expected", + ( + ( + "prob_cube", + [ + [0.1, 0.7, 1], + [0, 0, 0.35555556], + [0.2, 0.3, 0.4], + [0.23888889, 0.27777778, 0.9], + ], + ), + ( + "prob_cube_realizations", + [ + [ + [0.1, 0.7, 1], + [0, 0, 0.35555556], + [0.2, 0.3, 0.4], + [0.23888889, 0.27777778, 0.9], + ], + [ + [0.1, 0.7, 1], + [0, 0, 0.35555556], + [0.2, 0.3, 0.4], + [0.23888889, 0.27777778, 0.9], + ], + ], + ), + ( + "realization_cube", + [ + [1000, 3671.91601, -200.0, 428.08399], + [3000.0, 9671.91601, 0.0, 528.08399], + [4000.0, 10671.91601, 30.0, 478.08399], + ], + ), + ( + "percentile_cube", + [[1400, 3900, -500, 200], [2000, 9900, 0, 200], [3000, 10900, 30, 250]], + ), + ), +) +def test_cube(cube_name, expected, neighbour_cube, order, request): + """Tests plugin produces correct metadata and results for different input cubes""" + cube = request.getfixturevalue(cube_name) + coords_original = [c.name() for c in cube.dim_coords] + if order: + enforce_coordinate_ordering( + cube, [name, "realization", "percentile", "spot_index"] + ) + + coords_enforced = [c.name() for c in cube.dim_coords] + cube_units = cube.units + cube_title = cube.name() + try: + threshold_coord = cube.coord(name) + except CoordinateNotFoundError: + pass + result = SpotHeightAdjustment()(cube, neighbour_cube) + coords_result = [c.name() for c in result.dim_coords] + + if order: + enforce_coordinate_ordering(result, coords_original) + assert coords_result == coords_enforced + assert result.units == cube_units + assert result.name() == cube_title + if cube_name == ("prob_cube" or "prob_cube_realizations"): + assert result.coord(name) == threshold_coord + np.testing.assert_allclose(result.data, expected) + + +def test_prob_cube_threshold_unit(prob_cube, neighbour_cube): + """Tests that correct units and data are returned if probability cubes threshold is converted + to metres""" + prob_cube.coord(name).units = "km" + expected = [ + [0.1, 0.7, 1.0], + [0, 0, 0.3999556], + [0.2, 0.3, 0.4], + [0.1002, 0.2000778, 0.9], + ] + cube_units = prob_cube.coord(name).units + + result = SpotHeightAdjustment()(prob_cube, neighbour_cube) + assert result.coord(name).units == cube_units + assert result.coord(name).units == "km" + np.testing.assert_almost_equal(result.data, expected) + + +def test_insufficient_thresholds(prob_cube, neighbour_cube): + """Tests an error is raised if there are insufficient thresholds""" + cube = next(prob_cube.slices_over(name)) + with pytest.raises(ValueError, match="There are fewer than 2 thresholds"): + SpotHeightAdjustment()(cube, neighbour_cube)