Coverage for tests/continuous/test_murphy.py: 99%

156 statements  

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

1"""Tests for murphy metrics and thetas generation code.""" 

2import pathlib 

3import re 

4from datetime import datetime 

5from unittest.mock import Mock, patch 

6 

7import dask 

8import dask.array 

9import numpy as np 

10import pytest 

11import xarray as xr 

12 

13from scores.continuous import murphy_impl as murphy 

14from scores.continuous import murphy_thetas 

15from scores.continuous.murphy_impl import ( 

16 _expectile_thetas, 

17 _huber_thetas, 

18 _quantile_thetas, 

19) 

20 

21FCST = xr.DataArray( 

22 dims=("lead_day", "station_number", "valid_15z_date"), 

23 data=[[[0.0], [5.0]], [[10.0], [15.0]]], 

24 coords={ 

25 "lead_day": [1, 2], 

26 "station_number": [100001, 10000], 

27 "valid_15z_date": [datetime(2017, 1, 1, 15, 0)], 

28 }, 

29) 

30OBS = xr.DataArray( 

31 dims=("valid_15z_date", "station_number"), 

32 data=[[4.0, 2.0]], 

33 coords={ 

34 "station_number": [10000, 100001], # Different ordering from FCST 

35 "valid_15z_date": [datetime(2017, 1, 1, 15, 0)], 

36 }, 

37) 

38EXPECTED_QUANTILE = xr.Dataset( 

39 coords={"theta": [0.0, 2.0, 10.0], "lead_day": [1, 2]}, 

40 data_vars={ 

41 "total": xr.DataArray( 

42 dims=("lead_day", "theta"), 

43 data=[[0.25, 0.0, 0.0], [0.0, 0.25, 0.25]], 

44 ), 

45 "underforecast": xr.DataArray( 

46 dims=("lead_day", "theta"), 

47 data=[[0.25, 0.0, 0.0], [0.0, 0.0, 0.0]], 

48 ), 

49 "overforecast": xr.DataArray( 

50 dims=("lead_day", "theta"), 

51 data=[[0.0, 0.0, 0.0], [0.0, 0.25, 0.25]], 

52 ), 

53 }, 

54) 

55EXPECTED_HUBER = xr.Dataset( 

56 coords={"theta": [0.0, 2.0, 10.0], "lead_day": [1, 2]}, 

57 data_vars={ 

58 "total": xr.DataArray( 

59 dims=("theta", "lead_day"), 

60 data=[[0.5, 0.0], [0.0, 0.0], [0.0, 0.75]], 

61 ), 

62 "underforecast": xr.DataArray( 

63 dims=("theta", "lead_day"), 

64 data=[[0.5, 0.0], [0.0, 0.0], [0.0, 0.0]], 

65 ), 

66 "overforecast": xr.DataArray( 

67 dims=("theta", "lead_day"), 

68 data=[[0.0, 0.0], [0.0, 0.0], [0.0, 0.75]], 

69 ), 

70 }, 

71) 

72EXPECTED_EXPECTILE = xr.Dataset( 

73 coords={"theta": [0.0, 2.0, 10.0], "lead_day": [1, 2]}, 

74 data_vars={ 

75 "total": xr.DataArray( 

76 dims=("theta", "lead_day"), 

77 data=[[0.5, 0.0], [0.0, 0.0], [0.0, 1.5]], 

78 ), 

79 "underforecast": xr.DataArray( 

80 dims=("theta", "lead_day"), 

81 data=[[0.5, 0.0], [0.0, 0.0], [0.0, 0.0]], 

82 ), 

83 "overforecast": xr.DataArray( 

84 dims=("theta", "lead_day"), 

85 data=[[0.0, 0.0], [0.0, 0.0], [0.0, 1.5]], 

86 ), 

87 }, 

88) 

89 

90 

91def patch_scoring_func(monkeypatch, score_function, thetas): 

92 """Monkeypatch a scoring function, return the mock object.""" 

