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

1"""

2 Contains unit tests for scores.probability.brier_impl

3"""

4

5 import dask

6 import dask.array

7 import numpy as np

8 import pytest

9 import xarray as xr

10

11 from scores.probability import brier_score

12

13 FCST1 = 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)

18 OBS1 = 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)

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

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

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

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