Python numpy einsum用法及代码示例

本文简要介绍 python 语言中 numpy.einsum 的用法。

用法:

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

评估操作数上的 Einstein 求和约定。

使用爱因斯坦求和约定，可以以简单的方式表示许多常见的多维线性代数数组运算。在隐含的模式einsum计算这些值。

在明确的模式，einsum通过禁用或强制对指定的下标标签求和，为计算可能不被视为经典 Einstein 求和操作的其他数组操作提供了进一步的灵活性。

请参阅注释和示例以进行说明。

参数：

subscripts： str

将求和的下标指定为逗号分隔的下标标签列表。除非包含显式指示符“->”以及精确输出形式的下标标签，否则将执行隐式(经典爱因斯坦求和)计算。

operands： 数组 列表

这些是操作的数组。

out： ndarray，可选

如果提供，则计算将在此数组中完成。

dtype： {数据类型，无}，可选

如果提供，则强制计算使用指定的数据类型。请注意，您可能还必须提供更自由的强制转换参数以允许转换。默认为无。

order： {‘C’、‘F’、‘A’、‘K’}，可选

控制输出的内存布局。 “C”表示它应该是 C 连续的。 “F”表示它应该是 Fortran 连续的，“A”表示如果输入都是“F”，则它应该是“F”，否则为“C”。 “K”意味着它应该尽可能接近输入的布局，包括任意排列的轴。默认为“K”。

casting： {‘no’, ‘equiv’, ‘safe’, ‘same_kind’, ‘unsafe’}，可选

控制可能发生的数据类型转换。不建议将此设置为‘unsafe’，因为它会对累积产生不利影响。

• ‘no’ 表示根本不应该转换数据类型。

• ‘equiv’ 表示只允许更改字节顺序。

• ‘safe’ 意味着只允许可以保留值的强制转换。

• ‘same_kind’ 表示只允许安全类型转换或类型中的类型转换，如 float64 到 float32。

• ‘unsafe’ 表示可以进行任何数据转换。

默认为‘safe’。

optimize： {假，真，‘greedy’, ‘optimal’}，可选

控制是否应该进行中间优化。如果 False 和 True 将默认为 ‘greedy’ 算法，则不会发生优化。还接受来自np.einsum_path 函数的显式收缩列表。有关详细信息，请参阅np.einsum_path。默认为假。

返回：

output： ndarray

基于爱因斯坦求和约定的计算。

注意：

爱因斯坦求和约定可用于计算许多多维线性代数数组运算。 einsum 提供了一种简洁的方式来表示这些。

这些操作的非详尽列表可以通过 einsum 计算，如下所示以及示例：

• 数组的跟踪 numpy.trace 。

• 返回对角线 numpy.diag 。

• 数组轴总和， numpy.sum 。

• 换位和排列， numpy.transpose 。

• 矩阵乘法和点积， numpy.matmul numpy.dot 。

• 矢量内积和外积， numpy.inner numpy.outer 。

• 广播，逐元素和标量乘法， numpy.multiply 。

• 张量收缩， numpy.tensordot 。

• 链式数组操作，以高效的计算顺序 numpy.einsum_path 。

下标字符串是一个逗号分隔的下标标签列表，其中每个标签指的是相应操作数的一个维度。每当重复一个标签时，它就会被求和，因此 np.einsum('i,i', a, b) 等同于 np.inner(a,b) 。如果一个标签只出现一次，它不会被求和，所以np.einsum('i', a) 生成一个a 的视图，没有任何变化。另一个示例 np.einsum('ij,jk', a, b) 说明了传统的矩阵乘法并且等效于 np.matmul(a,b) 。一个操作数中的重复下标标签取对角线。例如，np.einsum('ii', a) 等价于 np.trace(a) 。

在隐式模式，选择的下标很重要，因为输出的轴按字母顺序重新排序。这意味着np.einsum('ij', a)不影响二维数组，而np.einsum('ji', a)取它的转置。此外，np.einsum('ij,jk', a, b)返回矩阵乘法，同时，np.einsum('ij,jh', a, b)返回乘法的转置，因为下标 ‘h’ 在下标 ‘i’ 之前。

在显式模式可以通过指定输出下标标签直接控制输出。这需要标识符“->”以及输出下标标签列表。此函数增加了函数的灵活性，因为可以在需要时禁用或强制求和。调用np.einsum('i->', a)就好像numpy.sum， 和np.einsum('ii->i', a)就好像numpy.diag.不同之处在于einsum默认情况下不允许广播。此外np.einsum('ij,jh->ih', a, b)直接指定输出下标标签的顺序，因此返回矩阵乘法，这与上面的隐式模式示例不同。

