Coverage for tests/test_utils.py: 100%

1"""

2 Containts tests for the scores.utils file

3"""

4

5 import pytest

6 import xarray as xr

7

8 from scores import utils

9 from scores.utils import DimensionError

10 from scores.utils import gather_dimensions as gd

11 from scores.utils import gather_dimensions2

12 from tests import utils_test_data

13

14

15 def test_dims_complement():

16 """Test `dims_complement` returns as expected."""

17 xr_data = utils_test_data.DA_RGB

18 expected_dims = sorted(["red", "green", "blue"])

19 complement = utils.dims_complement(xr_data)

20 assert complement == expected_dims

21

22 expected_dims = sorted(["red", "green"])

23 complement = utils.dims_complement(xr_data, ["blue"])

24 assert complement == expected_dims

25

26

27@pytest.mark.parametrize(

28 ("xr_data", "expected_dims", "mode"),

29 [

30 # 1-D DataArrays

31 (utils_test_data.DA_R, ["red"], "equal"),

32 (utils_test_data.DA_R, ["red"], None),

33 (utils_test_data.DA_R, ["red"], "subset"),

34 (utils_test_data.DA_R, ["red"], "superset"),

35 (utils_test_data.DA_R, ["green"], "disjoint"),

36 # 2-D DataArrays

37 (utils_test_data.DA_RG, ["red", "green"], "equal"),

38 (utils_test_data.DA_RG, ["red", "green"], None),

39 (utils_test_data.DA_RG, ["red", "green", "blue"], "subset"),

40 (utils_test_data.DA_RG, ["red", "green", "blue"], "proper subset"),

41 (utils_test_data.DA_RG, ["red"], "superset"),

42 (utils_test_data.DA_RG, ["red"], "proper superset"),

43 (utils_test_data.DA_RG, ["red", "green"], "superset"),

44 # 1-D Datasets

45 (utils_test_data.DS_R, ["red"], "equal"),

46 (utils_test_data.DS_R, ["red"], None),

47 (utils_test_data.DS_R, ["red"], "subset"),

48 (utils_test_data.DS_R, ["red"], "superset"),

49 (utils_test_data.DS_R, ["green"], "disjoint"),

50 # 2-D Datasets

51 (utils_test_data.DS_RG, ["red", "green"], "equal"),

52 (utils_test_data.DS_RG, ["red", "green"], None),

53 (utils_test_data.DS_RG, ["red", "green", "blue"], "subset"),

54 (utils_test_data.DS_RG, ["red", "green", "blue"], "proper subset"),

55 (utils_test_data.DS_RG, ["red"], "superset"),

56 (utils_test_data.DA_RG, ["red"], "proper superset"),

57 (utils_test_data.DS_RG, ["red", "green"], "superset"),

58 # Datasets with mutiple data variables

59 (utils_test_data.DS_RG_R, ["red", "green"], "subset"),

60 (utils_test_data.DS_RG_R, ["red"], "superset"),

61 (utils_test_data.DS_RG_RG, ["red", "green"], None),

62 (utils_test_data.DS_RG_RG, ["red", "green"], "subset"),

63 (utils_test_data.DS_RG_RG, ["red", "green"], "equal"),

64 (utils_test_data.DS_RG_RG, ["red", "green", "blue"], "subset"),

65 (utils_test_data.DS_RG_RG, ["red"], "superset"),

66 (utils_test_data.DS_RGB_GB, ["green", "blue"], "superset"),

67 # issue #162 - dims accepts any iterable

68 (utils_test_data.DS_RGB_GB, {"green", "blue"}, "superset"),

69 ],

70)

71 def test_check_dims(xr_data, expected_dims, mode):

72 """

73 Tests that check_dims passes when expected to

74 """

75 utils.check_dims(xr_data, expected_dims, mode=mode)

76

77

