NumPy・SciPyによる高速化の実力を見る

  • 19
    いいね
  • 0
    コメント

あらまし

以前の記事でNumPy・SciPyの高速化にまつわる事柄を書きました:

NumPyを用いた数値計算の高速化 : 基礎
NumPy・SciPyを用いた数値計算の高速化 : 応用その1
NumPy・SciPyを用いた数値計算の高速化 : 応用その2

ホントに早くなってるの?ちゃんと調べてみましょう.

調査方法

Pythonによるオレオレ実装と比較します. 速度よりシンプルさを重視した実装との比較なので正当な評価とは言い難いかもしれません. Pythonはanaconda3, 時間計測にはIPythonの%timeitを使用します.

--実行環境--
OS : Ubuntu16.04 LTS 64bit
Python : anaconda3-4.1.1
CPU : Intel Corei5 3550 (4-core / 4-thread)

リストの初期化

たとえば行列の初期化です.

before
N = 3000
a = [[i + j for i in range(N)] for j in range(N)]

こんなふうに内包表記で書いたとしましょう. 実行時間は661ms. 一方でNumPyのarangeでは

after
a = np.arange(0, N) + np.arange(0, N).reshape(N, 1)

15.5msでした. 100倍とは言わないですが, 40倍くらいは速いですね.

行列積

線形代数の演算はいろいろありますが, 簡単な行列積で試してみましょう. 200×200の行列を2つ用意して積を計算します:

setting
import numpy as np
from numpy.random import rand

N = 200
a = np.array(rand(N, N))
b = np.array(rand(N, N))
c = np.array([[0] * N for _ in range(N)])

行列積の定義のまま実装すると

befor
for i in range(N):
    for j in range(N):
        for k in range(N):
            c[i][j] = a[i][k] * b[k][j]

でしょうか. 実行時間は7.7s.

NumPyのユニバーサル関数を用いると

after
c = np.dot(a, b)

速度以前に単純なのがいいですね. 実行時間はなんと202us. 何倍とか言うのもおこがましいくらい一瞬ですね. これがMKLのBLASによるマルチコア処理の威力です. 1000×1000でも22.2ms. ここまでくるとforループの実装では手に負えなくなります.

微分

$\sin x$を微分してみましょう. 空間の分割数は100000にしています:

setting
import math as m

def f(x):
    return m.sin(x)

def g(x):
    return np.sin(x)

N, L = 100000, 2 * m.pi
x, dx = np.linspace(0, L, N), L / N

微分の定義をそのまま. 精度はこの際気にせずいきましょう:

before
diff = []
for i in x:
    diff.append((f(i + dx) - f(i)) / dx)

実行時間は91.9ms. もしかしたらappendが遅い原因かもしれないですね. 一方NumPyでは

after
diff = np.gradient(g(x), dx)

シンプルな上に, 戻り値の要素数も減りません(数値計算に際してはこれがどれだけ有り難いか!). 実行時間は8.25ms. 軽く10倍程度は速いですね. ちなみに

g = np.vectorize(f)

とすることで, 引数や戻り値をndarray仕様に変換することもできます.

積分

ちょっとコアな積分を用意しましょう:

\int_{-\infty}^{\infty} dx\int_{-\infty}^{x} dy\; e^{-(x^2 + y^2)}

$y$の積分区間に$x$が含まれているような二重積分です. 被積分関数と空間の分割数を

setting
def h(x, y):
    return np.exp(-(x**2 + y**2))

N, L = 2000, 100
x, dx = np.linspace(-L/2, L/2, N), L / N
y, dy = np.linspace(-L/2, L/2, N), L / N

のようにセットします. 二重積分を定義通り実装すると

before
ans = 0
for index, X in enumerate(x):
    for Y in y[:index+1]:
        ans += dx * dy * h(X, Y)

でしょうか. 実行時間は5.89s. しかしこれはあまりにも精度が悪いです. 自前で書くならSimpson積分にすべきですが, その場合はコードがそこそこ煩雑になります. また, Simpson積分でも広義積分には「空間を十分広く取る」という対応しかできません. しかし差分を細かくすると実行時間は$O(N^2)$のオーダーで増加していきます. 一方SciPyのdblquadではその全ての問題を解決できています:

