科学計算における NumPy の einsum

はじめに

科学計算において、効率的な線形代数演算を行うことは不可欠です。数値計算のための基本的な Python ライブラリである NumPy は、これらの演算に必要な多数のツールを提供しています。これらのツールの中でも、einsum（アインシュタイン縮約）は、複雑な配列演算を簡潔に表現するための特に強力な関数として際立っています。

このチュートリアルでは、einsumが何であるか、そして Python コードでどのように効果的に使用するかを理解するためのガイドを提供します。この実験の終わりまでに、einsum関数を使用して様々な配列演算を行うことができ、従来の NumPy 関数に比べたその利点を理解することができるようになります。

NumPy の einsum の基本を理解する

アインシュタイン縮約 (einsum) は、NumPy の強力な関数で、簡潔な表記を使って多くの配列演算を表現することができます。これは、物理学や数学で複雑な方程式を簡略化するために一般的に使用されるアインシュタイン縮約規則に従っています。

Python シェルを開く

まずは Python シェルを開きましょう。デスクトップ上でターミナルを開き、次のコマンドを入力します。

python3

Python のプロンプト (>>>) が表示されれば、Python の対話型シェルに入ったことを示します。

NumPy をインポートする

最初に、NumPy ライブラリをインポートする必要があります。

import numpy as np

einsum とは何か？

NumPy のeinsum関数を使うと、配列のどのインデックス（次元）に対して演算を行うかを記述した文字列表記を使って、配列演算を指定することができます。

einsum演算の基本的な形式は次の通りです。

np.einsum('notation', array1, array2, ...)

ここで、表記文字列は実行する演算を記述します。

簡単な例：ベクトルの内積

簡単な例から始めましょう。2 つのベクトルの内積を計算します。数学的な表記では、2 つのベクトル u と v の内積は次のようになります。

$$\sum_i u_i \times v_i$$

einsumを使ってこれを計算する方法は次の通りです。

## Create two random vectors
u = np.random.rand(5)
v = np.random.rand(5)

## Print the vectors to see their values
print("Vector u:", u)
print("Vector v:", v)

## Calculate dot product using einsum
dot_product = np.einsum('i,i->', u, v)
print("Dot product using einsum:", dot_product)

## Verify with NumPy's dot function
numpy_dot = np.dot(u, v)
print("Dot product using np.dot:", numpy_dot)

表記 'i,i->' は次の意味を持ちます。

• i は最初の配列 (u) のインデックスを表します。
• 2 番目の i は 2 番目の配列 (v) のインデックスを表します。
• 矢印 -> の後に何もないことは、スカラー結果（すべてのインデックスにわたる和）が欲しいことを示します。

einsum 表記を理解する

einsum表記は次の一般的なパターンに従います。

'index1,index2,...->output_indices'

• index1, index2：各入力配列の次元のラベル
• output_indices：出力配列の次元のラベル
• 入力配列間で繰り返されるインデックスは合計されます。
• 出力に現れるインデックスは結果に残ります。

例えば、表記 'ij,jk->ik' では：

• i, j は最初の配列の次元です。
• j, k は 2 番目の配列の次元です。
• j は両方の入力配列に現れるので、この次元にわたって合計します。
• i, k は出力に現れるので、これらの次元は残ります。

これはまさに行列積の公式です！

一般的な einsum 演算

これでeinsumの基本を理解したので、この強力な関数を使って実行できる一般的な演算をいくつか見てみましょう。

行列の転置

行列を転置するとは、その行と列を入れ替えることを意味します。数学的な表記では、行列 A の転置行列 A^T は次のように定義されます。

$$A^T_{ij} = A_{ji}$$

einsumを使って行列の転置を行う方法を見てみましょう。

## Create a random matrix
A = np.random.rand(3, 4)
print("Original matrix A:")
print(A)
print("Shape of A:", A.shape)  ## Should be (3, 4)

## Transpose using einsum
A_transpose = np.einsum('ij->ji', A)
print("\nTransposed matrix using einsum:")
print(A_transpose)
print("Shape of transposed A:", A_transpose.shape)  ## Should be (4, 3)