78@pytest.mark.parametrize(

79 ("xr_data", "expected_dims", "mode", "error_class", "error_msg_snippet"),

80 [

81 # check 'equal'

82 (

83 utils_test_data.DA_R,

84 ["green"],

85 "equal",

86 DimensionError,

87 "equal to the dimensions ['green'].",

88 ),

89 (

90 utils_test_data.DA_RG,

91 ["blue"],

92 "equal",

93 DimensionError,

94 "equal to the dimensions ['blue'].",

95 ),

96 (

97 utils_test_data.DS_R,

98 ["blue"],

99 "equal",

100 DimensionError,

101 "equal to the dimensions ['blue'].",

102 ),

103 (

104 utils_test_data.DS_RG,

105 ["red", "blue"],

106 "equal",

107 DimensionError,

108 "equal to the dimensions ['blue', 'red'].",

109 ),

110 (

111 utils_test_data.DS_RG_R,

112 ["red", "green"],

113 "equal",

114 DimensionError,

115 "['red'] of data variable 'DA_R' are not equal to the dimensions ['green', 'red']",

116 ),

117 (

118 utils_test_data.DS_RG_R,

119 ["red", "green"],

120 None,

121 DimensionError,

122 "['red'] of data variable 'DA_R' are not equal to the dimensions ['green', 'red']",

123 ),

124 (

125 utils_test_data.DS_RGB_GB,

126 ["red", "green", "blue"],

127 "equal",

128 DimensionError,

129 "['green', 'blue'] of data variable 'DA_GB' are not equal to the dimensions ['blue', 'green', 'red']",

130 ),

131 # check 'subset'

132 (

133 utils_test_data.DA_R,

134 ["green"],

135 "subset",

136 DimensionError,

137 "subset to the dimensions ['green'].",

138 ),

139 (

140 utils_test_data.DA_RG,

141 ["blue"],

142 "subset",

143 DimensionError,

144 "subset to the dimensions ['blue'].",

145 ),

146 (

147 utils_test_data.DS_R,

148 ["blue"],

149 "subset",

150 DimensionError,

151 "subset to the dimensions ['blue'].",

152 ),

153 (

154 utils_test_data.DS_RG,

155 ["red", "blue"],

156 "subset",

157 DimensionError,

158 "subset to the dimensions ['blue', 'red'].",

159 ),

160 (

161 utils_test_data.DS_RGB,

162 ["red", "green"],

163 "subset",

164 DimensionError,

165 "subset to the dimensions ['green', 'red'].",

166 ),

167 # check 'superset'

168 (

169 utils_test_data.DA_R,

170 ["green"],

171 "superset",

172 DimensionError,

173 "superset to the dimensions ['green'].",

174 ),

175 (

176 utils_test_data.DA_RG,

177 ["blue", "red", "green"],

178 "superset",

179 DimensionError,

180 "superset to the dimensions ['blue', 'green', 'red'].",

181 ),

182 (

183 utils_test_data.DS_R,

184 ["blue"],

185 "superset",

186 DimensionError,

187 "superset to the dimensions ['blue'].",

188 ),

189 (

190 utils_test_data.DS_RG,

191 ["red", "blue"],

192 "superset",

193 DimensionError,

194 "superset to the dimensions ['blue', 'red'].",

195 ),

196 (

197 utils_test_data.DS_RGB,

198 ["red", "green", "blue", "pink"],

199 "superset",

200 DimensionError,

201 "superset to the dimensions ['blue', 'green', 'pink', 'red'].",

202 ),

203 (

204 utils_test_data.DS_RG_R,

205 ["red", "green"],

206 "superset",

207 DimensionError,

208 "['red'] of data variable 'DA_R' are not superset to the dimensions ['green', 'red']",

209 ),

210 # check 'proper subset'

211 # these are the same as subset tests

212 (

213 utils_test_data.DA_R,

214 ["green"],

215 "proper subset",

216 DimensionError,

217 "proper subset to the dimensions ['green'].",

218 ),

219 (

220 utils_test_data.DA_RG,

221 ["blue"],

222 "proper subset",

223 DimensionError,

224 "proper subset to the dimensions ['blue'].",

225 ),

226 (

227 utils_test_data.DS_R,

228 ["blue"],

229 "proper subset",

230 DimensionError,

231 "proper subset to the dimensions ['blue'].",

232 ),

233 (

234 utils_test_data.DS_RG,

235 ["red", "blue"],

236 "proper subset",

237 DimensionError,

238 "proper subset to the dimensions ['blue', 'red'].",

239 ),

240 (

241 utils_test_data.DS_RGB,

242 ["red", "green"],

243 "proper subset",

244 DimensionError,

245 "proper subset to the dimensions ['green', 'red'].",

246 ),

247 # specifically for proper subset

248 (

249 utils_test_data.DS_RGB,

250 ["red", "green", "blue"],

251 "proper subset",

252 DimensionError,

253 "proper subset to the dimensions ['blue', 'green', 'red'].",

254 ),

255 (

256 utils_test_data.DA_R,

257 ["red"],

258 "proper subset",

259 DimensionError,

260 "proper subset to the dimensions ['red'].",

261 ),

262 # check 'proper superset'

263 # these are the same as superset tests

264 (

265 utils_test_data.DA_R,

266 ["green"],

267 "proper superset",

268 DimensionError,

269 "superset to the dimensions ['green'].",

270 ),

271 (

272 utils_test_data.DA_RG,

273 ["blue", "red", "green"],

274 "proper superset",

275 DimensionError,

276 "superset to the dimensions ['blue', 'green', 'red'].",

277 ),

278 (

279 utils_test_data.DS_R,

280 ["blue"],

281 "proper superset",

282 DimensionError,

283 "superset to the dimensions ['blue'].",

284 ),

285 (

286 utils_test_data.DS_RG,

287 ["red", "blue"],

288 "proper superset",

289 DimensionError,

290 "superset to the dimensions ['blue', 'red'].",

291 ),

292 (

293 utils_test_data.DS_RGB,

294 ["red", "green", "blue", "pink"],

295 "proper superset",

296 DimensionError,

297 "superset to the dimensions ['blue', 'green', 'pink', 'red'].",

298 ),

299 (

300 utils_test_data.DS_RG_R,

301 ["red", "green"],

302 "proper superset",

303 DimensionError,

304 "superset to the dimensions ['green', 'red'].",

305 ),

306 # specifically for proper superset

307 (

308 utils_test_data.DA_R,

309 ["red"],

310 "proper superset",

311 DimensionError,

312 "superset to the dimensions ['red'].",

313 ),

314 (

315 utils_test_data.DS_RG_R,

316 ["red"],

317 "proper superset",

318 DimensionError,

319 "['red'] of data variable 'DA_R' are not proper superset to the dimensions ['red']",

320 ),

321 (

322 utils_test_data.DS_RGB,

323 ["red", "green", "blue"],

324 "proper superset",

325 DimensionError,

326 "superset to the dimensions ['blue', 'green', 'red'].",

327 ),

328 (

329 utils_test_data.DS_RGB_GB,

330 ["green", "blue"],

331 "proper superset",

332 DimensionError,

333 "['green', 'blue'] of data variable 'DA_GB' are not proper superset to the dimensions ['blue', 'green']",

334 ),

335 # check 'disjoint'

336 (

337 utils_test_data.DA_R,

338 ["red"],

339 "disjoint",

340 DimensionError,

341 "disjoint to the dimensions ['red'].",

342 ),

343 (

344 utils_test_data.DA_R,

345 ["red", "green"],

346 "disjoint",

347 DimensionError,

348 "disjoint to the dimensions ['green', 'red'].",

349 ),

350 (

351 utils_test_data.DS_R,

352 ["red"],

353 "disjoint",

354 DimensionError,

355 "disjoint to the dimensions ['red'].",

356 ),

357 (

358 utils_test_data.DS_R,

359 ["red", "green"],

360 "disjoint",

361 DimensionError,

362 "disjoint to the dimensions ['green', 'red'].",

363 ),

364 (

365 utils_test_data.DA_RG,

366 ["red", "blue"],

367 "disjoint",

368 DimensionError,

369 "disjoint to the dimensions ['blue', 'red'].",

370 ),

371 (

372 utils_test_data.DS_RG,

373 ["red", "blue"],

374 "disjoint",

375 DimensionError,

376 "disjoint to the dimensions ['blue', 'red'].",

377 ),

378 (

379 utils_test_data.DS_RG_R,

380 ["green"],

381 "disjoint",

382 DimensionError,

383 "disjoint to the dimensions ['green'].",

384 ),

385 # check the modes

386 (utils_test_data.DA_R, ["red"], "frog", ValueError, "No such mode frog,"),

387 # check if a non data object is passed

388 (

389 [5],

390 ["red"],

391 "equal",

392 DimensionError,

393 "Supplied object has no dimensions",

394 ),

395 # duplicate values in dims

396 (

397 utils_test_data.DA_R,

398 ["red", "blue", "red"],

399 "equal",

400 ValueError,

401 "Supplied dimensions ['red', 'blue', 'red'] contains duplicate values.",

402 ),

403 # can't convert into a set

404 (

405 utils_test_data.DA_R,

406 [["red", "blue", "red"]],

407 "subset",

408 ValueError,

409 "Cannot convert supplied dims [['red', 'blue', 'red']] into a set. ",

410 ),

411 # if a string is passed

412 (

413 utils_test_data.DA_R,

414 "red",

415 "equal",

416 TypeError,

417 "'red' must be an iterable of strings",

418 ),

419 ],

420)