after
from scipy.integrate import dblquad
ans = dblquad(h, a = -np.inf, b = np.inf, gfun = lambda x : -np.inf, hfun = lambda x : x)[0]

$x$の積分区間が$[a, b]$, $y$の積分区間が$[{\rm gfun}, {\rm hfun}]$です. 適応型積分で, 誤差の範囲も指定できます. dblquadの戻り値はtupleで, 絶対誤差も一緒に返って来るようです. 実行時間は51.1ms. もはや何も言うことはありません. 最高です.

固有値方程式

(おまけです. 固有値方程式の数値解法は少々マニアックなので...)

調和振動子系のSchroedinger方程式を解きましょう. 自前で実装するとなるとJacobi法でしょうか. 興味があれば調べてみてください:

before
I = np.eye(N)
H = [[0 for i in range(N)] for j in range(N)]
for i in range(N):
    H[i][i] = 2.0/(2*dx**2) + 0.5*(-L/2+dx*i)**2
    if(0 <= i+1 < N):
        H[i][i+1] = -1.0/(2*dx**2)
    if(0 <= i-1 < N):
        H[i][i-1] = -1.0/(2*dx**2)
H = np.array(H)

# Jacobi法
flag = True
while(flag):
    # 非対角成分の最大値及びインデックスを調べる
    maxValue = 0
    cI, rI = None, None
    for j in range(N):
        for i in range(j):
            if(maxValue < abs(H[i][j])):
                maxValue = abs(H[i][j])
                rI, cI = i, j

    # 収束判定
    if(maxValue < 1e-4):
        flag = False
    # print(maxValue)

    # 回転行列の用意
    theta = None
    if(H[cI][cI] == H[rI][rI]):
        theta = m.pi/4
    else:
        theta = 0.5*m.atan(2.0*H[rI][cI]/(H[cI][cI]-H[rI][rI]))
        J = np.eye(N)
        J[rI][rI] = m.cos(theta)
        J[cI][cI] = m.cos(theta)
        J[rI][cI] = m.sin(theta)
        J[cI][rI] = -m.sin(theta)

    # 行列演算
    H = np.array(np.matrix(J.T)*np.matrix(H)*np.matrix(J))
    I = np.array(np.matrix(I)*np.matrix(J))

# 固有値・固有ベクトルの格納
v, w = I.transpose(), []
for i in range(N):
    w.append([H[i][i], i])
w.sort()

非対角項の最大値が十分小さくなったところで収束です. 固有値が昇順になっていなかったりと, あまり便利ではありません. 実行時間は15.6s. これ以上分割数を大きくするともう苦しいです. NumPyでは

after
# 系の設定
L, N = 10, 80
x, dx = np.linspace(-L/2, L/2, N), L / N

# 運動項K
K = np.eye(N, N)
K_sub = np.vstack((K[1:], np.array([0] * N)))
K = dx**-2 * (2 * K - K_sub - K_sub.T)

# ポテンシャル項
V = np.diag(np.linspace(-L/2, L/2, N)**2)

# エルミート行列の固有値方程式
# wが固有値, vが固有ベクトル
H = (K + V) / 2
w, v = np.linalg.eigh(H)

実行時間は1.03ms. 固有値方程式くらい複雑なアルゴリズムになると, もう色んな意味で勝ち目が無いですね. これもBLASがさぞ大暴れしていることでしょう. CだとLapackかGSLに頼ることになるのでしょうが, それも大変ですね.

おわりに

このくらいで十分でしょう. NumPy・SciPyがいかに高速に動作するか, またシンプルに記述できるかがわかると思います. もしC/C++で上のような計算を外部ライブラリに頼らずに書こうと思うと, それぞれの章の最初で書いたようなコードになります. また, C/C++の外部ライブラリは経験上, 触りにくいものが多い印象です.

あくまで参考ですが速度差をまとめると以下のような具合です:

リストの初期化 : 661ms15.5ms
行列積 : 7.7s202us
微分 : 91.9ms8.25ms
(2重)積分 : 5.89s51.1ms
固有値方程式 : 15.6s1.03ms

圧倒的です. 計算屋さんがPythonを使う価値は十二分にあると思います.