Корреляция Пирсона двух трехмерных массивов с NaN в питоне - PullRequest
Новая
       14

Корреляция Пирсона двух трехмерных массивов с NaN в питоне

0 голосов
/ 09 февраля 2019

У меня есть два трехмерных массива a и b с [time, lat, lon].Я хочу соотнести временные ряды каждой ячейки сетки, как correlate(a[:,0,0],b[:,0,0]), correlate(a[:,0,1],b[:,0,1]), ....Я стремлюсь к двум соотношениям.Один со всеми временными рядами и один только, где массив превышает определенный порог.Наборы данных также содержат некоторые пропущенные значения во временном ряду, и я прочитал оба набора данных с помощью Xarray.Корреляции и маскирование выполняются с использованием NumPy.В настоящий момент я прохожу каждую широту и долготу, собирая временные ряды, маскирую их, чтобы учесть nan и порог, и сопоставить их.Мой код выглядит так: 

def correlate(A, B, var1, var2, TH):

   name = "corr_"+var1+"_"+var2+"_TH_"+str(TH)+".nc"
   a = xr.open_dataset(A).sel(time=slice('1950-03','2013-12'))
   b = xr.open_dataset(B).sel(time=slice('1950-03','2013-12'))

   corr = np.empty([a[var1].shape[1],a[var1].shape[2]],dtype=float)
   corr_TH = corr
   varname_TH = "r_TH_"+str(TH)

   for lt in range(corr.shape[0]):
      for ln in range(corr.shape[1]):
         corr[lt,ln]         = np.ma.corrcoef(a[var1][:,lt,ln],b[var2][:,lt,ln], rowvar=True)[0,1]
         corr_TH[lt,ln] = np.ma.corrcoef(np.ma.masked_greater(a[var1][:,lt,ln],TH),b[var2][:,lt,ln], rowvar=True)[0,1]

   # save whole correlations
   ds = xr.Dataset({'r': (['lat', 'lon'],  corr),varname_TH: (['lat', 'lon'], corr_TH)},coords={'lon': a['lon'],'lat': a['lat']})

   return ds

Это работает в целом, но очень медленно.Я нашел функцию Xarray array.stack (), чтобы сгладить массивы, и попробовал что-то вроде: 

A_stack = A.var1.stack(z=('lat','lon'))
B_stack = B.var2.stack(z=('lat','lon'))

cov = ((A_stack - A_stack.mean(axis=0))* (B_stack - B_stack.mean(axis=0))).mean(axis=0)
corr = cov  / (A_stack.std(axis=0) * B_stack.std(axis=0))

Мультииндекс 'z', над которым массив находится в стеке, сохраняется в процессе, однако корреляциямассив в конце пуст.Я полагаю, это из-за Нанс.

У кого-нибудь есть идея сделать это?

спасибо

...