概要

• PythonコードをCythonで高速化して、Fortran90のコードに速度を近づけてみた。
• 最終的に、今回はFortran90の計算時間の2倍程度で済むようになった。(もちろん、計算課題や環境にも依る)

Cythonとは

詳しい解説は多方面でされているので、詳しいことは割愛するが、Cythonとは

• Pythonをベースにしたプログラミング言語で、ざっくり言うと「Cのデータ型を持ったPython」らしい。
• .pyxの拡張子を持つpythonコード内で、Cのデータ型をもたせて、Cythonで高速な.c形式のコードに変換してくれる。
• そして、.cのコードをCコンパイラでコンパイルし、Pythonインタプリタで呼び出して使える。

Cythonの書き方

Cythonの基本的な書き方として、
1. 変数の型宣言
2. 関数と関数の引数に型宣言
がある。

1.に関しては、.pyxコード内で、

cdef int N = 100000
cdef double bf0 = 1.
cdef double[3] bf = [0., 0., bf0]
cdef double[4][3]

のように記述する(後述で紹介するソースコードより抜粋)。cdefの後ろにコロン:を付けてインデントすることで、cdefの記述を省略して一気に型宣言することも可能。

2.に関しては、.pyxコード内で、

cdef double* rkfd(double *v):

    res[0] = e/m * (ef[0] + (v[1]*bf[2] - v[2]*bf[1]))
    res[1] = e/m * (ef[1] + (v[2]*bf[0] - v[0]*bf[2]))
    res[2] = e/m * (ef[2] + (v[0]*bf[1] - v[1]*bf[0]))

    return res

のように記述する(同じく抜粋)。関数の引数等に*が付いているのは、参照受け渡しの意味。

計算課題

今回の計算課題は、プラズマの運動方程式

\begin{align}
m\frac{d\boldsymbol{v}}{dt} &= q(\boldsymbol{E}+\boldsymbol{v}\times \boldsymbol{B}) \\
\frac{d\boldsymbol{r}}{dt} &= \boldsymbol{v}
\end{align}

とした。パラメータは、それぞれ、m: プラズマ質量、q: プラズマの電荷、r: プラズマ位置、v: プラズマ速度、E: 電場、B: 磁場、t:時間、である。r、v、E、Bについては、空間3次元のベクトル量。
これを差分化すると、

\begin{align}
\frac{\boldsymbol{v}^{n+1/2} - \boldsymbol{v}^{n-1/2}}{\Delta t} &= \frac{q}{m}(\boldsymbol{E}+\frac{\boldsymbol{v}^{n+1/2} + \boldsymbol{v}^{n-1/2}}{2}\times \boldsymbol{B}) \\
\boldsymbol{r}^{n+1} &= \boldsymbol{r}^{n} + \boldsymbol{v}^{n+1/2}\Delta t
\end{align}

となり、これを数値的に解いていくことになる。時間発展を逐次的に解いていくことになるため、Pythonには不向きなforループが必ず出てくることになる。
今回はこの逐次計算を、4次ルンゲクッタ法を用いて解いた。

ベンチマークテスト

今回は、以下の4種類の実装方法でベンチマークを行い、計算時間を比較した。
1. Fortran90のコード
2. 純粋なPythonのみのコード
3. 変数の型宣言のみを行うCythonコード
4. 変数の型宣言に加え、関数の引数と戻り値の型宣言も行うCythonコード
(Qiita上ではソースコードを割愛しますが、興味のある方は、Github[https://github.com/Hirotoshi-Uchino/cython_benchmark_gyro] に今回のソースコードがありますので、動かしてみて下さい。コードに関するコメント等も歓迎します。)

なお、今回のベンチマークテストの環境は以下の通り、
　・ PC: MacBook Pro (Retina, 13-inch, Early 2015)
　・ OS: macOS Sierra
　・ CPU: 2.9 GHz Intel Core i5
　・ メモリ 8 GB 1867MHz DDR3
　・ Fortanコンパイラ: gfortran (コンパイル時、-O3最適化オプションを使用)
　・ Python 3.5

ベンチマークテスト結果

　今回の逐次計算を10万ループした場合の結果を以下の表にまとめた。

数字は、10回の試行の平均と標準偏差から割り出したもの。

結果を見ると、純粋なPythonコードだと、300倍以上かかっていた計算時間が2倍程度まで抑えられていることがわかる。また、今回のコードは、関数呼び出しが非常に多いものとなっており、関数の引数や戻り値の型宣言まで行うことが重要であることが分かる。

まとめ

Cythonを用いることで、PythonのコードをFortran90の速度まで近づけるができた。Cythonの利用で発生した作業は、.pyから.pyxに書き換える際に、変数と関数の引数や戻り値に型宣言を行ったことと、.c → .pyxへの変換とCプログラムのコンパイル用のセットアップスクリプトを書いただけであったので、慣れると手早く利用できそうだと感じた。
Cythonを使うメリットとしては、高速化のためにわざわざ他言語に書き換えるための時間を削減できることであろう。うまく使えば、工数をかけずにPythonコードを高速化できそうである。

予告

今回の計算時間はわかったけど、計算結果(プラズマの運動はどうなるの？)って思った方もいるかと思う。これについては、可視化したものを今後紹介したい。