[ Python  Numpy : 3D matrix * 2D vector fast calculation ]
Hello everyone,
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 oneline 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 kith element of rotated_vectors contains the product of the kith rotation matrix with the kith 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.
Jagaral
Answer 1
For such reductions that require alignment along one or more axes, one can use np.einsum

rotated_vector = np.einsum('ijk,jk>ki',rotation_matrices,vectors)
Please note that the output would be of shape (N,2)
, where 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 ki
.
Runtime test 
In [24]: def org_app(rotation_matrices,vectors):
...: nb_of_vectors = vectors.shape[1]
...: 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 [25]: # Input arrays
...: rotation_matrices = np.random.rand(2,2,50)
...: vectors = np.random.rand(2,50)
...:
In [26]: out1 = org_app(rotation_matrices,vectors)
In [27]: out2 = np.einsum('ijk,jk>ki',rotation_matrices,vectors)
In [28]: np.allclose(out1,out2) # Verify results
Out[28]: True
In [29]: %timeit org_app(rotation_matrices,vectors)
10000 loops, best of 3: 196 µs per loop
In [30]: %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!
Answer 2
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]
The [...,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
Answer 3
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