将大小为n的3x3旋转矩阵数组与大小为3的3D向量数组相乘

Question

我有一个三维的位置向量数组 p [np.shape(p) 得到 (Nx, Ny, Nz, 3)]，以及一个旋转矩阵的 n 数组 Rn [np.shape(R) 得到 (n, 3, 3)]。

我想获得一个形状为 (n, Nx, Ny, Nz, 3) 的数组 PR，其中第 i 个 (0 < i < n) 在第 0 维上的条目是由数组 Rn 中索引为 i 的 3x3 旋转矩阵对位置向量数组 p 进行旋转后得到的三维数组。

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中保持向量化。

- camelCaseInPython

Answer 1

有两种可能的解决方案 -

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 的解决方案似乎非常慢，我不确定为什么，因为我本来以为它在内部使用了 @。

- Ananda

---

这是那些情况之一，其中einsum确实非常出色！ - Frank Yellin

你能把我的答案加入到你的计时中吗？ - Mad Physicist

@MadPhysicist，已完成。但是在我的系统中似乎非常慢，我不确定原因。 - Ananda

这很奇怪，速度如此缓慢。 - Mad Physicist

Answer 2

由于np.dot已经处理了维度的乘积（不像np.matmul会广播前导维度），因此您无需进行任何花哨的打包。

这里有两个额外的步骤：

1. 您需要添加一个尾随维度到p，使其成为3x3和3x1矩阵的乘积。
2. 结果将具有形状(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)

- Mad Physicist

Answer 3

让我们试一试：

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

- Quang Hoang