PythonのNumPy einsum()を使ってアインシュタインの縮約記法でテンソルの縮約を行う方法

Pythonでアインシュタインの縮約記法（Einstein summation convention）を用いたテンソルの縮約（Tensor contraction）を行うには、numpy.einsum() メソッドを使用します。第1引数は「subscript（添字）」で、縮約に使用する添字ラベルをカンマ区切りのリストとして指定します。第2引数は「operands（オペランド）」で、演算の対象となる配列を渡します。

einsum() メソッドは、オペランドに対してアインシュタインの縮約記法を評価します。この記法を活用すると、多次元配列に対する線形代数の代表的な演算の多くを、非常にシンプルな形式で表現できます。暗黙モードでは、einsum がこれらの計算を自動的に処理します。

一方、明示モードでは、特定の添字ラベルについての総和を無効化したり強制したりできるため、古典的なアインシュタインの縮約記法には分類されないような、より柔軟な配列演算も実行可能です。

手順

まず、必要なライブラリをインポートします。

import numpy as np

array() メソッドを使って2つのNumPy配列を作成します。

arr1 = np.arange(60.).reshape(3,4,5)
arr2 = np.arange(24.).reshape(4,3,2)

作成した配列を表示します。

print("Array1...\n",arr1)
print("\nArray2...\n",arr2)

両方の配列の次元数を確認します。

print("\nDimensions of Array1...\n",arr1.ndim)
print("\nDimensions of Array2...\n",arr2.ndim)

両方の配列の形状を確認します。

print("\nShape of Array1...\n",arr1.shape)
print("\nShape of Array2...\n",arr2.shape)

アインシュタインの縮約記法でテンソルの縮約を実行するには、numpy.einsum() メソッドを次のように呼び出します。

print("\nResult (Tensor contraction)...\n",np.einsum('ijk,jil->kl', arr1, arr2))

ここで指定している添字 'ijk,jil->kl' では、arr1 の i・j と arr2 の j・i が共通のラベルになっているため、これらの軸に沿って総和が計算されます。その結果、残された k と l の組み合わせごとの値を持つ配列が出力されます。

コード例

import numpy as np

# array()メソッドを使って2つのNumPy配列を作成
arr1 = np.arange(60.).reshape(3,4,5)
arr2 = np.arange(24.).reshape(4,3,2)

# 配列を表示
print("Array1...\n",arr1)
print("\nArray2...\n",arr2)

# 両方の配列の次元数を確認
print("\nDimensions of Array1...\n",arr1.ndim)
print("\nDimensions of Array2...\n",arr2.ndim)

# 両方の配列の形状を確認
print("\nShape of Array1...\n",arr1.shape)
print("\nShape of Array2...\n",arr2.shape)

# アインシュタインの縮約記法でテンソルの縮約を実行
print("\nResult (Tensor contraction)...\n",np.einsum('ijk,jil->kl', arr1, arr2))

出力結果

Array1...
[[[ 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. 27. 28. 29.]
[30. 31. 32. 33. 34.]
[35. 36. 37. 38. 39.]]

[[40. 41. 42. 43. 44.]
[45. 46. 47. 48. 49.]
[50. 51. 52. 53. 54.]
[55. 56. 57. 58. 59.]]]

Array2...
[[[ 0. 1.]
[ 2. 3.]
[ 4. 5.]]

[[ 6. 7.]
[ 8. 9.]
[10. 11.]]

[[12. 13.]
[14. 15.]
[16. 17.]]

[[18. 19.]
[20. 21.]
[22. 23.]]]

Dimensions of Array1...
3

Dimensions of Array2...
3

Shape of Array1...
(3, 4, 5)

Shape of Array2...
(4, 3, 2)

Result (Tensor contraction)...
[[4400. 4730.]
[4532. 4874.]
[4664. 5018.]
[4796. 5162.]
[4928. 5306.]]