Rotation of an array of vectors by an array of rotation matrices

Question

If we have a 3 x 3 rotation matrix R , it can be multiplied with v , a 3 x N array - an array of N column vectors - to produce a new 3 x N array of rotated vectors, like this:

v_rotated = R.dot(v)

Now suppose we have a N x M x 3 array, N times M vectors, which I want to rotate with N different 3 x 3 rotation matrices (one rotation matrix for each "row" of vectors). This is straightforward to do with a loop, but is there a faster and more compact (vectorized) way to do it, eg with numpy 's dot or tensorproduct ?

Example code for loop implementation:

from numpy import cos, sin, array, pi, linspace, random

# 100 different rotation matrices:
R = [array([[1, 0, 0], [0, cos(theta), -sin(theta)], [0, sin(theta), cos(theta)]]) for theta in linspace(0, pi, 100)]
# 100 x 200 random vectors:
v = random.random((100, 200, 3))

# rotate vectors in loop:
rotated_v = array([R_.dot(v_.T).T for R_, v_ in zip(R, v)])

Answer 1

let's assume that v.shape is (N, M, 3) and R.shape is (N, 3, 3) , you can use np.einsum

import numpy as np
rotated_v = np.einsum('lij, lkj->lki', R, v)

where l is the index on N , i and j are the indexes on 3x3 rotation dimension, and k is the index on M .

I matched my result with your as follow:

>>> print np.allclose(my_rotated_v, your_rotated_v)
True