Python で波動方程式の数値計算と動画 gif の書き出しをやらせてみよう

2次元正方形領域における波動方程式の初期値境界値問題を考えます。
こんな感じで波が伝播していく様子を気軽にシミュレーションできます。
f:id:wakabame:20180307201626g:plain
必要なものは imagemagick です。インストール方法は
kiito.hatenablog.com
を参考にしました。

単純な差分法として、

#%matplotlib nbagg #jupyter 
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
dx = 0.1
dy = dx
dt = 0.05
tmin = 0.0 
tmax = 6.0 #simulate time

#rectangle domain
xmin = 0.0
xmax = np.pi
ymin = 0.0
ymax = np.pi

c = 1.0 #propagation speed
rsq = (c*dt/dx)**2 #constant

nx = int((xmax-xmin)/dx) + 1
ny = int((ymax-ymin)/dy) + 1
nt = int((tmax-tmin)/dt) + 2

#mesh
X = np.linspace(xmin, xmax, nx)
Y = np.linspace(ymin, ymax, ny)
X, Y = np.meshgrid(Y, X)

#solution
u = np.zeros((nt,nx,ny))

#initial data
u_0 = np.exp(-((X-2)**2)*10)*np.exp(-((Y-2)**2)*10)*2
u_1 = np.zeros((nx,ny))
#np.sin(X)*np.sin(Y)
#np.exp(-(X**2)*10)*np.exp(-(Y**2)*10)
#np.exp(-((X-2)**2)*10)*np.exp(-((Y-2)**2)*10)

u[0] = u_0
u[1] = u[0] + dt * u_1

#simulation
for t in range(1,nt-1):
    for x in range(1,nx-1):
        for y in range(1,ny-1):
            u[t+1,x,y] = 2*(1-2*rsq)*u[t,x,y]-u[t-1,x,y]+rsq*(u[t,x-1,y]+u[t,x+1,y]+u[t,x,y-1]+u[t,x,y+1])
    #Neumann condition
    for x in range(1,nx-1):
        u[t+1,x,0] = u[t+1,x,1]
        u[t+1,x,ny-1] = u[t+1,x,ny-2]
    for y in range(1,ny-1):
        u[t+1,0,y] = u[t+1,1,y]
        u[t+1,nx-1,y] = u[t+1,nx-2,y]
    u[t,0,0] = (u[t,1,0]+u[t,0,1])/2
    u[t,nx-1,0] = (u[t,nx-2,0]+u[t,nx-1,1])/2
    u[t,0,ny-1] = (u[t,1,ny-1]+u[t,0,ny-2])/2
    u[t,nx-1,ny-1] = (u[t,nx-2,ny-1]+u[t,nx-1,ny-2])/2

fig = plt.figure()
fig.set_dpi(100)
ax = Axes3D(fig)

def animate(i):
    ax.clear()
    ax.plot_surface(X, Y, u[i], rstride=1, cstride=1, cmap=plt.cm.coolwarm,vmax=1,vmin=-1)
    ax.set_zlim(-2, 2)
    
anim = animation.FuncAnimation(fig,animate,frames=nt-1,interval=10,repeat=False)
#anim.save("wave2D-N2.gif", writer="imagemagick")
plt.show()

コンパイルすればできます。
動画を書きだす場合は下から2行目のコメントアウトを外しましょう。

初期条件を

u_0 = np.zeros((nx,ny))
u_1 = np.exp(-(X**2)*10)*np.exp(-(Y**2)*10)*20

に変更すると、初期速度場を与えて解くことができます。
f:id:wakabame:20180307204738g:plain

また、ディリクレ境界条件にしたい場合は、シミュレーションスキームにおいて #Neumann condition 以下を削除ないしコメントアウトをすればそのまま計算を行うことができます。
f:id:wakabame:20180307202245g:plain

さらに、正方形領域  \Omega = \left\{ (x,y)\in \mathbb{R}^2\, |\, 0 < x < \pi, 0 < y < \pi \right\} におけるラプラシアンの第1固有関数は
$$
\sin(x)\sin(y)
$$
でしたから、初期値として

u_0 = np.sin(X)*np.sin(Y)
u_1 = np.zeros((nx,ny))

を与えるとちゃんと固有振動が起きています。
f:id:wakabame:20180307203729g:plain

1次元の波動方程式も同じような方法で計算することができますね。

#%matplotlib nbagg
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

dx = 0.05
dt = 0.025
tmin = 0.0 
tmax = 12.0 #simulate time

#line domain
xmin = 0.0
xmax = np.pi

c = 1.0 #propagation speed
rsq = (c*dt/dx)**2 #constant

nx = int((xmax-xmin)/dx) + 1
nt = int((tmax-tmin)/dt) + 2
print(nx,nt)

X = np.linspace(xmin, xmax, nx)
#solution
u = np.zeros((nt,nx))

#initial data
u_0 = np.exp(-((X-1)**2)*10)
u_1 = np.zeros(nx)
u[0] = u_0
u[1] = u[0] + dt * u_1

#simulation
for t in range(1,nt-1):
    for x in range(1,nx-1):
        u[t+1,x] = 2*(1-rsq)*u[t,x]-u[t-1,x]+rsq*(u[t,x-1]+u[t,x+1])
    #Neumann condition
    u[t+1,0] = u[t+1,1]
    u[t+1,nx-1] = u[t+1,nx-2]

fig = plt.figure()
fig.set_dpi(100)
ax1 = fig.add_subplot(1,1,1)
def animate(i):
    ax1.clear()
    plt.ylim([-1.2,1.2])
    p, = plt.plot(u[i,:])
    
anim = animation.FuncAnimation(fig,animate,frames=nt-1,interval=50,repeat=False)
#anim.save("wave1dim21.gif", writer="imagemagick")
plt.show()

f:id:wakabame:20180307203946g:plain