Search code examples
pythonperformancenumpyvectorizationtensordot

Python vectorization of matrix-vector operation

I have a Matrix A with shape (2,2,N) and a Matrix V with shape (2,N)

I want to vectorize the following:

F = np.zeros(N)
for k in xrange(N):
    F[k] = np.dot( A[:,:,k], V[:,k] ).sum()

Any way this can be done with either tensordot or any other numpy function without explicit looping?

Solution

  • With np.einsum -

    F = np.einsum('ijk,jk->k',A,V)

    We can optimize it further with optimize flag (check docs) set as True.