Coverage for tests/test_weights.py: 100%
73 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# pylint: disable=missing-function-docstring
2import numpy as np
3import xarray as xr
5import scores.continuous
6import scores.functions
8ZERO = np.array([[1 for i in range(10)] for j in range(10)])
11# Standard forecast and observed test data which is static and can be used
12# across tests
13np.random.seed(0)
14LATS = [50, 51, 52, 53]
15LONS = [30, 31, 32, 33]
16fcst_temperatures_2d = 15 + 8 * np.random.randn(1, 4, 4)
17obs_temperatures_2d = 15 + 6 * np.random.randn(1, 4, 4)
18FCST_2D = xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[LATS, LONS])
19OBS_2D = xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[LATS, LONS])
20IDENTITY = np.ones((4, 4))
21ZEROS = np.zeros((4, 4))
24def simple_da(data):
25 """
26 Helper function for making a DataArray with a latitude coordinate variable roughly
27 over the Australian region
28 """
29 lats = np.arange(-5, -45, (-45 / len(data)))
30 arr = xr.DataArray.from_dict(
31 {
32 "coords": {
33 "lat": {"data": lats, "dims": "lat"},
34 },
35 "data": data,
36 }
37 )
38 return arr
41# These scores will be tested for valid processing of weights
42all_scores = [scores.continuous.mse, scores.continuous.mae]
45def test_weights_identity():
46 for score in all_scores:
47 unweighted = score(FCST_2D, OBS_2D)
48 weighted = score(FCST_2D, OBS_2D, weights=IDENTITY)
49 assert unweighted == weighted
52def test_weights_zeros():
53 for score in all_scores:
54 unweighted = score(FCST_2D, OBS_2D)
55 weighted = score(FCST_2D, OBS_2D, weights=ZEROS)
57 assert unweighted != weighted
58 assert weighted.sum() == 0
61def test_weights_latitude():
62 """
63 Tests the use of latitude weightings
64 """
66 lat_weightings_values = scores.functions.create_latitude_weights(OBS_2D.latitude)
68 for score in all_scores:
69 unweighted = score(FCST_2D, OBS_2D)
70 weighted = score(FCST_2D, OBS_2D, weights=lat_weightings_values)
71 assert unweighted != weighted
73 # Latitudes in degrees, tested to 8 decimal places
74 latitude_tests = [
75 (90, 0),
76 (89, 0.017452),
77 (45, 0.707107),
78 (22.5, 0.92388),
79 (0, 1),
80 (-22.5, 0.92388),
81 (-45, 0.707107),
82 (-89, 0.017452),
83 (-90, 0),
84 ]
85 latitudes, expected = zip(*latitude_tests)
86 latitudes = xr.DataArray(list(latitudes)) # Will not work from a tuple
87 expected = xr.DataArray(list(expected)) # Will not work from a tuple
89 found = scores.functions.create_latitude_weights(latitudes)
90 decimal_places = 6
91 found = found.round(decimal_places)
92 expected = expected.round(decimal_places)
93 assert found.equals(expected)
96def test_weights_NaN_matching():
97 da = xr.DataArray
99 fcst = da([np.nan, 0, 1, 2, 7, 0, 7, 1])
100 obs = da([np.nan, np.nan, 0, 1, 7, 0, 7, 0])
101 weights = da([1, 1, 1, 1, 1, 1, 0, np.nan])
102 expected = da([np.nan, np.nan, 1, 1, 0, 0, 0, np.nan])
104 result = scores.continuous.mae(fcst, obs, weights=weights, preserve_dims="all")
105 assert isinstance(result, xr.DataArray)
106 assert isinstance(expected, xr.DataArray)
108 assert result.equals(expected)
111def test_weights_add_dimension():
112 """
113 Test what happens when additional dimensions are added into weights which are not present in
114 fcst or obs. Repeats some of the NaN matching but the focus is really on the dimensional
115 expansion, using the same data to slowly build up the example and establish confidence.
116 """
118 da = simple_da # Make a DataArray with a latitude dimension
120 fcst = da([np.nan, 0, 1, 2, 7, 0, 7, 1])
121 obs = da([np.nan, np.nan, 0, 1, 7, 0, 7, 0])
122 simple_weights = [1, 1, 1, 1, 1, 1, 0, np.nan]
123 double_weights = [2, 2, 2, 2, 2, 2, 0, np.nan]
124 simple_expect = [np.nan, np.nan, 1, 1, 0, 0, 0, np.nan]
125 double_expect = [np.nan, np.nan, 2, 2, 0, 0, 0, np.nan]
127 simple = scores.continuous.mae(fcst, obs, weights=da(simple_weights), preserve_dims="all")
128 doubled = scores.continuous.mae(fcst, obs, weights=da(double_weights), preserve_dims="all")
130 assert simple.equals(da(simple_expect)) # type: ignore # Static analysis mireports this
131 assert doubled.equals(da(double_expect)) # type: ignore # Static analysis mireports this
133 composite_weights_data = [simple_weights, double_weights]
134 composite_expected_data = [simple_expect, double_expect]
136 composite_weights = xr.DataArray.from_dict(
137 {
138 "coords": {
139 "method": {"data": ["simpleweight", "doubleweight"], "dims": "method"},
140 "lat": {"data": list(fcst.lat), "dims": "lat"},
141 },
142 "data": composite_weights_data,
143 }
144 )
146 composite_expected = xr.DataArray.from_dict(
147 {
148 "coords": {
149 "method": {"data": ["simpleweight", "doubleweight"], "dims": "method"},
150 "lat": {"data": list(fcst.lat), "dims": "lat"},
151 },
152 "data": composite_expected_data,
153 }
154 )
156 composite = scores.continuous.mae(fcst, obs, weights=composite_weights, preserve_dims="all").transpose()
157 composite.broadcast_equals(composite_expected) # type: ignore # Static analysis mireports this