`numpy.einsum`是如何工作的？

Question

根据爱因斯坦求和来编写求和的正确方法对我来说是一个难题,所以我想在我的代码中尝试它.我在一些案例中取得了成功,但大多数都是经过反复试验.

现在有一个我无法弄清楚的案例.首先,一个基本问题.对于两个矩阵A和B是Nx1和 1xN分别,AB是NxN但BA是1x1.当我想NxN用np.einsum我可以做的计算案例时:

import numpy as np

a = np.asarray([[1,2]])
b = np.asarray([[2,3]])
print np.einsum('ij,ji->ij', a, b)

最后一个数组是2x2.然而

a = np.asarray([[1,2]])
b = np.asarray([[2,3]])
print np.einsum('ij,ij->ij', a, b)

返回一个1x2数组.我不太明白为什么这不能给出正确的结果.例如,对于上述案例numpy的指南说,箭头可用于强制求和或阻止它发生.但这对我来说似乎很模糊; 在上面的例子中,我不明白numpy如何根据索引的顺序(显然发生变化)决定输出数组的最终大小.

形式上我知道如下:当箭头右侧没有任何内容时,可以用数学方式将求和写为$\sum\limits_ {i = 0} ^ {N}\sum\limits_ {j = 0} ^ { M} A_ {ij} B_ {ij} $ for np.einsum('ij,ij',A,B),但是当有箭头时,我无法用正式的数学表达式来解释它.

Answer 1

In [22]: a
Out[22]: array([[1, 2]])
In [23]: b
Out[23]: array([[2, 3]])
In [24]: np.einsum('ij,ij->ij',a,b)
Out[24]: array([[2, 6]])
In [29]: a*b
Out[29]: array([[2, 6]])

这里,包括输出在内的所有部分中的索引的重复被解释为逐个元素的乘法.什么都没有总结. a[i,j]*b[i,j] = c[i,j]为了所有人i,j.

In [25]: np.einsum('ij,ji->ij',a,b)
Out[25]: 
array([[2, 4],
       [3, 6]])
In [28]: np.dot(a.T,b).T
Out[28]: 
array([[2, 4],
       [3, 6]])
In [38]: np.outer(a,b)
Out[38]: 
array([[2, 3],
       [4, 6]])

再次没有求和,因为相同的索引出现在左侧和右侧. a[i,j]*b[j,i] = c[i,j], 换一种说法:

[[1*2, 2*2],
 [1*3, 2*3]]

实际上是一种外在产品.看看a广播的方式b.T可能会有所帮助:

In [69]: np.broadcast_arrays(a,b.T)
Out[69]: 
[array([[1, 2],
        [1, 2]]), 
 array([[2, 2],
        [3, 3]])]

在陈述的左侧,重复的指数表明哪些维度相乘.匹配左侧和右侧确定它们是否相加.

np.einsum('ij,ji->j',a,b) # array([ 5, 10]) sum on i only
np.einsum('ij,ji->i',a,b) # array([ 5, 10]) sum on j only
np.einsum('ij,ji',a,b) # 15 sum on i and j

前段时间我制作了一个纯粹的Python,相当于einsum它解析字符串的方式.目标是创建一个nditer产品计算总和.但即使在Python中,它也不是一个简单易懂的脚本:

https://github.com/hpaulj/numpy-einsum/blob/master/einsum_py.py

显示这些求和规则的更简单序列:

In [53]: c=np.array([[1,2],[3,4]])
In [55]: np.einsum('ij',c)
Out[55]: 
array([[1, 2],
       [3, 4]])
In [56]: np.einsum('ij->i',c)
Out[56]: array([3, 7])
In [57]: np.einsum('ij->j',c)
Out[57]: array([4, 6])
In [58]: np.einsum('ij->',c)
Out[58]: 10

使用没有1维度的数组会消除广播复杂性:

In [71]: b2=np.arange(1,7).reshape(2,3)
In [72]: np.einsum('ij,ji',a2,b2)
...
ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (2,3)->(2,3) (2,3)->(3,2) 

或者我应该说,它暴露了广播的尝试.

省略号为einsum解释增加了一定程度的复杂性.当我解决了使用中的一个bug时,我开发了上面提到的github代码....但我没有花太多精力来完善文档.

省略号广播在numpy.einsum

当您需要一个可以处理各种大小数组的表达式时,省略号非常有用.如果你的数组总是2D,它不会做任何额外的事情.

举例来说,考虑一个概括dot,即将最后一个维度乘以A第二维到最后一维B.使用省略号,我们可以编写一个表达式,可以处理2d,3D和更大的数组:

np.einsum('...ij,...jk',np.ones((2,3)),np.ones((3,4)))  # (2,4)
np.einsum('...ij,...jk',np.ones((5,2,3)),np.ones((3,4)))  # (5,2,4)
np.einsum('...ij,...jk',np.ones((5,2,3)),np.ones((5,3,4))) # (5,2,4)
np.einsum('...ij,...jk',np.ones((5,2,3)),np.ones((7,5,3,4)))  # (7,5,2,4)
np.einsum('...ij,...jk->...ik',np.ones((5,2,3)),np.ones((7,5,3,4)) # (7, 5, 2, 4)

最后一个表达式使用默认的右侧索引...ik,省略号加上非求和索引.

您的原始示例可以写成

np.einsum('...j,j...->...j',a,b)

实际上,它填充i(或更多维度)以匹配数组的维度.

如果a或是b1d 也可以工作:

np.einsum('...j,j...->...j',a,b[0,:])

np.dot 推广到更大尺寸的方式是不同的

dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])

以einsum表示为:

np.einsum('ijo,kom->ijkm',np.ones((2,3,4)),np.ones((3,4,2)))

可以概括为

np.einsum('...o,kom->...km',np.ones((4,)),np.ones((3,4,2)))

要么

np.einsum('ijo,...om->ij...m',np.ones((2,3,4)),np.ones((3,4,2)))

但我认为我不能完全复制它einsum.也就是说,我不能告诉它填写索引A,然后是不同的索引B.