Coverage for src/CSET/operators/dfss.py: 56%

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"""A module containing the dFSS diagnostic.""" 

16 

17import multiprocessing as mp 

18import os 

19from functools import partial 

20from typing import List, Union 

21 

22import cartopy.crs as ccrs 

23import iris 

24import iris.coords as icoords 

25import iris.cube 

26import numpy as np 

27import numpy.ma as ma 

28from iris import coord_systems 

29 

30from .improver.nbhood import NeighbourhoodProcessing 

31 

32 

33def init_worker(): 

34 """Set env vars to ensure no over subscription.""" 

35 os.environ["OMP_NUM_THREADS"] = "1" 

36 os.environ["MKL_NUM_THREADS"] = "1" 

37 os.environ["OPENBLAS_NUM_THREADS"] = "1" 

38 

39 

40def _dfss_on_slice( 

41 slice, 

42 neighbourhood_lengths, 

43 centile_or_threshold, 

44 centile, 

45 threshold, 

46): 

47 time_point = slice.coord("time") 

48 dfss_cube, dfss_stdev_cube = _calc_dfss( 

49 slice, 

50 neighbourhood_lengths, 

51 time_point, 

52 centile_or_threshold, 

53 centile, 

54 threshold, 

55 ) 

56 

57 return dfss_cube, dfss_stdev_cube 

58 

59 

60def _parallel_calculate_dfss( 

61 cube_xy: iris.cube.Cube, 

62 neighbourhood_lengths: List[int], 

63 centile_or_threshold: str = "centile", 

64 centile: float = None, 

65 threshold: float = None, 

66): 

67 

68 time_slices = list(cube_xy.slices_over("time")) 

69 

70 with mp.Pool(initializer=init_worker) as pool: 

71 worker = partial( 

72 _dfss_on_slice, 

73 neighbourhood_lengths=neighbourhood_lengths, 

74 centile_or_threshold=centile_or_threshold, 

75 centile=centile, 

76 threshold=threshold, 

77 ) 

78 dfss_cubes, dfss_stdev_cubes = zip(*pool.imap(worker, time_slices), strict=True) 

79 

80 cube_list_dfss = iris.cube.CubeList(dfss_cubes) 

81 cube_list_dfss_stdev = iris.cube.CubeList(dfss_stdev_cubes) 

82 

83 forecast_period = cube_xy.coord("forecast_period") 

84 merged_cube_dfss = cube_list_dfss.merge_cube() 

85 

86 merged_cube_dfss_stdev = cube_list_dfss_stdev.merge_cube() 

87 

88 merged_cube_dfss.add_aux_coord(forecast_period, data_dims=(0)) 

89 merged_cube_dfss_stdev.add_aux_coord(forecast_period, data_dims=(0)) 

90 

91 out_cube_list = iris.cube.CubeList() 

92 

93 out_cube_list.append(merged_cube_dfss) 

94 out_cube_list.append(merged_cube_dfss_stdev) 

95 

96 return out_cube_list 

97 

98 

99def _serial_calculate_dfss( 

100 cube_xy: iris.cube.Cube, 

101 neighbourhood_lengths: List[int], 

102 centile_or_threshold: str = "centile", 

103 centile: float = None, 

104 threshold: float = None, 

105): 

106 

107 cube_list_dfss = iris.cube.CubeList() 

108 cube_list_dfss_stdev = iris.cube.CubeList() 

109 

110 for time_slice in cube_xy.slices_over("time"): 

111 time_point = time_slice.coord("time") 

112 dfss_cube, dfss_stdev_cube = _calc_dfss( 

113 time_slice, 

114 neighbourhood_lengths, 

115 time_point, 

116 centile_or_threshold, 

117 centile, 

118 threshold, 

119 ) 

120 cube_list_dfss.append(dfss_cube) 

121 cube_list_dfss_stdev.append(dfss_stdev_cube) 

122 

123 forecast_period = cube_xy.coord("forecast_period") 

124 merged_cube_dfss = cube_list_dfss.merge_cube() 