421 def test_check_dims_raises(xr_data, expected_dims, mode, error_class, error_msg_snippet):

422 """

423 Tests that check_dims correctly raises the correct error

424 """

425

426 with pytest.raises(error_class) as excinfo:

427 utils.check_dims(xr_data, expected_dims, mode=mode)

428 assert error_msg_snippet in str(excinfo.value)

429

430

431 def test_gather_dimensions_examples():

432 """

433 Test the logic for dimension handling with some examples

434 """

435

436 fcst_dims_conflict = set(["base_time", "lead_time", "lat", "lon", "all"])

437 fcst_dims = set(["base_time", "lead_time", "lat", "lon"])

438 obs_dims = []

439

440 # Basic tests on reduction

441 assert gd(fcst_dims, obs_dims, reduce_dims="lat") == set(["lat"])

442 assert gd(fcst_dims, obs_dims, reduce_dims=["lat", "lon"]) == set(["lat", "lon"])

443 assert gd(fcst_dims, obs_dims, reduce_dims=["lat", "lat", "lon"]) == set(["lat", "lon"])

444

445 # Tests if reduce_dims and preserve_dims are both None

446 assert gd(fcst_dims, obs_dims) == fcst_dims

447

448 # Reduce every dimension if the string "all" is specified

449 assert gd(fcst_dims, obs_dims, reduce_dims="all") == fcst_dims

450

451 # Reduce "all" as a named dimension explicitly

452 assert gd(fcst_dims_conflict, obs_dims, reduce_dims=["all"]) == set(["all"])

453

454 # Basic tests on preservation

455 assert gd(fcst_dims, obs_dims, preserve_dims="lat") == set(["base_time", "lead_time", "lon"])

456 assert gd(fcst_dims, obs_dims, preserve_dims=["lat", "lon"]) == set(["base_time", "lead_time"])

457 assert gd(fcst_dims, obs_dims, preserve_dims=["lat", "lat", "lon"]) == set(["base_time", "lead_time"])

458

459 # Preserve every dimension if the string "all" is specified

460 assert gd(fcst_dims, obs_dims, preserve_dims="all") == set([])

461

462 # Preserve "all" as a named dimension explicitly

463 assert gd(fcst_dims_conflict, obs_dims, preserve_dims=["all"]) == set(["base_time", "lead_time", "lat", "lon"])

464

465 # Test that preserve is the inverse of reduce

466 preserve_all = gd(fcst_dims, obs_dims, preserve_dims="all")

467 reduce_empty = gd(fcst_dims, obs_dims, reduce_dims=[])

468

469 assert preserve_all == reduce_empty

470 assert preserve_all == set([])

471

472 # Single dimensions specified as a string will be packed into a list

473 assert gd(fcst_dims, obs_dims, reduce_dims="lead_time") == set(["lead_time"])

474

475

476 def test_gather_dimensions_exceptions():

477 """

478 Confirm an exception is raised when both preserve and reduce arguments are specified

479 """

480

481 fcst_dims_conflict = set(["base_time", "lead_time", "lat", "lon", "all"])

482 fcst_dims = set(["base_time", "lead_time", "lat", "lon"])

483 obs_dims = []

484

485 # Confirm an exception if both preserve and reduce are specified

486 with pytest.raises(ValueError):

487 gd(fcst_dims, obs_dims, preserve_dims=[], reduce_dims=[])

488

489 # Attempt to reduce a non-existent dimension

490 with pytest.raises(ValueError) as excinfo:

491 assert gd(fcst_dims_conflict, obs_dims, reduce_dims="nonexistent") == []

