[ Python - Numpy : 3D matrix * 2D vector fast calculation ]
Here is what I want to do. I've got two arrays :
- rotation_matrices that contains 50 two dimensional rotation matrices. Each rotation matrix is shaped as (2,2). Thus, rotation_matrices is shaped as (2,2,50).
- vectors that contains 50 two dimensional vectors. Thus, it is shaped as (2,50).
I want (if it exists) a one-line numpy operation that gives me the (2,50) array that contains the rotated vectors, let's call it rotated_vectors. What I mean is that the k-ith element of rotated_vectors contains the product of the k-ith rotation matrix with the k-ith vector.
For the moment, I have come up with the loop that follows :
for ind,elt in enumerate(np.arange(nb_of_vectors)): rotated_vector[ind] = np.dot( rotation_matrices[:,:,ind], vectors[:,ind] )
I think there is room for improvement. If you have any suggestion, you are welcome.
Thank you for your time.
For such reductions that require alignment along one or more axes, one can use
rotated_vector = np.einsum('ijk,jk->ki',rotation_matrices,vectors)
Please note that the output would be of shape
N is the number of vectors. If instead, you were looking to have an output of shape
(2,N) and which would have required the original code to be :
rotated_vector[:,ind] = np.dot(...) instead, just edit the output string notation to
ik instead of
Runtime test -
In : def org_app(rotation_matrices,vectors): ...: nb_of_vectors = vectors.shape ...: r = np.zeros((nb_of_vectors,2)) ...: for ind,elt in enumerate(np.arange(nb_of_vectors)): ...: r[ind] = np.dot( rotation_matrices[:,:,ind], vectors[:,ind] ) ...: return r ...: In : # Input arrays ...: rotation_matrices = np.random.rand(2,2,50) ...: vectors = np.random.rand(2,50) ...: In : out1 = org_app(rotation_matrices,vectors) In : out2 = np.einsum('ijk,jk->ki',rotation_matrices,vectors) In : np.allclose(out1,out2) # Verify results Out: True In : %timeit org_app(rotation_matrices,vectors) 10000 loops, best of 3: 196 µs per loop In : %timeit np.einsum('ijk,jk->ki',rotation_matrices,vectors) 100000 loops, best of 3: 5.12 µs per loop
That proves again why iterating in NumPy is basically just bad!
Your axes are in an unusual order. First you'll want to put the matrix axes last:
rotation_matrices = np.rollaxis(rotation_matrices, -1) # shape (50, 2, 2) vectors = np.rollaxis(vectors, -1) # shape (50, 2)
Which would allow you to make your existing loop more readable:
for ind in np.arange(nb_of_vectors): rotated_vector[ind] = np.dot(rotation_matrices[ind], vectors[ind])
But instead, you can use the matrix multiply operator (or
np.matmul in python < 3.5)
rotated_vectors = (a @ vectors[...,None])[...,0] # rotated_vectors = np.matmul(a, vectors[...,None])[...,0]
[...,None] converts the array of vectors (shape
(n,) into an array of column matrices (shape
(n, 1)), and the trailing
[...,0] converts the column matrices back into vectors
This is an explicit formula version
result = np.array([vectors[0,:]*rotation_matrices[0,0,:] + vectors[1,:]*rotation_matrices[0,1,:], vectors[0,:]*rotation_matrices[1,0,:] + vectors[1,:]*rotation_matrices[1,1,:]]).transpose()
much (14x) faster than your original code but slower (2.6x) than
einsum on my machine