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
« 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"""
5import dask
6import dask.array
7import numpy as np
8import pytest
9import xarray as xr
11from scores.probability import brier_score
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)
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.
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)
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))
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}))
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)
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}))