Coverage for tests/continuous/quantile_loss_test_data.py: 100%

25 statements  

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

1""" 

2Generation of the test data used to test scores.continuous.quantile_score 

3""" 

4 

5import numpy as np 

6import pandas as pd 

7import xarray as xr 

8 

9FCST_ARRAY = [ 

10 [[0, 0, 0], [1, 1, 1]], 

11 [[0, 0, 0], [0.5, 0.5, 0.5]], 

12 [[0.5, 0.5, 0.5], [0, 0, 0]], 

13 [[1, 1, 1], [0, 0, 0]], 

14] 

15FCST_COORDS = { 

16 "lead_time": [15, 39, 63, 87], 

17 "valid_start": pd.date_range("2000-01-01 15:00:00", periods=2), 

18 "station_index": [10000, 10001, 10002], 

19} 

20FCST1 = xr.DataArray(data=FCST_ARRAY, coords=FCST_COORDS) 

21FCST2 = xr.DataArray([0.0, 1, 2, np.nan], coords=[[15, 39, 63, 87]], dims=["lead_time"]) 

22OBS_COORDS = {"valid_start": pd.date_range("2000-01-02 15:00:00", periods=2), "station_index": [10001, 10002, 10003]} 

23OBS_ARRAY = np.tile(np.array([1.0]), (2, 3)) 

24OBS1 = xr.DataArray(data=OBS_ARRAY, coords=OBS_COORDS) 

25OBS_DS = xr.Dataset( 

26 data_vars={ 

27 "temperature_1": (["valid_start", "station_index"], OBS_ARRAY), 

28 "temperature_2": (["valid_start", "station_index"], OBS_ARRAY), 

29 }, 

30 coords=OBS_COORDS, 

31) 

32FCST_DS = xr.Dataset( 

33 data_vars={ 

34 "temperature_1": (["lead_time", "valid_start", "station_index"], FCST_ARRAY), 

35 "temperature_2": (["lead_time", "valid_start", "station_index"], FCST_ARRAY), 

36 }, 

37 coords=FCST_COORDS, 

38) 

39DA1_2X2 = xr.DataArray([[5, 3], [7, 2]], coords=[[0, 1], [0, 1]], dims=["i", "j"]) 

40DA1_2X2X2 = xr.DataArray( 

41 [[[5, 3], [7, 2]], [[7, 4], [8, 2]]], 

42 coords=[[0, 1], [0, 1], [0, 1]], 

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

44) 

45DA2_2X2 = xr.DataArray([[2, 3], [8, 4]], coords=[[0, 1], [0, 1]], dims=["i", "j"]) 

46EXPECTED1 = xr.DataArray( 

47 [[[0.0, 0.0]], [[0.1, 0.1]], [[0.2, 0.2]], [[0.2, 0.2]]], 

48 coords=[[15, 39, 63, 87], [pd.to_datetime("2000-01-02 15:00:00")], [10001, 10002]], 

49 dims=["lead_time", "valid_start", "station_index"], 

50) 

51EXPECTED2 = xr.DataArray([0.0, 0.05, 0.1, 0.1], coords=[[15, 39, 63, 87]], dims=["lead_time"]) 

52EXPECTED3 = xr.DataArray([0.0, 0, 0, np.nan], coords=[[15, 39, 63, 87]], dims=["lead_time"]) 

53EXPECTED4 = xr.DataArray([0.0, 0.0, 0.1, 0.1], coords=[[15, 39, 63, 87]], dims=["lead_time"]) 

54EXPECTED_DS1 = xr.Dataset( 

55 data_vars={ 

56 "temperature_1": (["lead_time"], [0.0, 0.05, 0.1, 0.1]), 

57 "temperature_2": (["lead_time"], [0.0, 0.05, 0.1, 0.1]), 

58 }, 

59 coords={"lead_time": [15, 39, 63, 87]}, 

60) 

61EXPECTED_DS2 = xr.Dataset( 

62 data_vars={ 

63 "temperature_1": (["lead_time"], [0.0, 0.0, 0.1, 0.1]), 

64 "temperature_2": (["lead_time"], [0.0, 0.05, 0.1, 0.1]), 

65 }, 

66 coords={"lead_time": [15, 39, 63, 87]}, 

67) 

68WEIGHTS_ARRAY1 = [ 

69 [[0, 0, 0], [0, 0, 0]], 

70 [[0, 0, 0], [0, 0, 0]], 

71 [[1, 1, 1], [1, 1, 1]], 

72 [[1, 1, 1], [1, 1, 1]], 

73] 

74WEIGHTS_ARRAY2 = [ 

75 [[1, 1, 1], [1, 1, 0]], 

76 [[1, 1, 1], [1, 1, 1]], 

77 [[1, 1, 1], [1, 1, 1]], 

78 [[1, 1, 1], [1, 1, 1]], 

79] 

80WEIGHTS = xr.DataArray(data=WEIGHTS_ARRAY1, coords=FCST_COORDS) 

81WEIGHTS_DS = xr.Dataset( 

82 data_vars={ 

83 "temperature_1": (["lead_time", "valid_start", "station_index"], WEIGHTS_ARRAY1), 

84 "temperature_2": (["lead_time", "valid_start", "station_index"], WEIGHTS_ARRAY2), 

85 }, 

86 coords=FCST_COORDS, 

87)