Pythonでforをまわすような数値計算を行うとmatlabなどと同じで非常に遅い。
それを克服するためにpypyやcythonなどがある。しかし基本的にコードを書き換える必要があるので、もっと楽をしたい。
最近numbaというものを知った。ここのサンプルにある例だと、
from numba import autojit
@autojit
のおまじないだけで関数が劇的に早くなる。
少し、自分で例で実験してみた。
全球のSSTに対して、NINO3の相関をとるということをやってみる(これはnumbaの実験のために作った例であり、相関をとるだけならもっとエレガントな方法があるだろう)。
そのノートブックが
http://nbviewer.ipython.org/5396344
説明
(1) 図をプロットするためのおまじない。
(2) SSTアノマリとNINO3を読む関数を定義。
元データは NOAA OISST v2でmonthly dataのsst.mnmean.ncを使う。
データはPython Interface to GrADSを使って読む。
1982から2010年のデータから気候値からのanomaly(Bin Guan's GrADS Script Libraryのdeseasonを使った。古いバージョンを使ったので、今は使い方が違うかも知れない)とNINO3 SSTアノマリを定義し、返す。
ここらへんは以前の記事も参照。
/usr/local/lib/python2.7 なんちゃらというメッセージは私の環境でインストール上のゴミなので気にしないで。
(3) 上で定義した関数を呼んでいる。
(4) 時系列xと二次元時系列(水平2次元と時間の3次元)から2次元の相関係数を返す関数。
from numba import autojit
@autojit
の部分がnumba.
(5) ベンチマーク
この場合、5.7秒かかっている。
もし上で@autojitを付けない場合、7.34秒だった。
つまりnumbaを使っていることで、確かに速くなっているが、劇的にではなかった。
これは、forでまわしている中身がプリミティブな式ではなくて、numpyの関数corrcoefであるからだろうか。
ちなみに、欠損値処理は何もやっていないので、warningが出ている。この場合はプロットするときに陸地のデータはマスクするので、これで構わない。
欠損値処理もできるma.corrcoefを使うとかなり遅くなる。
numbaを使わない場合で50.2秒、numbaを47.5秒だった。
(6) 実際に相関係数を求める。
(7) land sea maskを読む関数
(8) land sea maskを読む
(9) プロット
エルニーニョに典型的な馬蹄形が見られる。インド洋ダイポールへの相関も見られる。
結論
今回試したケースでは、劇的というほどではないが、ちょっとした追加で、確かに速くはなった。場合によっては劇的に速くなる場合もあるだろう。。今後の進展に期待したい技術である。