Coverage for tests/continuous/test_isoreg.py: 100%

70 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-02-28 12:51 +1100

1""" 

2Tests for `scores.isoreg_impl`. 

3""" 

4import numpy as np 

5import pytest 

6import xarray as xr 

7from numpy import nan 

8 

9from scores.continuous.isoreg_impl import ( 

10 _bootstrap_ir, 

11 _confidence_band, 

12 _contiguous_ir, 

13 _do_ir, 

14 _iso_arg_checks, 

15 _nanquantile, 

16 _tidy_ir_inputs, 

17 _xr_to_np, 

18 isotonic_fit, 

19) 

20from tests.continuous import isoreg_test_data as itd 

21 

22 

23@pytest.mark.parametrize( 

24 ("fcst", "obs", "weight", "expected"), 

25 [ 

26 (itd.FCST_XRTONP, itd.OBS_XRTONP, None, itd.EXP_XRTONP1), 

27 (itd.FCST_XRTONP, itd.OBS_XRTONP, itd.WT_XRTONP, itd.EXP_XRTONP2), 

28 ], 

29) 

30def test__xr_to_np(fcst, obs, weight, expected): 

31 """Tests that `_xr_to_np` gives results as expected.""" 

32 result = _xr_to_np(fcst, obs, weight) 

33 for i in range(3): 

34 np.testing.assert_array_equal(result[i], expected[i]) 

35 

36 

37@pytest.mark.parametrize( 

38 ("fcst", "obs", "weight", "error_msg_snippet"), 

39 [ 

40 ( 

41 itd.FCST_XRTONP, 

42 itd.OBS_XRTONP2, 

43 None, 

44 "`fcst` and `obs` must have same dimensions.", 

45 ), 

46 ( 

47 itd.FCST_XRTONP, 

48 itd.OBS_XRTONP, 

49 itd.OBS_XRTONP2, 

50 "`fcst` and `weight` must have same dimensions.", 

51 ), 

52 ], 

53) 

54def test__xr_to_np_raises(fcst, obs, weight, error_msg_snippet): 

55 """Tests that `_xr_to_np` raises as expected.""" 

56 with pytest.raises(ValueError, match=error_msg_snippet): 

57 _xr_to_np(fcst, obs, weight) 

58 

59 

60@pytest.mark.parametrize( 

61 ( 

62 "fcst", 

63 "obs", 

64 "weight", 

65 "functional", 

66 "quantile_level", 

67 "solver", 

68 "bootstraps", 

69 "confidence_level", 

70 "error_msg_snippet", 

71 ), 

72 [ 

73 ( 

74 np.array([1, 2]), 

75 np.array([1]), 

76 None, 

77 "mean", 

78 None, 

79 None, 

80 None, 

81 0.9, 

82 "`fcst` and `obs` must have same shape.", 

83 ), 

84 ( 

85 np.array([1, 2, "string"]), 

86 np.array([1, 2, 3]), 

87 None, 

88 "mean", 

89 None, 

90 None, 

91 None, 

92 0.9, 

93 "`fcst` must be an array of floats or integers.", 

94 ), 

95 ( 

96 np.array([1, 2, 3]), 

97 np.array([1, 2, "category"]), 

98 None, 

99 "mean", 

100 None, 

101 None, 

102 None, 

103 0.9, 

104 "`obs` must be an array of floats or integers.", 

105 ), 

106 ( 

107 np.array([1, 2]), 

108 np.array([10, 20]), 

109 np.array([1]), 

110 "mean", 

111 None, 

112 None, 

113 None, 

114 0.9, 

115 "`fcst` and `weight` must have same shape.", 

116 ), 

117 ( 

118 xr.DataArray(data=[1, 2.1, 3]), 

119 xr.DataArray(data=[1, 2, nan]), 

120 xr.DataArray(data=[1, 2, "category"]), 

121 "mean", 

122 None, 

123 None, 

124 None, 

125 0.9, 

126 "`weight` must be an array of floats or integers, or else `None`.", 

127 ), 

128 ( 

129 np.array([1, 2, 3]), 

130 np.array([10, 20, 30]), 

131 np.array([1, -1, nan]), 

132 "mean", 

133 None, 

134 None, 

135 None, 

136 0.9, 

137 "`weight` values must be either positive or NaN.", 

138 ), 

139 ( 

140 np.array([1, 2, 3]), 

141 np.array([10, 20, 30]), 

142 np.array([1, 1, nan]), 

143 "median", 

144 None, 

145 None, 

146 None, 

147 0.9, 

148 "`functional` must be one of 'mean', 'quantile' or `None`.", 

149 ), 

150 ( 

151 np.array([1, 2, 3]), 

152 np.array([10, 20, 30]), 

153 None, 

154 "quantile", 

155 1.0, 

156 None, 

157 None, 

158 0.9, 

159 "`quantile_level` must be strictly between 0 and 1.", 

160 ), 

161 ( 

162 np.array([1, 2, 3]), 

163 np.array([10, 20, 30]), 

164 np.array([1, 1, nan]), 

165 "quantile", 

166 0.3, 

167 None, 

168 None, 

169 0.9, 

170 "Weighted quantile isotonic regression has not been implemented.", 

171 ), 

172 ( 

173 np.array([1, 2, 3]), 

174 np.array([10, 20, 30]), 

175 None, 

176 None, 

177 0.5, 

178 None, 

179 None, 

180 0.9, 

181 "`functional` and `solver` cannot both be `None`.", 

182 ), 

183 ( 

184 np.array([1, 2, 3]), 

185 np.array([10, 20, 30]), 

186 None, 

187 "quantile", 

188 0.5, 

189 np.quantile, 

190 None, 

191 0.9, 

192 "One of `functional` or `solver` must be `None`.", 

193 ), 

194 ( 

195 np.array([1, 2, 3]), 

196 np.array([10, 20, 30]), 

197 None, 

198 "quantile", 

199 0.5, 

200 None, 

201 13.6, 

202 0.9, 

203 "`bootstraps` must be a positive integer.", 

204 ), 

205 ( 

206 np.array([1, 2, 3]), 

207 np.array([10, 20, 30]), 

208 None, 

209 "quantile", 

210 0.5, 

211 None, 

212 -2, 

213 0.9, 

214 "`bootstraps` must be a positive integer.", 

215 ), 

216 ( 

217 np.array([1, 2, 3]), 

218 np.array([10, 20, 30]), 

219 None, 

220 "quantile", 

221 0.5, 

222 None, 

223 5000, 

224 0.0, 

225 "`confidence_level` must be strictly between 0 and 1.", 

226 ), 

227 ], 

228) 

