Question

我正在寻找一种快速（即矢量化）的方法来替换以下循环。我有2个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数组。 np.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,:]

Answer 1

只需使用广播：

v.shape
# (20738, 14)
MU.shape
# (31, 14)

v_mu_3d = v[:, None, :] * MU
p2_3d = v[:, None, :] / MU

v_mu_3d.shape
# (20738, 31, 14)
p2_3d.shape
# (20738, 31, 14)

2d numpy数组的乘法/除法产生3d数组

1 个答案: