技術系サラリーマン勉強記

数学、物理、プログラミングなど日々勉強した内容を取り扱っていきます。

Pythonで波動方程式の数値計算 & 結果のアニメーション

今回は両端を固定した弦の振動を表す一次元の波動方程式数値計算を行い結果をアニメーションで表示してみました。数値計算アルゴリズム数値計算 (理工系の基礎数学 8)を参考にしています。

f:id:mysciphmt:20191207180238g:plain

グラフのタイトルに時間経過(秒)を表示させていて、周期T=2sでアニメーションを繰り返しています。良い感じに振動を再現できています。

このアニメーションはJupyter notebook上で以下のコードを書くと表示できます。GIF形式で保存する際は、下から4行目のani.saveコメントアウトを外します。

コード

# matplotlibのアニメーション表示用
%matplotlib nbagg

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# 各種定数の設定
t0 = 0
tM = 2
M = 100
dt = (tM-t0)/M  # 時間間隔
x0 = 0
xN = 1
N = 20
dx = (xN-x0)/N # 空間間隔
a = (dt/dx)**2 # 差分方程式に出てくる係数
ux0 = 0 # 境界条件
uxN = 0 # 境界条件
A = 2 # 初期条件の式の振幅

# 変数の初期化
u = np.zeros([M+1, N+1])  # 行:時間、列:x

# 初期条件1 (t=0でのuの分布) u(x)=A*x(1-x) (0<=x<=1)
for xi in range(N+1):
    u[0,xi] = A*(xi*dx)*(1-xi*dx)
    
# 初期条件2 (t=dtでのuの分布) 
u[1,0] = 0
u[1,N] = 0
for xi in range(1,N):
    u[1,xi] = u[0,xi] + dt*0 + a/2*(u[0,xi+1] -2*u[0,xi] + u[0,xi-1])

# 差分解 (t=2*dt~M*dt)
for ti in range(1,M):
    # 両端の弦の変位は常に0
    u[ti+1,0] = ux0
    u[ti+1,N] = uxN
    
    # 両端以外を差分方程式により更新 
    for xj in range(1,N):
        u[ti+1, xj] = 2*u[ti,xj] -u[ti-1,xj] + a*(u[ti,xj+1] -2*u[ti,xj] + u[ti,xj-1])


# matplotlibで表示
x = np.linspace(x0, xN, N+1)  # maplotlibで表示するためのx軸の設定

fig = plt.figure(figsize=(6,4))
ax = fig.add_subplot(1,1,1)

# アニメ更新用の関数
def update_func(i):
    # 前のフレームで描画されたグラフを消去
    ax.clear()
    
    ax.plot(x, u[i,:], color='blue')
    ax.scatter(x, u[i,:], color='blue')
    # 軸の設定
    ax.set_ylim(-0.52, 0.52)
    # 軸ラベルの設定
    ax.set_xlabel('x', fontsize=12)
    ax.set_ylabel('u', fontsize=12)
    # サブプロットタイトルの設定
    ax.set_title('Time: ' + '{:.2f}'.format(tM*i/M))
    
ani = animation.FuncAnimation(fig, update_func, frames=M, interval=100,repeat=True)
# アニメーションの保存
#ani.save('test.gif', writer='imagemagick')

# 表示
plt.show()

参考文献