如何从角度向量向量化创建 N 个旋转矩阵？

Question

数据： theta 是 N 个角的向量。

问题：如何在矢量化时从 theta 创建 N 个 2D 旋转矩阵的矢量？

如果没有向量化，我可以想到一个 for 循环：

import numpy as np

N = 100
theta = np.random.rand(N)*2*np.pi

def R(theta_v):
    Rotation = np.empty((len(theta_v), 2, 2))
    for k, theta in enumerate(theta_v):
        Rotation[k] = np.array([ [np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)] ])
    return Rotation

Rotation = R(theta)

有没有办法避免 for loop 以实现更高效的代码？

Answer 1

您可以使用 cos 和 sin 的矢量化版本来矢量化您的函数，然后重新排列结果：

def R_vec(theta):
    c, s = np.cos(theta), np.sin(theta)
    return np.array([c, -s, s, c]).T.reshape(len(theta),2,2)

对于 N=100，矢量化版本比我计算机上的原始版本快 110 倍左右。