I am looking for a fast (ie vectorized) method to replace the following loop. I have 2 numpy arrays with dimensions: (20738,14) and (31,14). I need to multiply them element-wise to get an array of (20738,31,14). I've been experimenting with various configurations of broadcasting, but can't seem to get the desired result.
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)
A similar operation is the following:
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
I can replace the bottom portion with:
p2_3d = v_mu_3d.sum(axis=1) / v_mu_3d.sum(axis=2).sum(axis=1)[:,None]
but can't seem to figure out:
num / MU
which should take the (20738,14) num array and divide element-wise by the (31,14) MU array. np.repeat() might work, but speed is critical because the containing function is run several thousand times in a minimization routine.
UPDATE Based on the answer by Nils Werner, these loops can be reduced to: 1)
p1_3d = v_mu_3d / v_mu_3d.sum(axis=1)[:,None,:]
and 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,:]
Just use broadcasting:
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)
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.