Coverage for tests/probabilty/test_brier.py: 100%

30 statements  

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

1""" 

2Contains unit tests for scores.probability.brier_impl 

3""" 

4 

5import dask 

6import dask.array 

7import numpy as np 

8import pytest 

9import xarray as xr 

10 

11from scores.probability import brier_score 

12 

13FCST1 = xr.DataArray( 

14 [[[0.5, 0], [0, 0.5], [1, 0]], [[0.5, 0], [0.5, 0], [0, np.nan]]], 

15 dims=["a", "b", "c"], 

16 coords={"a": [0, 1], "b": [0, 1, 2], "c": [0, 1]}, 

17) 

18OBS1 = xr.DataArray( 

19 [[[1, 0], [0, 0], [np.nan, 0]], [[0, 0], [1, 0], [0, 1]]], 

20 dims=["a", "b", "c"], 

21 coords={"a": [0, 1], "b": [0, 1, 2], "c": [0, 1]}, 

22) 

23 

24 

25@pytest.mark.parametrize( 

26 ("fcst", "obs", "preserve_dims", "reduce_dims", "expected_bs"), 

27 [ 

28 # Reduce all dims 

29 ( 

30 FCST1, 

31 OBS1, 

32 None, 

33 None, 

34 xr.DataArray(0.1), 

35 ), 

36 # preserve dim "a" 

37 ( 

38 FCST1, 

39 OBS1, 

40 ["a"], 

41 None, 

42 xr.DataArray([0.1, 0.1], dims=["a"], coords={"a": [0, 1]}), 

43 ), 

44 # doesn't break with all NaNs 

45 ( 

46 FCST1, 

47 OBS1 * np.nan, 

48 ["a"], 

49 None, 

50 xr.DataArray([np.nan, np.nan], dims=["a"], coords={"a": [0, 1]}), 

51 ), 

52 # Check it works with DataSets 

53 ( 

54 xr.Dataset({"1": FCST1, "2": 0.5 * FCST1}), 

55 xr.Dataset({"1": OBS1, "2": OBS1}), 

56 ["a"], 

57 None, 

58 xr.Dataset( 

59 { 

60 "1": xr.DataArray([0.1, 0.1], dims=["a"], coords={"a": [0, 1]}), 

61 "2": xr.DataArray([0.125, 0.125], dims=["a"], coords={"a": [0, 1]}), 

62 } 

63 ), 

64 ), 

65 ], 

66) 

67def test_brier_score(fcst, obs, preserve_dims, reduce_dims, expected_bs): 

68 """ 

69 Tests brier_score. 

70 

71 Note that the underlying MSE function is tested more thoroughly. 

72 """ 

73 calculated_bs = brier_score(fcst, obs, preserve_dims=preserve_dims, reduce_dims=reduce_dims) 

74 xr.testing.assert_equal(calculated_bs, expected_bs) 

75 

76 

77def test_brier_score_dask(): 

78 """ 

79 Tests that the Brier score works with dask 

80 """ 

81 result = brier_score(FCST1.chunk(), OBS1.chunk()) 

82 assert isinstance(result.data, dask.array.Array) 

83 result = result.compute() 

84 assert isinstance(result.data, np.ndarray) 

85 xr.testing.assert_equal(result, xr.DataArray(0.1)) 

86 

87 

88@pytest.mark.parametrize( 

89 ("fcst", "obs", "error_msg_snippet"), 

90 [ 

91 # Fcst > 1 

92 (FCST1 + 0.0000001, OBS1, r"`fcst` contains values outside of the range \[0, 1\]"), 

93 # Fcst < 0 

94 (FCST1 - 0.0000001, OBS1, r"`fcst` contains values outside of the range \[0, 1\]"), 

95 # Obs = 1/2 

96 (FCST1, OBS1 / 2, "`obs` contains values that are not in the set {0, 1, np.nan}"), 

97 ], 

98) 

99def test_brier_score_raises(fcst, obs, error_msg_snippet): 

100 """ 

101 Tests that the Brier score raises the correct errors. 

102 """ 

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

104 brier_score(fcst, obs) 

105 # Check again but with input data as a DataSet 

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

107 brier_score(xr.Dataset({"x": fcst}), xr.Dataset({"x": obs})) 

108 

109 

110@pytest.mark.parametrize( 

111 ("fcst", "obs", "expected"), 

112 [ 

113 # FCST doubled 

114 (FCST1 * 2, OBS1, xr.DataArray(0.2)), 

115 # OBS halved 

116 (FCST1, OBS1 / 2, xr.DataArray(0.05)), 

117 ], 

118) 

119def test_brier_doesnt_raise(fcst, obs, expected): 

120 """ 

121 Tests that the Brier score doesn't raise an error when check_args=False 

122 """ 

123 result = brier_score(fcst, obs, check_args=False) 

124 xr.testing.assert_equal(result, expected) 

125 

126 # Check again but with input data as a DataSet 

127 result = brier_score(xr.Dataset({"x": fcst}), xr.Dataset({"x": obs}), check_args=False) 

128 xr.testing.assert_equal(result, xr.Dataset({"x": expected}))