Coverage for src/CSET/operators/radar_filter.py: 95%
111 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-20 14:18 +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 iris
18import iris.cube
19import iris.exceptions
20import numpy as np
22from CSET._common import iter_maybe
23from CSET.operators.filters import apply_mask, generate_mask
26def mask_list(model_names: list[str]) -> list[str]:
27 """Determine the Nimrod weights files to use.
29 Parameters
30 ----------
31 model_names: list[str]
32 A list containing model names and at least one Nimrod hourly
33 rainfall accumulation source.
34 Possible radar sources are:
35 "Nimrod2km", "Nimrod_2km".
36 "Nimrodxkm", "Nimrod_xkm".
37 "Nimrod1km", "Nimrod_1km".
39 Returns
40 -------
41 list[str]
42 A list of the Nimrod weights files to use with each of the input
43 models and radar sources.
45 Notes
46 -----
47 At least one of the entries in the input list must be a Nimrod hourly
48 rainfall accumulation source.
50 If just one Nimrod source is specified, then the weights file associated
51 with this source is used.
53 If more than one Nimrod source is in the input list, then each of the these
54 Nimrod sources is associated with its own weights file e.g. if the input list
55 contains ["Nimrod1km", "Nimrod2km"] then the weights files for these will
56 be ["Nimrod1km_weights", "Nimrod2km_weights"]. Any model fields in the input
57 list will be allocated a weights file according to the order of preference
58 specified in the list nimrod_preference e.g. if the input list is
59 ["UM_model", "Nimrod1km", "Nimrod2km"] then the output weights files list will
60 be ["Nimrod2km_weights", "Nimrod1km_weights", "Nimrod2km_weights"] as the Nimrod
61 weights for 2km data are preferred over those for 1km.
63 Examples
64 --------
65 >>> list_weights = mask_list( ["UM_model", "Nimrod1km", "Nimrod2km"] )
66 >>> print(list_weights)
67 ["Nimrod2km_weights", "Nimrod1km_weights", "Nimrod2km_weights"]
69 """
70 # Set the preference order for choosing a Nimrod radar weights source
71 # in order of most to least preferred.
72 nimrod_preference = [
73 "Nimrod2km",
74 "Nimrod_2km",
75 "Nimrodxkm",
76 "Nimrod_xkm",
77 "Nimrod1km",
78 "Nimrod_1km",
79 ]
81 # Define the string that helps form a Nimrod weights file.
82 wei = "_weights"
84 # Determine the preferred Nimrod mask to use.
85 empty_string = ""
86 preferred_nimrod = empty_string
87 for prefer in reversed(nimrod_preference):
88 if any(prefer in model for model in model_names):
89 preferred_nimrod = prefer
91 # Create the list of the required Nimrod masks.
92 mask_names_list = []
93 if preferred_nimrod != empty_string:
94 # Loop over the input model_names.
95 for model in model_names:
96 if any(model in nimrod for nimrod in nimrod_preference):
97 nimrod_mask = f"{model}{wei}"
98 else:
99 nimrod_mask = f"{preferred_nimrod}{wei}"
100 mask_names_list.append(nimrod_mask)
102 return mask_names_list
105def mask_by_weights(
106 cubes: iris.cube.CubeList,
107 model_names: list[str],
108 weights_names: list[str],
109 **kwargs,
110) -> iris.cube.CubeList:
111 """Filter a field using a radar weights field as a mask.
113 Parameters
114 ----------
115 cubes: iris.cube.CubeList
116 CubeList containing fields to mask and radar weights fields to use as the masks.
117 model_names: list[str]
118 A list of the model_names or radar sources to mask.
119 weights_names: list[str]
120 A list of radar weights sources to use as masks. There should be an entry
121 in weights_names to correspond with every entry in model_names.
123 Returns
124 -------
125 CubeList:
126 A CubeList of masked fields.
128 Examples
129 --------
130 >>> field_filtered = mask_by_weights(cubelist, model_names)
132 """
133 # Check the input unfiltered cubes and the mask cubes are both cubelists
134 # with the same number of cubes. If not, then add extra mask cubes.
135 if len(model_names) != len(weights_names): 135 ↛ 136line 135 didn't jump to line 136 because the condition on line 135 was never true
136 weights_names = mask_list(model_names)
138 # Create an empty cubelist to hold the filtered fields.
139 filtered_list = iris.cube.CubeList([])
141 # Loop over the fields to filter.
142 for model, mask in zip(
143 iter_maybe(model_names),
144 iter_maybe(weights_names),
145 strict=True,
146 ):
147 # Grab the field to filter.
148 model_constraint = iris.AttributeConstraint(model_name=model)
149 unfiltered_field = cubes.extract_cube(model_constraint)
151 # Select the field to use as the mask.
152 # Nice to do - put in support for a static mask.
153 mask_constraint = iris.AttributeConstraint(model_name=mask)
154 mask_field = cubes.extract_cube(mask_constraint)
156 # Create the mask - note that the condition e.g. "ge" can be set by a loader
157 # as can the threshold value.
158 mask_radar_wts = generate_mask(mask_field, "ge", 11)
160 # Apply the mask.
161 masked_radar_obs = apply_mask(unfiltered_field, mask_radar_wts)
163 # Put the filtered cube into the list of filtered cubes.
164 filtered_list.append(masked_radar_obs)
166 # Preserve returning a cube if only a cube has been supplied to filter.
167 if len(filtered_list) == 1: 167 ↛ 170line 167 didn't jump to line 170 because the condition on line 167 was always true
168 return filtered_list[0]
169 else:
170 return filtered_list
173def radar_apply_mask(
174 original_field: iris.cube.Cube | iris.cube.CubeList,
175 mask: iris.cube.Cube | iris.cube.CubeList,
176 boundary_margin: int = 8,
177) -> iris.cube.Cube | iris.cube.CubeList:
178 """Apply a mask to given field as a masked array.
180 Parameters
181 ----------
182 original_field: iris.cube.Cube | iris.cube.CubeList
183 The field(s) to be masked.
184 mask: iris.cube.Cube | iris.cube.CubeList
185 The mask(s) being applied to the original field(s).
186 boundary_margin: int, optional
187 Number of grid points from the domain boundary considered "unreliable".
188 Defaults to 8.
190 Returns
191 -------
192 masked_field: iris.cube.Cube | iris.cube.CubeList
193 A cube or CubeList of the masked field(s).
195 Notes
196 -----
197 The mask is first converted to 1s and NaNs before multiplication with
198 the original data.
200 As discussed in filters.generate_mask, you can combine multiple masks in a
201 recipe using other functions before applying the mask to the data.
203 Examples
204 --------
205 >>> radar_domain_only = radar_apply_mask( surface_microphysical_rainfall_rate, Nimrod2km_wts)
206 """
207 # Create an empty cubelist to hold the filtered fields.
208 masked_fields = iris.cube.CubeList([])
210 # Loop over the input mask and field cubes.
211 for M, F in zip(iter_maybe(mask), iter_maybe(original_field), strict=True):
212 masked_field = F.copy()
214 # Set the model perimeter to NaN as these gridpoints contain no useful data.
215 # c.f. boundary_margin in regrid.py
216 _mask_boundary(masked_field.data, boundary_margin)
218 # If the field and mask are on different grids, then regrid the field.
219 if M[0].shape != masked_field[0].shape: 219 ↛ 220line 219 didn't jump to line 220 because the condition on line 219 was never true
220 scheme = iris.analysis.Linear(extrapolation_mode="nan")
221 masked_field = masked_field.regrid(M, scheme)
223 # Apply the mask.
224 min_timesteps = min(M.shape[0], masked_field.shape[0])
225 masked_field = apply_mask(masked_field[0:min_timesteps], M[0:min_timesteps])
227 # Attach an attribute to the masked field detailing the mask used.
228 masked_field.attributes["mask"] = f"mask_of_{F.name()}"
230 # Append the masked field to the output list of masked fields.
231 masked_fields.append(masked_field)
233 # Return either a single cube or a cubelist.
234 if len(masked_fields) == 1:
235 return masked_fields[0]
236 else:
237 # return masked_fields
238 return masked_fields.merge()
241def _mask_boundary(data: np.ndarray, boundary_margin: int) -> None:
242 """Put a NaN strip along margin of width boundary_margin.
244 Parameters
245 ----------
246 data: array
247 The data array to be masked.
248 boundary_margin
249 Number of points to set as NaN at each edge.
250 """
251 if boundary_margin < 0:
252 raise ValueError(f"boundary_margin must be >= 0, got {boundary_margin}")
253 if boundary_margin == 0:
254 return
256 m = boundary_margin
257 data[..., :m, :] = np.nan # first rows
258 data[..., -m:, :] = np.nan # last rows
259 data[..., :m] = np.nan # first columns
260 data[..., -m:] = np.nan # last columns
263def radar_mask(
264 model_field: iris.cube.Cube | iris.cube.CubeList,
265 nimrod_field: iris.cube.Cube | iris.cube.CubeList,
266 nimrod_mask: iris.cube.Cube | iris.cube.CubeList,
267 boundary_margin: int = 8,
268 required_outputs: str = "radar",
269) -> iris.cube.Cube | iris.cube.CubeList:
270 """Apply a mask to given fields using a masked array.
272 Parameters
273 ----------
274 model_field: iris.cube.Cube | iris.cube.CubeList
275 The model field(s) to be masked.
276 nimrod_field: iris.cube.Cube | iris.cube.CubeList
277 The Nimrod field(s) to be masked.
278 nimrod_mask: iris.cube.Cube | iris.cube.CubeList
279 The Nimrod mask(s) to use. These are normally Nimrod wts fields.
280 boundary_margin: int, optional
281 Number of grid points from the domain boundary considered "unreliable".
282 Defaults to 8.
283 required_outputs: str, optional
284 Specifies which outputs are required:
285 "radar" outputs masked Nimrod rainfall field(s)
286 "model" outputs masked model field(s)
287 "all" outputs both masked model and masked Nimrod field(s).
289 Returns
290 -------
291 masked_field: iris.cube.Cube | iris.cube.CubeList
292 A cube or CubeList of the masked field(s).
294 Examples
295 --------
296 To mask both model and Nimrod fields using the Nimrod weights in nimrod_wts:
297 >>> masked_fields = radar_mask( model_fields, nimrod_fields, nimrod_wts, required_outputs="all")
299 """
300 # Create an empty cubelist to hold the filtered fields.
301 filtered_fields = iris.cube.CubeList([])
302 filtered_radar = iris.cube.CubeList([])
303 filtered_model = iris.cube.CubeList([])
305 # Loop over nimrod_field, model_field and nimrod_mask.
306 for M, F, N in zip(
307 iter_maybe(nimrod_mask),
308 iter_maybe(model_field),
309 iter_maybe(nimrod_field),
310 strict=True,
311 ):
312 # Apply the function radar_apply_mask to generate the re-gridded
313 # and masked model field.
314 masked_model_field = radar_apply_mask(F, M, boundary_margin=boundary_margin)
316 # Use the masked model field as the mask for the Nimrod field.
317 # Note: no re-gridding required.
318 min_timesteps = min(N.shape[0], masked_model_field.shape[0])
320 temp_mask = masked_model_field[0:min_timesteps].copy()
321 temp_mask.data[~np.isnan(temp_mask.data)] = 1.0
323 masked_nimrod_field = N[0:min_timesteps].copy()
324 masked_nimrod_field.data *= temp_mask.data
326 # Append the masked field to the output list of masked fields.
327 filtered_model.append(masked_model_field)
328 filtered_radar.append(masked_nimrod_field)
330 # Return either the masked model or Nimrod fields, or both.
331 match required_outputs:
332 case "radar":
333 filtered_fields.append(filtered_radar.merge_cube())
334 case "model":
335 filtered_fields.append(filtered_model.merge_cube())
336 case "all": 336 ↛ 341line 336 didn't jump to line 341 because the pattern on line 336 always matched
337 filtered_fields.append(filtered_model.merge_cube())
338 filtered_fields.append(filtered_radar.merge_cube())
340 # Return either a single cube or a cubelist.
341 if len(filtered_fields) == 1:
342 return filtered_fields[0]
343 else:
344 return filtered_fields
347def radar_mask_loop(
348 model_field: iris.cube.Cube | iris.cube.CubeList,
349 nimrod_field: iris.cube.Cube | iris.cube.CubeList,
350 nimrod_mask: iris.cube.Cube | iris.cube.CubeList,
351 boundary_margin: int = 8,
352 required_outputs: str = "all",
353) -> iris.cube.Cube | iris.cube.CubeList:
354 """Find common domains between a list of models and radar sources.
356 Parameters
357 ----------
358 model_field: iris.cube.Cube | iris.cube.CubeList
359 The model field(s) to be masked.
360 nimrod_field: iris.cube.Cube | iris.cube.CubeList
361 The Nimrod field(s) to be masked.
362 nimrod_mask: iris.cube.Cube | iris.cube.CubeList
363 The Nimrod mask(s) to use. These are normally Nimrod wts fields.
364 boundary_margin: int, optional
365 Number of grid points from the domain boundary considered "unreliable".
366 Defaults to 8.
367 required_outputs: str, optional
368 Specifies which outputs are required:
369 "radar" outputs masked Nimrod rainfall field(s)
370 "model" outputs masked model field(s)
371 "all" outputs both masked model and masked Nimrod field(s).
373 Returns
374 -------
375 masked_field: iris.cube.Cube | iris.cube.CubeList
376 A cube or CubeList of the masked field(s).
378 Examples
379 --------
380 To mask both model and Nimrod fields using the Nimrod weights in nimrod_wts:
381 >>> masked_fields = radar_mask_loop( model_fields, nimrod_fields, nimrod_wts, required_outputs="all")
383 """
384 # Create an empty cubelist to hold the filtered fields.
385 filtered_cubes = iris.cube.CubeList([])
387 use_nimrod_field = nimrod_field
388 use_nimrod_mask = nimrod_mask
390 # Loop over the models.
391 # using required_outputs="all" returns both
392 # masked model and radar fields
393 for model in model_field:
394 filtered_model = radar_mask(
395 model,
396 use_nimrod_field,
397 use_nimrod_mask,
398 boundary_margin=boundary_margin,
399 required_outputs=required_outputs,
400 )
401 filtered_cubes.append(filtered_model)
403 return filtered_cubes
406def match_varname_and_units(cubes: iris.cube.Cube | iris.cube.CubeList):
407 """Match the varname and units of a cube list.
409 Arguments
410 ---------
411 cubes: iris.cube.Cube | iris.cube.CubeList
412 A Cube or CubeList of a field to be matched.
414 Returns
415 -------
416 iris.cube.Cube | iris.cube.CubeList
417 The matched cubes.
420 Notes
421 -----
422 This function converts the names and units of a cube list to match
423 the first cube in the list. If just one cube is input, then this is
424 returned.
425 """
426 # If just one cube, then no need to match so return.
427 if isinstance(cubes, iris.cube.Cube):
428 cubes_in = iris.cube.CubeList([cubes])
429 else:
430 cubes_in = cubes
431 if len(cubes_in) == 1:
432 return cubes
434 # Initialise the list of matched cubes.
435 new_cubelist = iris.cube.CubeList([])
437 # Use the first cube in the CubeList as the base cube.
438 base_cube = cubes_in[0]
439 new_cubelist.append(base_cube)
441 # Loop over the cubes matching each to the base cube.
442 for cube in cubes_in[1:]:
443 new_cube = cube.copy()
445 # Match the cube varname.
446 new_cube.rename(base_cube.long_name)
447 new_cube.long_name = base_cube.long_name
448 new_cube.var_name = base_cube.var_name
450 # Match the cube units.
451 new_cube.units = base_cube.units
453 # Append the matched cube to the output cube list.
454 new_cubelist.append(new_cube)
456 return new_cubelist