Numpy中einsum函数用法
【导读】 einsum全称Einstein summation convention(爱因斯坦求和约定),又称为爱因斯坦标记法。能够计算任何维度的张量收缩。einsum的写法省去了求和符号,显得更加简洁。
一、一维张量收缩
对于一维张量,也即向量(Vector)。其收缩为零维张量,也即标量(Scalar):
c
=
∑
i
a
i
b
i
c =\sum_i a_i b_i
c=i∑aibi所以einsum 的写法就是:
c
=
a
i
b
i
c= a_i b_i
c=aibi用代码表示为:
>>> 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 i j a i j = ∑ i , j a i j a i j c=a_{ij }a_{ij}= \sum_{i,j} a_{ij }a_{ij} c=aijaij=i,j∑aijaij也即求矩阵内积
array([[0, 1],
[2, 3],
[4, 5]])
>>> np.einsum('ij,ij', a, b)
70
2.2 收缩到一维张量
c
i
=
∑
j
A
i
j
B
i
j
c_i = \sum_j A_{ij} B_{ij}
ci=j∑AijBij或
c
j
=
∑
i
A
i
j
B
i
j
c_j = \sum_i A_{ij} B_{ij}
cj=i∑AijBij
其中,
c
i
c_i
ci为两矩阵对应位置元素相乘后再把列元素相加;
c
j
c_j
cj为两矩阵对应位置元素相乘后把行相加
>>> 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
i=1\;,j=2\;,k=3\;,l=4
i=1,j=2,k=3,l=4,可知会产生
1
×
4
1\times4
1×4的二维张量(矩阵)。
由公式:
C
i
,
l
=
a
i
j
k
b
j
k
l
=
∑
j
=
0
1
∑
k
=
0
2
a
i
j
k
b
j
k
l
C_{i,l} = a_{ijk}b_{jkl}=\sum_{j=0}^1 \sum_{k=0}^2a_{ijk}b_{jkl}
Ci,l=aijkbjkl=j=0∑1k=0∑2aijkbjkl
不失一般性,我们验证:
C
0
,
0
=
220
C_{0,0} =220
C0,0=220
C
0
,
0
=
a
000
∗
b
000
+
a
001
∗
b
010
+
a
002
∗
b
020
+
a
010
∗
b
100
+
a
011
∗
b
110
+
a
012
∗
b
120
\begin{aligned}C_{0,0}=a_{000}*b_{000}+a_{001}*b_{010}+a_{002}*b_{020}\\+a_{010}*b_{100}+a_{011}*b_{110}+a_{012}*b_{120}\end{aligned}
C0,0=a000∗b000+a001∗b010+a002∗b020+a010∗b100+a011∗b110+a012∗b120
即
C
0
,
0
=
0
∗
0
+
1
∗
4
+
2
∗
8
+
3
∗
12
+
4
∗
16
+
5
∗
20
=
220
C_{0,0}=0*0+1*4+2*8+3*12+4*16+5*20=220
C0,0=0∗0+1∗4+2∗8+3∗12+4∗16+5∗20=220
矩阵形式的计算过程:
[
0
1
2
]
[
0
1
2
3
4
5
6
7
8
9
10
11
]
=
[
20
23
26
29
]
\begin{bmatrix}0&1&2\end{bmatrix}\begin{bmatrix} 0&1&2&3\\4&5&6&7\\8&9&10&11\end{bmatrix}=\begin{bmatrix}20&23&26&29\end{bmatrix}
[012]⎣
⎡04815926103711⎦
⎤=[20232629]
[
3
4
5
]
[
12
13
14
15
16
17
18
19
20
21
22
23
]
=
[
200
212
224
236
]
\begin{bmatrix}3&4&5\end{bmatrix}\begin{bmatrix} 12&13&14&15\\16&17&18&19\\20&21&22&23\end{bmatrix}=\begin{bmatrix}200&212&224& 236\end{bmatrix}
[345]⎣
⎡121620131721141822151923⎦
⎤=[200212224236]
[
220
235
250
265
]
=
[
20
23
26
9
]
+
[
200
212
224
236
]
\begin{bmatrix}220&235&250&265\end{bmatrix} =\begin{bmatrix}20&23&26&9\end{bmatrix} +\begin{bmatrix}200&212&224& 236\end{bmatrix}
[220235250265]=[2023269]+[200212224236]
对于更高维的张量收缩或大规模张量,不方便再用矩阵表示计算过程表示,建议直接由公式计算相关元素。
再给出一个复杂的例子:
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
i=3\;,j=4\;,k=5\;,l=2
i=3,j=4,k=5,l=2,可知会产生
5
×
2
5\times2
5×2的二维张量
由公式:
C
k
,
l
=
a
i
j
k
b
j
i
l
=
∑
i
=
0
2
∑
j
=
0
3
a
i
j
k
b
j
i
l
C_{k,l} = a_{ijk}b_{jil}=\sum_{i=0}^2 \sum_{j=0}^3a_{ijk}b_{jil}
Ck,l=aijkbjil=i=0∑2j=0∑3aijkbjil
同样,不失一般性验证
C
0
,
0
=
4400
C_{0,0}=4400
C0,0=4400
C
0
,
0
=
a
000
∗
b
000
+
a
010
∗
b
100
+
a
020
∗
b
200
+
a
030
∗
b
300
+
a
100
∗
b
010
+
a
110
∗
b
110
+
a
120
∗
b
210
+
a
130
∗
b
310
+
a
200
∗
b
020
+
a
210
∗
b
120
+
a
220
∗
b
220
+
a
230
∗
b
320
\begin{aligned}C_{0,0}=a_{000}*b_{000}+a_{010}*b_{100}+a_{020}*b_{200}+a_{030}*b_{300}\\+a_{100}*b_{010}+a_{110}*b_{110}+a_{120}*b_{210}+a_{130}*b_{310}\\+a_{200}*b_{020}+a_{210}*b_{120}+a_{220}*b_{220}+a_{230}*b_{320}\end{aligned}
C0,0=a000∗b000+a010∗b100+a020∗b200+a030∗b300+a100∗b010+a110∗b110+a120∗b210+a130∗b310+a200∗b020+a210∗b120+a220∗b220+a230∗b320
也即
C
0
,
0
=
0
∗
0
+
5
∗
6
+
10
∗
12
+
15
∗
18
+
20
∗
2
+
25
∗
8
+
30
∗
14
+
35
∗
20
+
40
∗
4
+
45
∗
10
+
50
∗
16
+
55
∗
22
=
4400
\begin{aligned}C_{0,0} &=0*0+5*6+10*12+15*18\\&+20*2+25*8+30*14+35*20\\&+40*4+45*10+50*16+55*22\\&=4400\end{aligned}
C0,0=0∗0+5∗6+10∗12+15∗18+20∗2+25∗8+30∗14+35∗20+40∗4+45∗10+50∗16+55∗22=4400
四、其他功能介绍(次要)
功能一:求矩阵的迹
>>> 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]])