93 if isinstance(thetas, xr.DataArray): 

94 thetas = thetas.values 

95 over = _rel_test_array(np.array([[999, 999, np.nan, np.nan]] * len(thetas)).T, theta=thetas) 

96 under = _rel_test_array(np.array([[np.nan, 888, np.nan, 888]] * len(thetas)).T, theta=thetas) 

97 mock_rel_fc_func = Mock(return_value=[over, under]) 

98 monkeypatch.setattr( 

99 murphy, 

100 score_function, 

101 mock_rel_fc_func, 

102 ) 

103 return mock_rel_fc_func 

104 

105 

106thetas_nc = xr.open_dataarray(pathlib.Path(__file__).parent / "test_data/thetas.nc") 

107thetas_list = [0.0, 2.0, 10.0] 

108 

109 

110@pytest.mark.parametrize( 

111 ("functional", "score_function", "thetas", "daskinput"), 

112 [ 

113 (murphy.QUANTILE, "_quantile_elementary_score", thetas_list, False), 

114 (murphy.HUBER, "_huber_elementary_score", thetas_list, False), 

115 (murphy.EXPECTILE, "_expectile_elementary_score", thetas_list, False), 

116 (murphy.EXPECTILE, "_expectile_elementary_score", thetas_nc, False), 

117 (murphy.EXPECTILE, "_expectile_elementary_score", thetas_nc, True), 

118 ], 

119) 

120def test_murphy_score_operations(functional, score_function, monkeypatch, thetas, daskinput): 

121 """murphy_score makes the expected operations on the scoring function output.""" 

122 fcst = _test_array([1.0, 2.0, 3.0, 4.0]) 

123 obs = _test_array([0.0, np.nan, 0.6, 137.4]) 

124 if daskinput: 

125 fcst = fcst.chunk() 

126 obs = obs.chunk() 

127 mock_rel_fc_func = patch_scoring_func(monkeypatch, score_function, thetas) 

128 result = murphy.murphy_score( 

129 fcst=fcst, 

130 obs=obs, 

131 thetas=thetas, 

132 functional=functional, 

133 alpha=0.5, 

134 huber_a=3, 

135 decomposition=True, 

136 preserve_dims=fcst.dims, 

137 ) 

138 if isinstance(thetas, xr.DataArray): 

139 thetas = thetas.values 

140 expected = xr.Dataset.from_dict( 

141 { 

142 "dims": ("station_number", "theta"), 

143 "data_vars": { 

144 "total": { 

145 "data": [ 

146 [999.0, 999.0, 999.0], 

147 [np.nan, np.nan, np.nan], 

148 [0.0, 0.0, 0.0], 

149 [888.0, 888.0, 888.0], 

150 ], 

151 "dims": ("station_number", "theta"), 

152 }, 

153 "underforecast": { 

154 "data": [ 

155 [0.0, 0.0, 0.0], 

156 [np.nan, np.nan, np.nan], 

157 [0.0, 0.0, 0.0], 

158 [888.0, 888.0, 888.0], 

159 ], 

160 "dims": ("station_number", "theta"), 

161 }, 

162 "overforecast": { 

163 "data": [ 

164 [999.0, 999.0, 999.0], 

165 [np.nan, np.nan, np.nan], 

166 [0.0, 0.0, 0.0], 

167 [0.0, 0.0, 0.0], 

168 ], 

169 "dims": ("station_number", "theta"), 

170 }, 

171 }, 

172 "coords": { 

173 "station_number": { 

174 "dims": ("station_number",), 

175 "data": [46012, 46126, 46128, 46129], 

176 }, 

177 "theta": { 

178 "dims": ("theta",), 

179 "data": thetas, 

180 }, 

181 }, 

182 } 

183 ) 

184 if daskinput: 

185 assert isinstance(result.total.data, dask.array.Array) 

186 result = result.compute() 

187 assert isinstance(result.total.data, np.ndarray) 

188 xr.testing.assert_identical(result, expected) 

