Coverage for src / CSET / operators / filters.py: 98%
66 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-09 14:55 +0000
1# © Crown copyright, Met Office (2022-2026) 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 | iris.cube.CubeList,
29 mask: iris.cube.Cube | iris.cube.CubeList,
30) -> iris.cube.Cube | iris.cube.CubeList:
31 """Apply a mask to given data as a masked array.
33 Parameters
34 ----------
35 original_field: iris.cube.Cube | iris.cube.CubeList
36 The field(s) to be masked.
37 mask: iris.cube.Cube | iris.cube.CubeList
38 The mask(s) being applied to the original field(s).
40 Returns
41 -------
42 masked_field: iris.cube.Cube | iris.cube.CubeList
43 A cube or cubelist of the masked field(s).
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 masked_fields = iris.cube.CubeList([])
58 for M, F in zip(iter_maybe(mask), iter_maybe(original_field), strict=True):
59 # Ensure mask data are floats and only 1s or NaNs.
60 M.data = np.float64(M.data)
61 M.data[M.data == 0.0] = np.nan
62 M.data[~np.isnan(M.data)] = 1.0
63 logging.info(
64 "Mask set to 1 or 0s, if addition of multiple masks results"
65 "in values > 1 these are set to 1."
66 )
67 masked_field = F.copy()
68 masked_field.data *= M.data
69 masked_field.attributes["mask"] = f"mask_of_{F.name()}"
70 masked_fields.append(masked_field)
71 if len(masked_fields) == 1:
72 return masked_fields[0]
73 else:
74 return masked_fields
77def filter_cubes(
78 cube: iris.cube.Cube | iris.cube.CubeList,
79 constraint: iris.Constraint,
80 **kwargs,
81) -> iris.cube.Cube:
82 """Filter a CubeList down to a single Cube based on a constraint.
84 Arguments
85 ---------
86 cube: iris.cube.Cube | iris.cube.CubeList
87 Cube(s) to filter
88 constraint: iris.Constraint
89 Constraint to extract
91 Returns
92 -------
93 iris.cube.Cube
95 Raises
96 ------
97 ValueError
98 If the constraint doesn't produce a single cube.
99 """
100 filtered_cubes = cube.extract(constraint)
101 # Return directly if already a cube.
102 if isinstance(filtered_cubes, iris.cube.Cube):
103 return filtered_cubes
104 # Check filtered cubes is a CubeList containing one cube.
105 if isinstance(filtered_cubes, iris.cube.CubeList) and len(filtered_cubes) == 1:
106 return filtered_cubes[0]
107 else:
108 raise ValueError(
109 f"Constraint doesn't produce single cube. Constraint: {constraint}"
110 f"\nSource: {cube}\nResult: {filtered_cubes}"
111 )
114def filter_multiple_cubes(
115 cubes: iris.cube.Cube | iris.cube.CubeList,
116 **kwargs,
117) -> iris.cube.CubeList:
118 """Filter a CubeList on multiple constraints, returning another CubeList.
120 Arguments
121 ---------
122 cube: iris.cube.Cube | iris.cube.CubeList
123 Cube(s) to filter
124 constraint: iris.Constraint
125 Constraint to extract. This must be a named argument. There can be any
126 number of additional constraints, they just need unique names.
128 Returns
129 -------
130 iris.cube.CubeList
132 Raises
133 ------
134 ValueError
135 Some constraints don't produce any cubes.
136 """
137 # Ensure input is a CubeList.
138 if isinstance(cubes, iris.cube.Cube):
139 cubes = iris.cube.CubeList((cubes,))
140 if len(kwargs) < 1:
141 raise ValueError("Must have at least one constraint.")
142 # Switch to extract due to lack of instance requiring one cube per
143 # constraint.
144 try:
145 filtered_cubes = cubes.extract(kwargs.values())
146 except iris.exceptions.ConstraintMismatchError as err:
147 raise ValueError("The constraints don't produce a cube or cubelist.") from err
148 if len(filtered_cubes) == 0:
149 raise ValueError("No cubes loaded. Please check your constraints.")
150 return filtered_cubes
153def generate_mask(
154 mask_field: iris.cube.Cube | iris.cube.CubeList,
155 condition: str,
156 value: float,
157) -> iris.cube.Cube | iris.cube.CubeList:
158 """Generate a mask to remove data not meeting conditions.
160 Parameters
161 ----------
162 mask_field: iris.cube.Cube | iris.cube.CubeList
163 The field(s) to be used for creating the mask.
164 condition: str
165 The type of condition applied, six available options:
166 'eq','ne','lt','le','gt', and 'ge'. The condition is consistent
167 regardless of whether mask_field is a cube or CubeList.
168 The conditions are as follows
169 eq: equal to,
170 ne: not equal to,
171 lt: less than,
172 le: less than or equal to,
173 gt: greater than,
174 ge: greater than or equal to.
175 value: float
176 The value on the right hand side of the condition. The value is
177 consistent regardless of whether mask_field is a cube or CubeList.
179 Returns
180 -------
181 mask: iris.cube.Cube | iris.cube.CubeList
182 Mask(s) meeting the condition applied.
184 Raises
185 ------
186 ValueError: Unexpected value for condition. Expected eq, ne, gt, ge, lt, le.
187 Got {condition}.
188 Raised when condition is not supported.
190 Notes
191 -----
192 The mask is created in the opposite sense to numpy.ma.masked_arrays. This
193 method was chosen to allow easy combination of masks together outside of
194 this function using misc.addition or misc.multiplication depending on
195 applicability. The combinations can be of any fields such as orography >
196 500 m, and humidity == 100 %.
198 The conversion to a masked array occurs in the apply_mask routine, which
199 should happen after all relevant masks have been combined.
201 Examples
202 --------
203 >>> land_mask = generate_mask(land_sea_mask,'gt',1)
204 """
205 mask_list = iris.cube.CubeList([])
206 for cube in iter_maybe(mask_field):
207 mask = cube.copy()
208 mask.data[:] = 0.0
209 match condition:
210 case "eq":
211 mask.data[cube.data == value] = 1.0
212 case "ne":
213 mask.data[cube.data != value] = 1.0
214 case "gt":
215 mask.data[cube.data > value] = 1.0
216 case "ge":
217 mask.data[cube.data >= value] = 1.0
218 case "lt":
219 mask.data[cube.data < value] = 1.0
220 case "le":
221 mask.data[cube.data <= value] = 1.0
222 case _:
223 raise ValueError("""Unexpected value for condition. Expected eq, ne,
224 gt, ge, lt, le. Got {condition}.""")
225 mask.attributes["mask"] = f"mask_for_{cube.name()}_{condition}_{value}"
226 mask.rename(f"mask_for_{cube.name()}_{condition}_{value}")
227 mask.units = "1"
228 mask_list.append(mask)
230 if len(mask_list) == 1:
231 return mask_list[0]
232 else:
233 return mask_list