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

1# pylint: disable=missing-function-docstring 

2import numpy as np 

3import xarray as xr 

4 

5import scores.continuous 

6import scores.functions 

7 

8ZERO = np.array([[1 for i in range(10)] for j in range(10)]) 

9 

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

22 

23 

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 

39 

40 

41# These scores will be tested for valid processing of weights 

42all_scores = [scores.continuous.mse, scores.continuous.mae] 

43 

44 

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 

50 

51 

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) 

56 

57 assert unweighted != weighted 

58 assert weighted.sum() == 0 

59 

60 

61def test_weights_latitude(): 

62 """ 

63 Tests the use of latitude weightings 

64 """ 

65 

66 lat_weightings_values = scores.functions.create_latitude_weights(OBS_2D.latitude) 

67 

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 

72 

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 

88 

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) 

94 

95 

96def test_weights_NaN_matching(): 

97 da = xr.DataArray 

98 

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]) 

103 

104 result = scores.continuous.mae(fcst, obs, weights=weights, preserve_dims="all") 

105 assert isinstance(result, xr.DataArray) 

106 assert isinstance(expected, xr.DataArray) 

107 

108 assert result.equals(expected) 

109 

110 

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 """ 

117 

118 da = simple_da # Make a DataArray with a latitude dimension 

119 

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] 

126 

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") 

129 

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 

132 

133 composite_weights_data = [simple_weights, double_weights] 

134 composite_expected_data = [simple_expect, double_expect] 

135 

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 ) 

145 

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 ) 

155 

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