## Verify with NumPy's transpose function
numpy_transpose = A.T
print("\nTransposed matrix using A.T:")
print(numpy_transpose)

表記 'ij->ji' は次の意味を持ちます。

• ij は入力行列のインデックスを表します（i は行、j は列）
• ji は出力行列のインデックスを表します（j は行、i は列）
• 実質的にインデックスの位置を入れ替えています

行列の乗算

行列の乗算は線形代数における基本的な演算です。2 つの行列 A と B について、それらの積 C は次のように定義されます。

$$C_{ik} = \sum_j A_{ij} \times B_{jk}$$

einsumを使って行列の乗算を行う方法は次の通りです。

## Create two random matrices
A = np.random.rand(3, 4)  ## 3x4 matrix
B = np.random.rand(4, 2)  ## 4x2 matrix

print("Matrix A shape:", A.shape)
print("Matrix B shape:", B.shape)

## Matrix multiplication using einsum
C = np.einsum('ij,jk->ik', A, B)
print("\nResult matrix C using einsum:")
print(C)
print("Shape of C:", C.shape)  ## Should be (3, 2)

## Verify with NumPy's matmul function
numpy_matmul = np.matmul(A, B)
print("\nResult matrix using np.matmul:")
print(numpy_matmul)

表記 'ij,jk->ik' は次の意味を持ちます。

• ij は行列 A のインデックスを表します（i は行、j は列）
• jk は行列 B のインデックスを表します（j は行、k は列）
• ik は出力行列 C のインデックスを表します（i は行、k は列）
• 繰り返されるインデックス j は合計されます（行列の乗算）

要素ごとの乗算

要素ごとの乗算は、2 つの配列の対応する要素同士を乗算することを意味します。同じ形状の 2 つの行列 A と B について、それらの要素ごとの積 C は次のようになります。

$$C_{ij} = A_{ij} \times B_{ij}$$

einsumを使って要素ごとの乗算を行う方法は次の通りです。

## Create two random matrices of the same shape
A = np.random.rand(3, 3)
B = np.random.rand(3, 3)

print("Matrix A:")
print(A)
print("\nMatrix B:")
print(B)

## Element-wise multiplication using einsum
C = np.einsum('ij,ij->ij', A, B)
print("\nElement-wise product using einsum:")
print(C)

## Verify with NumPy's multiply function
numpy_multiply = A * B
print("\nElement-wise product using A * B:")
print(numpy_multiply)

表記 'ij,ij->ij' は次の意味を持ちます。

• ij は行列 A のインデックスを表します
• ij は行列 B のインデックスを表します
• ij は出力行列 C のインデックスを表します
• インデックスは合計されず、対応する要素同士を単に乗算しています

高度な einsum 演算

これで基本的なeinsum演算に慣れたので、もっと高度な応用例をいくつか見てみましょう。これらの演算は、einsum関数の真の力と柔軟性を示しています。

対角成分の抽出

行列の対角成分を抽出することは、線形代数で一般的な演算です。行列 A について、その対角成分はベクトル d を形成し、次のようになります。

$$d_i = A_{ii}$$

einsumを使って対角成分を抽出する方法は次の通りです。

## Create a random square matrix
A = np.random.rand(4, 4)
print("Matrix A:")
print(A)

## Extract diagonal using einsum
diagonal = np.einsum('ii->i', A)
print("\nDiagonal elements using einsum:")
print(diagonal)

## Verify with NumPy's diagonal function
numpy_diagonal = np.diagonal(A)
print("\nDiagonal elements using np.diagonal():")
print(numpy_diagonal)

表記 'ii->i' は次の意味を持ちます。

• ii は A の対角成分の繰り返しインデックスを表します
• i はこれらの要素を 1 次元配列に抽出することを意味します

行列の跡

行列の跡は、その対角成分の和です。行列 A について、その跡は次のようになります。

$$\text{trace}(A) = \sum_i A_{ii}$$

einsumを使って跡を計算する方法は次の通りです。

