Using numpy matmul in row-wise manner with broadcasting

Question

I have an array of 3D points (n,3) that are to be rotated about the origin using a 3x3rotation matrix that is stored in the form of a nx3x3 array.

At the moment I'm simply doing this in a for loop with matmul but I'm figuring this is pointless as there must be a much faster broadcasting way of doing it.

Current code

n = 10
points = np.random.random([10,3])
rotation_matrices = np.tile(np.random.random([3,3]), (n,1,1))

result = []

for point in range(len(points)):
    rotated_point = np.matmul(rotation_matrices[point], points[point])

    result.append(rotated_point)

result = np.asarray(result)

Note: in this example I've just tiled the same rotation matrix but in my real case every 3x3 rotation matrix is different.

What I want to do

I'm guessing there must be some way of broadcasting this as the for loop becomes very slow when the point-cloud becomes very large. I'd like to do this:

np.matmul(rotation_matrices, points)

where each row in points is multiplied by it's corresponding rotation matrix. There is probably a way to do this with np.einsum but I couldn't figure out the signature.

Answer 1

If you see the doc, np.einsum('ij,jk', a, b) is the signature for matmul.

So you can try np.einsum with the signature:

np.einsum('kij,kj->ki', rotation_matrices, points)

Test:

einsum = np.einsum('kij,kj->ki', rotation_matrices, points)
manual = np.array([np.matmul(x,y) for x,y in zip (rotation_matrices, points)])
np.allclose(einsum, manual)
# True