我正在寻找一种快速(即矢量化)的方法来替换下面的循环。我有两个numpy数组,尺寸分别为:(20738,14)和(31,14)。我需要将它们元素相乘得到一个数组(20738,31,14)。我一直在尝试各种各样的广播配置,但似乎无法得到想要的结果。
v_mu_3d = np.zeros((v.shape[0], ALT_LEN, NEST_LEN))
for k in range(NEST_LEN):
v_mu_3d[:,:,k] = v * MU[:,k]
v_mu_3d = np.exp(v_mu_3d)
类似的操作如下:
p2_3d = np.zeros((v.shape[0], ALT_LEN, NEST_LEN))
for j in range(ALT_LEN):
num = v_mu_3d[:,:,:].sum(axis=1)
temp_MU = MU[j,:]
num = ne.evaluate('where(temp_MU >0, num / temp_MU, 0)')
denom = num.sum(axis=1)
denom = denom[:, np.newaxis]
p2_3d[:, j, :] = num / denom
我可以将底部替换为:
p2_3d = v_mu_3d.sum(axis=1) / v_mu_3d.sum(axis=2).sum(axis=1)[:,None]
但似乎不知道:
num / MU
它应该取(20738,14)num数组,并将元素按(31,14)mu数组进行划分。repeat()可能有效,但速度很关键,因为包含函数在最小化例程中运行了几千次。
更新
根据Nils-Werner的回答,这些循环可以简化为:
1)
p1_3d = v_mu_3d / v_mu_3d.sum(axis=1)[:,None,:]
2)
num = v_mu_3d.sum(axis=1)
num = np.where(MU>0, num[:,None,:] / MU, 0)
p2_3d = num / num.sum(axis=1)[:,None,:]