189 mock_rel_fc_func.assert_called_once() 

190 

191 

192@pytest.mark.parametrize( 

193 ("functional", "expected"), 

194 [ 

195 (murphy.QUANTILE, EXPECTED_QUANTILE), 

196 (murphy.HUBER, EXPECTED_HUBER), 

197 (murphy.EXPECTILE, EXPECTED_EXPECTILE), 

198 ], 

199) 

200def test_murphy_score(functional, expected): 

201 """murphy_score returns the expected object.""" 

202 thetas = [0.0, 2.0, 10.0] 

203 

204 result = murphy.murphy_score( 

205 fcst=FCST, 

206 obs=OBS, 

207 thetas=thetas, 

208 functional=functional, 

209 alpha=0.5, 

210 huber_a=3.0, 

211 decomposition=True, 

212 preserve_dims=["lead_day"], 

213 ) 

214 

215 xr.testing.assert_identical(result, expected) 

216 

217 

218def test_murphy_score_mean(monkeypatch): 

219 """ 

220 murphy_score returns the mean of the result if both reduce_dims and 

221 preserve_dims are None. 

222 """ 

223 fcst = _test_array([1.0, 2.0, 3.0, 4.0]) 

224 obs = _test_array([0.0, np.nan, 0.6, 137.4]) 

225 thetas = [0.0, 2.0, 10.0] 

226 _ = patch_scoring_func(monkeypatch, "_quantile_elementary_score", thetas) 

227 

228 result = murphy.murphy_score( 

229 fcst=fcst, 

230 obs=obs, 

231 thetas=thetas, 

232 functional=murphy.QUANTILE, 

233 alpha=0.5, 

234 huber_a=3, 

235 decomposition=True, 

236 ) 

237 

238 expected = xr.Dataset.from_dict( 

239 { 

240 "dims": ("theta"), 

241 "data_vars": { 

242 "total": { 

243 "data": [629.0, 629.0, 629.0], 

244 "dims": ("theta"), 

245 }, 

246 "underforecast": { 

247 "data": [296.0, 296.0, 296.0], 

248 "dims": ("theta"), 

249 }, 

250 "overforecast": { 

251 "data": [333.0, 333.0, 333.0], 

252 "dims": ("theta"), 

253 }, 

254 }, 

255 "coords": { 

256 "theta": { 

257 "dims": ("theta",), 

258 "data": thetas, 

259 }, 

260 }, 

261 } 

262 ) 

263 xr.testing.assert_identical(result, expected) 

264 

265 

266def test_murphy_score_no_decomposition(monkeypatch): 

267 """murphy_score returns only the total score if decomposition is False.""" 

268 fcst = _test_array([1.0, 2.0, 3.0, 4.0]) 

269 obs = _test_array([0.0, np.nan, 0.6, 137.4]) 

270 thetas = [0.0, 2.0, 10.0] 

271 _ = patch_scoring_func(monkeypatch, "_quantile_elementary_score", thetas) 

272 

273 result = murphy.murphy_score( 

274 fcst=fcst, 

275 obs=obs, 

276 thetas=thetas, 

277 functional=murphy.QUANTILE, 

278 alpha=0.5, 

279 huber_a=3, 

280 decomposition=False, 

281 ) 

282 

283 assert list(result.variables.keys()) == ["theta", "total"] 

284 

285 

286def _test_array(data): 

287 """Return a test array for a forecast or obs input.""" 

288 assert len(data) <= 4 

289 return xr.DataArray.from_dict( 

290 { 

291 "dims": ("station_number"), 

292 "data": data, 

293 "coords": { 

294 "station_number": { 

295 "dims": ("station_number",), 

296 "data": [46012, 46126, 46128, 46129][0 : len(data)], 

297 }, 

298 }, 

299 } 

300 ) 

301 

302 

303def _rel_test_array(data, theta): 

304 """Return a test array for *_elementary_score function inputs.""" 

305 assert len(data) <= 4 

