Search code examples
pythonnumpyvectorizationmultiplication

vectorize numpy array multiplication

a = np.array([[20, 12,  6],
              [12, 24, 18],
              [ 0, 14, 30]])
b = np.array([1,0.5])
c = np.array([b ** i for i in range(0, 3)][::-1])
array([[1.  , 0.25],
       [1.  , 0.5 ],
       [1.  , 1.  ]])

using the second column in c, I want to get

[[20*0.25, 12*0.25, 6*0.25],
 [12*0.5,  24*0.5,  18*0.5],
 [ 0*1,    14*1,    30*1]
].sum(axis=0)

I want to do this for each column in c. For 1-d array of c, I can do (a * c[:,None]).sum(axis=0). How can I vectorize the multiplication here without using loop?

Solution

  • You are looking for np.eisum

    np.einsum('ij,ik', a, c)

array([[32. , 11. ],
       [50. , 29. ],
       [54. , 40.5]])

    Note that you could use broadcasting although not as efficient as eisum:

    (a[:,None] * c[:, :,None]).sum(0)

array([[32. , 50. , 54. ],
       [11. , 29. , 40.5]])