简体   繁体   中英

Multiplication/Division of 2d numpy arrays to produce 3d array

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.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM