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

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. 

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 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 

142 

143 

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. 

150 

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 '==','!=','<','<=','>', and '>='. The condition is consistent 

158 regardless of whether mask_field is a cube or CubeList. 

159 value: float 

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

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

162 

163 Returns 

164 ------- 

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

166 Mask(s) meeting the condition applied. 

167 

168 Raises 

169 ------ 

170 ValueError: Unexpected value for condition. Expected ==, !=, >, >=, <, <=. 

171 Got {condition}. 

172 Raised when condition is not supported. 

173 

174 Notes 

175 ----- 

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

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

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

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

180 500 m, and humidity == 100 %. 

181 

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

183 should happen after all relevant masks have been combined. 

184 

185 Examples 

186 -------- 

187 >>> land_mask = generate_mask(land_sea_mask,'==',1) 

188 """ 

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

190 for cube in iter_maybe(mask_field): 

191 mask = cube.copy() 

192 mask.data[:] = 0.0 

193 match condition: 

194 case "==": 

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

196 case "!=": 

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

198 case ">": 

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

200 case ">=": 

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

202 case "<": 

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

204 case "<=": 

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

206 case _: 

207 raise ValueError("""Unexpected value for condition. Expected ==, !=, 

208 >, >=, <, <=. Got {condition}.""") 

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

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

211 mask.units = "1" 

212 mask_list.append(mask) 

213 

214 if len(mask_list) == 1: 

215 return mask_list[0] 

216 else: 

217 return mask_list