在广播中以行方式使用 numpy matmul

Question

我有一个 3D 点 (n,3) 数组，将使用以 nx3x3 数组形式存储的 3x3 旋转矩阵围绕原点旋转。

目前我只是在matmul 的 for 循环中执行此操作，但我认为这是没有意义的，因为必须有更快的广播方式来执行此操作。

当前代码

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)

注意：在此示例中，我只是平铺了相同的旋转矩阵，但在我的实际情况中，每个 3x3 旋转矩阵都是不同的。

我想做什么

我猜一定有某种方式来广播这个，因为当点云变得非常大时，for 循环变得非常慢。我想这样做：

np.matmul(rotation_matrices, points)

points 中的每个row 都乘以它对应的旋转矩阵。 np.einsum 可能有一种方法可以做到这一点，但我无法弄清楚签名。

Answer 1

如果您看到the doc，则np.einsum('ij,jk', a, b) 是matmul 的签名。

所以你可以试试np.einsum的签名：

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

测试：

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