时间戳

首页 > 解决方案 > 将 n 3x3 旋转矩阵的数组与 3 向量的 3d 数组相乘

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

问题描述

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

我正在尝试获取形状为 (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.xml 中保持向量化。

标签： pythonnumpyvectorizationmatrix-multiplication

解决方案

两种可能的解决方案是 -

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解决方案似乎慢得不合理。我不知道为什么，因为我会想象它@在引擎盖下使用。

推荐阅读