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

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

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

以前実施した拡散方程式の数値計算結果をアニメーションにしてみました。 myscitech.hatenablog.com

f:id:mysciphmt:20191205235632g:plain:w450

グラフのタイトルに時間経過(秒)を表示させています。

やはりアニメーションにすると分かりやすいですね。

このアニメーションは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 = 1
M = 500
dt = (tM-t0)/M  # 時間間隔
x0 = 0
xN = 1
N = 10
dx = (xN-x0)/N # 空間間隔
a = dt/(dx**2) # 差分方程式に出てくる係数
ux0 = 0 # 境界条件
uxN = 0 # 境界条件
A = 2 # 初期条件の式の振幅

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

# 初期条件 (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)

# 差分解
for ti in range(M):
    # 両端の温度は常に固定
    u[ti,0] = ux0
    u[ti,N] = uxN
    
    # 両端以外を差分方程式により更新 
    for xj in range(1,N):
        u[ti+1, xj] = a*u[ti,xj+1] + (1-2*a)*u[ti,xj] + a*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.02, 0.52)
    # 軸ラベルの設定
    ax.set_xlabel('x', fontsize=12)
    ax.set_ylabel('u', fontsize=12)
    # サブプロットタイトルの設定
    ax.set_title('Time: ' + '{:.2f}'.format(i/M))
    
ani = animation.FuncAnimation(fig, update_func, frames=M, interval=100,repeat=True)
# アニメーションの保存
#ani.save('test.gif', writer='imagemagick')

# 表示
plt.show()

参考文献