numpy: broadcast matrix multiply accross array

Question

I have a 3xN array, conceptually an array of N 3-vectors, I want to construct the array which results from matrix multiplying a given 3x3 matrix with each column of the array. Is there a good way to do this in a vectorized manner?

Currently, my problem is 3xN, but I may need to consider 3xNxM (or more) in the future.

Loopy approach

U=numpy.rand( [3,24] )

R=numpy.eye(3) # placeholder

for i in xrange( U.shape[1]):
    U[:,i]=numpy.dot( R, U[:,i] )

Answer 1

Using np.einsum function you can do it even for the multi dimension problem:

U = np.random.rand(3,24,5) 
R = np.eye(3,3)
result = np.einsum( "ijk,il", U,R )

The notation is a little tricky: the string you give first states the indeces of the dimensions of the arrays; so for U the indeces ijk are the running indeces for each dimension. It follows the einstein summation convention so indeces that have the same letter in the string will be summed over. For details read ScipyDocs. I'm sure in your case the dot is faster, because there is less overhead and it will probably use some blas routine, but as you stated that you want to expand to even more dimensions this might be the way to go.

Answer 2

In this case you can simply call np.dot(R, U) and it will work:

import numpy as np

np.random.seed(0)

U = np.random.rand(3,24) 
R = np.random.rand(3,3)

result = np.empty_like(U)

for i in range( U.shape[1]): 
    result[:,i] = np.dot(R, U[:,i])

print result

print np.allclose(result, np.dot(R, U))

For the (3xNxM) case you can reshape to (3x(N.M)), dot and reshape the result back, similar to my answer here