1a1a11a 1a1a11a - 2 months ago 14
Python Question

numpy list of vectors multiplication

I have a matrix M1 of shape (N*2) and another matrix M2 (2*N), I want to obtain a result of (1*N), each element i is the product of ith row of M1 and ith column of M2.
I tried to use dot in numpy, but it can only give me the matrix multiplication result, which is (N*N), of course, I can take the diagonal which is what I want, I would like to know is there a better way to do this?

Answer

You can use np.einsum -

np.einsum('ij,ji->i',M1,M2)

Runtime test -

In [407]: # Setup inputs
     ...: N = 1000
     ...: M1 = np.random.rand(N,2)
     ...: M2 = np.random.rand(2,N)
     ...: 

In [408]: np.allclose(original_app(M1,M2),np.einsum('ij,ji->i',M1,M2))
Out[408]: True

In [409]: %timeit original_app(M1,M2)
100 loops, best of 3: 2.08 ms per loop

In [410]: %timeit np.einsum('ij,ji->i',M1,M2)
100000 loops, best of 3: 12.9 ┬Ás per loop

Massive speedup there!