## Using the same matrix A from above
trace = np.einsum('ii->', A)
print("Trace of matrix A using einsum:", trace)

## Verify with NumPy's trace function
numpy_trace = np.trace(A)
print("Trace of matrix A using np.trace():", numpy_trace)

表記 'ii->' は次の意味を持ちます。

• ii は対角成分の繰り返しインデックスを表します
• 空の出力インデックスは、すべての対角成分を合計してスカラーを得ることを意味します

バッチ行列乗算

einsumは、多次元配列に対する演算を行う際に本当に威力を発揮します。たとえば、バッチ行列乗算では、2 つのバッチから行列のペアを乗算します。

形状が (n, m, p) の行列のバッチ A と、形状が (n, p, q) の行列のバッチ B がある場合、バッチ行列乗算によって形状が (n, m, q) の結果 C が得られます。

$$C_{ijk} = \sum_l A_{ijl} \times B_{ilk}$$

einsumを使ってバッチ行列乗算を行う方法は次の通りです。

## Create batches of matrices
n, m, p, q = 5, 3, 4, 2  ## Batch size and matrix dimensions
A = np.random.rand(n, m, p)  ## Batch of 5 matrices, each 3x4
B = np.random.rand(n, p, q)  ## Batch of 5 matrices, each 4x2

print("Shape of batch A:", A.shape)
print("Shape of batch B:", B.shape)

## Batch matrix multiplication using einsum
C = np.einsum('nmp,npq->nmq', A, B)
print("\nShape of result batch C:", C.shape)  ## Should be (5, 3, 2)

## Let's check the first matrix multiplication in the batch
print("\nFirst result matrix from batch using einsum:")
print(C[0])

## Verify with NumPy's matmul function
numpy_batch_matmul = np.matmul(A, B)
print("\nFirst result matrix from batch using np.matmul:")
print(numpy_batch_matmul[0])

表記 'nmp,npq->nmq' は次の意味を持ちます。

• nmp はバッチ A のインデックスを表します（n はバッチ、m は行、p は列）
• npq はバッチ B のインデックスを表します（n はバッチ、p は行、q は列）
• nmq は出力バッチ C のインデックスを表します（n はバッチ、m は行、q は列）
• 繰り返されるインデックス p は合計されます（行列の乗算）

einsum を使う理由

NumPy がこれらの演算に対して専用の関数を提供しているのに、なぜeinsumを使うのか疑問に思うかもしれません。以下にいくつかの利点を挙げます。

1. 統一的なインターフェース：einsumは多くの配列演算に対して単一の関数を提供します。
2. 柔軟性：他の方法では複数のステップが必要な演算を表現できます。
3. 可読性：表記法を理解すれば、コードがより簡潔になります。
4. パフォーマンス：多くの場合、einsum演算は最適化されており、効率的です。

複雑なテンソル演算において、einsumはしばしば最も明確で直接的な実装を提供します。

まとめ

この実験では、NumPy の強力なeinsum関数を探索しました。この関数は、配列演算に対してアインシュタイン縮約記法を実装しています。学んだ内容を振り返ってみましょう。

1. einsum の基本概念：einsumの表記法を使って配列演算を表現する方法を学びました。インデックスは配列の次元を表し、繰り返されるインデックスは和を表します。

2. 一般的な演算：einsumを使っていくつかの基本的な演算を実装しました。

   • ベクトルの内積
   • 行列の転置
   • 行列の乗算
   • 要素ごとの乗算

3. 高度な応用：より複雑な演算を探索しました。

   • 対角成分の抽出
   • 行列の跡
   • バッチ行列乗算

einsum関数は、NumPy における配列演算に対して統一的で柔軟なアプローチを提供します。np.dot、np.matmul、np.transposeなどの専用関数が特定の演算に利用できますが、einsumは幅広い演算に対して一貫したインターフェースを提供します。これは、多次元配列を扱う際に特に価値があります。

科学計算やデータサイエンスの旅を続ける中で、einsumは複雑な配列演算を簡潔で読みやすいコードで実行するための強力なツールになります。