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

1# Copyright 2022 Met Office and 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 generate constraints to filter with.""" 

16 

17from collections.abc import Iterable 

18from datetime import datetime 

19 

20import iris 

21import iris.cube 

22import iris.exceptions 

23 

24 

25def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint: 

26 """Generate constraint from STASH code. 

27 

28 Operator that takes a stash string, and uses iris to generate a constraint 

29 to be passed into the read operator to minimize the CubeList the read 

30 operator loads and speed up loading. 

31 

32 Arguments 

33 --------- 

34 stash: str 

35 stash code to build iris constraint, such as "m01s03i236" 

36 

37 Returns 

38 ------- 

39 stash_constraint: iris.AttributeConstraint 

40 """ 

41 # At a later stage str list an option to combine constraints. Arguments 

42 # could be a list of stash codes that combined build the constraint. 

43 stash_constraint = iris.AttributeConstraint(STASH=stash) 

44 return stash_constraint 

45 

46 

47def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint: 

48 """Generate constraint from variable name. 

49 

50 Operator that takes a CF compliant variable name string, and uses iris to 

51 generate a constraint to be passed into the read operator to minimize the 

52 CubeList the read operator loads and speed up loading. 

53 

54 Arguments 

55 --------- 

56 varname: str 

57 CF compliant name of variable. Needed later for LFRic. 

58 

59 Returns 

60 ------- 

61 varname_constraint: iris.Constraint 

62 """ 

63 varname_constraint = iris.Constraint(name=varname) 

64 return varname_constraint 

65 

66 

67def generate_level_constraint( 

68 coordinate: str, levels: int | list[int], **kwargs 

69) -> iris.Constraint: 

70 """Generate constraint for particular levels on the specified coordinate. 

71 

72 Operator that generates a constraint to constrain to specific model or 

73 pressure levels. If no levels are specified then any cube with the specified 

74 coordinate is rejected. 

75 

76 Typically ``coordinate`` will be ``"pressure"`` or ``"model_level_number"`` 

77 for UM, or ``"full_levels"`` or ``"half_levels"`` for LFRic. 

78 

79 Arguments 

80 --------- 

81 coordinate: str 

82 Level coordinate name about which to constraint. 

83 levels: int | list[int] 

84 CF compliant levels. 

85 

86 Returns 

87 ------- 

88 constraint: iris.Constraint 

89 """ 

90 # Ensure is iterable. 

91 if not isinstance(levels, Iterable): 

92 levels = [levels] 

93 

94 # When no levels specified reject cube with level coordinate. 

95 if len(levels) == 0: 

96 

97 def no_levels(cube): 

98 # Reject cubes for which coordinate exists. 

99 return not bool(cube.coords(coordinate)) 

100 

101 return iris.Constraint(cube_func=no_levels) 

102 

103 # Filter the coordinate to the desired levels. 

104 # Dictionary unpacking is used to provide programmatic keyword arguments. 

105 return iris.Constraint(**{coordinate: levels}) 

106 

107 

108def generate_cell_methods_constraint(cell_methods: list, **kwargs) -> iris.Constraint: 

109 """Generate constraint from cell methods. 

110 

111 Operator that takes a list of cell methods and generates a constraint from 

112 that. 

113 

114 Arguments 

115 --------- 

116 cell_methods: list 

117 cube.cell_methods for filtering 

118 

119 Returns 

120 ------- 

121 cell_method_constraint: iris.Constraint 

122 """ 

123 

124 def check_cell_methods(cube: iris.cube.Cube): 

125 return cube.cell_methods == tuple(cell_methods) 

126 

127 cell_methods_constraint = iris.Constraint(cube_func=check_cell_methods) 

128 return cell_methods_constraint 

129 

130 

131def generate_time_constraint( 

132 time_start: str, time_end: str = None, **kwargs 

133) -> iris.AttributeConstraint: 

134 """Generate constraint between times. 

135 

136 Operator that takes one or two ISO 8601 date strings, and returns a 

137 constraint that selects values between those dates (inclusive). 

138 

139 Arguments 

140 --------- 

141 time_start: str | datetime.datetime 

142 ISO date for lower bound 

143 

144 time_end: str | datetime.datetime 

145 ISO date for upper bound. If omitted it defaults to the same as 

146 time_start 

147 

148 Returns 

149 ------- 

150 time_constraint: iris.Constraint 

151 """ 

152 if isinstance(time_start, str): 

153 time_start = datetime.fromisoformat(time_start) 

154 if time_end is None: 

155 time_end = time_start 

156 elif isinstance(time_end, str): 

157 time_end = datetime.fromisoformat(time_end) 

158 time_constraint = iris.Constraint(time=lambda t: time_start <= t.point <= time_end) 

159 return time_constraint 

160 

161 

162def generate_area_constraint( 

163 lat_start: float, lat_end: float, lon_start: float, lon_end: float, **kwargs 

164) -> iris.Constraint: 

165 """Generate an area constraint between latitude/longitude limits. 

166 

167 Operator that takes a set of latitude and longitude limits and returns a 

168 constraint that selects grid values only inside that area. Works with the 

169 data's native grid so is defined within the rotated pole CRS. 

170 

171 Arguments 

172 --------- 

173 lat_start: float 

174 Latitude value for lower bound 

175 lat_end: float 

176 Latitude value for top bound 

177 lon_start: float 

178 Longitude value for left bound 

179 lon_end: float 

180 Longitude value for right bound 

181 

182 Returns 

183 ------- 

184 area_constraint: iris.Constraint 

185 """ 

186 area_constraint = iris.Constraint( 186 ↛ exitline 186 didn't jump to the function exit

187 coord_values={ 

188 "grid_latitude": lambda cell: lat_start < cell < lat_end, 

189 "grid_longitude": lambda cell: lon_start < cell < lon_end, 

190 } 

191 ) 

192 return area_constraint 

193 

194 

195def combine_constraints( 

196 constraint: iris.Constraint = None, **kwargs 

197) -> iris.Constraint: 

198 """ 

199 Operator that combines multiple constraints into one. 

200 

201 Arguments 

202 --------- 

203 constraint: iris.Constraint 

204 First constraint to combine. 

205 additional_constraint_1: iris.Constraint 

206 Second constraint to combine. This must be a named argument. 

207 additional_constraint_2: iris.Constraint 

208 There can be any number of additional constraint, they just need unique 

209 names. 

210 ... 

211 

212 Returns 

213 ------- 

214 combined_constraint: iris.Constraint 

215 

216 Raises 

217 ------ 

218 TypeError 

219 If the provided arguments are not constraints. 

220 """ 

221 # If the first argument is not a constraint, it is ignored. This handles the 

222 # automatic passing of the previous step's output. 

223 if isinstance(constraint, iris.Constraint): 

224 combined_constraint = constraint 

225 else: 

226 combined_constraint = iris.Constraint() 

227 

228 for constr in kwargs.values(): 

229 combined_constraint = combined_constraint & constr 

230 return combined_constraint