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 <marcus.spelman@metoffice.gov.uk>

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):

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,

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