229def test__iso_arg_checks( # pylint: disable=too-many-locals, too-many-arguments 

230 fcst, 

231 obs, 

232 weight, 

233 functional, 

234 quantile_level, 

235 solver, 

236 bootstraps, 

237 confidence_level, 

238 error_msg_snippet, 

239): 

240 """Tests that `_iso_arg_checks` raises as expected.""" 

241 with pytest.raises(ValueError, match=error_msg_snippet): 

242 _iso_arg_checks( 

243 fcst=fcst, 

244 obs=obs, 

245 weight=weight, 

246 functional=functional, 

247 quantile_level=quantile_level, 

248 solver=solver, 

249 bootstraps=bootstraps, 

250 confidence_level=confidence_level, 

251 ) 

252 

253 

254@pytest.mark.parametrize( 

255 ("fcst", "obs", "weights", "expected"), 

256 [ 

257 (itd.FCST_TIDY1, itd.OBS_TIDY1, None, itd.EXP_TIDY1), 

258 (itd.FCST_TIDY1, itd.OBS_TIDY1, itd.WEIGHT_TIDY1, itd.EXP_TIDY2), 

259 ], 

260) 

261def test__tidy_ir_inputs(fcst, obs, weights, expected): 

262 """Tests that `_tidy_ir_inputs` gives results as expected.""" 

263 result = _tidy_ir_inputs(fcst, obs, weights) 

264 for i in range(3): 

265 np.testing.assert_array_equal(result[i], expected[i]) 

266 

267 

268def test__tidy_ir_inputs_raises(): 

269 """Tests that _tidy_ir_inputs raises as expected.""" 

270 with pytest.raises(ValueError, match="pairs remaining after NaNs removed."): 

271 _tidy_ir_inputs(np.array([0.0, nan, 4.1, 3]), np.array([nan, 0.0, nan, nan])) 

272 

273 

274@pytest.mark.parametrize( 

275 ("functional", "solver", "expected"), 

276 [ 

277 ("mean", None, np.array([2.0, 2.0, 2.0])), 

278 ("quantile", None, np.array([1.0, 1.0, 1.0])), 

279 ("none", np.max, np.array([5.0, 5.0, 5.0])), 

280 ], 

281) 

282def test__do_ir(functional, solver, expected): 

283 """ 

284 Tests that `_do_ir` gives results as expected. 

285 Simultaneously supplies simple confirmation tests for `_contiguous_mean_ir`, 

286 `_contiguous_quantile_ir` and `_contiguous_ir`. 

287 """ 

288 result = _do_ir(np.array([1, 2, 3]), np.array([5.0, 1, 0]), None, functional, 0.5, solver) 

289 np.testing.assert_array_equal(result, expected) 

290 

291 

292def _wmean_solver(y, weight): 

293 """solver for mean isotonic regression""" 

294 return np.average(y, weights=weight) 

295 

296 