要启用和控制广播，请使用省略号。默认 NumPy 风格的广播是通过在每个术语的左侧添加省略号来完成的，例如 np.einsum('...ii->...i', a) 。要沿着第一个和最后一个轴进行跟踪，您可以执行 np.einsum('i...i', a) ，或者要使用最左边的索引而不是最右边的索引执行 matrix-matrix 乘积，可以执行 np.einsum('ij...,jk...->ik...', a, b) 。

当只有一个操作数，没有轴相加，也没有提供输出参数时，返回操作数的视图而不是新数组。因此，将对角线作为np.einsum('ii->i', a) 会生成一个视图(在 1.10.0 版本中更改)。

einsum 还提供了另一种方法来提供下标和操作数为 einsum(op0, sublist0, op1, sublist1, ..., [sublistout]) 。如果未以这种格式提供输出形状，einsum 将以隐式模式计算，否则将显式执行。下面的示例具有对应的einsum 调用以及两个参数方法。

只要输入数组可写，从 einsum 返回的视图现在都是可写的。例如，np.einsum('ijk...->kji...', a) 现在将具有与 np.swapaxes(a, 0, 2) 相同的效果，而np.einsum('ii->i', a) 将返回二维数组对角线的可写视图。

添加了 optimize 参数，该参数将优化 einsum 表达式的收缩顺序。对于具有三个或更多操作数的收缩，这可以在计算期间以更大的内存占用为代价大大提高计算效率。

通常会应用 ‘greedy’ 算法，经验测试表明在大多数情况下返回最佳路径。在某些情况下，‘optimal’ 将通过更昂贵、更详尽的搜索返回最高级路径。对于迭代计算，最好计算一次最优路径并通过将其作为参数提供来重用该路径。下面给出一个例子。

有关详细信息，请参阅 numpy.einsum_path 。

例子：

>>> a = np.arange(25).reshape(5,5)
>>> b = np.arange(5)
>>> c = np.arange(6).reshape(2,3)

矩阵的迹：

>>> np.einsum('ii', a)
60
>>> np.einsum(a, [0,0])
60
>>> np.trace(a)
60

提取对角线(需要显式形式)：

>>> 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])

在轴上求和(需要显式形式)：

>>> 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])

对于高维数组求和单轴可以用省略号来完成：

>>> np.einsum('...j->...', a)
array([ 10,  35,  60,  85, 110])
>>> np.einsum(a, [Ellipsis,1], [Ellipsis])
array([ 10,  35,  60,  85, 110])

计算矩阵转置，或重新排序任意数量的轴：

>>> 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]])

矢量内积：

>>> np.einsum('i,i', b, b)
30
>>> np.einsum(b, [0], b, [0])
30
>>> np.inner(b,b)
30

矩阵向量乘法：

>>> np.einsum('ij,j', a, b)
array([ 30,  80, 130, 180, 230])
>>> np.einsum(a, [0,1], b, [1])
array([ 30,  80, 130, 180, 230])
>>> np.dot(a, b)
array([ 30,  80, 130, 180, 230])
>>> np.einsum('...j,j', a, b)
array([ 30,  80, 130, 180, 230])

广播和标量乘法：

>>> np.einsum('..., ...', 3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.einsum(',ij', 3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.einsum(3, [Ellipsis], c, [Ellipsis])
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.multiply(3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])

矢量外积：

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

张量收缩：

>>> 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.]])
>>> np.tensordot(a,b, axes=([1,0],[0,1]))
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

可写返回数组(从 1.10.0 版本开始)：

>>> a = np.zeros((3, 3))
>>> np.einsum('ii->i', a)[:] = 1
>>> a
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])

省略号使用示例：

>>> a = np.arange(6).reshape((3,2))
>>> b = np.arange(12).reshape((4,3))
>>> np.einsum('ki,jk->ij', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])
>>> np.einsum('ki,...k->i...', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])
>>> np.einsum('k...,jk', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])

链式数组操作。对于更复杂的收缩，可以通过重复计算 ‘greedy’ 路径或预先计算 ‘optimal’ 路径并使用 einsum_path 插入(自 1.12.0 版起)重复应用它来实现加速。对于较大的阵列，性能改进可能特别显著：

>>> a = np.ones(64).reshape(2,4,8)

基本einsum：~1520ms(在 3.1GHz Intel i5 上进行基准测试。)

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a)

次优einsum(由于重复路径计算时间)：~330ms

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal')

Greedy einsum(更快的最佳路径逼近)：~160ms

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='greedy')

最佳einsum(某些用例中的最佳使用模式)：~110ms

>>> path = np.einsum_path('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal')[0]
>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize=path)