Coverage for tests/test_functions.py: 100%

1"""

2 Tests scores.functions

3"""

4 import numpy as np

5 import pytest

6 import xarray as xr

7

8 from scores.functions import angular_difference

9

10 DA_DIR_A = xr.DataArray(

11 data=[[[0, 50], [340, 100]], [[300, np.nan], [90, 0]]], # include nan

12 dims=["i", "j", "k"],

13 coords={"i": [1, 2], "j": [10000, 10001], "k": [1, 2]},

14)

15 DA_DIR_B = xr.DataArray(

16 data=[[160, 359, 10], [100, np.nan, 20]],

17 dims=["k", "j"], # reversed dim order, missing 'i' dim

18 coords={"k": [1, 2], "j": [10001, 10000, 10002]}, # extra station, changed order

19)

20 DS_DIR_A = xr.Dataset({"a": DA_DIR_A, "b": DA_DIR_B})

21 DS_DIR_B = xr.Dataset({"a": DA_DIR_B, "b": DA_DIR_A})

22 DA_DIR_NAN = xr.DataArray( # array of only nans

23 data=[[[np.nan, np.nan], [np.nan, np.nan]], [[np.nan, np.nan], [np.nan, np.nan]]],

24 dims=["i", "j", "k"],

25 coords={"i": [1, 2], "j": [10000, 10001], "k": [1, 2]},

26)

27 EXP_AD = xr.DataArray(

28 data=[[[1, np.nan], [180, 0]], [[59, np.nan], [70, 100]]],

29 dims=["i", "j", "k"],

30 coords={"i": [1, 2], "j": [10000, 10001], "k": [1, 2]},

31)

32 EXP_AD_NAN = xr.DataArray(

33 data=[[[np.nan, np.nan], [np.nan, np.nan]], [[np.nan, np.nan], [np.nan, np.nan]]],

34 dims=["i", "j", "k"],

35 coords={"i": [1, 2], "j": [10000, 10001], "k": [1, 2]},

36)

37 EXP_DS = xr.Dataset({"a": EXP_AD, "b": EXP_AD})

38

39

40@pytest.mark.parametrize(

41 ("source_a", "source_b", "expected"),

42 [

43 (DA_DIR_A, DA_DIR_B, EXP_AD),

44 (DA_DIR_A, DA_DIR_NAN, EXP_AD_NAN),

45 (DS_DIR_A, DS_DIR_B, EXP_DS),

46 ],

47)

48 def test_angular_difference(source_a, source_b, expected):

49 """Tests that angular_difference returns correct values"""

50 output = angular_difference(source_a, source_b)

51 output.broadcast_equals(expected)