From e250bc07fc59c9dc723a04c248d1d6ee5fbd24d0 Mon Sep 17 00:00:00 2001 From: mspelman07 <99179165+mspelman07@users.noreply.github.com> Date: Wed, 24 Jul 2024 16:08:42 +0100 Subject: [PATCH] Adds masked_add to cube combiner (#2015) * Adds masked_add to cube combiner * Remove unnecessary import * formatting * add docstrings and simplifications * formatting --------- Co-authored-by: Marcus Spelman --- improver/cli/combine.py | 2 +- improver/cube_combiner.py | 31 +++++++++- .../cube_combiner/test_CubeCombiner.py | 57 ++++++++++++++++++- 3 files changed, 86 insertions(+), 4 deletions(-) diff --git a/improver/cli/combine.py b/improver/cli/combine.py index e297995473..0580b25103 100755 --- a/improver/cli/combine.py +++ b/improver/cli/combine.py @@ -32,7 +32,7 @@ def process( An iris CubeList to be combined. operation (str): An operation to use in combining input cubes. One of: - +, -, \*, add, subtract, multiply, min, max, mean + +, -, \*, add, subtract, multiply, min, max, mean, masked_add new_name (str): New name for the resulting dataset. broadcast (str): diff --git a/improver/cube_combiner.py b/improver/cube_combiner.py index f43326c9d5..6968ed8146 100644 --- a/improver/cube_combiner.py +++ b/improver/cube_combiner.py @@ -128,6 +128,32 @@ def process(self, *cubes: Union[Cube, CubeList]) -> Cube: return self.plugin(CubeList(filtered_cubes), self.new_name) +def masked_add( + masked_array: np.ma.MaskedArray, masked_array_2: np.ma.MaskedArray +) -> np.ma.MaskedArray: + """ + Operation to add two masked arrays treating masked points as 0. + + Args: + masked_array (numpy.ma.MaskedArray): + An array that may be masked. + masked_array_2 (numpy.ma.MaskedArray): + An array that may be masked. + + Returns: + numpy.ma.MaskedArray: + The sum of the two masked arrays with masked points treated as 0. + """ + new_array_1 = np.ma.filled(masked_array, 0) + new_array_2 = np.ma.filled(masked_array_2, 0) + + new_mask = np.ma.getmask(masked_array) * np.ma.getmask(masked_array_2) + + summed_cube = np.ma.MaskedArray(np.add(new_array_1, new_array_2), mask=new_mask) + + return summed_cube + + class CubeCombiner(BasePlugin): """Plugin for combining cubes using linear operators""" @@ -140,8 +166,9 @@ class CubeCombiner(BasePlugin): "multiply": np.multiply, "max": np.maximum, "min": np.minimum, - "mean": np.add, - } # mean is calculated in two steps: sum and normalise + "mean": np.add, # mean is calculated in two steps: sum and normalise + "masked_add": masked_add, # masked_add sums arrays but treats masked points as 0 + } def __init__( self, diff --git a/improver_tests/cube_combiner/test_CubeCombiner.py b/improver_tests/cube_combiner/test_CubeCombiner.py index da865bbfa6..baee8076f9 100644 --- a/improver_tests/cube_combiner/test_CubeCombiner.py +++ b/improver_tests/cube_combiner/test_CubeCombiner.py @@ -9,12 +9,13 @@ import iris import numpy as np +import pytest from iris.coords import CellMethod from iris.cube import Cube from iris.exceptions import CoordinateNotFoundError from iris.tests import IrisTest -from improver.cube_combiner import Combine, CubeCombiner +from improver.cube_combiner import Combine, CubeCombiner, masked_add from improver.synthetic_data.set_up_test_cubes import ( add_coordinate, set_up_probability_cube, @@ -497,3 +498,57 @@ def test_with_Combine(self): if __name__ == "__main__": unittest.main() + + +@pytest.fixture +def cube1(): + """Set up a probability cube with data for testing""" + data = np.full((1, 2, 2), 0.5, dtype=np.float32) + cube1 = set_up_probability_cube( + data, + np.array([0.001], dtype=np.float32), + variable_name="lwe_thickness_of_precipitation_amount", + time=datetime(2015, 11, 19, 0), + time_bounds=(datetime(2015, 11, 18, 23), datetime(2015, 11, 19, 0)), + frt=datetime(2015, 11, 18, 22), + ) + return cube1 + + +@pytest.fixture +def cube2(): + """Set up a second probability cube with data for testing""" + data = np.full((1, 2, 2), 0.6, dtype=np.float32) + cube2 = set_up_probability_cube( + data, + np.array([0.001], dtype=np.float32), + variable_name="lwe_thickness_of_precipitation_amount", + time=datetime(2015, 11, 19, 0), + time_bounds=(datetime(2015, 11, 18, 23), datetime(2015, 11, 19, 0)), + frt=datetime(2015, 11, 18, 22), + ) + return cube2 + + +@pytest.mark.parametrize("cube1_mask", [False, True]) +@pytest.mark.parametrize("cube2_mask", [False, True]) +def test_masked_add(cube1, cube2, cube1_mask, cube2_mask): + """Tests the plugin works with the masked_add option""" + mask = [[False, True], [False, False]] + expected_output = np.array(np.full((2, 2), 1.1, dtype=np.float32)) + expected_mask = [[False, False], [False, False]] + + if cube1_mask: + cube1.data = np.ma.MaskedArray(cube1.data, mask=mask) + expected_output[0][1] = 0.6 + if cube2_mask: + cube2.data = np.ma.MaskedArray(cube2.data, mask=mask) + if cube1_mask: + expected_mask = [[False, True], [False, False]] + expected_output[0][1] = 0.0 + else: + expected_output[0][1] = 0.5 + result = masked_add(cube1.data, cube2.data) + assert np.allclose(result.data, expected_output) + assert np.allclose(result.mask, expected_mask) + assert result.dtype == np.float32