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

2次元正方形領域における波動方程式の初期値境界値問題を考えます。
こんな感じで波が伝播していく様子を気軽にシミュレーションできます。
f:id:wakabame:20180307201626g:plain

単純な差分法として、

# %matplotlib inline
import numpy as np
from IPython.display import HTML
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="pillow")
# HTML(anim.to_jshtml())

を実行すればできます。
notebook で利用する場合は下から最初と最終行のコメントアウトを外しましょう。
f:id:wakabame:20210620012658p:plain

初期条件を

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固有関数は

\begin{equation}
\sin(x)\sin(y)
\end{equation}

でしたから、初期値として

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

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

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

# %matplotlib inline
import numpy as np
from IPython.display import HTML
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
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="pillow")
# HTML(anim.to_jshtml())

f:id:wakabame:20180307203946g:plain