Coverage for src/CSET/operators/filters.py: 100%
58 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-10-30 15:17 +0000
1# © Crown copyright, Met Office (2022-2025) and CSET contributors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15"""Operators to perform various kind of filtering."""
17import logging
19import iris
20import iris.cube
21import iris.exceptions
22import numpy as np
24from CSET._common import iter_maybe
27def apply_mask(
28 original_field: iris.cube.Cube,
29 mask: iris.cube.Cube,
30) -> iris.cube.Cube:
31 """Apply a mask to given data as a masked array.
33 Parameters
34 ----------
35 original_field: iris.cube.Cube
36 The field to be masked.
37 mask: iris.cube.Cube
38 The mask being applied to the original field.
40 Returns
41 -------
42 masked_field: iris.cube.Cube
43 A cube of the masked field.
45 Notes
46 -----
47 The mask is first converted to 1s and NaNs before multiplication with
48 the original data.
50 As discussed in generate_mask, you can combine multiple masks in a
51 recipe using other functions before applying the mask to the data.
53 Examples
54 --------
55 >>> land_points_only = apply_mask(temperature, land_mask)
56 """
57 # Ensure mask is only 1s or NaNs.
58 mask.data[mask.data == 0] = np.nan
59 mask.data[~np.isnan(mask.data)] = 1
60 logging.info(
61 "Mask set to 1 or 0s, if addition of multiple masks results"
62 "in values > 1 these are set to 1."
63 )
64 masked_field = original_field.copy()
65 masked_field.data *= mask.data
66 masked_field.attributes["mask"] = f"mask_of_{original_field.name()}"
67 return masked_field
70def filter_cubes(
71 cube: iris.cube.Cube | iris.cube.CubeList,
72 constraint: iris.Constraint,
73 **kwargs,
74) -> iris.cube.Cube:
75 """Filter a CubeList down to a single Cube based on a constraint.
77 Arguments
78 ---------
79 cube: iris.cube.Cube | iris.cube.CubeList
80 Cube(s) to filter
81 constraint: iris.Constraint
82 Constraint to extract
84 Returns
85 -------
86 iris.cube.Cube
88 Raises
89 ------
90 ValueError
91 If the constraint doesn't produce a single cube.
92 """
93 filtered_cubes = cube.extract(constraint)
94 # Return directly if already a cube.
95 if isinstance(filtered_cubes, iris.cube.Cube):
96 return filtered_cubes
97 # Check filtered cubes is a CubeList containing one cube.
98 if isinstance(filtered_cubes, iris.cube.CubeList) and len(filtered_cubes) == 1:
99 return filtered_cubes[0]
100 else:
101 raise ValueError(
102 f"Constraint doesn't produce single cube. Constraint: {constraint}"
103 f"\nSource: {cube}\nResult: {filtered_cubes}"
104 )
107def filter_multiple_cubes(
108 cubes: iris.cube.Cube | iris.cube.CubeList,
109 **kwargs,
110) -> iris.cube.CubeList:
111 """Filter a CubeList on multiple constraints, returning another CubeList.
113 Arguments
114 ---------
115 cube: iris.cube.Cube | iris.cube.CubeList
116 Cube(s) to filter
117 constraint: iris.Constraint
118 Constraint to extract. This must be a named argument. There can be any
119 number of additional constraints, they just need unique names.
121 Returns
122 -------
123 iris.cube.CubeList
125 Raises
126 ------
127 ValueError
128 The constraints don't produce a single cube per constraint.
129 """
130 # Ensure input is a CubeList.
131 if isinstance(cubes, iris.cube.Cube):
132 cubes = iris.cube.CubeList((cubes,))
133 if len(kwargs) < 1:
134 raise ValueError("Must have at least one constraint.")
135 try:
136 filtered_cubes = cubes.extract_cubes(kwargs.values())
137 except iris.exceptions.ConstraintMismatchError as err:
138 raise ValueError(
139 "The constraints don't produce a single cube per constraint."
140 ) from err
141 return filtered_cubes
144def generate_mask(
145 mask_field: iris.cube.Cube | iris.cube.CubeList,
146 condition: str,
147 value: float,
148) -> iris.cube.Cube | iris.cube.CubeList:
149 """Generate a mask to remove data not meeting conditions.
151 Parameters
152 ----------
153 mask_field: iris.cube.Cube | iris.cube.CubeList
154 The field(s) to be used for creating the mask.
155 condition: str
156 The type of condition applied, six available options:
157 'eq','ne','lt','le','gt', and 'ge'. The condition is consistent
158 regardless of whether mask_field is a cube or CubeList.
159 The conditions are as follows
160 eq: equal to,
161 ne: not equal to,
162 lt: less than,
163 le: less than or equal to,
164 gt: greater than,
165 ge: greater than or equal to.
166 value: float
167 The value on the right hand side of the condition. The value is
168 consistent regardless of whether mask_field is a cube or CubeList.
170 Returns
171 -------
172 mask: iris.cube.Cube | iris.cube.CubeList
173 Mask(s) meeting the condition applied.
175 Raises
176 ------
177 ValueError: Unexpected value for condition. Expected eq, ne, gt, ge, lt, le.
178 Got {condition}.
179 Raised when condition is not supported.
181 Notes
182 -----
183 The mask is created in the opposite sense to numpy.ma.masked_arrays. This
184 method was chosen to allow easy combination of masks together outside of
185 this function using misc.addition or misc.multiplication depending on
186 applicability. The combinations can be of any fields such as orography >
187 500 m, and humidity == 100 %.
189 The conversion to a masked array occurs in the apply_mask routine, which
190 should happen after all relevant masks have been combined.
192 Examples
193 --------
194 >>> land_mask = generate_mask(land_sea_mask,'gt',1)
195 """
196 mask_list = iris.cube.CubeList([])
197 for cube in iter_maybe(mask_field):
198 mask = cube.copy()
199 mask.data[:] = 0.0
200 match condition:
201 case "eq":
202 mask.data[cube.data == value] = 1
203 case "ne":
204 mask.data[cube.data != value] = 1
205 case "gt":
206 mask.data[cube.data > value] = 1
207 case "ge":
208 mask.data[cube.data >= value] = 1
209 case "lt":
210 mask.data[cube.data < value] = 1
211 case "le":
212 mask.data[cube.data <= value] = 1
213 case _:
214 raise ValueError("""Unexpected value for condition. Expected eq, ne,
215 gt, ge, lt, le. Got {condition}.""")
216 mask.attributes["mask"] = f"mask_for_{cube.name()}_{condition}_{value}"
217 mask.rename(f"mask_for_{cube.name()}_{condition}_{value}")
218 mask.units = "1"
219 mask_list.append(mask)
221 if len(mask_list) == 1:
222 return mask_list[0]
223 else:
224 return mask_list