将 n 个 3x3 旋转矩阵的数组与 3 向量的 3d 数组相乘

Question

我有一个 3d 位置向量数组 p [np.shape(p) yields (Nx, Ny, Nz, 3)] 和一个包含 n 个旋转矩阵的数组 Rn [np.shape(R ) 产生 (n, 3, 3)].

我正在尝试获取形状为 (n, Nx, Ny, Nz, 3) 的数组 PR，其中维度 0 的第 i 个 (0 < i < n) 条目是位置向量 p 的 3d 数组由数组 Rn 的索引 i 处的 3x3 旋转矩阵旋转。

theta = np.arange(0, 2*np.pi, np.pi/50)
phi = np.arange(0, np.pi, np.pi/100)

a = np.arange(100)
b = np.arange(50)
p = np.array(np.meshgrid(a, b, a, indexing="xy"))
p = np.moveaxis(p, 1, 2)
p = np.moveaxis(p, 0, 3)
# np.shape(p) => (100,50,100,3) 
Rn = np.array([np.array([np.cos(theta)*np.cos(phi), np.cos(theta)*np.sin(phi), -np.sin(theta)]),
               np.array([-np.sin(phi),              np.cos(phi),               np.zeros(np.shape(phi))]),
               np.array([np.cos(phi)*np.sin(theta), np.sin(theta)*np.sin(phi), np.cos(theta)])])
Rn = np.moveaxis(Rn , 1, 2)
Rn = np.moveaxis(Rn , 0, 1)
# np.shape(Rn) => (100, 3, 3)

到目前为止，我已经尝试了以下操作，但未成功。

PR= np.matmul(Rn, p)

执行此操作最有效的方法是什么？我知道如何使用 For 循环执行此操作，但为了提高效率，我一直试图在 numpy 中保持向量化。

Answer 1

让我们试试：

out = ((Rn @ p.reshape(-1,3).T)
          .reshape(Rn.shape[0],3,-1) 
          .swapaxes(1,2)
          .reshape(-1, *p.shape)
      )

Answer 2

两种可能的解决方案是 -

np.einsum("ijkl,nal->nijka", p, Rn, optimize=True)
td = np.moveaxis(np.tensordot(p, Rn, axes=((-1), (-1))), 3, 0)

我还会将这些解决方案与该线程中的其他答案进行比较。

p = np.random.rand(10, 20, 30, 3)
Rn = np.random.rand(100, 3, 3)

es = np.einsum("ijkl,nal->nijka", p, Rn, optimize=True)
td = np.moveaxis(np.tensordot(p, Rn, axes=((-1), (-1))), 3, 0)
d = np.squeeze(np.moveaxis(np.dot(Rn, p[..., None]), 1, -2), -1)
out = ((Rn @ p.reshape(-1,3).T)
          .reshape(Rn.shape[0],3,-1) 
          .swapaxes(1,2)
          .reshape(-1, *p.shape)
      )

print(np.allclose(es, out))
print(np.allclose(td, out))
print(np.allclose(d, out))

全部给出True.

如果您尝试对他们的表现进行基准测试，

%timeit np.einsum("ijkl,nal->nijka", p, Rn, optimize=True)
%timeit np.moveaxis(np.tensordot(p, Rn, axes=((-1), (-1))), 3, 0)
%timeit ((Rn @ p.reshape(-1,3).T).reshape(Rn.shape[0],3,-1) .swapaxes(1,2).reshape(-1, *p.shape))
%timeit np.moveaxis(np.squeeze(np.dot(Rn, p[..., None]), -1), 1, -1)

给予，

3.91 ms ± 129 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.15 ms ± 168 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
2.45 ms ± 29.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
29.1 ms ± 98.9 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

对于我系统上给定大小的数组。

einsum 和 tensordot 似乎具有相当的性能，而 @ 解决方案似乎是最快的。 dot 解决方案似乎慢得不合理。我不确定为什么，因为我会想象它在引擎盖下使用 @。

Answer 3

你不需要做任何花哨的包装，因为 np.dot 已经处理了维度的乘积（不像 np.matmul，它一起广播主要维度）。

还有两个额外的步骤：

您需要向 p 添加尾随维度，使其成为 3x3 x 3x1 矩阵的乘积。
由于产品的原因，结果将具有 (n, 3, Nx, Ny, Nz, 1) 的形状。你会想把第二个维度移到倒数第二个，挤出最后一个：

np.moveaxis(np.squeeze(np.dot(Rn, p[..., None]), -1), 1, -1)

或

np.squeeze(np.moveaxis(np.dot(Rn, p[..., None]), 1, -2), -1)

将 n 个 3x3 旋转矩阵的数组与 3 向量的 3d 数组相乘

Multiply array of n 3x3 rotation matrices with 3d array of 3-vectors

python

numpy

vectorization

matrix-multiplication