はじめに
今回は、Pythonの数値計算モジュールNumPyが提供している関数einsumの使い方を解説する。einsumはアインシュタイン(Einstein)の縮約記法を実装した関数である。テンソル間の複雑な計算を1行で書くことができる。前回実装した変分推論においても用いたので、この機会に紹介したい。
アインシュタインの縮約記法とは
いま、ベクトルとの内積を考える。
(1)
ここで、上の式の和の記号を省略し、同じ文字の添え字が2つ続く場合はそれらについて和を取ると約束する。
(2)
左辺の簡略化した書き方をアインシュタインの縮約記法と呼ぶ。この記法を最初に考案したのはアインシュタインである。厳密な話をすると、アインシュタインが一般相対論を定式化する際に用いた数学(リーマン幾何学)では2種類の添え字を扱う。上付き添え字と下付き添え字である。前者を反変ベクトル、後者を共変ベクトルと呼び、これらは区別される。そして、同じ上付き添え字と下付き添え字が続く場合の和を簡略化する記法として、上の縮約記法が導入された。
(3)
本解説では簡単のため、最初に説明した方(下付き添え字だけの場合)の記法もアインシュタインの縮約記法と呼ぶことにして、話を進める。
NumPy.einsumの使い方
最初に、上で見た2つのベクトルの内積の場合を考える。コードは以下の通り。
1 2 3 4 5 6 7 8 9 10 11 |
# 2つのベクトルを定義する。 x = np.array([1, 2, 3]) y = np.array([3, 2, 1]) # 内積を計算する。dotを使用した場合。 z = x.dot(y) assert(z == 10) # einsumの場合。 u = np.einsum("i,i", x, y) assert(u == z) |
10行目でeinsum
に渡している第1引数の文字列"i,i"
は、式(2)の添え字に一致している。すなわち、式(2)のの添え字との添え字とを対応させ和を取ることをeinsum
に教えている。
次に、2つの行列の積を考える。縮約記法で書くと以下のようになる。
(4)
ここで、は行列の成分を表す。コードは以下の通りである。
1 2 3 4 5 6 7 8 9 10 11 |
# 2つの行列を定義する。 A = np.array([[1, 2], [2, 1]]) B = np.array([[0, 3], [3, 0]]) # 積を計算する。matmulを使用した場合。 C = np.matmul(A, B) assert(np.all(C == np.array([[6, 3], [3, 6]]))) # einsumの場合。 D = np.einsum("in,nj->ij", A, B) assert(np.all(C == D)) |
10行目でeinsum
の第1引数に渡している文字列"in,nj->ij"
は式(4)の添え字と一致している。すなわち、式(4)の右辺にあるの添え字との添え字から、左辺のの添え字を作る操作であることをeinsum
に教えている。
今度は、2つのベクトルから行列を作る演算を考える。
(5)
(6)
となる。この式には和は出現しないが、以下のようにeinsum
を用いて書くことができる。
1 2 3 4 5 6 |
# 2つのベクトルを定義する。 x = np.array([1, 2, 3]) y = np.array([1, 3]) z = np.einsum("i,j->ij", x, y) assert(np.all(z == np.array([[1, 3], [2, 6], [3, 9]]))) |
5行目のeinsum
に渡している引数"i,j->ij"
は式(6)の添え字と対応している。すなわち、右辺にあるの添え字との添え字から、左辺のの添え字を作る操作であることをeinsum
に教えている。本来の縮約記法は和記号を省略することを意味したが、NumPy.einsumの適用範囲は拡張されている。さらに、次の例を考えてみたい。
(7)
ここで、はベクトルである。成分表示すると
(8)
となる。ここで、添え字について縮約記法を用いた。einsum
を用いると次のよう書ける。
1 2 3 4 5 6 |
# 2つの行列を定義する。 x = np.array([[1, 2, 3], [1, 2, 3]]) y = np.array([[1, 3], [1, 3]]) z = np.einsum("ni,nj->ij", x, y) assert(np.all(z == np.array([[2, 6], [4, 12], [6, 18]]))) |
もう、説明は必要ないであろう。
最後に、複雑な例を挙げて終わりにしたい。
(9)
上の式は次式の縮約記法である。
(10)
コードは以下の通り。
1 2 3 4 5 |
x = np.arange(16).reshape(2, 2, 2, 2) y = np.arange(8).reshape(2, 2, 2) z = np.arange(16).reshape(2, 2, 2, 2) w = np.einsum("iabc,abc,abcj->ij", x, y, z) assert((2, 2) == w.shape) |
まとめ
今回は、PyNum.einsumの使い方を説明した。通常であればループの入れ子を用いて書くことになる計算を1行で書くことができる。しかし、上のサンプルコードでも示したが、内積にはdot関数が、行列同士の積にはmatmulが用意されている。同じことをする関数がすでに存在するのであれば、既存関数を使った方が高速であるようだ。速度比較についてはこちらのサイトが詳しいので参照してほしい。