306 return xr.DataArray.from_dict( 

307 { 

308 "dims": ("station_number", "theta"), 

309 "data": data, 

310 "coords": { 

311 "station_number": { 

312 "dims": ("station_number",), 

313 "data": [46012, 46126, 46128, 46129][0 : len(data)], 

314 }, 

315 "theta": { 

316 "dims": ("theta",), 

317 "data": theta, 

318 }, 

319 }, 

320 } 

321 ) 

322 

323 

324def test__quantile_elementary_score(): 

325 """_quantile_elementary_score returns the expected values.""" 

326 fcst = _rel_test_array(data=np.array([[0, 1, 4]] * 2).T, theta=[0, 2]) 

327 obs = _rel_test_array(data=np.array([[1, 3, 1]] * 2).T, theta=[0, 2]) 

328 theta = _rel_test_array(data=[[0, 2]] * 3, theta=[0, 2]) 

329 alpha = 0.1 

330 

331 result = murphy._quantile_elementary_score(fcst, obs, theta, alpha) 

332 

333 assert len(result) == 2 

334 np.testing.assert_equal(result[0], np.array([[np.nan, np.nan], [np.nan, np.nan], [np.nan, 0.9]])) 

335 np.testing.assert_equal(result[1], np.array([[0.1, np.nan], [np.nan, 0.1], [np.nan, np.nan]])) 

336 

337 

338def test__huber_elementary_score(): 

339 """_huber_elementary_score returns the expected values.""" 

340 fcst = _rel_test_array(data=np.array([[0, 1, 4]] * 2).T, theta=[0, 2]) 

341 obs = _rel_test_array(data=np.array([[1, 3, 1]] * 2).T, theta=[0, 2]) 

342 theta = _rel_test_array(data=[[0, 2]] * 3, theta=[0, 2]) 

343 alpha = 0.1 

344 

345 result = murphy._huber_elementary_score(fcst, obs, theta, alpha, huber_a=0.5) 

346 

347 assert len(result) == 2 

348 np.testing.assert_equal(result[0], np.array([[np.nan, np.nan], [np.nan, np.nan], [np.nan, 0.45]])) 

349 np.testing.assert_equal(result[1], np.array([[0.05, np.nan], [np.nan, 0.05], [np.nan, np.nan]])) 

350 

351 

352def test__expectile_elementary_score(): 

353 """_expectile_elementary_score returns the expected values.""" 

354 fcst = _rel_test_array(data=np.array([[0, 1, 4]] * 2).T, theta=[0, 2]) 

355 obs = _rel_test_array(data=np.array([[1, 3, 1]] * 2).T, theta=[0, 2]) 

356 theta = _rel_test_array(data=[[0, 2]] * 3, theta=[0, 2]) 

357 alpha = 0.1 

358 

359 result = murphy._expectile_elementary_score(fcst, obs, theta, alpha) 

360 

361 assert len(result) == 2 

362 np.testing.assert_equal(result[0], np.array([[np.nan, np.nan], [np.nan, np.nan], [np.nan, 0.9]])) 

363 np.testing.assert_equal(result[1], np.array([[0.1, np.nan], [np.nan, 0.1], [np.nan, np.nan]])) 

364 

365 

366@pytest.mark.parametrize( 

367 ("new_kwargs", "expected_exception_msg"), 

368 ( 

369 [ 

370 {"alpha": 0}, 

371 "alpha (=0) argument for Murphy scoring function should be strictly " "between 0 and 1.", 

372 ], 

373 [ 

374 {"alpha": 1}, 

375 "alpha (=1) argument for Murphy scoring function should be strictly " "between 0 and 1.", 

376 ], 

377 [ 

378 {"functional": "?"}, 

379 "Functional option '?' for Murphy scoring function is unknown, should be " 

380 "one of ['quantile', 'huber', 'expectile'].", 

381 ], 

382 [ 

383 {"functional": "huber", "huber_a": 0}, 

384 "huber_a (=0) argument should be > 0 when functional='huber'.", 

385 ], 

386 [ 

387 {"functional": "huber", "huber_a": None}, 

388 "huber_a (=None) argument should be > 0 when functional='huber'.", 

389 ], 

390 ), 

391) 

