numpy で2次元の行列計算と for とのパフォーマンス比較

この記事は、次のような2次元の行列計算を Python の numpy ではどう書くの？というものである。

for i in range(0, ni-1):
    for j in range(0, nj-1):
        y[i, j] = x[i, j] + x[i+1, j] - x[i+1, j+1] * x[i, j+1]

後で、パフォーマンスの比較もする。Python のバージョンは、3.7.4。anaconda3-2019.10 を利用している。

前置き

COVID-19 + 転職に伴い、東京の自宅にスタックし続ける日々に、暇を持て余し、新しい言語を覚えることにした。次の職場では Python がメインのようなので、Python を。身銭を切らないと本腰が入らないたちなので、Udemy でコースを取って1週間くらいかけて一通りの基礎を学んだ。長らくプログラミングはやってきたが、Python は、経験 2 週間程度のビギナーだ。

一通りやってみて、Python で Web 開発はできそうな気がするが、せっかく Python をやるので、データ分析や機械学習などをやってみたい。というわけで、はるか昔にやった1次元移流方程式の解法を Python で書き直してみた。

▶ 1 Dimensional Flow Calculation by Python

細かいところはチェックしておらず、諸々手抜きだが、望む結果が得られたので、初めてにしてはと満足した。しかし、何というか達成感？みたいなものが感じられなかったので、2次元の計算に手を出してみたが、行列計算が極めてダルい。これ以上このまま書きすすめるべきではないと私の Ghost が囁く。そして、
Python では、

• for を書いたら負け
• for をネストすると実行速度が遅くなる。残酷なまでに。

ともっぱらの噂だ。
解決方法として、numpy 等を使うとのことだったので、使ってみたが、混乱するばかりで、すぐには身につかなさそうなので、メモとして残しておく。

この記事を読んで得られるのは「行列の計算は numpy を使え（for を使ってはいけない）」という読まなくても分かることだけだ。それ以外については、読んでもピンとこないと思うので、理解するのであれば、同じように実際に手を動かして確認して、手に馴染むまで繰り返してみると良い。

list と numpy array の書き方の比較

Python には配列を表すものとして、list, tuple, dict などがあり、2次元の配列の場合、他の言語からやってきたものとしては、list で、x[i][j]などとなっていると分かりやすい。が、数学的表記からすると x[i, j]の方が分かりやすい。numpy の array では、後者のような書き方になっている。実際に書いてみれば分かるが、何度も x[i][j] と書くのは面倒臭いし、見にくいので、x[i, j] のように書けるのはありがたい。

では早速。例えば、ndarray 形式の x, y を計算する、次のような式があったとする。

y[i, j] = x[i, j] + x[i+1, j] - x[i+1, j+1] * x[i, j+1]

特に意味のある式ではないが、数物系に限らず数値計算処理では、これ的な計算をよく目にするはず。これを i と j のループでネストしてみると、

for i in range(0, x.shape[0]-1):
    for j in range(0, x.shape[1]-1):
        y[i, j] = x[i, j] + x[i+1, j] - x[i+1, j+1] * x[i, j+1]

となる。非常に分かりやすい。最大分割数を ni, nj などとするなら、

for i in range(0, ni-1):
    for j in range(0, nj-1):
        y[i, j] = x[i, j] + x[i+1, j] - x[i+1, j+1] * x[i, j+1]

という意味だ。

さて、numpy では、この計算は以下のように書くのがスマートらしい。

y[0:-1, 0:-1] = x[0:-1, 0:-1] + x[1:, 0:-1] - x[1:, 1:] * x[0:-1, 1:]

え？なに？

じっくり見てみれば、まぁ何？スライスされている範囲と式が対応している？そんな気がしなくもない。
自由自在に書ける気はしない。

やってみる

いったん、本当に合ってるか確認してみる。

import numpy as np


x = np.random.randint(10, size=(100, 50))
y = np.zeros((100, 50))

for i in range(0, x.shape[0]-1):
    for j in range(0, x.shape[1]-1):
        y[i, j] = x[i, j] + x[i+1, j] - x[i+1, j+1] * x[i, j+1]

z = np.zeros_like(y)
z[0:-1, 0:-1] = x[0:-1, 0:-1] + x[1:, 0:-1] - x[1:, 1:] * x[0:-1, 1:]

if (z == y).all():
    print('合ってる')
else:
    print('違う')

実行結果は、……合ってる。range の 0 は必要ないが、理解のためあえて。
他にも

