numpy.linalg.solve vs. scipy.sparse.linalg.bicgstab

Introduction

量子コンピュータの流行に触発され、量子力学をもう一度勉強することに決めた私。しかし「シュレディンガー方程式を数値的に解いて、その結果の確率分布を可視化したい」なんて考えたときに、ふと思ったのです。連立一次方程式のソルバーはどれが良いのかと。

NumPyのsolve vs. SciPyのbicgstab

というわけで両者の速度を比較してみました。NumPyからは.solveメソッド, SciPyからは反復解法であるBiCGstab法をチョイスしました。
(チョイス理由は「私が過去に使ったことがあるから」程度にしか考えておらず、他により高速な解法があれば、今後そちらを試したいと思います)

実行スクリプトと各々の計測結果

numpy.linalg.solve()

以下にNumPyのsolveメソッドを利用した計算スクリプトを示します。

solve_numpy.py

'''
    solve_numpy.py
    purpose: solve Ax = b with numpy
'''

from time import time
import numpy as np


# make matrix A
a1_lis = [1, 0, 0]
a2_lis = [0, 2, 0]
a3_lis = [0, 0, 4] 
a_matrix=np.array([a1_lis, a2_lis, a3_lis])

# set b vector
b_vec = np.array([4, 5, 6])

t1 = time()

# solve Ax = b
for i in range(0, 1000):
    np.linalg.solve(a_matrix, b_vec)

t2 = time()

print(t2-t1)

一回の実行では瞬殺なので、1000回繰り返したときに要した時間を計測しています。
結果は大体0.0067でした。

scipy.sparse.linalg.bicgstab()

続いてSciPyのbicgstabメソッドを使用して同じことをしてみます。

solve_scipy.py

'''
    solve_scipy.py
    purpose: solve Ax = b with BiCGStab in scipy
'''

from time import time
import numpy as np
from scipy.sparse.linalg import bicgstab


# set matrix A
a1_lis = [1, 0, 0]
a2_lis = [0, 2, 0]
a3_lis = [0, 0, 4] 
a_matrix=np.array([a1_lis, a2_lis, a3_lis])

# set b vector
b_vec = np.array([4, 5, 6]) 

t1 = time()

# solve Ax = b
for i in range(0, 1000):
    #x_vec = bicgstab(a_matrix, b_vec)
    bicgstab(a_matrix, b_vec)

t2 = time()

print(t2-t1)

こちらも同じく1000回繰り返したときに要した時間を計測しています。
結果は大体0.16でした。

結果としてははNumPyのsolveメソッドを今後使った方が良さそうな気がしますね。

なぜNumPyのsolveの方が早いのか？

調べてみるとNumPyのsolveメソッドはLAPACKを使っているらしいです。なるほど速いわけですね。

結論と今後

今回の結論としてはNumPyのsolveメソッドに軍配が上がりました。しかし行列の大きさなどによっては、処理速度が変わってくるかもしれません。実際に量子力学の適当な問題を解いたときに実行速度の測定をしてみる必要がありそうです。