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 

26logger = logging.getLogger(__name__) 

27 

28 

29def apply_mask( 

30 original_field: iris.cube.Cube | iris.cube.CubeList, 

31 mask: iris.cube.Cube | iris.cube.CubeList, 

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

33 """Apply a mask to given data as a masked array. 

34 

35 Parameters 

36 ---------- 

37 original_field: iris.cube.Cube | iris.cube.CubeList 

38 The field(s) to be masked. 

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

40 The mask(s) being applied to the original field(s). 

41 

42 Returns 

43 ------- 

44 masked_field: iris.cube.Cube | iris.cube.CubeList 

45 A cube or cubelist of the masked field(s). 

46 

47 Notes 

48 ----- 

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

50 the original data. 

51 

52 As discussed in generate_mask, you can combine multiple masks in a 

53 recipe using other functions before applying the mask to the data. 

54 

55 Examples 

56 -------- 

57 >>> land_points_only = apply_mask(temperature, land_mask) 

58 """ 

59 masked_fields = iris.cube.CubeList([]) 

60 for M, F in zip(iter_maybe(mask), iter_maybe(original_field), strict=True): 

61 # Ensure mask data are floats and only 1s or NaNs. 

62 M.data = np.float64(M.data) 

63 M.data[M.data == 0.0] = np.nan 

64 M.data[~np.isnan(M.data)] = 1.0 

65 logger.info( 

66 "Mask set to 1 or 0s, if addition of multiple masks results" 

67 "in values > 1 these are set to 1." 

68 ) 

69 masked_field = F.copy() 

70 masked_field.data *= M.data 

71 masked_field.attributes["mask"] = f"mask_of_{F.name()}" 

72 masked_fields.append(masked_field) 

73 if len(masked_fields) == 1: 

74 return masked_fields[0] 

75 else: 

76 return masked_fields 

77 

78 

79def filter_cubes( 

80 cube: iris.cube.Cube | iris.cube.CubeList, 

81 constraint: iris.Constraint, 

82 **kwargs, 

83) -> iris.cube.Cube: 

84 """Filter a CubeList down to a single Cube based on a constraint. 

85 

86 Arguments 

87 --------- 

88 cube: iris.cube.Cube | iris.cube.CubeList 

89 Cube(s) to filter 

90 constraint: iris.Constraint 

91 Constraint to extract 

92 

93 Returns 

94 ------- 

95 iris.cube.Cube 

96 

97 Raises 

98 ------ 

99 ValueError 

100 If the constraint doesn't produce a single cube. 

101 """ 

102 filtered_cubes = cube.extract(constraint) 

103 # Return directly if already a cube. 

104 if isinstance(filtered_cubes, iris.cube.Cube): 

105 return filtered_cubes 

106 # Check filtered cubes is a CubeList containing one cube. 

107 if isinstance(filtered_cubes, iris.cube.CubeList) and len(filtered_cubes) == 1: 

108 return filtered_cubes[0] 

109 else: 

110 raise ValueError( 

111 f"Constraint doesn't produce single cube. Constraint: {constraint}" 

112 f"\nSource: {cube}\nResult: {filtered_cubes}" 

113 ) 

114 

115 

116def filter_multiple_cubes( 

117 cubes: iris.cube.Cube | iris.cube.CubeList, 

118 **kwargs, 

119) -> iris.cube.CubeList: 

120 """Filter a CubeList on multiple constraints, returning another CubeList. 

121 

122 Arguments 

123 --------- 

124 cube: iris.cube.Cube | iris.cube.CubeList 

125 Cube(s) to filter 

126 constraint: iris.Constraint 

127 Constraint to extract. This must be a named argument. There can be any 

128 number of additional constraints, they just need unique names. 

129 

130 Returns 

131 ------- 

132 iris.cube.CubeList 

133 

134 Raises 

135 ------ 

136 ValueError 

137 Some constraints don't produce any cubes. 

138 """ 

139 # Ensure input is a CubeList. 

140 if isinstance(cubes, iris.cube.Cube): 

141 cubes = iris.cube.CubeList((cubes,)) 

142 if len(kwargs) < 1: 

143 raise ValueError("Must have at least one constraint.") 

144 # Switch to extract due to lack of instance requiring one cube per 

145 # constraint. 

146 try: 

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

148 except iris.exceptions.ConstraintMismatchError as err: 

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

150 if len(filtered_cubes) == 0: 

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

152 return filtered_cubes 

153 

154 

155def generate_mask( 

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

157 condition: str, 

158 value: float, 

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

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

161 

162 Parameters 

163 ---------- 

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

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

166 condition: str 

167 The type of condition applied, six available options: 

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

169 regardless of whether mask_field is a cube or CubeList. 

170 The conditions are as follows 

171 eq: equal to, 

172 ne: not equal to, 

173 lt: less than, 

174 le: less than or equal to, 

175 gt: greater than, 

176 ge: greater than or equal to. 

177 value: float 

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

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

180 

181 Returns 

182 ------- 

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

184 Mask(s) meeting the condition applied. 

185 

186 Raises 

187 ------ 

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

189 Got {condition}. 

190 Raised when condition is not supported. 

191 

192 Notes 

193 ----- 

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

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

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

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

198 500 m, and humidity == 100 %. 

199 

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

201 should happen after all relevant masks have been combined. 

202 

203 Examples 

204 -------- 

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

206 """ 

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

208 for cube in iter_maybe(mask_field): 

209 mask = cube.copy() 

210 mask.data[:] = 0.0 

211 match condition: 

212 case "eq": 

213 mask.data[cube.data == value] = 1.0 

214 case "ne": 

215 mask.data[cube.data != value] = 1.0 

216 case "gt": 

217 mask.data[cube.data > value] = 1.0 

218 case "ge": 

219 mask.data[cube.data >= value] = 1.0 

220 case "lt": 

221 mask.data[cube.data < value] = 1.0 

222 case "le": 

223 mask.data[cube.data <= value] = 1.0 

224 case _: 

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

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

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

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

229 mask.units = "1" 

230 mask_list.append(mask) 

231 

232 if len(mask_list) == 1: 

233 return mask_list[0] 

234 else: 

235 return mask_list