Coverage for src/CSET/operators/aggregate.py: 96%

1# © Crown copyright, Met Office (2022-2024) 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 aggregate across either 1 or 2 dimensions.""" 

16 

17import logging 

18 

19import iris 

20import iris.analysis 

21import iris.coord_categorisation 

22import iris.cube 

23import iris.exceptions 

24import iris.util 

25import isodate 

26import numpy as np 

27 

28from CSET._common import iter_maybe 

29from CSET.operators._utils import is_time_aggregatable 

30 

31logger = logging.getLogger(__name__) 

32 

33 

34def time_aggregate( 

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

36 method: str, 

37 interval_iso: str, 

38 **kwargs, 

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

40 """Aggregate cube/cubes by its time coordinate. 

41 

42 Aggregates fields in a cube or cube list for the specified coordinate and 

43 using the method supplied. The aggregated cube/cubelist will keep the coordinate and 

44 add a coordinate with the aggregated end time points. 

45 

46 Can also handle multiple forecast reference times. 

47 

48 Examples include generating hourly or 6-hourly precipitation accumulations for precipitation, or maximum screen level temperature every 3 hours. 

49 

50 We use the isodate class to convert ISO 8601 durations into time intervals 

51 for creating a new time coordinate for aggregation. 

52 

53 We use the lambda function to pass coord and interval into the callable 

54 category function in add_categorised to allow users to define their own 

55 sub-daily intervals for the new time coordinate. 

56 

57 Arguments 

58 --------- 

59 cubes: iris.cube.Cube | iris.cube.CubeList 

60 Cube or CubeList to aggregate and iterate over one dimension 

61 method: str 

62 Type of aggregate i.e. method: 'SUM', getattr creates 

63 iris.analysis.SUM, etc. 

64 interval_iso: isodate timedelta ISO 8601 object i.e PT6H (6 hours), PT30M (30 mins) 

65 Interval to aggregate over. 

66 interval_iso: str 

67 A string containing a datetime timedelta for resampling over in hours, i.e. PT3H, PT24H. 

68 

69 Returns 

70 ------- 

71 resampled_cubes: iris.cube.Cube | iris.cube.CubeList 

72 Cube or CubeList containing aggregated cubes. 

73 

74 Raises 

75 ------ 

76 ValueError 

77 If the constraint doesn't produce a single cube containing a field. 

78 """ 

79 # Return unchanged cubes if interval_iso is 0, to allow this operator to be used across multiple fields where only some will need resampling. 

80 if interval_iso == "0": 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true

81 return cubes 

82 

83 cubes = iter_maybe(cubes) 

84 

85 resampled_cubes = iris.cube.CubeList() 

86 

87 timedelta = isodate.parse_duration(interval_iso) 

88 interval = int(timedelta.total_seconds() / 3600) 

89 

90 for cube in cubes: 

91 if cube.coord("forecast_reference_time").shape[0] > 1: 

92 # Handle cubes with multiple forecast cycles. 

93 aggregated_cube = _aggregate_in_time_multiple_frt(cube, method, interval) 

94 else: 

95 aggregated_cube = _aggregate_in_time_single_frt(cube, method, interval) 

96 

97 resampled_cubes.append(aggregated_cube) 

98 if len(resampled_cubes) == 1: 98 ↛ 100line 98 didn't jump to line 100 because the condition on line 98 was always true

99 return resampled_cubes[0] 

100 return resampled_cubes 

101 

102 

103def ensure_aggregatable_across_cases( 

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

105) -> iris.cube.CubeList: 

106 """Ensure a Cube or CubeList can be aggregated across multiple cases. 

107 

108 The cubes are grouped into buckets of compatible cubes, then each bucket is 

109 converted into a single aggregatable cube with ``forecast_period`` and 

110 ``forecast_reference_time`` dimension coordinates. 

111 

112 Arguments 

113 --------- 

114 cubes: iris.cube.Cube | iris.cube.CubeList 

115 Each cube is checked to determine if it has the necessary 

116 dimensional coordinates to be aggregatable, being processed if needed. 

117 

118 Returns 

119 ------- 

120 cubes: iris.cube.CubeList 

121 A CubeList of time aggregatable cubes. 

122 

123 Raises 

124 ------ 

125 ValueError 

126 If any of the provided cubes cannot be made aggregatable. 

127 

128 Notes 

129 ----- 

130 This is a simple operator designed to ensure that a Cube is aggregatable 

131 across cases. If a CubeList is presented it will create an aggregatable Cube 

132 from that list. Its functionality is for case study (or trial) aggregation 

133 to ensure that the full dataset can be loaded as a single cube. This 

134 functionality is particularly useful for percentiles, Q-Q plots, and 

135 histograms. 

136 

137 The necessary dimension coordinates for a cube to be aggregatable are 

138 ``forecast_period`` and ``forecast_reference_time``. 

139 """ 

140 

141 # Group compatible cubes. 

142 class Buckets: 

143 def __init__(self): 

144 self.buckets = [] 

145 

146 def add(self, cube: iris.cube.Cube): 

147 """Add a cube into a bucket. 

148 

149 If the cube is compatible with an existing bucket it is added there. 

150 Otherwise it gets its own bucket. 

151 """ 

152 for bucket in self.buckets: 

153 if bucket[0].is_compatible(cube): 

154 bucket.append(cube) 

155 return 

156 self.buckets.append(iris.cube.CubeList([cube])) 

157 

158 def get_buckets(self) -> list[iris.cube.CubeList]: 

159 return self.buckets 

160 

161 b = Buckets() 

162 for cube in iter_maybe(cubes): 

163 b.add(cube) 

164 buckets = b.get_buckets() 

165 

166 logger.debug("Buckets:\n%s", "\n---\n".join(str(b) for b in buckets)) 

167 

168 # Ensure each bucket is a single aggregatable cube. 

169 aggregatable_cubes = iris.cube.CubeList() 

170 for bucket in buckets: 

171 # Single cubes that are already aggregatable won't need processing. 

172 if len(bucket) == 1 and is_time_aggregatable(bucket[0]): 

173 aggregatable_cube = bucket[0] 

174 aggregatable_cube = _add_nref(aggregatable_cube) 

175 aggregatable_cubes.append(aggregatable_cube) 

176 continue 

177 

178 # Create an aggregatable cube from the provided CubeList. 

179 to_merge = iris.cube.CubeList() 

180 for cube in bucket: 

181 try: 

182 to_merge.extend( 

183 cube.slices_over(["forecast_period", "forecast_reference_time"]) 

184 ) 

185 except iris.exceptions.CoordinateNotFoundError as err: 

186 raise ValueError( 

187 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.", 

188 cube, 

189 ) from err 

190 aggregatable_cube = to_merge.merge_cube() 

191 

192 # Add attribute on number of forecast_reference_times 

193 aggregatable_cube = _add_nref(aggregatable_cube) 

194 

195 # Verify cube is now aggregatable. 

196 if not is_time_aggregatable(aggregatable_cube): 196 ↛ 197line 196 didn't jump to line 197 because the condition on line 196 was never true

197 raise ValueError( 

198 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.", 

199 aggregatable_cube, 

200 ) 

201 aggregatable_cubes.append(aggregatable_cube) 

202 

203 return aggregatable_cubes 

204 

205 

206def add_hour_coordinate( 

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

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

209 """Add a category coordinate of hour of day to a Cube or CubeList. 

210 

211 Arguments 

212 --------- 

213 cubes: iris.cube.Cube | iris.cube.CubeList 

214 Cube of any variable that has a time coordinate. 

215 Note input Cube or CubeList items should only have 1 time dimension. 

216 

217 Returns 

218 ------- 

219 cube: iris.cube.Cube 

220 A Cube with an additional auxiliary coordinate of hour. 

221 

222 Notes 

223 ----- 

224 This is a simple operator designed to be used prior to case aggregation for 

225 histograms, Q-Q plots, and percentiles when aggregated by hour of day. 

226 """ 

227 new_cubelist = iris.cube.CubeList() 

228 for cube in iter_maybe(cubes): 

229 # Add a category coordinate of hour into each cube. 

230 iris.util.promote_aux_coord_to_dim_coord(cube, "time") 

231 iris.coord_categorisation.add_hour(cube, "time", name="hour") 

232 cube.coord("hour").units = "hours" 

233 new_cubelist.append(cube) 

234 

235 if len(new_cubelist) == 1: 

236 return new_cubelist[0] 

237 else: 

238 return new_cubelist 

239 

240 

241def rolling_window_time_aggregation( 

242 cubes: iris.cube.Cube | iris.cube.CubeList, method: str, window: int 

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

244 """Aggregate a cube along the time dimension using a rolling window. 

245 

246 Arguments 

247 --------- 

248 cubes: iris.cube.Cube | iris.cube.CubeList 

249 Cube or Cubelist of any variable to be aggregated over a rolling window 

250 in time. 

251 method: str 

252 Type of aggregate i.e. method: 'MAX', getattr creates 

253 iris.analysis.MAX, etc. 

254 window: int 

255 The rolling window size. 

256 

257 Returns 

258 ------- 

259 cube: iris.cube.Cube | iris.cube.CubeList 

260 A Cube or Cubelist of the rolling window aggregate. The Cubes will have 

261 a time dimension that is reduced in size to the original cube by the 

262 window size. 

263 

264 Notes 

265 ----- 

266 This operator is designed to be used to help create daily maxima and minima 

267 for any variable. 

268 """ 

269 new_cubelist = iris.cube.CubeList() 

270 for cube in iter_maybe(cubes): 

271 # Use a rolling window in time to applied specified aggregation method 

272 # over a specified window length. 

273 window_cube = cube.rolling_window( 

274 "time", getattr(iris.analysis, method), window 

275 ) 

276 new_cubelist.append(window_cube) 

277 

278 if len(new_cubelist) == 1: 

279 return new_cubelist[0] 

280 else: 

281 return new_cubelist 

282 

283 

284def _add_nref(cube: iris.cube.Cube): 

285 """Retain information on number of forecast_reference_time inputs. 

286 

287 This preserves information on number of aggregated cases that can 

288 otherwise be lost on subsequent calls to collapse functions. 

289 """ 

290 nref = np.size(cube.coord("forecast_reference_time").points) 

291 cube.coord("time").attributes["number_reference_times"] = nref 

292 return cube 

293 

294 

295def _aggregate_in_time_multiple_frt( 

296 cube: iris.cube.Cube, method: str, interval: int 

297) -> iris.cube.Cube: 

298 """Aggregate a cube with multiple forecast reference times. 

299 

300 Aggregates each forecast cycle separately, then concatenates the results 

301 back into a single cube along forecast_reference_time. 

302 """ 

303 aggregated_cycles = iris.cube.CubeList() 

304 for frt_cube in cube.slices_over("forecast_reference_time"): 

305 iris.coord_categorisation.add_categorised_coord( 

306 frt_cube, 

307 "interval", 

308 "time", 

309 lambda coord, cell: cell // interval * interval, 

310 ) 

311 agg = frt_cube.aggregated_by( 

312 "interval", 

313 getattr(iris.analysis, method), 

314 ) 

315 agg.remove_coord("interval") 

316 agg = iris.util.new_axis( 

317 agg, 

318 agg.coord("forecast_reference_time"), 

319 ) 

320 aggregated_cycles.append(agg) 

321 # Current approach is to remove time 2d auxcoord as it causes issues concatenating into a single cube. 

322 # Time can be reconstructed from forecast_reference_time and forecast_period later on. 

323 # as we can construct it if needed. 

324 for cb in aggregated_cycles: 

325 cb.remove_coord("time") 

326 return aggregated_cycles.concatenate_cube() 

327 

328 

329def _aggregate_in_time_single_frt(cube: iris.cube.Cube, method: str, interval: int): 

330 """Aggregate a cube with one forecast reference time.""" 

331 iris.coord_categorisation.add_categorised_coord( 

332 cube, 

333 "interval", 

334 "time", 

335 lambda coord, cell: cell // interval * interval, 

336 ) 

337 aggregated_cube = cube.aggregated_by( 

338 "interval", 

339 getattr(iris.analysis, method), 

340 ) 

341 aggregated_cube.remove_coord("interval") 

342 return aggregated_cube