[FIXED] Is there a vectorized way to apply a transformation matrix to a group of vectors?

Issue

Imagine you have a group of vectors, e.g. in the form of a trajectory. Is there a vectorized way of applying the transformation matrix to all data points at once, or are you stuck with a for-loop? Here is some sample code:

import numpy as np


angle = np.deg2rad(90)
rotM = np.array(
    [
        [np.cos(angle), -np.sin(angle), 0],
        [np.sin(angle),  np.cos(angle), 0],
        [            0,              0, 1],
    ]
)

# trajectory with columns t, x, y, z
trajectory = np.array(
    [
        [1, 1, 0, 0],
        [2, 2, 1, 0],
        [3, 3, 2, 0],
        [4, 4, 3, 1],
        [5, 6, 4, 2],
        [6, 9, 5, 3],
    ]
)

# transform coordinates
for i in range(len(trajectory)):
    trajectory[i][1:] = np.dot(rotM, trajectory[i][1:])

All I found so far is numpy.linalg.multi_dot, and these two posts (one, two), none of which seem to apply to my case.

Solution

For this case, use broadcasting along with np.matmul/@. You can multiply a 3x3 martrix by an Nx3x1 array of vectors:

trajectory[:, 1:] = rotM @ trajectory[:, 1:, None]

A cleaner and more flexible solution might be to use scipy.spatial.transform.Rotation objects instead of hand-crafting the matrix yourself:

rotM = Rotation.from_euler('z', angle)
trajectory[:, 1:] = rotM.apply(trajectory[:, 1:])

No need to add shim dimensions in this case.

Answered By - Mad Physicist