Julia tips #3: ODE solverの速度比較 with Python

JuliaではSundials.jlをインストールするだけで、SundialsのODEソルバ、CVODEが使えます！！（マニアックすぎ？）

未知システムのパラメータ推定をする際、数値積分の区間をうまいこと刻まないと値が飛びすぎて計算できなかったり、刻む必要のないところで刻んでしまって計算不能に陥ることがあります。

そこでCVODEなどのAdaptive stepなODEソルバが必要なわけです。

何個か種類がありますが、有名どころでは

• LSODA
• CVODE

が2強といったところで、Pythonでよく使われるodeintもadaptive stepです（内部的にはLSODEとかいうやつらしいけど...良く知りません）。

じゃあ、JuliaでCVODEを呼び出して使うのと、Pythonでodeint使うのはどっちがどんくらい速いのよ？と思ったので比べてみました。

なお、Juliaのソルバについては比較してくれている方がいました。

CVODEの圧勝です。

Juliaのコードもこちらから借りてきて、Pythonと比較したいと思います。

qiita.com

実装

---

Julia（上記リンク先から部分的に拝借）

using Sundials

# 非線形パラメータ
mu = 0.2
# 関数定義
function vdp_sun(t, u, du)
    # dx/dt
    du[1] = u[2]
    # dy/dt
    du[2] = mu * (1.0 - u[1]^2.0) * u[2] - u[1]
    du
end
# パラメータ設定
u0 = [0.2; 0.2]
t = [0.0:0.1:10.0;]

Python (3.6)

import numpy as np
from scipy.integrate import odeint

# 非線形パラメータ
mu = 0.2
# 関数定義
def vdp_sun(u, t):
    # dx/dt
    dxdt = u[1]
    # dy/dt
    dydt = mu * (1.0 - u[0]**2.0) * u[1] - u[0]
    return [dxdt, dydt]
    
# パラメータ設定
u0 = [0.2, 0.2]
t = np.arange(0,10.1,0.1)
# メインの計算
res = odeint(vdp_sun, u0, t) 

速度比較

Juliaは1ループごとに出力、Pythonのtimeitではループにかかった平均時間を出力していますが、大雑把な比較はこれで良いでしょう。
見やすさのため10回だけ回しています。

Julia

res = Sundials.cvode(vdp_sun, u0, t);
for i in 1:10
    @time res = Sundials.cvode(vdp_sun, u0, t);
end

Python (3.6)

%timeit -n 10 odeint(vdp_sun, u0, t)

結果

Julia

  0.000206 seconds (1.81 k allocations: 43.734 KiB)
  0.000132 seconds (1.81 k allocations: 43.734 KiB)
  0.000130 seconds (1.81 k allocations: 43.734 KiB)
  0.000117 seconds (1.81 k allocations: 43.734 KiB)
  0.000118 seconds (1.81 k allocations: 43.734 KiB)
  0.000115 seconds (1.81 k allocations: 43.734 KiB)
  0.000115 seconds (1.81 k allocations: 43.734 KiB)
  0.000115 seconds (1.81 k allocations: 43.734 KiB)
  0.000116 seconds (1.81 k allocations: 43.734 KiB)
  0.000124 seconds (1.81 k allocations: 43.734 KiB)

Python (3.6)

716 µs ± 5.63 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

まとめると以下のようになります。

• Julia: 128.8 µs ± 26.43 µs
• Python: 716 µs ± 5.63 µs

Juliaの方が5倍ほど速い！

シミュレーションを何千回、何万回と繰り返すものにとっては5倍の高速化は結構大きいと思います。

予想ではもっと変わるかなと思っていたのですがこんなものでした。

Pythonを高速化

Pythonを高速化したら、Sundialsに迫るのでは？？

と考えたので、numbaで高速化したPythonで戦ってみた。

odeintに渡す微分方程式をjit+型指定すると速くなる。

from scipy.integrate import odeint
import numpy as np
from numba import jit

# 型指定jit, tが逐次的に渡されることに注意
@jit("float64[:](float64[:], float64)", nopython=True)  
def vdp_sun_jit(u, t):
    # dx/dt
    dxdt = u[1]
    # dy/dt
    dydt = mu * (1.0 - u[0]**2.0) * u[1] - u[0]
    return np.array([dxdt, dydt]).astype(np.float64)
    
# パラメータ設定
u0 = np.array([0.2, 0.2]).astype(np.float64)   # 渡す型に合わせる
t = np.arange(0,10.1,0.1).astype(np.float64)   # 渡す型に合わせる
# メインの計算
res = odeint(vdp_sun, u0, t)

%timeit -n 10 odeint(vdp_sun_jit, u0, t)

結果

387 µs ± 11.6 µs

うーん、速くはなりましたが2倍行かないくらいでしたね。

まとめ

• Julia: 128.8 µs ± 26.43 µs
• Python: 716 µs ± 5.63 µs
• Python+numba: 387 µs ± 11.6 µs

やっぱSundials凄い。

数値計算ではJuliaを使おう！