Numpy库中einsum函数用法

【导读】 einsum全称Einstein summation convention（爱因斯坦求和约定），又称为爱因斯坦标记法。能够计算任何维度的张量收缩。einsum的写法省去了求和符号，显得更加简洁。

一、一维张量收缩

对于一维张量，也即向量（Vector）。其收缩为零维张量，也即标量（Scalar）：

$$c=\sum _i{a}_i{b}_i$$所以einsum 的写法就是：
$$c=a_i{b}_i$$用代码表示为：

>>> a = np.arange(10)
>>> b = np.arange(10)+1
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> b
array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])
>>> np.einsum('i,i', a, b)
330

二、二维张量收缩

同样对于二维张量，也即矩阵（Matrix）。其收缩为一维张量或零维张量：

2.1 收缩到零维张量

$$c=a_{ij}{a}_{ij}=\sum _{i,j}{a}_{ij}{a}_{ij}$$也即求矩阵内积

array([[0, 1],
       [2, 3],
       [4, 5]])
>>> np.einsum('ij,ij', a, b)
70

2.2 收缩到一维张量

$$c_i=\sum _j{A}_{ij}{B}_{ij}$$或$$c_j=\sum _i{A}_{ij}{B}_{ij}$$
其中，$c_i$为两矩阵对应位置元素相乘后再把列元素相加；$c_j$为两矩阵对应位置元素相乘后把行相加

>>> np.einsum('ij,ij->i', a, b)
array([ 8, 62])
>>> np.einsum('ij,ij->j', a, b)
array([12, 22, 36])

再给出一个例子

>>> a1 = np.array([[0,1],[2,3],[4,5]])
>>> a1
array([[0, 1],
       [2, 3],
       [4, 5]])
>>> np.einsum('ji,ij->i', a1, a)
array([10, 40])

再给出一个矩阵与向量的例子

>>> a2 = np.array([[1,2],[3,4],[5,6]])
>>> b2 = np.array([1,2])
>
>>> np.einsum('ij,j', a2, b2)
array([ 5, 11, 17])
>>> np.einsum('ij,j->...', a2, b2)	# 收缩到零维
33

再给出一个多个向量的例子

>>> c2 = np.array([1,2,3])
>>> np.einsum('ij,j,i->i', a2, b2, c2)
array([ 5, 22, 51])
>>> np.einsum('ij,j,i->...', a2, b2, c2)  # 收缩到零维
78

三、三维张量收缩（重难点）

由于其组合非常多样，如三维张量与矩阵、向量的爱因斯坦求和运算，收缩到二维、一维、零维，我们不一一讨论。这里我们只给出三维张量之间进行爱因斯坦求和、收缩到二维张量的例子

3.1 例1

>>> a = np.arange(6).reshape(1,2,3)
>>> b = np.arange(24).reshape(2,3,4)
>>> a
array([[[0, 1, 2],
        [3, 4, 5]]])
>>> b
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [20, 21, 22, 23]]])
        
>>> np.einsum('ijk,jkl->il', a, b)
array([[220, 235, 250, 265]])

其中，$i=1,j=2,k=3,l=4$，可知会产生$1\times 4$的二维张量（矩阵）。
由公式：
$$C_{i,l}=a_{ijk}{b}_{jkl}=\sum _{j=0}^1\sum _{k=0}^2{a}_{ijk}{b}_{jkl}$$
不失一般性，我们验证：$C_{0,0}=220$
$$\begin{array}{r}{C}_{0,0}=a_{000}\ast b_{000}+a_{001}\ast b_{010}+a_{002}\ast b_{020}\\ +a_{010}\ast b_{100}+a_{011}\ast b_{110}+a_{012}\ast b_{120}\end{array}$$
即
$$C_{0,0}=0\ast 0+1\ast 4+2\ast 8+3\ast 12+4\ast 16+5\ast 20=220$$
矩阵形式的计算过程：
$$\left[\begin{array}{ccc}0& 1& 2\end{array}\right]\left[\begin{array}{cccc}0& 1& 2& 3\\ 4& 5& 6& 7\\ 8& 9& 10& 11\end{array}\right]=\left[\begin{array}{cccc}20& 23& 26& 29\end{array}\right]$$
$$\left[\begin{array}{ccc}3& 4& 5\end{array}\right]\left[\begin{array}{cccc}12& 13& 14& 15\\ 16& 17& 18& 19\\ 20& 21& 22& 23\end{array}\right]=\left[\begin{array}{cccc}200& 212& 224& 236\end{array}\right]$$

$$\left[\begin{array}{cccc}220& 235& 250& 265\end{array}\right]=\left[\begin{array}{cccc}20& 23& 26& 9\end{array}\right]+\left[\begin{array}{cccc}200& 212& 224& 236\end{array}\right]$$
对于更高维的张量收缩或大规模张量，不方便再用矩阵表示计算过程表示，建议直接由公式计算相关元素。

