通过广播以行方式使用 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

通过广播以行方式使用 numpy matmul

Using numpy matmul in row-wise manner with broadcasting

python

numpy

matrix-multiplication

array-broadcasting

numpy-einsum

当前代码

我想做什么