297@pytest.mark.parametrize( 

298 ("y", "solver", "weight", "expected"), 

299 [ 

300 (itd.Y1, np.median, None, itd.EXP_IR_MEDIAN), 

301 (itd.Y1, np.mean, None, itd.EXP_IR_MEAN), 

302 (itd.Y1, _wmean_solver, itd.W1, itd.EXP_IR_WMEAN), 

303 (np.ndarray(0), np.mean, None, 0), 

304 ], 

305) 

306def test__contiguous_ir(y, solver, weight, expected): 

307 """Tests that `_contiguous_ir` gives results as expected.""" 

308 result = _contiguous_ir(y, solver, weight) 

309 np.testing.assert_array_equal(result, expected) 

310 

311 

312def test__contiguous_ir_raises(): 

313 """Tests that `_contiguous_ir` raises as expected.""" 

314 with pytest.raises(ValueError, match="`y` and `weight` must have same length."): 

315 _contiguous_ir(np.array([1.0, 2, 3]), np.mean, weight=np.array([1, 1])) 

316 

317 

318@pytest.mark.parametrize( 

319 ("weight", "functional", "q_level", "solver", "bootstrap", "expected"), 

320 [ 

321 (itd.BS_WT, None, None, _wmean_solver, 3, itd.BS_EXP1), 

322 (None, "mean", None, None, 3, itd.BS_EXP2), 

323 (None, "quantile", 0.5, None, 1, itd.BS_EXP3), 

324 ], 

325) 

326def test__bootstrap_ir( # pylint: disable=too-many-locals, too-many-arguments 

327 weight, functional, q_level, solver, bootstrap, expected 

328): 

329 """Tests that `_contiguous_ir` gives results as expected.""" 

330 np.random.seed(seed=1) 

331 result = _bootstrap_ir( 

332 fcst=itd.BS_FCST, 

333 obs=itd.BS_OBS, 

334 weight=weight, 

335 functional=functional, 

336 quantile_level=q_level, 

337 solver=solver, 

338 bootstraps=bootstrap, 

339 ) 

340 np.testing.assert_array_equal(result, expected) 

341 

342 

343def test__confidence_band(): 

344 """Tests that `_confidence_band` gives results as expected.""" 

345 result = _confidence_band(itd.CB_BOOT_INPUT, 0.5, 4) 

346 for i in range(2): 

347 np.testing.assert_array_equal(result[i], itd.EXP_CB[i]) 

348 

349 

350def test__confidence_band_nan(): 

351 """_confidence_band returns expected objects with all NaN input""" 

352 result = _confidence_band(np.where(False, itd.CB_BOOT_INPUT, np.nan), 0.5, 4) 

353 

354 expected = np.array([np.nan] * 7) 

355 np.testing.assert_array_equal(result[0], expected) 

356 np.testing.assert_array_equal(result[1], expected) 

357 

358 

359def test__nanquantile(): 

360 """_nanquantile returns expected results.""" 

361 for quant in [0.75, 0.25]: 

362 result = _nanquantile(itd.CB_BOOT_INPUT, quant) 

363 np.testing.assert_array_equal(result, itd.EXP_CB2[quant]) 

364 

365 

366@pytest.mark.parametrize( 

367 ("fcst", "obs", "bootstraps", "report_bootstrap_results", "expected"), 

368 [ 

369 (itd.FCST_XR, itd.OBS_XR, None, None, itd.EXP_IF1), 

370 (itd.FCST_ARRAY, itd.OBS_ARRAY, None, None, itd.EXP_IF1), 

371 (itd.FCST_ARRAY, itd.OBS_ARRAY, 3, False, itd.EXP_IF2), 

372 (itd.FCST_ARRAY, itd.OBS_ARRAY, 3, True, itd.EXP_IF3), 

373 ], 

374) 

375def test_isotonic_fit(fcst, obs, bootstraps, report_bootstrap_results, expected): 

376 """Tests that `isotonic_fit` gives results as expected.""" 

377 np.random.seed(seed=1) 

378 result = isotonic_fit( 

379 fcst, 

380 obs, 

381 functional=None, 

382 solver=np.mean, 

383 bootstraps=bootstraps, 

384 confidence_level=0.5, 

385 min_non_nan=3, 

386 report_bootstrap_results=report_bootstrap_results, 

387 ) 

388 # check arrays are equal 

389 array_keys = [ 

390 "fcst_sorted", 

391 "fcst_counts", 

392 "regression_values", 

393 "confidence_band_lower_values", 

394 "confidence_band_upper_values", 

395 ] 

396 if report_bootstrap_results: 

397 array_keys.append("bootstrap_results") 

398 for key in array_keys: 

399 np.testing.assert_array_equal(result[key], expected[key]) 

400 

401 assert result["confidence_band_levels"] == expected["confidence_band_levels"] 

402 

403 for key in [ 

404 "regression_func", 

405 "confidence_band_lower_func", 

406 "confidence_band_upper_func", 

407 ]: 

408 np.testing.assert_array_equal(result[key](itd.TEST_POINTS), expected[key](itd.TEST_POINTS))