pythonを使って網羅的計算 〜高速計算を求めて

というなかなか自己矛盾なタイトルでの投稿。 pythonはスクリプトの書く時間を短縮するのが利点であって、pythonに計算をやらせるのはそんなに早くないことが多いので。
だからと言ってふと思い立ったときにCを書くかと言うとちょっと手が渋るのも事実。 この投稿ではそんな大量計算をpythonで頑張ってやってみた時の経験について書きたい。

私的結論から書いておくと道は二つ。 一つはやはり面倒臭がらずにC系統で書く。 pythonを使うことにこだわるならpython的な高速化は求めずに、愚直にスクリプトを書いてからcythonに書き直したりnumba@jitで高速化を目指す。 このどちらかの方が結果として早くなるんじゃないかなと思った。

しかしこの投稿で書くのはそう言う方法ではなく、なるべくpython的に頑張ろうとした時の試行錯誤をいくつか書いていきたい。 もちろんpythonianな人たちはもっとエクセレントな方法を知っているだろう。 気づいたことがあったら教えていただけるととても助かる。

pythonでどのような網羅的計算をしたかったか

ちょっとここの説明が長いので、この章を飛ばしてもらってもなんとなくわかると思う。
4つのパラメータa,b,c,dがある式。 この式は3つの整数変数h,k,lの総組み合わせに対してそれぞれ解を吐き出すような式。 例えば

s = h*a+k*b+c*l+d

といった感じで。 実際はもう少し複雑な計算をさせているが。 いずれにしろ10個ずつ整数変数があるとしたら一つのa,b,c,dに対して1000個のs値が計算されることになる。

一方でこの時の実験値が8個。 この実験値と上記で計算された値sとが良く合うパラメータa,b,c,dのセットを探したいというのがやりたかったこと(どのsと実験値を合わせても良い)。
ある程度情報があって最小二乗法などで解けるなら何も問題がない。 そうでない場合はちょっと考えないといけない。

いくつかある方法の中で一番愚直な方法は、とりうるパラメータa,b,c,dを細かい刻みで動かして全ての値のを網羅的に計算する。 その中で一番実験値と合う値を引き算で算出。 その値を叩き出したパラメータa,b,c,dを求めるという方法だ。

プログラムの構成

a,b,c,dのパラメータ変数を適当な間隔で定義する。 各a,b,c,dから式の値を計算して、実験値との差分を計算する関数を回す。 差分が例えば0.001以下のパラメータセットを出力する。 そんな単純プログラム。 しかしa,b,c,dを増やせば増やすほどガンガンと計算量が増えるプログラムだった。

for文を避けてnumpyで頑張る

上記のようなプログラムを書きたい時はCだったらfor文をガスガス回すと言うのは正しい戦略だと思う。 わかりやすいし早いし。
ただpythonはfor文が割と遅いという特徴がある。 for文を如何に避けるかはpythonで重たい計算をするときの一つの方針にはなる。
と言うわけでまず4つのパラメータに対して式を計算するのにfor文を回すのではなく、numpyのarrayに収納してみることをトライ。 どの程度速度が変わるかをまずは単純計算で比較。 4つのパラメータを10個ずつ取ったとき。 まずは単なる網羅的足し算での比較を。

for文:

import time
import numpy as np

def main():
 start = time.time()
 for a in range(10):
 for b in range(10):
 for c in range(10):
 for d in range(10):
 a+b+c+d
 elapsed_time=time.time()-start
 print ('elapsed_time',elapsed_time)
 return
結果
elapsed_time 0.00111 [s]

array文:

import time
import numpy as np

def main():
 start = time.time()
 a = np.linspace(0,9,10)
 b = np.linspace(0,9,10)
 c = np.linspace(0,9,10)
 d = np.linspace(0,9,10)

 ar1 = np.array(np.meshgrid(a,b,c,d)).T.reshape(-1,4)
 np.sum(ar1,1)

 elapsed_time=time.time()-start
 print ('elapsed_time',elapsed_time)
 return
結果
elapsed_time 0.00107 [s]

