python中带有NaN的两个3D数组的Pearson相关性

时间:2019-02-09 06:14:53

标签: python numpy multidimensional-array correlation python-xarray

我有两个带有[time,lat,lon]的三维数组a和b。我想关联每个网格单元的时间序列,例如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”,但是最后相关数组为空。我想那是因为Nans。

有人知道这样做吗?

谢谢

0 个答案:

没有答案