python - 使用 numpy 处理点矩阵乘法组

Question

给定两个平行数组，一个是旋转矩阵数组，另一个是一组 3D 点数组，我正在寻找将每个子组与其对应矩阵相乘的最快方法。

我能够通过使用 numpy.einsum 遍历每个组来实现我想要的。我希望有一种方法可以在没有循环的情况下做到这一点。这是我到目前为止的代码:

import numpy as np
N_GROUPS = 10
N_SUBGROUPS = 4

p = np.random.random((N_SUBGROUPS,N_GROUPS,3,)) # N_GROUPS of N_SUBGROUPS of 3D points
M = np.random.random((N_GROUPS,3,3,))           # N_GROUPS of rotation matrices
I = np.linalg.inv(M)                            # Inverse of M for testing purposes

# Use a loop to transform every subgroup.
for i in xrange(N_SUBGROUPS):
    p_ = np.einsum('ij,ijk->ik', p[i], M)

    # test
    p__= np.einsum('ij,ijk->ik', p_, I)
    print np.allclose(p[i],p__)# Returns True

重写 einsum 表达式以处理我的情况的任何帮助，或有关如何使用其他方法的建议将不胜感激。

Answer 1

这真的很简单:您自己完成了大部分工作!

只需取与子组对应的索引并将其放在 einsum 方程的两边:这将为您提供所需的维度数组 (N_SUBGROUPS, N_GROUPS, 3)。

假设我们称子群索引为l，那么:

p_ = np.einsum('lij,ijk->lik', p, M)

# I've changed the subgroup range index for clarity
for l in range(N_SUBGROUPS):
    # test
    p__= np.einsum('ij,ijk->ik', p_[l], I)
    print(np.allclose(p[l],p__))  # Returns True