python - 将每列的点积与自身取并存储在新的数组维度中

Question

我需要在给定的二维数组中将每列的点积与自身的转置相结合。目前，我正在使用循环，它既慢又丑，所以我想知道如何使用纯数组操作（如果可能）以及例如使用einsum和类似工具来做到这一点。

MWE：

import numpy as np
A = np.random.randint(2,15,size=(2,3))

>>>A
array([[14,  9,  4],
       [12,  8,  2]])

这就是我想要的，包括我目前的做法：

AA = [A[:,i].reshape(-1,1).dot(A[:,i].reshape(-1,1).T) for i in range(3)]
>>>AA
[array([[196, 168],
        [168, 144]]),
 array([[81, 72],
        [72, 64]]),
 array([[16,  8],
        [ 8,  4]])]

最后

>>>np.stack(AA,axis=0)
array([[[196, 168],
        [168, 144]],

       [[ 81,  72],
        [ 72,  64]],

       [[ 16,   8],
        [  8,   4]]])

有形状(3,2,2)。

现在显然在我真正的问题中，A是巨大的，所以目前的方法是不可行的。因此，如何才能做得更好？

标签： pythonnumpymultidimensional-arraylinear-algebra

Answer 1

您可以使用numpy.einsum：

import numpy as np
A = np.array([
       [14,  9,  4],
       [12,  8,  2]
    ])

np.einsum('ik,jk->kij', A, A)
array([[[196, 168],
        [168, 144]],

       [[ 81,  72],
        [ 72,  64]],

       [[ 16,   8],
        [  8,   4]]])