Pythonのeinsum()でアインシュタイン縮約記法により行列のトレース（対角和）を求める方法

einsum()メソッドは、アインシュタインの縮約記法（Einstein summation convention）に基づいて演算を評価します。この記法を活用すると、多次元配列に対する線形代数のさまざまな操作を、非常にシンプルな形で表現できます。暗黙モードでは、einsumがこれらの値を自動的に計算してくれます。

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

アインシュタイン縮約記法を使って行列のトレース（対角成分の総和）を求めるには、Pythonのnumpy.einsum()メソッドを使用します。

• 第1引数（subscript）：総和を指定するための添字ラベルを、カンマ区切りの文字列で渡します。トレースを求める場合は「'ii'」のように同じラベルを2回指定します。
• 第2引数（operands）：演算の対象となる配列を渡します。

手順

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

import numpy as np

arange()メソッドとreshape()メソッドを使って、NumPy配列を作成します。

arr = np.arange(16).reshape(4,4)

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

print("Our Array...\n",arr)

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

print("\nDimensions of our Array...\n",arr.ndim)

配列のデータ型を取得します。

print("\nDatatype of our Array object...\n",arr.dtype)

配列の形状を取得します。

print("\nShape of our Array object...\n",arr.shape)

アインシュタイン縮約記法で行列のトレースを求めるには、numpy.einsum()メソッドに「'ii'」を指定します。これは「行と列が同じ位置にある要素（対角成分）をすべて足し合わせる」ことを意味します。

print("\nResult (trace)...\n",np.einsum('ii', arr))

完全なコード例

import numpy as np

# arange()とreshape()を使ってNumPy配列を作成
arr = np.arange(16).reshape(4,4)

# 配列を表示
print("Our Array...\n",arr)

# 次元数を確認
print("\nDimensions of our Array...\n",arr.ndim)

# データ型を取得
print("\nDatatype of our Array object...\n",arr.dtype)

# 形状を取得
print("\nShape of our Array object...\n",arr.shape)

# アインシュタイン縮約記法でトレースを求める
print("\nResult (trace)...\n",np.einsum('ii', arr))

出力結果

Our Array...
[[ 0 1 2 3]
[ 4 5 6 7]
[ 8 9 10 11]
[12 13 14 15]]

Dimensions of our Array...
2

Datatype of our Array object...
int64

Shape of our Array object...
(4, 4)

Result (trace)...
30

この例では、4×4の行列の対角成分は0、5、10、15であり、その合計である30がトレースとして返されます。なお、同様の結果はnp.trace(arr)でも得られますが、einsum()を使うことで他の複雑なテンソル演算にも応用できる点が大きなメリットです。