• 0 から ni, nj (x.shape[0], x.shape[1])
• 1 から ni, nj (x.shape[0], x.shape[1])
• 1 から ni-1, nj-1 (x.shape[0]-1, x.shape[1]-1)
• 2 から ni-3, nj-3 (x.shape[0]-3, x.shape[1]-3)

などをやってみる。

# 0 から ni, nj (x.shape[0], x.shape[1])
for i in range(0, x.shape[0]):
    for j in range(x.shape[1]):
        y[i, j] = x[i, j] + x[i, j]

z[0:, 0:] = x[0:, 0:] + x[0:, 0:]

# 1 から ni, nj (x.shape[0], x.shape[1])
for i in range(1, x.shape[0]):
    for j in range(1, x.shape[1]):
        y[i, j] = x[i, j] + x[i-1, j] - x[i-1, j-1] * x[i, j-1]

z[1:, 1:] = x[1:, 1:] + x[0:-1, 1:] - x[0:-1, 0:-1] * x[1:, 0:-1]

# 1 から ni-1, nj-1 (x.shape[0]-1, x.shape[1]-1)
for i in range(1, x.shape[0]-1):
    for j in range(1, x.shape[1]-1):
        y[i, j] = x[i, j] + x[i-1, j] - x[i+1, j-1] * x[i, j+1]

z[1:-1, 1:-1] = x[1:-1, 1:-1] + x[0:-2, 1:-1] - x[2:, 0:-2] * x[1:-1, 2:]

# 2 から ni-3, nj-3 (x.shape[0]-3, x.shape[1]-3)
for i in range(2, x.shape[0]-3):
    for j in range(2, x.shape[1]-3):
        y[i, j] = x[i, j] + x[i-1, j] - x[i+1, j-1] * x[i, j+1]

z[2:-3, 2:-3] = x[2:-3, 2:-3] + x[1:-4, 2:-3] - x[3:-2, 1:-4] * x[2:-3, 3:-2]

書き方のまとめ

んー、つまりこういうことか？

# i, j を n から ni-m, nj-m までループ
# x[i, j], x[i-ln, j], x[i+ln, j-ln], x[i, j+ln] などを計算

for i in range(n, x.shape[0]-m):
    for j in range(n, x.shape[1]-m):
        y[i, j] = x[i, j] + x[i-ln, j] - x[i+ln, j-ln] * x[i, j+ln]

z[n:-m, n:-m] = x[n:-m, n:-m] + x[n-ln:-m-ln, n:-m] - \
    x[n+ln:-m+ln, n-ln:-m-ln] * x[n:-m, n+ln:-m+ln

まとめの散らかりっぷり。読みづらいこと極まりないので、表にしてみる。

start	end	i	i-1	i+1	i-ln
0	ni	0:	-	-	-
0	ni-1	0:-1	-	1:	-
1	ni	1:	0:-1	-	-
1	ni-1	1:-1	0:-2	2:	-
n	ni-m	n:-m	n-1:-m-1	n+1:-m+1	n-ln:-m-ln

※ ni = x.shape[0]

あってるかな？どう書いても分かりにくい。Python 初心者の私としてみては、メンテナンス性の観点より利用すべきでないというのが率直な感想だ。少なくともこの時点では、そう思ってました。

処理速度の比較

計算が1億回回るようにループして計測してみた。最初にやった1次元の計算ではメッシュを細かくして、2,000万回程度計算したので、1億回は多いように見えて、実際の計算では全然多くはない。
端数を合わせるのは面倒だったので、そのまま。計測結果はそれぞれ 10 回ずつ実行した平均。

loop	for (sec)	slice (sec)	ave: s/f %	max: s/f %	min: s/f %	median: s/f %
100,007,840	148.3013	1.3558	0.91%	0.95%	0.88%	0.91%
10,020,160	13.4111	0.1179	0.88%	1.00%	0.72%	0.86%
1,024,160	1.4228	0.0146	1.03%	1.20%	0.84%	1.04%
110,720	0.1536	0.0017	1.10%	1.35%	0.96%	1.08%

圧倒的じょのいこ！！！ぺっぺっ

メンテナンス性の観点云々のレベルではない。ざっくり1％程度で終わるということは、1時間かかる計算なら、36秒で終わるということだ。1日なら？1週間なら？？1ヶ月なら？？？読みやすさのために、for を使うなんて選択肢は絶対にない。コメントにでも書いておいてくれ。

ただし、この numpy による計算が速いかどうかは別の話だ。numpy は C言語および Fortran で実装されているらしいが、Fortran で do ループをネストした方が速ぇーしっ！てことはあると思う。検証はしていない。

以上。もっとこうやった方が良いよなどはあるかと思うが、いったんこれにて。