PythonのNumPy einsum()でアインシュタインの縮約記法を使い行列の対角成分を抽出する方法

NumPyのeinsum()メソッドは、オペランド（演算対象の配列）に対してアインシュタインの縮約記法（Einstein summation convention）を適用して評価を行います。この記法を活用すると、多次元配列に対する線形代数の一般的な操作の多くを、非常にシンプルな形式で表現できます。

暗黙モードでは、einsumはこれらの値をそのまま計算します。一方、明示モードでは、特定の添字ラベルへの総和を無効化または強制できるため、古典的なアインシュタインの縮約操作とは見なされないような配列操作も柔軟に計算可能です。

Pythonでアインシュタインの縮約記法を使って行列の対角成分を抽出するには、numpy.einsum()メソッドを使用します。

第1引数は「subscript（添字）」で、総和に使用する添字をカンマ区切りのラベルリストとして指定します。第2引数は「operands（オペランド）」で、演算の対象となる配列を渡します。

今回使用する'ii->i'という指定では、入力側に同じ添字「i」が2回現れることで対角要素が選択され、矢印の右辺で出力の添字を「i」と指定しているため、総和は行われず対角成分だけが抽出されます。

手順

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

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()メソッドを使用します。

print("\nResult...\n",np.einsum('ii->i', 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...\n",np.einsum('ii->i', 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...
[ 0 5 10 15]