【何倍？】Numba CUDAをつかって超高速でFitz-Hugh 南雲方程式の位相場をかく【万倍】

概要

数値積分による常微分方程式の数値的な求解を500×500個の初期値それぞれにたいして約6万ステップで行った。
通常のPythonは論外として、NumbaをつかったCPUレベルでの高速化をまず行ったところ、一晩かかった。しかしNumba CUDAを使って10秒ほどで計算が完了した。CUDAに入門してはや1週間ほど、めざましい成果が出たのはこれが初めてだったので記録がてら。

#結果
CPUでひと晩かけてじっくりと計算した位相場

GPUで10秒で計算した位相場

#結論
無数の初期値にたいして全く同じ演算を行う場合GPUの処理方式がベストマッチしていた。つかうべし。

#ソースコード
CUDAツールキットとかは一通り入っている前提

1. 方程式のリミットサイクルを固定刻み幅の4段Runge-Kutta法でもとめる
2. リミットサイクル上の点の値(x,y)が求まる。適当な点をえらび、またその点と(ほぼ)同じ点に帰ってくるまでの経過時間を位相とする。刻み幅は変えていないので点の個数が経過時間とそのまま対応する。
3. リミットサイクルを含んだ平面の格子点からn周期ぶんRunge-Kutta法をまわす
4. それぞれの格子点がたどりついた点と、2で求めたリミットサイクル上の点とをコサイン類似度で比較
5. もっとも類似度が高い点と同じ位相をその格子点の位相として、色で対応付ける。

理論についてはよくわからん場合詳しくはこちら
https://www.notion.so/4-1-9bead7cd44fb443f8da408e9bc11c0bc

3-4のところの処理

cpu.py

import numpy as np
import matplotlib.pyplot as plt
from numba import jit
r=500
@jit('f8(f8, f8)')
def fhn_xdot(x, y):
    return x*(x+0.1)*(1-x)-y
@jit('f8(f8, f8)')
def fhn_ydot(x, y):
    return (x-0.5*y)/100
#Tは周期(今回は12648×0.01秒), lcvecotrはリミットサイクル上の点の座標がはいったndarray(T×2),lcnormはその点の原点からの距離(T×1) 
for i in range(r):
    x = np.linspace(-0.5, 1, r)[i]
    for j in range(r):
        y = np.linspace(-0.1, 0.3,r)[j]
        x2, y2, _ = rk4th_2d(fhn_xdot, fhn_ydot, x, y, T)#適宜積分する。
        col = np.argmax([e-mail]([x2[-1], y2[-1]])/(lcnorm*np.sqrt(x2[-1]**2+y2[-1]**2)))/T#コサイン類似度
        pf[i][j] = col

forは入れ子にしなくともいいのかもしれない。詳しい人教えて下さい

gpu.py

from numba import cuda

#デバイス関数を定義しておく
@cuda.jit('f8(f8, f8)', device=True)
def fhn_xdot(x, y):
    return x*(x+0.1)*(1-x)-y
@cuda.jit('f8(f8, f8)', device=True)
def fhn_ydot(x, y):
    return (x-0.5*y)/100

@cuda.jit
def rk4_gpu_fhn(pf, mx, my):
    #x→何番目の振動子か?、y→状態量はどれか?
    k, j = cuda.grid(2)
    x, y = mx[k], my[j]
    dt = 0.01
    for _ in range(12648*5):
        k_1x = fhn_xdot(x,y)
        k_1y = fhn_ydot(x,y)
        k_2x = fhn_xdot(x+k_1x/2*dt,y+k_1y/2*dt)
        k_2y = fhn_ydot(x+k_1x/2*dt,y+k_1y/2*dt)
        k_3x = fhn_xdot(x+k_2x/2*dt,y+k_2y/2*dt)
        k_3y = fhn_ydot(x+k_2x/2*dt,y+k_2y/2*dt)
        k_4x = fhn_xdot(x+k_3x*dt,y+k_3y*dt)
        k_4y = fhn_ydot(x+k_3x*dt,y+k_3y*dt)
       
        x+=((k_1x+2*k_2x+2*k_3x+k_4x)/6*dt)
        y+=((k_1y+2*k_2y+2*k_3y+k_4y)/6*dt)
    pf[k, j] = x, y
    

mx1 = np.linspace(-0.5, 1, r)
my1 = np.linspace(-0.1, 0.3, r)
griddim=r
blockdim=1, r
pf = np.empty((r, r,2))
rk4_gpu_fhn[griddim, blockdim](pf, mx1, my1)

#いろつける
pfc = np.empty((r, r))
for i in range(r):
    for j in range(r):
        col = np.argmax(lcvector@pf[i][j]/(lcnorm*np.linalg.norm(pf[i][j])))/T#コサイン類似度
        pfc[i][j] = col

cpuでは、添字をつかってじゅんぐりに積分していくようなことしかできなかったのだけど、CUDAのプログラミングモデルでは｢任意の添字にたいする一般的な処理｣という形で書くことができて、blockdim×griddimぶん同時にその処理を走らせることができる。｢任意の添字｣は｢cuda.grid(2)｣で取得している。

可視化も結構めんどくさかったので参考で書いておく

visualize.py

cm = plt.get_cmap("jet")
fig=plt.figure(dpi=200)
ax = fig.add_subplot(1, 1, 1)
ax.plot(px1, py1)
ax.set_xlim(-0.5, 1)
ax.set_ylim(-0.1,0.3)
ax.set_xlabel("u")
ax.set_ylabel("v")

xlim = ax.get_xlim()
ylim = ax.get_ylim()
mp=ax.imshow(pfc.T*2*np.pi, extent=[*xlim, *ylim], aspect='auto',vmin=0, vmax=2*np.pi, origin="lower", cmap=cm)
cbar = fig.colorbar(mp, ticks=[0, np.pi, 2*np.pi])
cbar.ax.set_yticklabels([r'0', r'$\pi$', r'$2\pi$'])