diff --git a/src/CSET/operators/filters.py b/src/CSET/operators/filters.py index 298e83d6e..da99bfd7f 100644 --- a/src/CSET/operators/filters.py +++ b/src/CSET/operators/filters.py @@ -14,11 +14,59 @@ """Operators to perform various kind of filtering.""" +import logging from typing import Union import iris import iris.cube import iris.exceptions +import numpy as np + + +def apply_mask( + original_field: iris.cube.Cube, + mask: iris.cube.Cube, +) -> iris.cube.Cube: + """Apply a mask to given data as a masked array. + + Parameters + ---------- + original_field: iris.cube.Cube + The field to be masked. + mask: iris.cube.Cube + The mask being applied to the original field. + + Returns + ------- + A masked field. + + Return type + ----------- + iris.cube.Cube + + Notes + ----- + The mask is first converted to 1s and NaNs before multiplication with + the original data. + + As discussed in generate_mask, you can combine multiple masks in a + recipe using other functions before applying the mask to the data. + + Examples + -------- + >>> land_points_only = apply_mask(temperature, land_mask) + """ + # Ensure mask is only 1s or NaNs. + mask.data[mask.data == 0] = np.nan + mask.data[~np.isnan(mask.data)] = 1 + logging.info( + "Mask set to 1 or 0s, if addition of multiple masks results" + "in values > 1 these are set to 1." + ) + masked_field = original_field.copy() + masked_field.data *= mask.data + masked_field.attributes["mask"] = f"mask_of_{original_field.name()}" + return masked_field def filter_cubes( @@ -93,3 +141,68 @@ def filter_multiple_cubes( "The constraints don't produce a single cube per constraint." ) from err return filtered_cubes + + +def generate_mask( + mask_field: iris.cube.Cube, + condition: str, + value: float, +) -> iris.cube.Cube: + """Generate a mask to remove data not meeting conditions. + + Parameters + ---------- + mask_field: iris.cube.Cube + The field to be used for creating the mask. + condition: str + The type of condition applied, six available options: + '==','!=','<','<=','>', and '>='. + value: float + The value on the right hand side of the condition. + + Returns + ------- + mask: iris.cube.Cube + Mask meeting the condition applied. + + Raises + ------ + ValueError: Unexpected value for condition. Expected ==, !=, >, >=, <, <=. + Got {condition}. + Raised when condition is not supported. + + Notes + ----- + The mask is created in the opposite sense to numpy.ma.masked_arrays. This + method was chosen to allow easy combination of masks together outside of + this function using misc.addition or misc.multiplication depending on + applicability. The combinations can be of any fields such as orography > + 500 m, and humidity == 100 %. + + The conversion to a masked array occurs in the apply_mask routine, which + should happen after all relevant masks have been combined. + + Examples + -------- + >>> land_mask = generate_mask(land_sea_mask,'==',1) + """ + mask = mask_field.copy() + mask.data = np.zeros(mask.data.shape) + match condition: + case "==": + mask.data[mask_field.data == value] = 1 + case "!=": + mask.data[mask_field.data != value] = 1 + case ">": + mask.data[mask_field.data > value] = 1 + case ">=": + mask.data[mask_field.data >= value] = 1 + case "<": + mask.data[mask_field.data < value] = 1 + case "<=": + mask.data[mask_field.data <= value] = 1 + case _: + raise ValueError("""Unexpected value for condition. Expected ==, !=, + >, >=, <, <=. Got {condition}.""") + mask.attributes["mask"] = f"mask_for_{mask_field.name()}_{condition}_{value}" + return mask diff --git a/src/CSET/recipes/example_combined_mask_addition.yaml b/src/CSET/recipes/example_combined_mask_addition.yaml new file mode 100644 index 000000000..9e14eb9af --- /dev/null +++ b/src/CSET/recipes/example_combined_mask_addition.yaml @@ -0,0 +1,50 @@ +category: Quick Look +title: Example adding masks for 'OR' type filters +description: | + Generates and applies two masks, and adds them together for more complex stratified analysis. + Examples would include masking for variables outside a range, e.g. temperature < 273 K or temperature + > 313 K. Addition of masks would be akin to finding a probability of two events happening at the same time in different places. + +steps: + - operator: read.read_cubes + + - operator: filters.apply_mask + original_field: + operator: filters.filter_cubes + constraint: + operator: constraints.combine_constraints + variable_constraint: + operator: constraints.generate_var_constraint + varname: surface_air_temperature + cell_method_constraint: + operator: constraints.generate_cell_methods_constraint + cell_methods: [] + masks: + operator: misc.addition + addend_1: + operator: filters.generate_mask + mask_field: + operator: filters.filter_cubes + constraint: + operator: constraints.generate_var_constraint + varname: surface_altitude + condition: '<' + value: 500 + addend_2: + operator: filters.generate_mask + mask_field: + operator: filters.filter_cubes + constraint: + operator: constraints.generate_var_constraint + varname: surface_altitude + condition: '>' + value: 200 + + - operator: collapse.collapse + coordinate: [grid_latitude, grid_longitude] + method: MEAN + + - operator: write.write_cube_to_nc + overwrite: True + + - operator: plot.plot_line_series diff --git a/src/CSET/recipes/example_combined_mask_multiplication.yaml b/src/CSET/recipes/example_combined_mask_multiplication.yaml new file mode 100644 index 000000000..e7c725cde --- /dev/null +++ b/src/CSET/recipes/example_combined_mask_multiplication.yaml @@ -0,0 +1,50 @@ +category: Quick Look +title: Example adding masks for 'AND' type filters +description: | + Generates and applies two masks, and multiplies them together for more complex stratified analysis. + Examples would include masking for temperature > 273 K and orography > 500 m. Multiplication of masks + would be akin to finding a probability of two events happening at the same time in the same place. + +steps: + - operator: read.read_cubes + + - operator: filters.apply_mask + original_field: + operator: filters.filter_cubes + constraint: + operator: constraints.combine_constraints + variable_constraint: + operator: constraints.generate_var_constraint + varname: surface_air_temperature + cell_method_constraint: + operator: constraints.generate_cell_methods_constraint + cell_methods: [] + masks: + operator: misc.multiplication + multiplicand: + operator: filters.generate_mask + mask_field: + operator: filters.filter_cubes + constraint: + operator: constraints.generate_var_constraint + varname: surface_altitude + condition: '>' + value: 500 + multiplier: + operator: filters.generate_mask + mask_field: + operator: filters.filter_cubes + constraint: + operator: constraints.generate_var_constraint + varname: surface_air_temperature + condition: '<=' + value: 273 + + - operator: collapse.collapse + coordinate: [grid_latitude, grid_longitude] + method: MEAN + + - operator: write.write_cube_to_nc + overwrite: True + + - operator: plot.plot_line_series diff --git a/src/CSET/recipes/example_simple_mask.yaml b/src/CSET/recipes/example_simple_mask.yaml new file mode 100644 index 000000000..d39d8d683 --- /dev/null +++ b/src/CSET/recipes/example_simple_mask.yaml @@ -0,0 +1,36 @@ +category: Quick Look +title: Example simple mask application +description: Generates and applies a simple mask to a field for stratified analysis. + +steps: + - operator: read.read_cubes + + - operator: filters.apply_mask + original_field: + operator: filters.filter_cubes + constraint: + operator: constraints.combine_constraints + variable_constraint: + operator: constraints.generate_var_constraint + varname: surface_air_temperature + cell_method_constraint: + operator: constraints.generate_cell_methods_constraint + cell_methods: [] + masks: + operator: filters.generate_mask + mask_field: + operator: filters.filter_cubes + constraint: + operator: constraints.generate_var_constraint + varname: surface_altitude + condition: '>=' + value: 500 + + - operator: collapse.collapse + coordinate: [grid_latitude, grid_longitude] + method: MEAN + + - operator: write.write_cube_to_nc + overwrite: True + + - operator: plot.plot_line_series diff --git a/src/CSET/recipes/example_spatial_plot_of_mask.yaml b/src/CSET/recipes/example_spatial_plot_of_mask.yaml new file mode 100644 index 000000000..c2e7b8883 --- /dev/null +++ b/src/CSET/recipes/example_spatial_plot_of_mask.yaml @@ -0,0 +1,21 @@ +category: Surface Spatial Plot +title: Example plot of a mask +description: | + Generates a mask and then provides a spatial map of the mask. + +steps: + - operator: read.read_cubes + + - operator: filters.generate_mask + mask_field: + operator: filters.filter_cubes + constraint: + operator: constraints.generate_var_constraint + varname: surface_altitude + condition: '>=' + value: 500 + + - operator: plot.spatial_pcolormesh_plot + + - operator: write.write_cube_to_nc + overwrite: True diff --git a/tests/operators/test_filters.py b/tests/operators/test_filters.py index 8eefe3b78..a35352c75 100644 --- a/tests/operators/test_filters.py +++ b/tests/operators/test_filters.py @@ -15,6 +15,7 @@ """Test filter operators.""" import iris.cube +import numpy as np import pytest from CSET.operators import constraints, filters @@ -176,3 +177,103 @@ def test_generate_level_constraint_returnallmodlev(load_verticalcoord_cubelist): assert expected_coordstr in repr( load_verticalcoord_cubelist.extract(constraint_1)[0].coords ) + + +def test_generate_mask_fail_wrong_condition(cube): + """Fails for a wrong condition.""" + with pytest.raises(ValueError): + filters.generate_mask(cube, "!<", 10) + + +def test_generate_mask_equal_to(cube): + """Generates a mask with values equal to a specified value.""" + mask = cube.copy() + mask.data = np.zeros(mask.data.shape) + mask.data[cube.data == 276] = 1 + assert np.allclose( + filters.generate_mask(cube, "==", 276).data, + mask.data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_generate_mask_not_equal_to(cube): + """Generates a mask with values not equal to a specified value.""" + mask = cube.copy() + mask.data = np.zeros(mask.data.shape) + mask.data[cube.data != 276] = 1 + assert np.allclose( + filters.generate_mask(cube, "!=", 276).data, + mask.data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_generate_mask_greater_than(cube): + """Generates a mask with values greater than a specified value.""" + mask = cube.copy() + mask.data = np.zeros(mask.data.shape) + mask.data[cube.data > 276] = 1 + assert np.allclose( + filters.generate_mask(cube, ">", 276).data, + mask.data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_generate_mask_greater_equal_to(cube): + """Generates a mask with values greater than or equal to a specified value.""" + mask = cube.copy() + mask.data = np.zeros(mask.data.shape) + mask.data[cube.data >= 276] = 1 + assert np.allclose( + filters.generate_mask(cube, ">=", 276).data, + mask.data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_generate_mask_less_than(cube): + """Generates a mask with values less than a specified value.""" + mask = cube.copy() + mask.data = np.zeros(mask.data.shape) + mask.data[cube.data < 276] = 1 + assert np.allclose( + filters.generate_mask(cube, "<", 276).data, + mask.data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_generate_mask_less_equal_to(cube): + """Generates a mask with values less than or equal to a specified value.""" + mask = cube.copy() + mask.data = np.zeros(mask.data.shape) + mask.data[cube.data <= 276] = 1 + assert np.allclose( + filters.generate_mask(cube, "<=", 276).data, + mask.data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_apply_mask(cube): + """Apply a mask to a cube.""" + mask = filters.generate_mask(cube, "==", 276) + mask.data[mask.data == 0] = np.nan + mask.data[~np.isnan(mask.data)] = 1 + test_data = cube.copy() + test_data.data *= mask.data + assert np.allclose( + filters.apply_mask(cube, mask).data, + test_data.data, + rtol=1e-06, + atol=1e-02, + equal_nan=True, + )