125 merged_cube_dfss_stdev = cube_list_dfss_stdev.merge_cube() 

126 

127 merged_cube_dfss.add_aux_coord(forecast_period, data_dims=(0)) 

128 merged_cube_dfss_stdev.add_aux_coord(forecast_period, data_dims=(0)) 

129 

130 out_cube_list = iris.cube.CubeList() 

131 

132 out_cube_list.append(merged_cube_dfss) 

133 out_cube_list.append(merged_cube_dfss_stdev) 

134 

135 return out_cube_list 

136 

137 

138def calculate_dfss( 

139 cube_xy: iris.cube.Cube, 

140 neighbourhood_lengths: List[int], 

141 centile_or_threshold: str = "centile", 

142 centile: float = None, 

143 threshold: float = None, 

144 run_parallel: bool = True, 

145): 

146 """Do the dfss calculation. 

147 

148 Args: 

149 cube_xy: 

150 Original xy cube 

151 neighbourhood_lengths: list of neighbourhood lengths 

152 centile_or_threshold: centile or threshold method 

153 centile: centile value 

154 threshold: threshold value 

155 run_parallel: run the multiple time points in parallel 

156 

157 

158 Returns 

159 ------- 

160 dfss_cube 

161 cube with dfss variable 

162 dfss_stdev_cube 

163 cube with dfss standard deviation variable 

164 

165 """ 

166 if len(cube_xy.coord("realization").points) == 1: 

167 raise ValueError("dFSS is only valid for an ensemble") 

168 

169 # force serial run if running in workflow 

170 if "CYLC_RUN_DIR" in os.environ: 170 ↛ 171line 170 didn't jump to line 171 because the condition on line 170 was never true

171 run_parallel = False 

172 

173 if run_parallel: 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true

174 out_cube_list = _parallel_calculate_dfss( 

175 cube_xy, neighbourhood_lengths, centile_or_threshold, centile, threshold 

176 ) 

177 else: 

178 out_cube_list = _serial_calculate_dfss( 

179 cube_xy, neighbourhood_lengths, centile_or_threshold, centile, threshold 

180 ) 

181 

182 for cubes in out_cube_list: 

183 cubes.attributes["method"] = centile_or_threshold 

184 if centile_or_threshold == "centile": 184 ↛ 186line 184 didn't jump to line 186 because the condition on line 184 was always true

185 cubes.attributes["centile"] = centile 

186 if centile_or_threshold == "threshold": 186 ↛ 187line 186 didn't jump to line 187 because the condition on line 186 was never true

187 cubes.attributes["threshold"] = threshold 

188 

189 return out_cube_list 

190 

191 

192def _calc_dfss( 

193 cube_xy: iris.cube.Cube, 

194 neighbourhood_lengths: Union[List[int]], 

195 time_point, 

196 centile_or_threshold: str = "centile", 

197 centile: float = None, 

198 threshold: float = None, 

199): 

200 _ = ( 

201 cube_xy.data 

202 ) # NOTE: without realising the data, dask is very slow to run this code 

203 del _ 

204 ens_members = cube_xy.coord("realization").points 

205 number_of_members = len(ens_members) 

206 dfss = np.zeros(len(neighbourhood_lengths)) 

207 dfss_stdev = np.zeros_like(dfss) 

208 

209 for i, neighbourhood_length in enumerate(neighbourhood_lengths): 

210 fss_array = np.full((number_of_members, number_of_members), np.nan) 

211 iu1 = np.triu_indices(number_of_members, 1) 

212 

213 fss_array[iu1] = 0.0 

214 for i_a, memb_a in enumerate(ens_members): 

215 ens_member_a = cube_xy.extract(iris.Constraint(realization=memb_a)) 

216 for i_b, memb_b in enumerate(ens_members): 

217 if not np.isnan(fss_array[i_a, i_b]): 

218 ens_member_b = cube_xy.extract(iris.Constraint(realization=memb_b)) 

219 

220 fss_array[i_a, i_b] = _calc_fss( 

221 ens_member_a, 

222 ens_member_b, 

223 neighbourhood_length, 

224 centile_or_threshold=centile_or_threshold, 

225 centile=centile, 

226 threshold=threshold, 

227 ) 

