基于 numpy 的爱因斯坦求和约定（einsum）初探

0.爱因斯坦求和约定（Einstein Notation）

    在数学里，特别是将线性代数套用到物理时，爱因斯坦求和约定（Einstein summation convention）是一种标记的约定，又称为爱因斯坦标记法（Einstein notation），在处理关于坐标的方程式时非常有用。这约定是由阿尔伯特·爱因斯坦于1916年提出的。其计算目标和方式如下：

$$a_i{b}_i=\sum _{i=1}^N{a}_i{b}_i=a_1{b}_1+a_2{b}_2+a_3{b}_3+...$$

1.numpy 爱因斯坦求和约定函数

numpy.einsum(subscripts, *operands, out=None, dtype=None, order=‘K’, casting=‘safe’, optimize=False)

功能： 计算操作数组的爱因斯坦求和约定。

通常情况下，执行同样的操作时，einsum 效率远高于 sum 等其他计算函数。

参数：
  subscripts: str。将求和的下标指定为以逗号分隔的下标标签列表。除非包含显式指示符“->”以及精确输出形式的下标标签，否则将执行隐式（经典爱因斯坦求和）计算。
  operands:array。操作数组。

可选参数：
  out，dtype，order，casting，optimiz：与 numpy 其他函数的对应参数一致，这里不做探索。

简单示例：
  假定我们有两个 3 元素 的一维数组 A、B，我们希望两个数组对应位置的元素相乘并相加，其表达式可以为：
$$S=\sum _{i=1}^3{A}_i{B}_i$$
  利用 einsum 函数解析上述表达式：

import numpy as np
C = np.array([1,2,3])
D = np.array([4,5,6])
S = np.einsum('i,i->', C, D)
print(S)

  输出结果为：

32

  subscripts 标签中第一个 i 代表第一个操作数组的下标 i ，即 Ai；第二个 i 为第二个操作数组的下标，即 Bi。-> 代表求和（此处可以隐去）。

  可见构造 subscripts 标签是 einsum 函数基础。

2.subscripts 标签构建

    subscripts 字符串指导 operands 数组的计算方式。因此，了解并编写 subscripts 字符串是使用 einsum 函数的基础。

2.1 数组分隔

    subscripts 求和式每个数组变量之间用 , 隔开。（参考简单示例）

2.2 计算维度

    subscripts 求和式每组数的下标的顺序代表了维度位置，例如：

一维数据可以表示 i（参考简单示例）； 二维数据可以表示为 ij；

若求两个二维数组对应位置的元素相乘并相加的结果，例如将 简单示例 A、B变为两个二维数组，则 subscripts 可改写为 ij,ij->，其他不变。

多维数据可按照二维方式扩展。即：ijklmn,ijklmn-> 等的形式。

忽略维度。若维度过多，可使用 ... 省略前后端多余的维度。 即：...ij,...ij->... 等的形式。但计算结果必须设置到忽略维度上。即 -> 后的 ... 不能省略。

    目前，numpy 最多支持 32 维数据。einsum 下标可用 a-z 和 A-Z共计 52 个字母。

3.其他方法的简单示例

import numpy as np
A = np.array([[1,2,3], 
			  [1,1,1]])
B = np.array([[4,5,6], 
			  [2,2,2]])
C = np.array([1,2,3])
D = np.array([4,5,6])
E = np.array([[1,2,3],
              [1,1,1],
              [7,8,9]])
F = np.linspace(0,1000,180,dtype = np.int32).reshape(4,5,3,3)
G = np.linspace(1,1000, 9, dtype = np.int32).reshape(3,3)

3.1 原数组

einsum	numpy 对应函数操作	描述
(‘ij’, A)	A	返回自身。A 有多少个维度，subscripts 就应有多少个不相同的字母

S = np.einsum('ij', A) 
print(S)

[[1 2 3]
[1 1 1]]

3.2 求和

einsum	numpy 对应函数操作	描述
(‘ij->’, A)	np.sum(A)	所有值求和

S = np.einsum('ij->', A) 
print(S)

9

3.3 相乘

einsum	结果	描述
(‘ij,ij->ij’, A, B)	A * B	数组 A 和数组 B 对应相乘

S = np.einsum('ij,ij->ij', A, B) 
print(S)

[[ 4 10 18]
[ 2 2 2]]

3.4 内积

einsum	结果	描述
(‘i,i’, C, D)	np.inner(C, D)	数组 C 和数组 D 对应相乘并相加

    参考：简单示例。

3.5 外积

einsum	结果	描述
(‘i,j->ij’, C, D)	np.outer(C, D)	C 中的每一个元素与 D 中每一个元素相乘后组合成新数组

S = np.einsum('i,j->ij', C, D)
print(S)

[[ 4 5 6]
[ 8 10 12]
[12 15 18]]

3.6 转置

einsum	结果	描述
(‘ij->ji’, A)	A.T	行列互换

S = np.einsum('ij->ji', A)
print(S)

[[1 1]
[2 1]
[3 1]]

3.7 取对角

einsum	结果	描述
(‘ii->i’, E)	np.diag(E)	取数组左上至右下对角线上的元素

S = np.einsum('ii->i', A) 
print(S)

[1 1 9]

3.8 迹

迹（trace）表示主对角线所有元素的和。

einsum	结果	描述
(‘ii->’, E)	np.trace(E)	取数组左上至右下对角线上的元素之和

S = np.einsum('ii->', A) 
print(S)

11

3.9 按照维度求和

einsum	结果	描述
(‘ij->i’, E)	np.sum(E, axis = 1)	按照 axis = 1 这条轴求和，其他维度或类似

S = np.einsum('ij->i', E) 
print(S)

11

einsum	结果	描述
(‘ij,ij->i’, A, B)	np.sum(A * B, axis = 1)	A 和 B 求内积，然后按照 axis = 1 这条轴求和

S = np.einsum('ij,ij->i', A, B) 
print(S)

[32 6]

3.10 忽略一些维度

einsum	结果	描述
(‘…ij->…’, F)	np.sum(F, axis = (-2， -1))	按照后两维求和

S = print(np.einsum('...ij->...', F))
print(S)

[[ 197 649 1102 1554 2007]
[2459 2912 3364 3816 4270]
[4721 5175 5627 6079 6532]
[6984 7437 7889 8342 8795]]

einsum	结果	描述
(‘…ij,ij->…’, F, G)	np.sum(F * G, axis = (-2,-1))	按照后两维相乘后求和

S = np.einsum('...ij,ij->...', F, G)
print(S)

[[ 140125 366175 593100 818900 1045451]
[1271876 1498301 1724476 1950152 2177452]
[2403002 2629803 2856228 3082153 3308828]
[3534504 3761054 3987354 4213530 4440580]]

4.einsum 之于 gma

    最近在构建一些卷积函数。相对于成熟的scipy或者opencv的卷积，gma 相当于从头理了思路，并对其进行了实现。相较于普通的循环，gma实现卷积有数倍的速度提升，但对于scipy或者opencv任然要慢上不少。
    在此期间，偶然发现了einsum这个函数，其速度是相当之快。但其也有缺点，比如初次使用隐式（爱因斯坦求和约定）构建较为困难。当然，gma 并没有引用 这个函数，主要原因在：einsum 不支持传入掩膜数组，也就是说，如果传入的数组带有掩膜，那么掩膜会被直接忽略，这不符合我构建卷积的初衷，因此决定舍弃。但此函数依旧非常好用，说不定在不久的将来，我还会用它替代现有的方法。
  以上是仅限于我的理解对此函数的认知，欢迎各位参考。