Apply mask to cube (#2014)

* Add apply_mask cli and function

* remove print statement

* Updates to make sure plugin doesn't require cubelist input

* Update Docstring

---------

Co-authored-by: Marcus Spelman <marcus.spelman@metoffice.gov.uk>

improver/cli/apply_mask.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+# (C) Crown copyright, Met Office. All rights reserved.
+#
+# This file is part of IMPROVER and is released under a BSD 3-Clause license.
+# See LICENSE in the root of the repository for full licensing details.
+"""Script to apply provided mask to cube data."""
+
+from improver import cli
+
+
+@cli.clizefy
+@cli.with_output
+def process(*cubes: cli.inputcube, mask_name: str, invert_mask: bool = "False"):
+    """
+    Applies provided mask to cube data. The mask_name is used to extract the mask cube
+    from the input cubelist. The other cube in the cubelist is then masked using the
+    mask data. If invert_mask is True, the mask will be inverted before it is applied.
+
+    Args:
+        cubes (iris.cube.CubeList):
+            A list of iris cubes that should contain exactly two cubes: a mask to be applied
+            and a cube to apply the mask to. The cubes should have the same dimensions.
+        mask_name (str):
+            The name of the cube containing the mask data. This should match with exactly one
+            of the cubes in the input cubelist.
+        invert_mask (bool):
+            Use to select whether the mask should be inverted before being applied to the data.
+
+    Returns:
+        A cube with the mask applied to the data. The metadata will exactly match the input cube.
+    """
+    from improver.utilities.mask import apply_mask
+
+    return apply_mask(*cubes, mask_name=mask_name, invert_mask=invert_mask)

improver/utilities/mask.py

@@ -0,0 +1,72 @@
+# (C) Crown copyright, Met Office. All rights reserved.
+#
+# This file is part of IMPROVER and is released under a BSD 3-Clause license.
+# See LICENSE in the root of the repository for full licensing details.
+"""Module for applying mask to a cube."""
+
+from typing import Union
+
+import iris
+import numpy as np
+
+from improver.utilities.common_input_handle import as_cubelist
+from improver.utilities.cube_checker import find_dimension_coordinate_mismatch
+from improver.utilities.cube_manipulation import (
+    enforce_coordinate_ordering,
+    get_coord_names,
+)
+
+
+def apply_mask(
+    *cubes: Union[iris.cube.CubeList, iris.cube.Cube],
+    mask_name: str,
+    invert_mask: bool = False,
+) -> iris.cube.Cube:
+    """
+    Apply a provided mask to a cube. If invert_mask is True, the mask will be inverted.
+
+    Args:
+        cubes:
+            A list of iris cubes that should contain exactly two cubes: a mask and a cube
+            to apply the mask to. The cubes should have the same dimensions.
+        mask_name:
+            The name of the mask cube. It should match with exactly one of the cubes in
+            the input cubelist.
+        invert_mask:
+            If True, the mask will be inverted before it is applied.
+    Returns:
+        A cube with a mask applied to the data.
+
+    Raises:
+    ValueError: If the number of cubes provided is not equal to 2.
+    ValueError: If the input cube and mask cube have different dimensions.
+
+    """
+    cubes = as_cubelist(*cubes)
+    cube_names = [cube.name() for cube in cubes]
+    if len(cubes) != 2:
+        raise ValueError(
+            f"""Two cubes are required for masking, a mask and the cube to be masked.
+                         Provided cubes are {cube_names}"""
+        )
+
+    mask = cubes.extract_cube(mask_name)
+    cubes.remove(mask)
+    cube = cubes[0]
+
+    # Ensure mask is in a boolean form and invert if requested
+    mask.data = mask.data.astype(np.bool)
+    if invert_mask:
+        mask.data = ~mask.data
+
+    coord_list = get_coord_names(cube)
+    enforce_coordinate_ordering(mask, coord_list)
+
+    # This error is required to stop the mask from being broadcasted to a new shape by numpy. When
+    # the mask and cube have different shapes numpy will try to broadcast the mask to be the same
+    # shape as the cube data. This might succeed but masks unexpected data points.
+    if find_dimension_coordinate_mismatch(cube, mask):
+        raise ValueError("Input cube and mask cube must have the same dimensions")
+
+    cube.data = np.ma.array(cube.data, mask=mask.data)
+    return cube

improver_tests/acceptance/SHA256SUMS

@@ -70,6 +70,10 @@ c7eb9bab2ad43ac19ecc071730479a9f27a58992fcb22050743d34cdf2ad9639  ./apply-lapse-
 6eb89848a2a36007d9a2210e2ff4cb459921502329eb4bb6ef0bde74701d0c27  ./apply-lapse-rate/realizations/enukx_orography.nc
 25eeaa6b86b95354efd9406b1f9f1b5b43f0a7dd46f75fc5f3d144f52a30059e  ./apply-lapse-rate/realizations/enukx_temperature.nc
 a747cf91ea8a2b720e1a22a9569cc1a2c22bb7ad28282e661127413470ba0392  ./apply-lapse-rate/realizations/kgo.nc
+e455d391d15c30fe3cf3680fa0c355a3d4f0a9a84450f685f618655eb962ff52  ./apply-mask/kgo.nc
+a3aca2f841ff96a7780cb98a5ccf874001905f18db0c6166dffcf116a3f1b319  ./apply-mask/kgo_inverted.nc
+f33c0f3d10846ee69edd24fb318bebd4a2645363c89a434f56aeff24313fc886  ./apply-mask/mask.nc
+9fba2a41268f30c530e9362bf15ef34c781d24b5e580589cec7715b1a6d25b27  ./apply-mask/wind_speed.nc
 3a3d1b902e9e97f6ec7c74937e43dfce460c6599ef1890ac90f711959415f6ed  ./apply-night-mask/global_basic/input.nc
 65ba596b39ae589e1033b8e35fcb71f341b2fe64337fc2e4d88edff4331925e9  ./apply-night-mask/global_basic/kgo.nc
 c5860ce0801d097e1705130d6e3f49b19c43fa399041bd7726b7a1257335f2fb  ./apply-night-mask/uk_basic/input.nc

improver_tests/acceptance/test_apply_mask.py

@@ -0,0 +1,38 @@
+# (C) Crown copyright, Met Office. All rights reserved.
+#
+# This file is part of IMPROVER and is released under a BSD 3-Clause license.
+# See LICENSE in the root of the repository for full licensing details.
+"""
+Tests for the apply-mask 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("invert", [True, False])
+def test_apply_mask(tmp_path, invert):
+    """Test apply-mask CLI."""
+    kgo_dir = acc.kgo_root() / "apply-mask/"
+    kgo_path = kgo_dir / "kgo.nc"
+    if invert:
+        kgo_path = kgo_dir / "kgo_inverted.nc"
+
+    output_path = tmp_path / "output.nc"
+    args = [
+        kgo_dir / "wind_speed.nc",
+        kgo_dir / "mask.nc",
+        "--mask-name",
+        "land_binary_mask",
+        "--invert-mask",
+        f"{invert}",
+        "--output",
+        output_path,
+    ]
+    run_cli(args)
+    acc.compare(output_path, kgo_path)

improver_tests/utilities/test_mask.py

@@ -0,0 +1,78 @@
+# (C) Crown copyright, Met Office. All rights reserved.
+#
+# This file is part of IMPROVER and is released under a BSD 3-Clause license.
+# See LICENSE in the root of the repository for full licensing details.
+"""
+Unit tests for the function apply_mask.
+"""
+
+import iris
+import numpy as np
+import pytest
+
+from improver.synthetic_data.set_up_test_cubes import set_up_variable_cube
+from improver.utilities.cube_manipulation import enforce_coordinate_ordering
+from improver.utilities.mask import apply_mask
+
+
+@pytest.fixture
+def wind_gust_cube():
+    data = np.full((2, 3), 10)
+    return set_up_variable_cube(
+        data=data, attributes={"wind_gust_type": "10m_ratio"}, name="wind_gust"
+    )
+
+
+@pytest.fixture
+def mask():
+    data = np.array([[0, 0, 1], [1, 1, 0]])
+    return set_up_variable_cube(data=data, name="land_sea_mask")
+
+
+@pytest.mark.parametrize("invert_mask", [True, False])
+@pytest.mark.parametrize("switch_coord_order", [True, False])
+def test_basic(wind_gust_cube, mask, switch_coord_order, invert_mask):
+    """
+    Test the basic functionality of the apply_mask plugin. Checks that the
+    mask is correctly applied and inverted if requested. Also checks plugin
+    can cope with different orderings of coordinates on the input cubes."""
+
+    expected_data = np.full((2, 3), 10)
+    expected_mask = np.array([[False, False, True], [True, True, False]])
+    if switch_coord_order:
+        enforce_coordinate_ordering(wind_gust_cube, ["longitude", "latitude"])
+        expected_data = expected_data.transpose()
+        expected_mask = expected_mask.transpose()
+    if invert_mask:
+        expected_mask = np.invert(expected_mask)
+
+    input_list = [wind_gust_cube, mask]
+
+    result = apply_mask(
+        iris.cube.CubeList(input_list),
+        mask_name="land_sea_mask",
+        invert_mask=invert_mask,
+    )
+
+    assert np.allclose(result.data, expected_data)
+    assert np.allclose(result.data.mask, expected_mask)
+
+
+def test_different_dimensions(wind_gust_cube, mask):
+    """ Test that the function will raise an error if the mask cube has different
+    dimensions to other cube."""
+    mask = mask[0]
+    input_list = [wind_gust_cube, mask]
+    with pytest.raises(
+        ValueError, match="Input cube and mask cube must have the same dimensions"
+    ):
+        apply_mask(iris.cube.CubeList(input_list), mask_name="land_sea_mask")
+
+
+def test_too_many_cubes(wind_gust_cube, mask):
+    """
+    Test that the function will raise an error if more than two cubes are provided.
+    """
+    input_list = [wind_gust_cube, wind_gust_cube, wind_gust_cube]
+    with pytest.raises(ValueError, match="Two cubes are required for masking"):
+        apply_mask(iris.cube.CubeList(input_list), mask_name="land_sea_mask")