492 assert str(excinfo.value.args[0]) == utils.ERROR_SPECIFIED_NONPRESENT_REDUCE_DIMENSION

493

494 # Attempt to preserve a non-existent dimension

495 with pytest.raises(ValueError) as excinfo:

496 assert gd(fcst_dims_conflict, obs_dims, preserve_dims="nonexistent") == []

497 assert str(excinfo.value.args[0]) == utils.ERROR_SPECIFIED_NONPRESENT_PRESERVE_DIMENSION

498

499 # Preserve "all" as a string but named dimension present in data

500 with pytest.warns(UserWarning):

501 assert gd(fcst_dims_conflict, obs_dims, preserve_dims="all") == set([])

502

503 # Preserve "all" as a string but named dimension present in data

504 with pytest.warns(UserWarning):

505 assert gd(fcst_dims_conflict, obs_dims, reduce_dims="all") == fcst_dims_conflict

506

507

508@pytest.mark.parametrize(

509 ("fcst", "obs", "weights", "reduce_dims", "preserve_dims", "special_fcst_dims", "error_msg_snippet"),

510 [

511 (

512 utils_test_data.DA_RGB,

513 utils_test_data.DA_R,

514 None,

515 ["red"],

516 ["blue"],

517 None,

518 utils.ERROR_OVERSPECIFIED_PRESERVE_REDUCE,

519 ),

520 # checks for special_fcst_dims

521 (

522 utils_test_data.DA_RGB,

523 utils_test_data.DA_R,

524 None,

525 None,

526 None,

527 ["black"],

528 "`special_fcst_dims` must be a subset of `fcst` dimensions",

529 ),

530 (

531 utils_test_data.DA_RGB,

532 utils_test_data.DA_R,

533 None,

534 None,

535 None,

536 ["red"],

537 "`obs.dims` must not contain any `special_fcst_dims`",

538 ),

539 (

540 utils_test_data.DA_RGB,

541 utils_test_data.DA_R,

542 None,

543 None,

544 None,

545 "red",

546 "`obs.dims` must not contain any `special_fcst_dims`",

547 ),

548 (

549 utils_test_data.DA_RGB,

550 utils_test_data.DA_R,

551 utils_test_data.DA_G,

552 None,

553 None,

554 "green",

555 "`weights.dims` must not contain any `special_fcst_dims`",

556 ),

557 (

558 utils_test_data.DA_RGB,

559 utils_test_data.DA_R,

560 utils_test_data.DA_G,

561 "blue",

562 None,

563 "blue",

564 "`reduce_dims` and `preserve_dims` must not contain any `special_fcst_dims`",

565 ),

566 (

567 utils_test_data.DA_RGB,

568 utils_test_data.DA_R,

569 utils_test_data.DA_G,

570 None,

571 ["blue"],

572 "blue",

573 "`reduce_dims` and `preserve_dims` must not contain any `special_fcst_dims`",

574 ),

575 (

576 utils_test_data.DA_RGB,

577 utils_test_data.DA_R,

578 utils_test_data.DA_G,

579 None,

580 ["blue", "yellow"],

581 None,

582 utils.ERROR_SPECIFIED_NONPRESENT_PRESERVE_DIMENSION2,

583 ),

584 (

585 utils_test_data.DA_RGB,

586 utils_test_data.DA_R,

587 utils_test_data.DA_G,

588 "yellow",

589 None,

590 "blue",

591 utils.ERROR_SPECIFIED_NONPRESENT_REDUCE_DIMENSION2,

592 ),

593 ],

594)

595 def test_gather_dimensions2_exceptions(

596 fcst, obs, weights, reduce_dims, preserve_dims, special_fcst_dims, error_msg_snippet

597):

598 """

599 Confirm `gather_dimensions2` raises exceptions as expected.

600 """

601 with pytest.raises(ValueError) as excinfo:

602 gather_dimensions2(

603 fcst,

604 obs,

605 weights=weights,

606 reduce_dims=reduce_dims,

607 preserve_dims=preserve_dims,

608 special_fcst_dims=special_fcst_dims,

609 )

610 assert error_msg_snippet in str(excinfo.value)

611

612

613 def test_gather_dimensions2_warnings():