392def test_murphy_score_invalid_input(new_kwargs, expected_exception_msg): 

393 """murphy_score raises an exception for invalid inputs.""" 

394 fcst = _test_array([1.0, 2.0, 3.0, 4.0]) 

395 obs = _test_array([0.0, np.nan, 0.6, 137.4]) 

396 thetas = [0.0, 2.0, 10.0] 

397 kwargs = { 

398 "fcst": fcst, 

399 "obs": obs, 

400 "thetas": thetas, 

401 "functional": murphy.QUANTILE, 

402 "alpha": 0.5, 

403 "huber_a": 3, 

404 "decomposition": True, 

405 } 

406 

407 with pytest.raises(ValueError, match=re.escape(expected_exception_msg)): 

408 _ = murphy.murphy_score(**{**kwargs, **new_kwargs}) 

409 

410 

411@pytest.mark.parametrize( 

412 ("functional", "left_limit_delta", "expected"), 

413 [ 

414 [murphy.QUANTILE, 0.1, [0.0, 2.0, 4.0, 5.0, 10.0, 15.0]], 

415 [ 

416 murphy.HUBER, 

417 0.1, 

418 [-0.1, 0.0, 1.5, 2.0, 2.5, 3.5, 4.0, 4.5, 4.9, 5.0, 9.9, 10.0, 14.9, 15.0], 

419 ], 

420 [murphy.EXPECTILE, 0.1, [-0.1, 0.0, 2.0, 4.0, 4.9, 5.0, 9.9, 10.0, 14.9, 15.0]], 

421 [murphy.EXPECTILE, None, [0.0, 2.0, 4.0, 5.0, 10.0, 15.0]], 

422 ], 

423) 

424def test_murphy_thetas(functional, left_limit_delta, expected): 

425 """murphy_thetas returns the expected object.""" 

426 

427 result = murphy_thetas( 

428 forecasts=[FCST], 

429 obs=OBS, 

430 functional=functional, 

431 huber_a=0.5, 

432 left_limit_delta=left_limit_delta, 

433 ) 

434 

435 assert result == expected 

436 

437 

438@pytest.mark.parametrize( 

439 ("functional"), 

440 [murphy.QUANTILE, murphy.HUBER, murphy.EXPECTILE], 

441) 

442@patch("scores.continuous.murphy_impl._quantile_thetas", autospec=True) 

443@patch("scores.continuous.murphy_impl._huber_thetas", autospec=True) 

444@patch("scores.continuous.murphy_impl._expectile_thetas", autospec=True) 

445def test_murphy_thetas_calls(mock__expectile_thetas, mock__huber_thetas, mock__quantile_thetas, functional): 

446 """murphy_thetas makes the expected function call.""" 

447 result = murphy_thetas( 

448 forecasts=1, # type: ignore # due to mocking 

449 obs=2, # type: ignore # due to mocking 

450 functional=functional, 

451 huber_a=4, 

452 left_limit_delta=5, 

453 ) 

454 

455 expected_kwargs = { 

456 "forecasts": 1, 

457 "obs": 2, 

458 "huber_a": 4, 

459 "left_limit_delta": 5, 

460 } 

461 if functional == murphy.QUANTILE: 

462 assert result == mock__quantile_thetas.return_value 

463 mock__quantile_thetas.assert_called_once_with(**expected_kwargs) 

464 mock__huber_thetas.assert_not_called() 

465 mock__expectile_thetas.assert_not_called() 

466 elif functional == murphy.HUBER: 

467 assert result == mock__huber_thetas.return_value 

468 mock__quantile_thetas.assert_not_called() 

469 mock__huber_thetas.assert_called_once_with(**expected_kwargs) 

