Coverage for tests/continuous/test_standard.py: 97%

253 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2024-02-28 12:51 +1100

1""" 

2Contains unit tests for scores.continuous.standard 

3""" 

4# pylint: disable=missing-function-docstring 

5# pylint: disable=line-too-long 

6 

7import dask 

8import dask.array 

9import numpy as np 

10import numpy.random 

11import pandas as pd 

12import pytest 

13import xarray as xr 

14 

15import scores.continuous 

16 

17PRECISION = 4 

18 

19# Mean Squared Error 

20 

21 

22def test_mse_xarray_1d(): 

23 """ 

24 Test both value and expected datatype matches for xarray calculation 

25 """ 

26 fcst_as_xarray_1d = xr.DataArray([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

27 obs_as_xarray_1d = xr.DataArray([1, 1, 1, 2, 1, 2, 1, 1, 1, 3, 1]) 

28 result = scores.continuous.mse(fcst_as_xarray_1d, obs_as_xarray_1d) 

29 

30 expected = xr.DataArray(1.0909) 

31 assert isinstance(result, xr.DataArray) 

32 assert result.round(PRECISION) == expected.round(PRECISION) 

33 

34 

35def test_mse_pandas_series(): 

36 """ 

37 Test calculation works correctly on pandas series 

38 """ 

39 

40 fcst_pd_series = pd.Series([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

41 obs_pd_series = pd.Series([1, 1, 1, 2, 1, 2, 1, 1, 1, 3, 1]) 

42 expected = 1.0909 

43 result = scores.continuous.mse(fcst_pd_series, obs_pd_series) 

44 assert isinstance(result, float) 

45 assert round(result, 4) == expected 

46 

47 

48def test_mse_dataframe(): 

49 """ 

50 Test calculation works correctly on dataframe columns 

51 """ 

52 

53 fcst_pd_series = pd.Series([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

54 obs_pd_series = pd.Series([1, 1, 1, 2, 1, 2, 1, 1, 1, 3, 1]) 

55 df = pd.DataFrame({"fcst": fcst_pd_series, "obs": obs_pd_series}) 

56 expected = 1.0909 

57 result = scores.continuous.mse(df["fcst"], df["obs"]) 

58 assert isinstance(result, float) 

59 assert round(result, PRECISION) == expected 

60 

61 

62def test_mse_xarray_to_point(): 

63 """ 

64 Test MSE calculates the correct value for a simple 1d sequence 

65 Currently breaks type hinting but here for future pandas support 

66 """ 

67 fcst_as_xarray_1d = xr.DataArray([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

68 result = scores.continuous.mse(fcst_as_xarray_1d, 1) # type: ignore 

69 expected = xr.DataArray(1.45454545) 

70 assert isinstance(result, xr.DataArray) 

71 assert result.round(PRECISION) == expected.round(PRECISION) 

72 

73 

74def test_2d_xarray_mse(): 

75 """ 

76 Test MSE calculates the correct value on a 2d array 

77 """ 

78 numpy.random.seed(0) 

79 lats = [50, 51, 52, 53] 

80 lons = [30, 31, 32, 33] 

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

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

83 fcst_temperatures_xr_2d = xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

84 obs_temperatures_xr_2d = xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

85 

86 result = scores.continuous.mse(fcst_temperatures_xr_2d, obs_temperatures_xr_2d) 

87 expected = xr.DataArray(142.33162) 

88 assert isinstance(result, xr.DataArray) 

89 assert result.round(PRECISION) == expected.round(PRECISION) 

90 

91 

92def test_2d_xarray_mse_with_dimensions(): 

93 """ 

94 Assert that MSE can correctly calculate MSE along a specified dimension 

95 """ 

96 numpy.random.seed(0) 

97 lats = [50, 51, 52, 53] 

98 lons = [30, 31, 32, 33] 

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

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

101 fcst_temperatures_xr_2d = xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

102 obs_temperatures_xr_2d = xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

103 

104 result = scores.continuous.mse(fcst_temperatures_xr_2d, obs_temperatures_xr_2d, reduce_dims="latitude") 

105 

106 expected_values = [290.0929, 90.8107, 12.2224, 176.2005] 

107 expected_dimensions = ("longitude",) 

108 assert isinstance(result, xr.DataArray) 

109 assert all(result.round(4) == expected_values) # type: ignore # We don't want full xarray comparison, and static analysis is confused about types 

110 assert result.dims == expected_dimensions 

111 

112 

113# Root Mean Squared Error 

114@pytest.fixture 

115def rmse_fcst_xarray(): 

116 """Creates forecast Xarray for test.""" 

117 return xr.DataArray([-1, 3, 1, 3, 0, 2, 2, 1, 1, 2, 3]) 

118 

119 

120@pytest.fixture 

121def rmse_fcst_nan_xarray(): 

122 """Creates forecast Xarray containing NaNs for test.""" 

123 return xr.DataArray([-1, 3, 1, 3, np.nan, 2, 2, 1, 1, 2, 3]) 

124 

125 

126@pytest.fixture 

127def rmse_obs_xarray(): 

128 """Creates observation Xarray for test.""" 

129 return xr.DataArray([1, 1, 1, 2, 1, 2, 1, -1, 1, 3, 1]) 

130 

131 

132@pytest.fixture 

133def rmse_fcst_pandas(): 

134 """Creates forecast Pandas series for test.""" 

135 return pd.Series([-1, 3, 1, 3, 0, 2, 2, 1, 1, 2, 3]) 

136 

137 

138@pytest.fixture 

139def rmse_fcst_nan_pandas(): 

140 """Creates forecast Pandas series containing NaNs for test.""" 

141 return pd.Series([-1, 3, 1, 3, np.nan, 2, 2, 1, 1, 2, 3]) 

142 

143 

144@pytest.fixture 

145def rmse_obs_pandas(): 

146 """Creates observation Pandas series for test.""" 

147 return pd.Series([1, 1, 1, 2, 1, 2, 1, 1, -1, 3, 1]) 

148 

149 

150@pytest.mark.parametrize( 

151 "forecast, observations, expected, request_kwargs", 

152 [ 

153 ("rmse_fcst_xarray", "rmse_obs_xarray", xr.DataArray(1.3484), {}), 

154 ( 

155 "rmse_fcst_xarray", 

156 "rmse_obs_xarray", 

157 xr.DataArray([2, 2, 0, 1, 1, 0, 1, 2, 0, 1, 2]), 

158 {"preserve_dims": "all"}, 

159 ), 

160 ("rmse_fcst_nan_xarray", "rmse_obs_xarray", xr.DataArray(1.3784), {}), 

161 ("rmse_fcst_xarray", 1, xr.DataArray(1.3484), {}), 

162 ("rmse_fcst_nan_xarray", 1, xr.DataArray(1.3784), {}), 

163 ("rmse_fcst_pandas", "rmse_obs_pandas", 1.3484, {}), 

164 ("rmse_fcst_pandas", 1, 1.3484, {}), 

165 ("rmse_fcst_nan_pandas", "rmse_obs_pandas", 1.3784, {}), 

166 ], 

167 ids=[ 

168 "simple-1d", 

169 "preserve-1d", 

170 "simple-1d-w-nan", 

171 "to-point", 

172 "to-point-w-nan", 

173 "pandas-series-1d", 

174 "pandas-to-point", 

175 "pandas-series-nan-1d", 

176 ], 

177) 

178def test_rmse_xarray_1d(forecast, observations, expected, request_kwargs, request): 

179 """ 

180 Test RMSE for the following cases: 

181 * Calculates the correct value for a simple xarray 1d sequence 

182 * Calculates the correct value for an xarray 1d sequence preserving all 

183 * Calculates the correct value for a simple pandas 1d series 

184 * Calculates the correct value for an xarray 1d sequence comparing to a point 

185 """ 

186 if isinstance(forecast, str): 186 ↛ 188line 186 didn't jump to line 188, because the condition on line 186 was never false

187 forecast = request.getfixturevalue(forecast) 

188 if isinstance(observations, str): 

189 observations = request.getfixturevalue(observations) 

190 result = scores.continuous.rmse(forecast, observations, **request_kwargs) 

191 if not isinstance(result, float): 

192 assert (result.round(PRECISION) == expected).all() 

193 else: 

194 assert np.round(result, PRECISION) == expected 

195 

196 

197@pytest.fixture 

198def rmse_2d_fcst_xarray(): 

199 """Creates 2D forecast Xarray for test.""" 

200 numpy.random.seed(0) 

201 lats = [50, 51, 52, 53] 

202 lons = [30, 31, 32, 33] 

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

204 fcst_temperatures_2d[0, 2, 1] = np.nan 

205 return xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

206 

207 

208@pytest.fixture 

209def rmse_2d_obs_xarray(): 

210 """Creates 2D observation Xarray for test.""" 

211 numpy.random.seed(0) 

212 lats = [50, 51, 52, 53] 

213 lons = [30, 31, 32, 33] 

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

215 obs_temperatures_2d[0, 1, 2] = np.nan 

216 return xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

217 

218 

219@pytest.fixture 

220def rmse_2d_expected_xarray(): 

221 """Creates 2D forecast Xarray to be used as expected result for test.""" 

222 lons = [30, 31, 32, 33] 

223 exp_temperatures_2d = [2.6813, 1.2275, 1.252, 2.6964] 

224 return xr.DataArray(exp_temperatures_2d, dims=["longitude"], coords=[lons]) 

225 

226 

227@pytest.mark.parametrize( 

228 "forecast, observations, expected, request_kwargs, expected_dimensions", 

229 [ 

230 ("rmse_2d_fcst_xarray", "rmse_2d_obs_xarray", xr.DataArray(2.1887), {}, ()), 

231 ( 

232 "rmse_2d_fcst_xarray", 

233 "rmse_2d_obs_xarray", 

234 "rmse_2d_expected_xarray", 

235 {"reduce_dims": "latitude"}, 

236 ("longitude",), 

237 ), 

238 ( 

239 "rmse_2d_fcst_xarray", 

240 "rmse_2d_obs_xarray", 

241 "rmse_2d_expected_xarray", 

242 {"preserve_dims": "longitude"}, 

243 ("longitude",), 

244 ), 

245 ], 

246 ids=["simple-2d", "reduce-2d", "preserve-2d"], 

247) 

248def test_rmse_xarray_2d_rand( # pylint: disable=too-many-arguments 

249 forecast, observations, expected, request_kwargs, expected_dimensions, request 

250): 

251 """ 

252 Test RMSE for the following cases on 2d Data: 

253 * Calculates the correct value for a simple xarray 2d sequence 

254 * Calculates the correct value for an xarray 2d sequence reducing over set dim 

255 * Calculates the correct value for an xarray 2d sequence preserving set dim 

256 """ 

257 if isinstance(forecast, str): 257 ↛ 259line 257 didn't jump to line 259, because the condition on line 257 was never false

258 forecast = request.getfixturevalue(forecast) 

259 if isinstance(observations, str): 259 ↛ 261line 259 didn't jump to line 261, because the condition on line 259 was never false

260 observations = request.getfixturevalue(observations) 

261 if isinstance(expected, str): 

262 expected = request.getfixturevalue(expected) 

263 

264 result = scores.continuous.rmse(forecast, observations, **request_kwargs) 

265 xr.testing.assert_allclose(result.round(PRECISION), expected) 

266 assert result.dims == expected_dimensions 

267 

268 

269def create_xarray(data: list[list[float]]): 

270 """Creates an Xarray.DataArray to be used in the tests""" 

271 npdata = np.array(data) 

272 lats = list(np.arange(npdata.shape[0]) + 50) 

273 lons = list(np.arange(npdata.shape[1]) + 30) 

274 return xr.DataArray(npdata, dims=["latitude", "longitude"], coords=[lats, lons]) 

275 

276 

277@pytest.mark.parametrize( 

278 "forecast, observations, expected, request_kwargs,", 

279 [ 

280 (create_xarray([[0, 0], [1, 1]]), create_xarray([[0, 0], [1, 1]]), xr.DataArray(0), {}), 

281 (create_xarray([[-1, 0], [1, 1]]), create_xarray([[-1, 0], [1, 1]]), xr.DataArray(0), {}), 

282 (create_xarray([[1, 0], [1, 1]]), create_xarray([[-1, 0], [1, 1]]), xr.DataArray(1), {}), 

283 (create_xarray([[-1, 0], [1, 1]]), create_xarray([[1, 0], [1, 1]]), xr.DataArray(1), {}), 

284 (create_xarray([[np.nan, 0], [1, 1]]), create_xarray([[1, 0], [1, 1]]), xr.DataArray(0), {}), 

285 ( 

286 create_xarray([[np.nan, 1], [1, 1]]), 

287 create_xarray([[1, 0], [1, 1]]), 

288 xr.DataArray(np.sqrt(1 / 3).round(PRECISION)), 

289 {}, 

290 ), 

291 (create_xarray([[1, 0], [1, 1]]), create_xarray([[np.nan, 0], [1, 1]]), xr.DataArray(0), {}), 

292 ( 

293 create_xarray([[1, 0], [1, 1]]), 

294 create_xarray([[np.nan, 1], [1, 1]]), 

295 xr.DataArray(np.sqrt(1 / 3).round(PRECISION)), 

296 {}, 

297 ), 

298 ], 

299 ids=[ 

300 "perfect", 

301 "perfect-neg", 

302 "single", 

303 "single_neg", 

304 "perfect-nan", 

305 "single-nan", 

306 "perfect-nan-r", 

307 "single-nan-r", 

308 ], 

309) 

310def test_rmse_xarray_2d_defined(forecast, observations, expected, request_kwargs, request): 

311 """ 

312 Test RMSE Values for defined edge cases 

313 

314 """ 

315 if isinstance(forecast, str): 315 ↛ 316line 315 didn't jump to line 316, because the condition on line 315 was never true

316 forecast = request.getfixturevalue(forecast) 

317 if isinstance(observations, str): 317 ↛ 318line 317 didn't jump to line 318, because the condition on line 317 was never true

318 observations = request.getfixturevalue(observations) 

319 if isinstance(expected, str): 319 ↛ 320line 319 didn't jump to line 320, because the condition on line 319 was never true

320 expected = request.getfixturevalue(expected) 

321 

322 result = scores.continuous.rmse(forecast, observations, **request_kwargs) 

323 xr.testing.assert_allclose(result.round(PRECISION), expected) 

324 

325 

326# Mean Absolute Error 

327 

328 

329def test_mae_xarray_1d(): 

330 """ 

331 Test both value and expected datatype matches for xarray calculation 

332 """ 

333 fcst_as_xarray_1d = xr.DataArray([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

334 obs_as_xarray_1d = xr.DataArray([1, 1, 1, 2, 1, 2, 1, 1, 1, 3, 1]) 

335 result = scores.continuous.mae(fcst_as_xarray_1d, obs_as_xarray_1d) 

336 fcst_as_xarray_1d = xr.DataArray([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

337 

338 expected = xr.DataArray(0.7273) 

339 assert result.round(PRECISION) == expected.round(PRECISION) 

340 

341 

342def test_mae_pandas_series(): 

343 """ 

344 Test calculation works correctly on pandas series 

345 """ 

346 

347 fcst_pd_series = pd.Series([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

348 obs_pd_series = pd.Series([1, 1, 1, 2, 1, 2, 1, 1, 1, 3, 1]) 

349 expected = 0.7273 

350 result = scores.continuous.mae(fcst_pd_series, obs_pd_series) 

351 assert isinstance(result, float) 

352 assert round(result, 4) == expected 

353 

354 

355def test_mae_dataframe(): 

356 """ 

357 Test calculation works correctly on dataframe columns 

358 """ 

359 

360 fcst_pd_series = pd.Series([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

361 obs_pd_series = pd.Series([1, 1, 1, 2, 1, 2, 1, 1, 1, 3, 1]) 

362 df = pd.DataFrame({"fcst": fcst_pd_series, "obs": obs_pd_series}) 

363 expected = 0.7273 

364 result = scores.continuous.mae(df["fcst"], df["obs"]) 

365 assert isinstance(result, float) 

366 assert round(result, PRECISION) == expected 

367 

368 

369def test_mae_xarray_to_point(): 

370 """ 

371 Test MAE calculates the correct value for a simple sequence 

372 Tests unhinted types but this is useful 

373 """ 

374 fcst_as_xarray_1d = xr.DataArray([1, 3, 1, 3, 2, 2, 2, 1, 1, 2, 3]) 

375 result = scores.continuous.mae(fcst_as_xarray_1d, 1) # type: ignore 

376 expected = xr.DataArray(0.9091) 

377 assert result.round(PRECISION) == expected.round(PRECISION) 

378 

379 

380def test_2d_xarray_mae(): 

381 """ 

382 Test MAE calculates the correct value on a 2d array 

383 """ 

384 numpy.random.seed(0) 

385 lats = [50, 51, 52, 53] 

386 lons = [30, 31, 32, 33] 

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

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

389 fcst_temperatures_xr_2d = xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

390 obs_temperatures_xr_2d = xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

391 

392 result = scores.continuous.mae(fcst_temperatures_xr_2d, obs_temperatures_xr_2d) 

393 expected = xr.DataArray(8.7688) 

394 assert result.round(PRECISION) == expected.round(PRECISION) 

395 

396 

397def test_2d_xarray_mae_with_dimensions(): 

398 """ 

399 Test MAE is calculated correctly along a specified dimension 

400 """ 

401 numpy.random.seed(0) 

402 lats = [50, 51, 52, 53] 

403 lons = [30, 31, 32, 33] 

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

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

406 fcst_temperatures_xr_2d = xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

407 obs_temperatures_xr_2d = xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

408 

409 result = scores.continuous.mae(fcst_temperatures_xr_2d, obs_temperatures_xr_2d, reduce_dims="latitude") 

410 

411 expected_values = [13.2397, 9.0065, 2.9662, 9.8629] 

412 expected_dimensions = ("longitude",) 

413 assert all(result.round(4) == expected_values) # type: ignore # We don't want full xarray comparison, and static analysis is confused about types 

414 assert result.dims == expected_dimensions 

415 

416 

417def test_xarray_dimension_handling_with_arrays(): 

418 """ 

419 Test MAE is calculated correctly along a specified dimension 

420 """ 

421 numpy.random.seed(0) 

422 lats = [50, 51, 52, 53] 

423 lons = [30, 31, 32, 33] 

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

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

426 fcst_temperatures_xr_2d = xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

427 obs_temperatures_xr_2d = xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

428 

429 reduce_lat = scores.continuous.mae(fcst_temperatures_xr_2d, obs_temperatures_xr_2d, reduce_dims="latitude") 

430 

431 preserve_lon = scores.continuous.mae(fcst_temperatures_xr_2d, obs_temperatures_xr_2d, preserve_dims="longitude") 

432 

433 expected_values = [13.2397, 9.0065, 2.9662, 9.8629] 

434 expected_dimensions = ("longitude",) 

435 assert all(preserve_lon.round(4) == expected_values) # type: ignore # We don't want full xarray comparison, and static analysis is confused about types 

436 assert reduce_lat.dims == expected_dimensions 

437 assert preserve_lon.dims == expected_dimensions 

438 

439 

440def test_xarray_dimension_preservations_with_arrays(): 

441 """ 

442 Test MAE is calculated correctly along a specified dimension 

443 """ 

444 numpy.random.seed(0) 

445 lats = [50, 51, 52, 53] 

446 lons = [30, 31, 32, 33] 

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

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

449 fcst_temperatures_xr_2d = xr.DataArray(fcst_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

450 obs_temperatures_xr_2d = xr.DataArray(obs_temperatures_2d[0], dims=["latitude", "longitude"], coords=[lats, lons]) 

451 

452 reduce_empty = scores.continuous.mae(fcst_temperatures_xr_2d, obs_temperatures_xr_2d, reduce_dims=[]) 

453 preserve_all = scores.continuous.mae(fcst_temperatures_xr_2d, obs_temperatures_xr_2d, preserve_dims="all") 

454 

455 assert reduce_empty.dims == fcst_temperatures_xr_2d.dims # Nothing should be reduced 

456 assert (reduce_empty == preserve_all).all() 

457 

458 

459FCST_CHUNKED = xr.DataArray( 

460 data=np.array([[1, 2], [3, 10]]), dims=["dim1", "dim2"], coords={"dim1": [1, 2], "dim2": [1, 2]} 

461).chunk() 

462OBS_CHUNKED = xr.DataArray( 

463 data=np.array([[0, 0], [0, np.nan]]), dims=["dim1", "dim2"], coords={"dim1": [1, 2], "dim2": [1, 2]} 

464).chunk() 

465 

466 

467# Dask tests 

468def test_mse_with_dask(): 

469 """ 

470 Test that mse works with dask 

471 """ 

472 result = scores.continuous.mse(FCST_CHUNKED, OBS_CHUNKED, reduce_dims="dim1") 

473 assert isinstance(result.data, dask.array.Array) # type: ignore # Static analysis fails to recognise the type of 'result' correctly 

474 result = result.compute() # type: ignore # Static analysis thinks this is a float, but it's a dask array 

475 assert isinstance(result.data, np.ndarray) 

476 expected = xr.DataArray(data=[5, 4], dims=["dim2"], coords={"dim2": [1, 2]}) 

477 xr.testing.assert_equal(result, expected) 

478 

479 

480def test_mae_with_dask(): 

481 """ 

482 Test that mae works with dask 

483 """ 

484 result = scores.continuous.mae(FCST_CHUNKED, OBS_CHUNKED, reduce_dims="dim1") 

485 assert isinstance(result.data, dask.array.Array) # type: ignore 

486 result = result.compute() 

487 assert isinstance(result.data, np.ndarray) 

488 expected = xr.DataArray(data=[2, 2], dims=["dim2"], coords={"dim2": [1, 2]}) 

489 xr.testing.assert_equal(result, expected) 

490 

491 

492def test_rmse_with_dask(): 

493 """ 

494 Test that rmse works with dask 

495 """ 

496 result = scores.continuous.rmse(FCST_CHUNKED, OBS_CHUNKED, reduce_dims="dim1") 

497 assert isinstance(result.data, dask.array.Array) 

498 result = result.compute() 

499 assert isinstance(result.data, np.ndarray) 

500 expected = xr.DataArray(data=[np.sqrt(5), 2], dims=["dim2"], coords={"dim2": [1, 2]}) 

501 xr.testing.assert_equal(result, expected) 

502 

503 

504# Angular / directional tests 

505DA1_ANGULAR = xr.DataArray([[10, 10], [90, 90]], coords=[[0, 1], [0, 1]], dims=["i", "j"]) 

506DA2_ANGULAR = xr.DataArray([[350, 180], [270, 280]], coords=[[0, 1], [0, 1]], dims=["i", "j"]) 

507 

508 

509def test_mse_angular(): 

510 """Tests that `mse` returns the expected object with `angular` is True""" 

511 

512 expected = xr.DataArray( 

513 [[20**2, 170**2], [180**2, 170**2]], 

514 coords=[[0, 1], [0, 1]], 

515 dims=["i", "j"], 

516 name="mean_squared_error", 

517 ) 

518 

519 result = scores.continuous.mse(DA1_ANGULAR, DA2_ANGULAR, preserve_dims=["i", "j"], angular=True) 

520 

521 xr.testing.assert_equal(result, expected) 

522 

523 

524def test_mae_angular(): 

525 """Tests that `mae` returns the expected object with `angular` is True""" 

526 

527 expected = xr.DataArray( 

528 [[20, 170], [180, 170]], 

529 coords=[[0, 1], [0, 1]], 

530 dims=["i", "j"], 

531 name="mean_squared_error", 

532 ) 

533 

534 result = scores.continuous.mae(DA1_ANGULAR, DA2_ANGULAR, preserve_dims=["i", "j"], angular=True) 

535 

536 xr.testing.assert_equal(result, expected) 

537 

538 

539def test_rmse_angular(): 

540 """Tests that `rmse` returns the expected object with `angular` is True""" 

541 

542 expected = xr.DataArray( 

543 [((20**2 + 170**2) / 2) ** 0.5, ((180**2 + 170**2) / 2) ** 0.5], 

544 coords={"i": [0, 1]}, 

545 dims=["i"], 

546 name="mean_squared_error", 

547 ) 

548 

549 result = scores.continuous.rmse(DA1_ANGULAR, DA2_ANGULAR, preserve_dims=["i"], angular=True) 

550 

551 xr.testing.assert_equal(result, expected)