-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathindex_calculation.py
134 lines (107 loc) · 4.68 KB
/
index_calculation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#Date:2023-10-31
import numpy as np
def area_ave_rmse(data_real: np.array, data_predict: np.array, lat:np.array) -> np.array:
"""_summary_
Args:
data_real (np.array): real data of shape (time,day,lev,lat,lon)
data_predict (np.array): predict data of shape (time,day,lev,lat,lon)
lat (np.array): lat of shape (lat,) degree
Returns:
np.array: rmse of shape (time,day,lev)
"""
mask = np.array(np.isnan(data_real[0,0,0]))
angle_radians = np.radians(lat)
angle_radians = angle_radians[:, np.newaxis]
angle_radians = np.tile(angle_radians, [1, data_real.shape[-1]])
angle_radians_mask = angle_radians
angle_radians_mask[mask] = np.nan
weights = np.cos(angle_radians_mask)/ np.nansum(np.cos(angle_radians_mask))
data_diff = data_predict - data_real
data_diff2 = data_diff**2
data_diff2 = data_diff2*weights
data_diff2 = data_diff2.reshape(data_diff.shape[0], data_diff.shape[1], data_diff.shape[2], data_diff.shape[3]*data_diff.shape[4])
mse = np.nansum(data_diff2, axis=3)
rmse = np.sqrt(mse)
return rmse
def area_point_rmse(data_real: np.array, data_predict: np.array) -> np.array:
"""_summary_
Args:
data_real (np.array): real data of shape (time,day,lev,lat,lon)
data_predict (np.array): predict data of shape (time,day,lev,lat,lon)
Returns:
np.array: rmse of shape (day,lev,lat,lon)
"""
data_diff = data_predict - data_real
data_diff2 = data_diff**2
mse = np.nanmean(data_diff2,axis=0)
rmse = np.sqrt(mse)
return rmse
def area_point_r(data_real: np.array, data_predict: np.array) -> np.array:
"""_summary_
Args:
data_real (np.array): real data of shape (time,day,lev,lat,lon)
data_predict (np.array): predict data of shape (time,day,lev,lat,lon)
Returns:
np.array: r of shape (day,lev,lat,lon)
"""
r = np.zeros(shape=(data_real.shape[1], data_real.shape[2], data_real.shape[3], data_real.shape[4]))
for i in range(data_real.shape[1]):
print(f'lead {i} day')
for j in range(data_real.shape[2]):
for k in range(data_real.shape[3]):
for l in range(data_real.shape[4]):
sample_x = data_real[:, i, j, k, l]
sample_y = data_predict[:, i, j, k, l]
r[i,j,k,l] = np.corrcoef(sample_x, sample_y)[0,1]
return r
def area_ave_r(data_real: np.array, data_predict: np.array) -> np.array:
"""_summary_
Args:
data_real (np.array): real data of shape (time,day,lev,lat,lon)
data_predict (np.array): predict data of shape (time,day,lev,lat,lon)
Returns:
np.array: r of shape (time, day, lev)
"""
r = np.zeros(shape=(data_real.shape[0], data_real.shape[1], data_real.shape[2]))
for i in range(data_real.shape[0]):
print(f'lead {i} day')
for j in range(data_real.shape[1]):
for k in range(data_real.shape[2]):
data_a = data_real[i, j, k, :, :]
data_b = data_predict[i, j, k, :, :]
valid_indices = ~np.isnan(data_a) & ~np.isnan(data_b)
cleaned_data_a = data_a[valid_indices]
cleaned_data_b = data_b[valid_indices]
r[i,j,k] = np.corrcoef(cleaned_data_a, cleaned_data_b)[0,1]
return r
def area_ave_bias(data_real: np.array, data_predict: np.array, lat: np.array) -> np.array:
"""_summary_
Args:
data_real (np.array): real data of shape (time,day,lev,lat,lon)
data_predict (np.array): predict data of shape (time,day,lev,lat,lon)
Returns:
np.array: rmse of shape (time,day,lev)
"""
mask = np.array(np.isnan(data_real[0,0,0]))
angle_radians = np.radians(lat)
angle_radians = angle_radians[:, np.newaxis]
angle_radians = np.tile(angle_radians, [1, data_real.shape[-1]])
angle_radians_mask = angle_radians
angle_radians_mask[mask] = np.nan
weights = np.cos(angle_radians_mask)/ np.nansum(np.cos(angle_radians_mask))
data_diff = data_predict - data_real
data_diff = data_diff*weights
data_diff = data_diff.reshape(data_diff.shape[0],data_diff.shape[1],data_diff.shape[2],data_diff.shape[3]*data_diff.shape[4])
bias = np.nanmean(data_diff,axis=3)
return bias
def area_point_bias(data_real: np.array, data_predict: np.array) -> np.array:
"""_summary_
Args:
data_real (np.array): real data of shape (time,day,lev,lat,lon)
data_predict (np.array): predict data of shape (time,day,lev,lat,lon)
Returns:
np.array: rmse of shape (day,lev,lat,lon)
"""
data_diff = data_predict - data_real
bias = np.nanmean(data_diff,axis=0)
return bias