とつとつとしてろうとせず

ひまつぶしにどうぞ。

擬スペクトル法によるKdV方程式の数値計算

注意

以下の記述は私的な備忘録であり、内容の正確さに関しては一切保証いたしませんので、悪しからず。

概要

PODのテスト用に時系列データのサンプルが欲しかったので、擬スペクトル法を用いてKdV方程式を解いてみた(PODのサンプルとして適切な選択なのかはおいおい確認することにしよう……)。
実行環境は例によってColaboratoryでござい。
「擬スペクトル法ってみんな使ってるし簡単だよね!」と舐めてかかっていたら、何の報いか色々うまくいかずに引っかかってしまったので、つまずいたところなどをメモしておく。

KdV方程式

Korteweg-de Vries方程式は次式で与えられる非線形波動方程式である(係数の取り方によって異なる定義が複数ある)。

\begin{align}
\frac{\partial u}{\partial t} + \alpha u \frac{\partial u}{\partial x} + \beta \frac{\partial^3 u}{\partial x^3} = 0
\end{align}
ここで、$\alpha, \beta > 0$はパラメータである。今回は特に$\alpha = 1$、$\beta = 0.022^2$とおいて有名な論文の結果が得られるか確認する(Zabusky & Kruskal (1965))。また、領域$[0,L]$は$L=2$の周期領域とする。

擬スペクトル法

擬スペクトル法(pseudo spectral method)は偏微分方程式に対する数値計算法の一つで、以下のように離散化を行う。

  1. 空間について離散フーリエ変換(DFT)を計算し、波数空間で空間微分などを計算(波数をかけるだけ)
  2. 時間についてはそのまま離散化する

これにより、偏微分方程式が波数空間における常微分方程式に変換されるため、Runge-Kutta法などによって時間積分を行い、得られた結果を逆変換すれば求める解が得られる。差分化を行わないので数値拡散がなく、高精度の解を得られるが、非線形項の取扱が難しかったり(波数空間だと計算コストが高いので一度実空間に戻す)、任意の境界形状は表現できなかったりするらしい。
とはいえ、1次元問題であれば、DFTも使えるpythonでやるにはお手軽な手法なので、今回は擬スペクトル法でやってみる。

実行例

コードはこんな感じ。ライブラリのインポート部分は後述のように実行時には使ってない関数も含まれるので注意。

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from numpy import pi, sin, cos, exp, real, imag
from scipy.integrate import solve_ivp
from numpy.fft import fft, ifft, fftfreq, rfft, irfft
from scipy.fftpack import diff as psdiff

# 空間、時間、波数の設定
L = 2.0 # 空間の長さ
T = 10 # 計算時間
N = 256 # 空間分割数
m = N/2+1 # 波数空間のデータ数
dt =1.0e-2 # 時間刻み
x = np.linspace(0, L, N) #空間
t = np.arange(0, T, dt) # 時間
k = fftfreq(m)*m*2*pi/L #波数

# パラメータと初期条件
a = 1.0
b = 0.022**2
u0 = cos(pi*x)

# KdV方程式
def kdv(t,u) : 
  v = rfft(u)
  ux = irfft(1.0j*k*v) # gradient term
  conv = u*ux # convection term
  disp = irfft( (1.0j*k)**3*v ) # dispersion term
  return -a*conv -b*disp

# 数値積分の実行
sol = solve_ivp(kdv, [0, T], u0, method='Radau', t_eval=t)
u = sol.y # 解を代入

コードのコアの部分は右辺の関数の定義部分なわけだが、擬スペクトル法だと非常に簡潔に書けてしまうのが魅力。真面目に差分法のコードを書くとさすがにもう少し長くなるはず。また、今回は勾配の計算のみを波数空間で行う形にしているので、最終結果は実空間での値である。波数空間で時間積分して最後に逆変換という形でも問題ないはず。どの道、非線形項の計算のために一度実空間に戻さないといけないので、変換と逆変換は毎ステップ行うことになる(線形方程式ならこの手の処理は要らないけど)。

結果のアニメーション

せっかくなのでアニメーションを作ろうと思ったのだが、Colaboratory上では一個のセル辺りのデータ容量?の上限が決まっているようで(9M?)、フレーム数が増えると実行できないようだったので、アニメーションだけはローカルの環境で作成した。gifにしようかとも思ったんだが、容量オーバーなのかアップロードに失敗したので、無難にyoutube経由でリンクを貼る方法にした。


最初の1フレーム目にグラフを重ね書きしてしまう不具合が直せなかったのだけど、全体の挙動は見えるので今回はこれでよしとしよう。参考にしたサイト様やWikipediaの動画などと比べると時間進行の速さにずれがあるような気がするが……結果はあまり信用できないのかもしれない。波形の様子なんかは似ているので、大間違いとまではいかないと思うけど。

大雑把なコードの説明

今回使った、あるいは使用を検討した関数などについて、コメントがてらつまずいたところなどをメモしておく。

fftとrfft

numpyで離散フーリエ変換をやる関数は大きく分けて"fft"と"rfft"の2種類がある。コマンドリファレンスなどを見ると、両者の違いは波数空間での負の周波数成分を考慮するかどうからしい。そのため、係数に複素数が含まれるような場合には"fft"を使う。一方、入力が実数のみであるなら負の周波数を考慮しないrfftの方が便利、ということらしい。ただし、ここに一つ落とし穴があった。

擬スペクトル法では関数の勾配を計算するために、関数をフーリエ変換し、波数空間で$i k$を掛けて逆変換を行う。この波数は次のように計算できる。

N = 256 # 空間分割数
m = N/2+1 # 波数空間のデータ数
k = fftfreq(m)*m*2*pi/L #波数

