将n个3×3旋转矩阵的数组与3d向量的3d数组相乘
我有一个 3d 位置向量数组 p [np.shape(p) yields (Nx, Ny, Nz, 3)] 和一个包含 n 个旋转矩阵的数组 Rn [np.shape(R) yields (n, 3, 3) ]。
我正在尝试获得一个形状为 (n, Nx, Ny, Nz, 3) 的数组 PR，其中第 i 个 (0 < i < n) 在维度 0 的条目是由 3x3 旋转旋转的位置向量 p 的 3d 数组数组 Rn 的索引 i 处的矩阵。
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 中保持矢量化。
回答
两种可能的解决方案是 -
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解决方案似乎异常缓慢。我不知道为什么，因为我会想象它@在引擎盖下使用。
- This is one of those cases where einsum really does shine!