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)