Julia + CUDA で熱方程式を解く

Introduction

Juliaで一次元熱方程式のDirichlet問題

\begin{align}
u_t(x,t)-cu_{xx}(x,t)&=0,& \quad &0<x<1,\ t>0, \\
u(0,t)&=g_0(t),& \quad &t>0, \\ 
u(1,t)&=g_1(t),& \quad &t>0, \\ 
u(x,0)&=f(x),& \quad &0<x<1 \\ 
\end{align}

をCUDAを援用してCrank-Nicolson法で解くプログラムを作成した．以下，$f(x)=1$，$g_0(t)=g_1(t)=0$，$c=2.0$としてソースコードを紹介し，GPUを使わなかった場合との比較を紹介する．

CUDA

CUDAとはNVIDIAが開発・提供している、GPU向けの汎用並列コンピューティングプラットフォームである(CUDA)．GeForceとかで数値計算が可能となる．

CUDAをインストールする方法については，他の記事を参照のこと．

Crank-Nicolson法

差分法を参照のこと．ただし，行列を使って説明した方がJuliaとの相性がいいので，簡単に説明しておく．

$x$変数の微分を中心差分としたとき，Dirichlet境界条件で$x$の二階微分を与える差分化は行列を用いて

D
=
\begin{pmatrix}
-2 & 1 & 0 & 0 & \cdots & 0 & 0 & 0 & 0 \\
1 & -2 & 1 & 0 & \cdots & 0 & 0 & 0 & 0 \\
0 & 1 & -2 & 1 & \cdots & 0 & 0 & 0 & 0 \\
\vdots & \vdots & \vdots & \vdots & \cdots  & \vdots & \vdots & \vdots & \vdots\\
0 & 0 & 0 & 0 & \cdots & 1 & -2 & 1 & 0 \\
0 & 0 & 0 & 0 & \cdots & 0 & 1 & -2 & 1 \\
0 & 0 & 0 & 0 & \cdots & 0 & 0 & 1 & -2 \\
\end{pmatrix}

と書ける．そこで，$h$を空間刻み幅，$k$を時間刻み幅として$r=\frac{ck}{h^2}$として

\frac{\vec{u}_{n+1}-\vec{u}_{n}}{k}=\frac{c}{h^2}D\vec{u}_{n}+\frac{c}{h^2}\vec{b}_{\mathrm{f}}
\Leftrightarrow
\vec{u}_{n+1}
=(I+rD)\vec{u}_{n}
+r\vec{b}_{\mathrm{f}}

とするのが前進差分解法であり，

\frac{\vec{u}_{n}-\vec{u}_{n-1}}{k}=\frac{c}{h^2}D\vec{u}_{n}
+
\frac{c}{h^2}\vec{b}_{\mathrm{b}}
\Leftrightarrow
(I-rD)\vec{u}_{n+1}
=\vec{u}_{n}
+r\vec{b}_{\mathrm{b}}

とするのが後退差分解法である．ここで，$I$ は単位行列，$\vec{b}$ は境界条件に応じて設定する与えられたベクトルである．そこで$0\leq \theta\leq 1$を固定して，$\theta+\eta=1$として前進差分解法を$\theta$倍，後退差分解法を$\eta$倍して和をとると

(1-\theta r D)\vec{u}_{n+1}
=
(1+\eta r D)\vec{u}_{n}
+\eta r\vec{b}_{\mathrm{f}}
+\theta r\vec{b}_{\mathrm{b}}

となる．そこで，$D_{\mathrm{f}}=I+\eta r D$，$D_{\mathrm{b}}=I-\theta r D$とすれば

\vec{u}_{n+1}
=
D_{\mathrm{b}}^{-1}
(D_{\mathrm{f}}
\vec{u}_{n}
+\eta r\vec{b}_{\mathrm{f}}
+\theta r\vec{b}_{\mathrm{b}}
)

が得られる．$\vec{u}_0$を初期値を用いて定めれば，この漸化式により熱方程式の数値計算ができる．$\theta=0.5$としたものをCrank-Nicolson法という．

数値計算アイデアのJuliaによる実装

必要なパッケージの読み込み

線形代数(LinearAlgebra)とCUDAのパッケージを読み込む．

using LinearAlgebra
using CUDA

インストールされていない場合は

using Pkg
Pkg.add("CUDA")

のようにCUDAのパッケージをインストールする．

計算するだけなら，LinearAlgebraとCUDAだけでよいが，データを出力したり，グラフを表示させるためには

using Plots
using CSV
using DataFrames

も追加で入れておく．

初期関数の設定

