Coverage for src / CSET / operators / filters.py: 98%

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. 

14 

15"""Operators to perform various kind of filtering.""" 

16 

17import logging 

18 

19import iris 

20import iris.cube 

21import iris.exceptions 

22import numpy as np 

23 

24from CSET._common import iter_maybe 

25 

26 

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. 

32 

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. 

39 

40 Returns 

41 ------- 

42 masked_field: iris.cube.Cube 

43 A cube of the masked field. 

44 

45 Notes 

46 ----- 

47 The mask is first converted to 1s and NaNs before multiplication with 

48 the original data. 

49 

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. 

52 

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 

68 

69 

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. 

76 

77 Arguments 

78 --------- 

79 cube: iris.cube.Cube | iris.cube.CubeList 

80 Cube(s) to filter 

81 constraint: iris.Constraint 

82 Constraint to extract 

83 

84 Returns 

85 ------- 

86 iris.cube.Cube 

87 

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 ) 

105 

106 

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. 

112 

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. 

120 

121 Returns 

122 ------- 

123 iris.cube.CubeList 

124 

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 # Switch to extract due to lack of instance requiriing one cube per 

136 # constraint. 

137 try: 

138 filtered_cubes = cubes.extract(kwargs.values()) 

139 except iris.exceptions.ConstraintMismatchError as err: 

140 raise ValueError("The constraints don't produce a cube or cubelist.") from err 

141 if len(filtered_cubes) == 0: 

142 raise ValueError("No cubes loaded. Please check your constraints.") 

143 return filtered_cubes 

144 

145 

146def generate_mask( 

147 mask_field: iris.cube.Cube | iris.cube.CubeList, 

148 condition: str, 

149 value: float, 

150) -> iris.cube.Cube | iris.cube.CubeList: 

151 """Generate a mask to remove data not meeting conditions. 

152 

153 Parameters 

154 ---------- 

155 mask_field: iris.cube.Cube | iris.cube.CubeList 

156 The field(s) to be used for creating the mask. 

157 condition: str 

158 The type of condition applied, six available options: 

159 'eq','ne','lt','le','gt', and 'ge'. The condition is consistent 

160 regardless of whether mask_field is a cube or CubeList. 

161 The conditions are as follows 

162 eq: equal to, 

163 ne: not equal to, 

164 lt: less than, 

165 le: less than or equal to, 

166 gt: greater than, 

167 ge: greater than or equal to. 

168 value: float 

169 The value on the right hand side of the condition. The value is 

170 consistent regardless of whether mask_field is a cube or CubeList. 

171 

172 Returns 

173 ------- 

174 mask: iris.cube.Cube | iris.cube.CubeList 

175 Mask(s) meeting the condition applied. 

176 

177 Raises 

178 ------ 

179 ValueError: Unexpected value for condition. Expected eq, ne, gt, ge, lt, le. 

180 Got {condition}. 

181 Raised when condition is not supported. 

182 

183 Notes 

184 ----- 

185 The mask is created in the opposite sense to numpy.ma.masked_arrays. This 

186 method was chosen to allow easy combination of masks together outside of 

187 this function using misc.addition or misc.multiplication depending on 

188 applicability. The combinations can be of any fields such as orography > 

189 500 m, and humidity == 100 %. 

190 

191 The conversion to a masked array occurs in the apply_mask routine, which 

192 should happen after all relevant masks have been combined. 

193 

194 Examples 

195 -------- 

196 >>> land_mask = generate_mask(land_sea_mask,'gt',1) 

197 """ 

198 mask_list = iris.cube.CubeList([]) 

199 for cube in iter_maybe(mask_field): 

200 mask = cube.copy() 

201 mask.data[:] = 0.0 

202 match condition: 

203 case "eq": 

204 mask.data[cube.data == value] = 1 

205 case "ne": 

206 mask.data[cube.data != value] = 1 

207 case "gt": 

208 mask.data[cube.data > value] = 1 

209 case "ge": 

210 mask.data[cube.data >= value] = 1 

211 case "lt": 

212 mask.data[cube.data < value] = 1 

213 case "le": 

214 mask.data[cube.data <= value] = 1 

215 case _: 

216 raise ValueError("""Unexpected value for condition. Expected eq, ne, 

217 gt, ge, lt, le. Got {condition}.""") 

218 mask.attributes["mask"] = f"mask_for_{cube.name()}_{condition}_{value}" 

219 mask.rename(f"mask_for_{cube.name()}_{condition}_{value}") 

220 mask.units = "1" 

221 mask_list.append(mask) 

222 

223 if len(mask_list) == 1: 

224 return mask_list[0] 

225 else: 

226 return mask_list