使用numpy einsum计算矩阵列向量的内积

时间:2018-01-02 19:43:58

标签: python numpy matrix multidimensional-array numpy-einsum

假设我有一个像这样的numpy矩阵:

[[   1    2    3]
 [  10  100 1000]]

我想用自己计算每列的内积,结果是:

[1*1 + 10*10    2*2 + 100*100    3*3 + 1000*1000] == [101, 10004, 1000009]

我想知道是否可以使用einsum函数(并更好地理解它)。

到目前为止,我能得到的最接近的结果是:

import numpy as np

arr = np.array([[1, 2, 3], [10, 100, 1000]])

res = np.einsum('ij,ik->jk', arr, arr)

# [[    101    1002   10003]
#  [   1002   10004  100006]
#  [  10003  100006 1000009]]

对角线包含预期结果,但我想知道是否可以避免边缘计算。

2 个答案:

答案 0 :(得分:4)

使用np.einsum,就像这样 -

np.einsum('ij,ij->j',arr,arr)

示例运行 -

In [243]: np.einsum('ij,ij->j',arr,arr)
Out[243]: array([    101,   10004, 1000009])

np.sum -

In [244]: (arr**2).sum(0)
Out[244]: array([    101,   10004, 1000009])

numexpr module -

In [248]: import numexpr as ne

In [249]: ne.evaluate('sum(arr**2,0)')
Out[249]: array([    101,   10004, 1000009])

答案 1 :(得分:1)

通过Divakar's einsum answer的一个中间步骤,可以直观地理解您在此期待的内容。

In [19]: arr
Out[19]: 
array([[   1,    2,    3],
       [  10,  100, 1000]])

# simply take element-wise product with the array itself
In [20]: np.einsum('ij, ij -> ij', arr, arr)
Out[20]: 
array([[      1,       4,       9],
       [    100,   10000, 1000000]])

但是,这并没有给出你期望的结果。因此,如果您观察到上述结果,我们只需要沿第一维(即轴0)求和。因此,我们在einsum结果中省略了i之后的下标->,这意味着我们要求它沿着该轴求和,并产生预期的结果:

In [21]: np.einsum('ij, ij -> j', arr, arr)
Out[21]: array([    101,   10004, 1000009])

P.S。另外,有关np.einsum的一般理解,请参阅此处的详细讨论:understanding-numpy-einsum