228 

229 if fss_array[i_a, i_b] == np.nan: 229 ↛ 230line 229 didn't jump to line 230 because the condition on line 229 was never true

230 dfss[:] = np.nan 

231 dfss_stdev[:] = np.nan 

232 return dfss, dfss_stdev 

233 

234 dfss[i] = (ma.masked_invalid(fss_array)).mean() 

235 dfss_stdev[i] = (ma.masked_invalid(fss_array)).std() 

236 

237 neighbourhood_coord = icoords.DimCoord( 

238 neighbourhood_lengths, 

239 var_name="neighbourhoods", 

240 ) 

241 

242 dfss_cube = iris.cube.Cube( 

243 dfss, long_name="dfss", dim_coords_and_dims=[(neighbourhood_coord, 0)] 

244 ) 

245 dfss_cube.add_aux_coord(time_point) 

246 dfss_stdev_cube = iris.cube.Cube( 

247 dfss_stdev, 

248 long_name="dfss_stdev", 

249 dim_coords_and_dims=[(neighbourhood_coord, 0)], 

250 ) 

251 dfss_stdev_cube.add_aux_coord(time_point) 

252 return dfss_cube, dfss_stdev_cube 

253 

254 

255def _calc_fss( 

256 cube_a_in: iris.cube.Cube, 

257 cube_b_in: iris.cube.Cube, 

258 neighbourhood_length: int, 

259 centile_or_threshold: str = "centile", 

260 centile: float = None, 

261 threshold: float = None, 

262): 

263 # Set the threshold of interest 

264 

265 cube_a = cube_a_in.copy() 

266 cube_b = cube_b_in.copy() 

267 

268 if centile_or_threshold == "centile": 268 ↛ 280line 268 didn't jump to line 280 because the condition on line 268 was always true

269 cube_a_has_mask = isinstance(cube_a.data, np.ma.MaskedArray) 

270 cube_b_has_mask = isinstance(cube_b.data, np.ma.MaskedArray) 

271 if cube_a_has_mask and cube_b_has_mask: 271 ↛ 272line 271 didn't jump to line 272 because the condition on line 271 was never true

272 threshold_a = np.nanpercentile(cube_a.data.filled(np.nan), centile) 

273 threshold_b = np.nanpercentile(cube_b.data.filled(np.nan), centile) 

274 elif not cube_a_has_mask and not cube_b_has_mask: 274 ↛ 278line 274 didn't jump to line 278 because the condition on line 274 was always true

275 threshold_a = np.percentile(cube_a.data, centile) 

276 threshold_b = np.percentile(cube_b.data, centile) 

277 else: 

278 msg = "Mask status must be the same for both cubes" 

279 raise UserWarning(msg) 

280 elif centile_or_threshold == "threshold": 

281 threshold_a = threshold_b = threshold 

282 else: 

283 msg = ( 

284 "Function argument centile_or_threshold must equal one " 

285 "of [centile, threshold]" 

286 ) 

287 raise UserWarning(msg) 

288 

289 cube_a.data = np.ma.where(cube_a.data > threshold_a, 1, 0) 

290 cube_b.data = np.ma.where(cube_b.data > threshold_b, 1, 0) 

291 

292 n_a = np.count_nonzero(cube_a.data) 

293 n_b = np.count_nonzero(cube_b.data) 

294 

295 n_tot = np.size(cube_a.data) 

296 frac_cov_a = n_a / n_tot 

297 frac_cov_b = n_b / n_tot 

298 

299 if np.maximum(frac_cov_a, frac_cov_b) <= 0.002: 299 ↛ 303line 299 didn't jump to line 303 because the condition on line 299 was always true

300 fss = np.nan 

301 return fss 

302 

303 nbhooder = NeighbourhoodProcessing( 

304 neighbourhood_method="square", radii=neighbourhood_length 

305 ) 

306 

307 cube_a = _regrid_lat_lon_cube_to_xy_cube(cube_a) 

