Coverage for tests/test_weights.py: 100%

1# pylint: disable=missing-function-docstring

2 import numpy as np

3 import xarray as xr

4

5 import scores.continuous

6 import scores.functions

7

8 ZERO = 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

13 np.random.seed(0)

14 LATS = [50, 51, 52, 53]

15 LONS = [30, 31, 32, 33]

16 fcst_temperatures_2d = 15 + 8 * np.random.randn(1, 4, 4)

17 obs_temperatures_2d = 15 + 6 * np.random.randn(1, 4, 4)

18 FCST_2D = xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[LATS, LONS])

19 OBS_2D = xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[LATS, LONS])

20 IDENTITY = np.ones((4, 4))

21 ZEROS = np.zeros((4, 4))

22

23

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

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

43

44

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

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

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

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

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