再给出一个复杂的例子：

3.2 例2

>>> a = np.arange(60.).reshape(3,4,5)
>>> b = np.arange(24.).reshape(4,3,2)
>>> np.einsum('ijk,jil->kl', a, b)
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])
>>> np.einsum(a, [0,1,2], b, [1,0,3], [2,3])	# 掌握上一种形式就好
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

该例中$i=3,j=4,k=5,l=2$，可知会产生$5\times 2$的二维张量
由公式：
$$C_{k,l}=a_{ijk}{b}_{jil}=\sum _{i=0}^2\sum _{j=0}^3{a}_{ijk}{b}_{jil}$$
同样，不失一般性验证$C_{0,0}=4400$
$$\begin{array}{r}{C}_{0,0}=a_{000}\ast b_{000}+a_{010}\ast b_{100}+a_{020}\ast b_{200}+a_{030}\ast b_{300}\\ +a_{100}\ast b_{010}+a_{110}\ast b_{110}+a_{120}\ast b_{210}+a_{130}\ast b_{310}\\ +a_{200}\ast b_{020}+a_{210}\ast b_{120}+a_{220}\ast b_{220}+a_{230}\ast b_{320}\end{array}$$
也即
$$\begin{array}{rl}{C}_{0,0}& =0\ast 0+5\ast 6+10\ast 12+15\ast 18\\ & +20\ast 2+25\ast 8+30\ast 14+35\ast 20\\ & +40\ast 4+45\ast 10+50\ast 16+55\ast 22\\ & =4400\end{array}$$

四、其他功能介绍（次要）

功能一：求矩阵的迹

 >>> a = np.arange(25).reshape(5,5)
 >>> a
 >>> array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

 >>> b = np.arange(5)
 >>> b
 array([0, 1, 2, 3, 4])
 
 >>> c = np.arange(6).reshape(2,3)
 >>> c
 array([[0, 1, 2],
       [3, 4, 5]])

Trace of a matrix:
    
    >>> np.einsum('ii', a)
    60
    >>> np.einsum(a, [0,0])
    60
    >>> np.trace(a)
    60

功能二：取对角元素

Extract the diagonal (requires explicit form):
    
    >>> np.einsum('ii->i', a)
    array([ 0,  6, 12, 18, 24])
    >>> np.einsum(a, [0,0], [0])
    array([ 0,  6, 12, 18, 24])
    >>> np.diag(a)
    array([ 0,  6, 12, 18, 24])

功能三 ：对某个维度求和

  Sum over an axis (requires explicit form):
    
    >>> np.einsum('ij->i', a)
    array([ 10,  35,  60,  85, 110])
    >>> np.einsum(a, [0,1], [0])
    array([ 10,  35,  60,  85, 110])
    >>> np.sum(a, axis=1)
    array([ 10,  35,  60,  85, 110])

对于更高维度的矩阵，对某一维度求和可以使用省略号

>>> a = np.arange(27).reshape(3,3,3)
>>> a
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])
>>> np.einsum('...j->...', a)
array([[ 3, 12, 21],
       [30, 39, 48],
       [57, 66, 75]])

功能四：矩阵转置

>>> c = np.arange(6).reshape(2,3)
>>> c
array([[0, 1, 2],
       [3, 4, 5]])
>>> np.einsum('ji', c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> 
>>> np.einsum('ij->ji', c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.einsum(c, [1,0])
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.transpose(c)
array([[0, 3],
       [1, 4],
       [2, 5]])

功能五：矩阵或向量的内积

>>> b = np.arange(5)
>>> b
array([0, 1, 2, 3, 4])
>>> np.einsum('i,i', b, b)
    30
>>> np.einsum(b, [0], b, [0])
    30
>>> np.inner(b,b)
    30

功能六：矩阵和向量乘法

>>> a = np.arange(6).reshape(2,3)
>>> a
array([[0, 1, 2],
       [3, 4, 5]])
>>> b = np.arange(3)
>>> b
array([0, 1, 2])
>>> np.einsum('ij,j', a, b)
array([ 5, 14])
>>> np.dot(a,b)
array([ 5, 14])

为减少不必要的重复，接下来代码中不再说明同样功能的其他函数或enisum函数的另一种使用形式
如：

>>> np.einsum('...j->...', a)
array([ 3, 12])
>>> np.einsum(a, [Ellipsis,1], [Ellipsis])
array([ 3, 12])

功能七：向量外积

>>> a
array([0, 1, 2])
>>> 
>>> b
array([1, 2, 3])
>>> np.einsum('i,j', a, b)
array([[0, 0, 0],
       [1, 2, 3],
       [2, 4, 6]])