308 cube_b = _regrid_lat_lon_cube_to_xy_cube(cube_b) 

309 

310 fraction_fields_a = nbhooder.process(cube_a) 

311 fraction_fields_b = nbhooder.process(cube_b) 

312 

313 field_a = fraction_fields_a.data 

314 field_b = fraction_fields_b.data 

315 

316 fss = _calc_fss_two_fields(field_a.data, field_b.data) 

317 

318 return fss 

319 

320 

321def _calc_fss_two_fields(field_a, field_b): 

322 field_diff = field_a - field_b 

323 mse = np.sum(np.sum(field_diff**2)) 

324 abs_val = field_a**2 + field_b**2 

325 mse_ref = np.sum(np.sum(abs_val)) 

326 if mse_ref > 0: 

327 fss = 1 - (mse / mse_ref) 

328 else: 

329 fss = np.nan 

330 

331 return fss 

332 

333 

334def _get_spatial_coords(cube): 

335 """Return the x, y coordinates of an input :class:`iris.cube.Cube`.""" 

336 x_coord = cube.coord(axis="x") 

337 y_coord = cube.coord(axis="y") 

338 return [x_coord, y_coord] 

339 

340 

341def _regrid_lat_lon_cube_to_xy_cube(cube_latlon): 

342 """Regrid a lat-lon cube to xy. 

343 

344 Takes a cube specified on a lat-lon grid and re-grids 

345 it to an x-y grid 

346 

347 Args: 

348 cube_latlon: 

349 Original cube on a lat-lon grid 

350 

351 Returns 

352 ------- 

353 cube_xy: 

354 Original cube regridded onto an x-y grid 

355 

356 """ 

357 x_coord, y_coord = _get_spatial_coords(cube_latlon) 

358 

359 # Transform max and min lon and lat points to set new x,y array on 

360 # new coordinate system 

361 src_crs = y_coord.coord_system.as_cartopy_crs() 

362 trg_crs = ccrs.TransverseMercator(central_latitude=0, central_longitude=0) 

363 trg_crs_iris = coord_systems.TransverseMercator(0, 0, 0, 0, 1.0) 

364 

365 lons = [np.min(x_coord.points), np.max(x_coord.points)] 

366 lats = [np.min(y_coord.points), np.max(y_coord.points)] 

367 

368 x, y = [], [] 

369 for lon, lat in zip(lons, lats, strict=True): 

370 x_trg, y_trg = trg_crs.transform_point(lon, lat, src_crs) 

371 x.append(x_trg) 

372 y.append(y_trg) 

373 

374 numb_x_points = np.size(x_coord.points) 

375 numb_y_points = np.size(y_coord.points) 

376 total_numb_points = numb_x_points * numb_y_points 

377 

378 new_x = icoords.DimCoord( 

379 np.linspace(x[0], x[1], numb_x_points), 

380 standard_name="projection_x_coordinate", 

381 units="m", 

382 coord_system=trg_crs_iris, 

383 ) 

384 

385 new_y = icoords.DimCoord( 

386 np.linspace(y[0], y[1], numb_y_points), 

387 standard_name="projection_y_coordinate", 

388 units="m", 

389 coord_system=trg_crs_iris, 

390 ) 

391 

392 new_ens = icoords.DimCoord( 

393 cube_latlon.coord("realization").points, standard_name="realization" 

394 ) 

395 

396 number_members = len(cube_latlon.coord("realization").points) 

397 new_data = np.zeros(number_members * total_numb_points).reshape( 

398 number_members, numb_y_points, numb_x_points 

399 ) 

400 

401 # Create blank cube in new coordinate system 

402 new_cube = iris.cube.Cube( 

403 new_data, 

404 long_name=cube_latlon.name(), 

405 dim_coords_and_dims=[(new_ens, 0), (new_y, 1), (new_x, 2)], 

406 units=cube_latlon.units, 

407 ) 

408 

409 # Regrid original cube onto new cube 

410 cube_xy = cube_latlon.regrid(new_cube, iris.analysis.Nearest()) 

411 

412 return cube_xy