614 """Tests that gather_dimensions2 warns as expected with correct output."""

615 # Preserve "all" as a string but named dimension present in data

616 with pytest.warns(UserWarning):

617 result = gather_dimensions2(

618 utils_test_data.DA_R.rename({"red": "all"}), utils_test_data.DA_R, preserve_dims="all"

619 )

620 assert result == set([])

621

622 with pytest.warns(UserWarning):

623 result = gather_dimensions2(

624 utils_test_data.DA_R.rename({"red": "all"}), utils_test_data.DA_R, reduce_dims="all"

625 )

626 assert result == {"red", "all"}

627

628

629@pytest.mark.parametrize(

630 ("fcst", "obs", "weights", "reduce_dims", "preserve_dims", "special_fcst_dims", "expected"),

631 [

632 # test that fcst and obs dims are returned

633 (utils_test_data.DA_B, utils_test_data.DA_R, None, None, None, None, {"blue", "red"}),

634 # test that fcst, obs and weights dims are returned

635 (utils_test_data.DA_B, utils_test_data.DA_R, utils_test_data.DA_G, None, None, None, {"blue", "red", "green"}),

636 # two tests that fcst, obs and weights dims are returned, without the special fcst dim

637 (utils_test_data.DA_B, utils_test_data.DA_R, utils_test_data.DA_G, None, None, "blue", {"red", "green"}),

638 (utils_test_data.DA_B, utils_test_data.DA_R, utils_test_data.DA_G, None, None, ["blue"], {"red", "green"}),

639 # test that reduce_dims="all" behaves as expected

640 (utils_test_data.DA_B, utils_test_data.DA_R, utils_test_data.DA_G, "all", None, None, {"blue", "red", "green"}),

641 # three tests for reduce_dims

642 (utils_test_data.DA_B, utils_test_data.DA_R, utils_test_data.DA_G, "blue", None, None, {"blue"}),

643 (utils_test_data.DA_B, utils_test_data.DA_R, utils_test_data.DA_G, ["blue"], None, None, {"blue"}),

644 (utils_test_data.DA_RGB, utils_test_data.DA_R, utils_test_data.DA_G, ["green"], None, "blue", {"green"}),

645 # test for preserve_dims="all"

646 (utils_test_data.DA_RGB, utils_test_data.DA_B, None, None, "all", "red", set([])),

647 # three tests for preserve_dims

648 (utils_test_data.DA_RGB, utils_test_data.DA_R, utils_test_data.DA_G, None, "green", None, {"red", "blue"}),

649 (utils_test_data.DA_RGB, utils_test_data.DA_R, None, None, ["green"], None, {"red", "blue"}),

650 (utils_test_data.DA_RGB, utils_test_data.DA_B, None, None, ["green"], "red", {"blue"}),

651 ],

652)

653 def test_gather_dimensions2_examples(fcst, obs, weights, reduce_dims, preserve_dims, special_fcst_dims, expected):

654 """

655 Test that `gather_dimensions2` gives outputs as expected.

656 """

657 result = gather_dimensions2(

658 fcst,

659 obs,

660 weights=weights,

661 reduce_dims=reduce_dims,

662 preserve_dims=preserve_dims,

663 special_fcst_dims=special_fcst_dims,

664 )

665 assert result == expected

666

667

668 def test_tmp_coord_name_namecollision():

669 names = []

670 number_of_names = 3

671 data = xr.DataArray(data=[1, 2, 3])

672 names = utils.tmp_coord_name(data, count=number_of_names)

673 assert len(set(names)) == len(names)

674 assert len(names) == number_of_names

675

676

677 def test_tmp_coord_name():

678 """

679 Tests that `tmp_coord_name` returns as expected.

680 """

681 data = xr.DataArray(data=[1, 2, 3])

682 assert utils.tmp_coord_name(data) == "newdim_0"

683

684 data = xr.DataArray(data=[1, 2, 3], dims=["stn"], coords=dict(stn=[101, 202, 304]))

685 assert utils.tmp_coord_name(data) == "newstnstn"

686

687 data = xr.DataArray(data=[1, 2, 3], dims=["stn"], coords=dict(stn=[101, 202, 304], elevation=("stn", [0, 3, 24])))

688 assert utils.tmp_coord_name(data) == "newstnstnelevation"