"fftfreq"は離散フーリエ変換時のサンプリング周波数を返す関数であり、定義式上データ点の個数で割られているため、データ点の個数と$2\pi / L$をかけることで波数の配列が得られる。このとき、フーリエ変換後のデータ点の個数には注意が必要だ。fft関数は負の周波数も考慮するため、波数空間でのデータ点の個数は元の空間のデータ点の個数と等しい。一方、rfftは正の周波数のみを考慮するため、データ点の個数は元々の半分になる(正確には、元のデータ点数$n$が偶数なら$(n/2)+1$、奇数なら$(n+1)/2$になる)。したがって、波数のデータ点数とフーリエ変換後のデータ点数が合うように定義していないと勾配の計算時にエラーになる(後で気づいたが、事前に適当なフーリエ変換をしておいて、その長さを"len(v)"のように取得する方がスマートだな)。

solve_ivp

scioyのv1.0から?入ったらしいODEを解く新しい関数。最初は上述のfftを使わないといけないと思っており、計算途中に複素数が入ってくるため、複素数でも同じような使用感で使えるらしいこの関数を使ってみた。結果としてrfftを使うならodeintでええやんというオチになってしまったが、せっかくなので簡単に使い方をメモしておく。

solve_ivpの使い方はodeintとほとんど同じだが、細かいところが色々と変わっている。

from scipy.integrate import solve_ivp

sol = solve_ivp(kdv, [0, T], u0, method='Radau', t_eval=t)
u = sol.y

まず、地味にややこしいのが右辺関数の定義がodeintと逆で"f(t,y)"の順になる。また、関数の引数としてパラメータを追加するような使い方は基本できないらしい(ラムダ式を使えば云々という情報がネットにあったが、俺の場合は特に困らなかったのでスルー)。
積分時間の指定も変わっていて、基本的にはソルバー側が時間刻みを自動で決める方針なので、"[t_start, t_end]"みたいな形で積分する区間を始点と終点のタプルで与える。積分点を指定したい場合は"t_eval"オプションで与える。
初期値とスキームの指定はご覧の通り。また、時間刻みをソルバーが決める関係上、戻り値は解と時間の配列がセットで返ってくる("sol.y"、"sol.t"のようにしてアクセス)。

係数に複素数が含まれる場合でも使い方は同じだが(odeだと書き方が変わるっぽい?)、その場合は初期値を複素数型のデータとして与える必要がある(初期値が実数であったとしても、"0.0j"を足すみたいな感じで複素数型にしておく)。また、スキームによって複素数に対応しているものとそうでないものがあるので、場合によってはスキームを変更する必要が出てくる。

今回は使ってないが、"dense_output"をTrueにすると、結果を時間方向に補間した関数で返してくれるようだ(戻り値は"sol.sol"みたいな感じ)。補間関数として何が使われるかもスキームに依存しているので、使う際はその辺も確認しないといけないだろう。また、補間関数が個別の評価点を使うためなのか、評価点("t_eval")をこちらで指定するとエラーになる。

ちなみに、デフォルトのRunge-Kutta法('RK45')では、今回使った時間刻みの値だと結果にものすごい高周波ノイズが入ってしまったので、陰的Runge-Kutta法('Radau')にスキームを変えている("BDF"でも可。odeintと同じスキームにしたいなら"LSODA"とする)。時間刻みをかなり小さくとればRK45でも計算できるのかもしれないが、問題がstiffならそこは素直に陰的スキームを使う法が適切だろう。

psdiff

最初はscipyのcookbookをもとにpsdiff("scipy.fftpack.diff")を使おうとしていた。これは以下のような感じで擬スペクトル法をもとに勾配を簡単に計算してくれる関数だ。

from scipy.fftpack import diff as psdiff

u = cos(pi*x)
ux = psdiff(u, period=L)
uxxx = psdiff(u, period=L, order=3)

擬スペクトル法を使っているので、関数の周期を入力する必要がある(指定しない場合は$2\pi$)。微分階数はデフォルトで1、それ以上の場合は"order"オプションで指定する。

この関数、使いやすいしコードも見やすくなって擬スペクトル法の計算のためにあるような便利なものなんだが、試しに計算してみると特に高階の計算でなぜか区間の両端にノイズが入ることがあり、使い物にならなかった。正確には、これを使ってもKdV方程式の計算そのものはうまくいったんだけど、それはノイズが出る分散項の係数が小さく設定されているからじゃないのかという疑いが晴れず、一般のケースに適用するには不安が残るため、結局は使わなかったのだ。原因自体はよくわからず、何か使い方を間違えているのかもしれないが、定義通りに計算するとノイズも出ないので、今回は愚直な方法を選んだ次第だ。

まとめ

紆余曲折あったが、それらしい結果が得られたので今回はこれでよしとする。もちろん、真面目にやるなら厳密解とか文献ときちんと合っているのか比較したり、保存性のチェックなんかをしなきゃいけないんだろうけど、面倒だしそれが目的でもないのでやらない。

追記

確認したらpsdiffを使わない方法(定義通りの計算)でも両端にノイズ入ってましたわ……結果が同じなら簡単に書けるpsdiffの方がいいね。確認不足でパチこいてすんません。
でもそうだとすると、余計にシミュレーション結果が怪しくなってくるな。擬スペクトル法ならではの処置を施さないとやっぱりうまいかないんだろうか。だとすると、素人にはちょっと難しい問題だ。
試しにnumpy.gradientで計算したところ、よっぽどきれいな結果(勾配と分散)が得られた。ただし、やはり境界で値がおかしくなるけど(当然っちゃ当然か)。一々境界値修正するのも面倒だしなぁ。すぐにはどうにもなりそうにないので、不完全燃焼だが、ここいらで撤退することにする。

上記の結果はあまり鵜呑みにしないでくださいますよう。