我有一个位置向量的 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，其中维度 0 的第 i 个 (0
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解决方案似乎不合理地缓慢。我不确定为什么，因为我会想象它在引擎盖下使用@。
让我们试试：
out = ((Rn @ p.reshape(-1,3).T)
.reshape(Rn.shape[0],3,-1)
.swapaxes(1,2)
.reshape(-1, *p.shape)
)
你不需要做任何花哨的包装，因为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)