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

1"""

2 Tests for `scores.isoreg_impl`.

3"""

4 import numpy as np

5 import pytest

6 import xarray as xr

7 from numpy import nan

8

9 from 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)

20 from 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)

30 def 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)

54 def 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)

229 def 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)

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

268 def 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)

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

292 def _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)

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

312 def 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)

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

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

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

359 def 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)

375 def 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))