あれ? ちょっとだけ早いけど10000回の計算だと大して変わらない。 しかもこのnumpyのarrayの組み方は弱点がある。 網目を大きくするとarrayがとっても大きくなる。 その分だけメモリを食うという弱点。
というわけで網目を30x30x30x30にした場合。

for: elapsed_time 0.0681 [s]
array: elapsed_time 0.1256 [s]

網目を50x50x50x50にした場合(6250000回計算)。

for: elapsed_time 0.51124 [s]
array: elapsed_time 1.07741 [s]

というわけでこれだけだと実はforの方が良かった。

さてこれをちょっとだけ複雑な計算にしてあげるとどうなるか。 各30パラメータの例で、変更は二乗の和をとるだけ。 変更と結果は下記の通り。
for:
a**2+b**2+c**2+d**2
elapsed_time 0.90076 [s]

array:
ar1 = ar1*ar1
ar2 = np.sum(ar1,1)
elapsed_time 0.13798 [s]

forは二乗を取るだけでかなり遅くなる。 一方でarrayは大して変わらない。 というわけでここでだいぶarrayの方が早くなった。

とはいえこういう単純計算だったらもっと早く書く方法があるのかな。 pythonianな方に教えて頂きたいところ。

ループで使用する関数内での変数宣言を避ける

これは回数の少ないループだとむしろ中に書いておいた方がわかりやすいかもしれない。 しかし100万回計算するループの中に変数宣言が入っていたりするとpythonはとても遅くなる。
普段はあまり使いたくないgrobal変数だけど、この場合は使用を検討してみると早くなるかもしれない。 grobal変数は、関数の外でgrobalに続けて宣言すればよい。

grobal h
h = np.linspace(0,5,6,dtype=int)

など。

argminよりargpartitionを使う

本題に戻って私がしたかったことは、実験値と計算値を比べること。 これは単純に差分を取ってあげて、そこから最小値を抜き取ればよい。 ここでパッと考えるとargminを使いたくなるところ。 しかしargpartitionという関数があり、こちらの方が早い。

dif2 = np.absolute(slist - dobs)
min =np.argpartition(dif2,0,axis=0)

こんな感じで差分を取ったarrayの最小のインデックスが抜き出せる。

@numba.jit

時折とても強力に高速化してくれる。 けれどむしろ遅くなる場合もある。 全体にかけて効くときもあるし、関数1個だけにかけた方がよく効くときもある。 よくわからないけれど、効いてくれたら嬉しい機能。 試す価値はあり。

数学関数はmathかnumpyか

物によってはmathのが早いし、numpyの方が早かったこともある。 というかあまり変わらない。 最後の最後にどっちのが早いか見てみてもよいかも。 100万回も計算すれば違いが出るかもしれない。

anacondaのpythonを使う

私がanacondaのpython3に乗り換えたのはこれが原因。 なぜかちょっとだけ計算が速くなる。 もっと早いのがあるなら教えて欲しい。

計算量を減らす努力

小さな努力が大きな結果に跳ね返る場合もある。 計算に対称性などがあって計算が減らせる場合は減らしておく。 計算量が多い場合だとパラメータが減ったり式が単純化されるだけで結構な違いに跳ね返る。 なんだかんだでマニュアル努力が大切。

終わりに

というわけでいくつか私がpythonの高速化を計った時に使った方法を書いてみた。 私は結局目的に使い物になるくらいのところ(1回の試行が帰宅前に回し始めて朝には終わっている)で諦めてしまったのだけど。
最終盤のプログラムでもCで適当に書いた似たようなプログラムに数十倍の速度差をつけられていた。 本職pythonの人がどのくらい早くできるのかは気になるところではある。

cythonへの変換も試した。 しかしいろいろnumpyで書いていたせいか、変数の宣言などcythonを単純に使っても速度を上げることはできなかった。
というわけである程度計算量が必要なプログラムを書く時は、最初からcythonを意識した書き方をしておく。 もしくはcythonフォーマットで書き始めるというのが正しい道なのかな。

関連記事

pythonのまとめ

 

D