用numpy处理点矩阵乘法组

时间:2016-07-26 00:29:13

标签: python arrays numpy matrix

给定两个并行数组,一个是旋转矩阵数组,另一个是三维点数组,我正在寻找将每个子组乘以其相应矩阵的最快方法。

通过使用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表达式以处理我的情况的帮助,或者有关如何使用其他方法的建议都将非常感激。

1 个答案:

答案 0 :(得分: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