470 mock__expectile_thetas.assert_not_called() 

471 elif functional == murphy.EXPECTILE: 471 ↛ exitline 471 didn't return from function 'test_murphy_thetas_calls', because the condition on line 471 was never false

472 assert result == mock__expectile_thetas.return_value 

473 mock__quantile_thetas.assert_not_called() 

474 mock__huber_thetas.assert_not_called() 

475 mock__expectile_thetas.assert_called_once_with(**expected_kwargs) 

476 

477 

478def test__quantile_thetas(): 

479 """_quantile_thetas returns the expected values.""" 

480 forecasts = [_test_array([1.0, 2.0, 3.0]), _test_array([0.0, 10.0, np.nan])] 

481 obs = _test_array([0.0, 0.6, 137.4]) 

482 

483 result = _quantile_thetas(forecasts, obs) 

484 

485 assert result == [0.0, 0.6, 1.0, 2.0, 3.0, 10.0, 137.4] 

486 

487 

488@pytest.mark.parametrize( 

489 ("left_limit_delta", "expected"), 

490 [ 

491 (0.1, [-10.0, -0.1, 0.0, 0.9, 1.0, 1.9, 2.0, 10.0, 127.4, 137.4, 147.4]), 

492 (0, [-10.0, 0.0, 1.0, 2.0, 10.0, 127.4, 137.4, 147.4]), 

493 ], 

494) 

495def test__huber_thetas(left_limit_delta, expected): 

496 """_huber_thetas returns the expected values.""" 

497 forecasts = [_test_array([1.0, 2.0]), _test_array([0.0, np.nan])] 

498 obs = _test_array([0.0, 137.4]) 

499 

500 result = _huber_thetas(forecasts, obs, huber_a=10.0, left_limit_delta=left_limit_delta) 

501 

502 assert result == expected 

503 

504 

505@pytest.mark.parametrize( 

506 ("left_limit_delta", "expected"), 

507 [ 

508 (0.1, [-0.1, 0.0, 0.9, 1.0, 1.9, 2.0, 137.4]), 

509 (0, [0.0, 1.0, 2.0, 137.4]), 

510 ], 

511) 

512def test__expectile_thetas(left_limit_delta, expected): 

513 """_expectile_thetas returns the expected values.""" 

514 forecasts = [_test_array([1.0, 2.0]), _test_array([0.0, np.nan])] 

515 obs = _test_array([0.0, 137.4]) 

516 

517 result = _expectile_thetas(forecasts, obs, left_limit_delta=left_limit_delta) 

518 

519 assert result == expected 

520 

521 

522@pytest.mark.parametrize( 

523 ("new_kwargs", "expected_exception_msg"), 

524 ( 

525 [ 

526 {"functional": "?"}, 

527 "Functional option '?' for Murphy scoring function is unknown, should be " 

528 "one of ['quantile', 'huber', 'expectile'].", 

529 ], 

530 [ 

531 {"functional": "huber", "huber_a": 0}, 

532 "huber_a (=0) argument should be > 0 when functional='huber'.", 

533 ], 

534 [ 

535 {"functional": "expectile", "left_limit_delta": -0.1}, 

536 "left_limit_delta (=-0.1) argument should be >= 0.", 

537 ], 

538 ), 

539) 

540def test_murphy_thetas_invalid_inputs(new_kwargs, expected_exception_msg): 

541 """murphy_thetas raises an exception for invalid inputs.""" 

542 forecasts = [_test_array([1.0, 2.0]), _test_array([0.0, np.nan])] 

543 obs = _test_array([0.0, 137.4]) 

544 kwargs = { 

545 "forecasts": forecasts, 

546 "obs": obs, 

547 "functional": murphy.QUANTILE, 

548 "huber_a": 10.0, 

549 "left_limit_delta": 1.0, 

550 } 

551 

552 with pytest.raises(ValueError, match=re.escape(expected_exception_msg)): 

553 _ = murphy_thetas(**{**kwargs, **new_kwargs})