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

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

2 import pathlib

3 import re

4 from datetime import datetime

5 from unittest.mock import Mock, patch

6

7 import dask

8 import dask.array

9 import numpy as np

10 import pytest

11 import xarray as xr

12

13 from scores.continuous import murphy_impl as murphy

14 from scores.continuous import murphy_thetas

15 from scores.continuous.murphy_impl import (

16 _expectile_thetas,

17 _huber_thetas,

18 _quantile_thetas,

19)

20

21 FCST = 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)

30 OBS = 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)

38 EXPECTED_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)

55 EXPECTED_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)

72 EXPECTED_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

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

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

107 thetas_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)

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

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

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

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

286 def _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

303 def _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

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

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

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

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

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

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

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

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

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

540 def 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})