Julia + CUDA で熱方程式を解く

# Introduction
Juliaで一次元熱方程式のDirichlet問題
```math
\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](https://ja.wikipedia.org/wiki/CUDA))．GeForceとかで数値計算が可能となる．

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

## Crank-Nicolson法
[差分法](https://ja.wikipedia.org/wiki/%E5%B7%AE%E5%88%86%E6%B3%95)を参照のこと．ただし，行列を使って説明した方がJuliaとの相性がいいので，簡単に説明しておく．

$x$変数の微分を中心差分としたとき，Dirichlet境界条件で$x$の二階微分を与える差分化は行列を用いて
```math
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}$として
```math
\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}}
```
とするのが前進差分解法であり，
```math
\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$倍して和をとると
```math
(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$とすれば
```math
\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のパッケージを読み込む．
```julia
using LinearAlgebra
using CUDA
```

インストールされていない場合は
```julia
using Pkg
Pkg.add("CUDA")
```
のようにCUDAのパッケージをインストールする．

計算するだけなら，LinearAlgebraとCUDAだけでよいが，データを出力したり，グラフを表示させるためには
```julia
using Plots
using CSV
using DataFrames
```
も追加で入れておく．


## 初期関数の設定
$f(x)=1.0$を定義する．
```julia
function f(x)
    return 1.0
end
```

# 境界条件の設定
$g_0(t)=0.0$，$g_1(t)=0.0$を定義する．
```julia
function g_0(t)
    return 0.0
end
function g_1(t)
    return 0.0
end
```

## 定数の設定
いろいろな定数を設定する．
```julia
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 は拡散数と呼ばれている，空間と時間の刻み幅の比率である([差分法](https://ja.wikipedia.org/wiki/%E5%B7%AE%E5%88%86%E6%B3%95)，[拡散数](https://ja.wikipedia.org/wiki/%E6%8B%A1%E6%95%A3%E6%95%B0))．t は時間変数についての繰り返し回数である．

## Dirichlet Laplacianの離散化と前進差分行列の定義
行列の定義については，GPUを使用せずに作成する．
```julia
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で定める．
```julia
Db = I - r * theta * D
Db_inv = inv(Db)
```

## 境界条件を与えるベクトルの生成
bf と bb を境界条件を与えるベクトルとして初期化する．
```julia
bf = zeros(N-1)
bb = zeros(N-1)
```

## 差分法による次の時間ステップを計算するための関数生成
漸化式の右辺に対応する関数を生成する．この計算において，gpuを用いる．そのために，cu(Db_inv)やcu(Df)などにより，行列をGPU計算できる変数に変換する．[Julia+CUDAで気軽にGPU計算](https://qiita.com/Shoichiro-Tsutsui/items/b2ff07f62d1ca12122bc)を参考のこと．
```julia
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に初期値を設定する．
```julia
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計算できる変数を通常の配列に戻すコマンドである．
```julia

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なしの場合で実行時間を比較してみた．実時間で比較すると

<dl>
  <dt>GPUあり</dt>
  <dd>22.493秒</dd>
  <dt>GPUなし</dt>
  <dd>43.8835秒</dd>
</dl>

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

# あとがき

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