python - Python：在 4 维数组上向量化

Question

我有 4-d 风速阵列（时间、ht、x、y），我试图从中计算湍流动能。为了做到这一点，在每个时间步长和每个高度，我必须计算三个风分量的平均值，从每个网格点的速度中减去该平均值，计算每个网格点的 TKE，然后在那个处取平均值高度。正如您在下面的代码中看到的那样，这会导致很多循环。这个我跑了一次，跑了一个多星期，不利于快速分析。

我的问题：有没有办法对这些循环中正在进行的过程进行矢量化，以帮助我的代码运行得更快？我已经看到了一维和二维数组的示例，但没有达到我想要做的规模。

tkehgt = np.empty((289,39),np.float32)
for c in range(0,289,1):
    print count
    u_avg = np.empty((39),np.float32)
    v_avg = np.empty((39),np.float32)
    w_avg = np.empty((39),np.float32)
    for d in range(hgt_final.shape[1]):
        u_avg[d] = np.average(u[c,d,:,:])
        v_avg[d] = np.average(v[c,d,:,:])
        w_avg[d] = np.average(w[c,d,:,:])
    for e in range(hgt_final.shape[1]):
        tkegrid = np.empty((300,300),np.float32)
        for f in range(hgt_final.shape[2]):
            for g in range(hgt_final.shape[3]):
                uprime = u[c,e,f,g] - u_avg[e]
                vprime = v[c,e,f,g] - v_avg[e]
                wprime = w[c,e,f,g] - w_avg[e]
                tkegrid[f,g] = 0.5*(uprime**2+vprime**2+wprime**2)
        tkehgt[count,e] = np.average(tkegrid[:,:])

标签： pythonvectorization