$f(x)=1.0$を定義する．

function f(x)
    return 1.0
end

境界条件の設定

$g_0(t)=0.0$，$g_1(t)=0.0$を定義する．

function g_0(t)
    return 0.0
end
function g_1(t)
    return 0.0
end

定数の設定

いろいろな定数を設定する．

c = 2.0
theta = 0.5
eta = 1 - theta
r = 0.4
N = 10000
h = 1.0/N
k = (r*h^2)/c
t = 1000

r は拡散数と呼ばれている，空間と時間の刻み幅の比率である(差分法，拡散数)．t は時間変数についての繰り返し回数である．

Dirichlet Laplacianの離散化と前進差分行列の定義

行列の定義については，GPUを使用せずに作成する．

function build_D_dirichlet(N)
    D = (-2)*Matrix{Float64}(I, N-1, N-1)
    for i in 2:N-1
        D[i-1,i]=1
        D[i,i-1]=1
    end
    return D
end
D = build_D_dirichlet(N)
Df = I + r * eta * D

後退差分行列の定義

Dbで後退差分の行列を生成する．そして，Dbの逆行列をDb_invで定める．

Db = I - r * theta * D
Db_inv = inv(Db)

境界条件を与えるベクトルの生成

bf と bb を境界条件を与えるベクトルとして初期化する．

bf = zeros(N-1)
bb = zeros(N-1)

差分法による次の時間ステップを計算するための関数生成

漸化式の右辺に対応する関数を生成する．この計算において，gpuを用いる．そのために，cu(Db_inv)やcu(Df)などにより，行列をGPU計算できる変数に変換する．Julia+CUDAで気軽にGPU計算を参考のこと．

Db_inv_gpu = cu(Db_inv)
Df_gpu = cu(Df)
function next_step(v,time)
    bf[1] = r * eta * g_0(time)
    bf[end] = r * eta * g_1(time)
    bb[1] = -r * theta * g_0(time+k)
    bb[end] = -r * theta * g_1(time+k)
    bf_gpu = cu(bf)
    bb_gpu = cu(bb)
    return Db_inv_gpu * (Df_gpu * v + bf_gpu -bb_gpu)
end

なお，bf_gpu=cu(bf)とした後にbf_gpu[0]=r * eta * g_0(time)とするとエラーになった．要素にアクセスするときはgpuを使わないのがよいようだ．

初期値の生成

u を初期化して，uに初期値を設定する．

u = zeros(N-1)
for i in 1:N-1
    u[i] = f(h * i)
end

時間発展の計算

漸化式を用いて，繰り返しの計算を行う．LinRange(0,1,N+1)は$0\leq x\leq 1$を$(N+1)$等分して配列を作るものである．cu(u)でuをgpu計算できる変数に変換してから，next_stepでgpu計算を行う．Array(w_interior_gpu)はgpu計算できる変数を通常の配列に戻すコマンドである．

u_gpu = cu(u)

x = LinRange(0, 1, N+1)
for num in 1:t
    time = num * k

    w_interior_gpu = next_step(u_gpu, time)

    w = vcat(g_0(time), Array(w_interior_gpu), g_1(time))

# csvファイルにて，(x,w)座標のデータを出力するなら下記の4行のコメントを外す
#
#   mkpath("./csv")
#   df = DataFrame(x_data = x, y_data = w)
#   csv_filename = "./csv/" * lpad(string(num), 5, "0") * ".csv"
#   CSV.write(csv_filename, df)
#
# pngファイルにて，(x,w)データのplotを出力するなら，下記の5行のコメントを外す
#
#    mkpath("./png")
#    png_filename = "./png/" * lpad(string(num), 5, "0") * ".png"
#    plot(x, w, color = :black, xlims=(0,1), ylims=(0,1))
#    title!("t = $(round(time, digits=6))")
#    savefig(png_filename)    
#
    global u_gpu = w_interior_gpu
end

数値計算結果

上記の定数で，GPUありの場合とGPUなしの場合で実行時間を比較してみた．実時間で比較すると

GPUあり

22.493秒

GPUなし

43.8835秒

jupyterでステップごとに計算してみると，逆行列の計算に20秒程度を要しているため，時間発展の計算ではGPUありの場合はほぼすぐに終わっているようである．GPUなしだと20秒弱かかっているので，行列計算で速さが10倍程度違っているようである．

あとがき

ソースコードやアルゴリズムの精査はしていないので，修正すべきことは多々あると思われる．それでも，微分方程式を気軽に数値計算することができることを身につけて，いろいろと試してみる価値があると思われる．