Rearranging for explicit terms, we get:
import holoviews as hv
hv.extension('bokeh', logo=False)
import numba
import numpy as np
@numba.jit
def step(U, V, W, dt, dx, dy, g, t, f):
"""U = u^n+1, V = u^n, W = u^n-1."""
N, M = U.shape
for i in range(1, N-1):
for j in range(1, M-1):
U[i, j] = dt**2 / dx ** 2 * (V[i+1, j] - 2 * V[i, j] + V[i-1, j]) + \
dt**2 / dy ** 2 * (V[i, j+1] - 2 * V[i, j] + V[i, j-1]) + \
dt**2 * g(t) * f(i * dx, j * dy) + \
2 * V[i, j] - W[i, j]
return U
@numba.jit
def g(t):
return np.sin(2 * np.pi * 15e-2 * t)
@numba.jit
def f(x, y):
x0, y0 = (62.5, 32.5)
return np.exp(- np.abs((x - x0)**2 + (y - y0)**2))
def n_steps(U, V, W, dt, dx, dy, g, t0, f, nsteps=100):
for i in range(nsteps):
t = t0 + i * dt
U = step(U, V, W, dt, dx, dy, g, t, f)
W = V.copy()
V = U.copy()
return U, V, W, t
N, M = 250, 260
U = np.zeros((N, M))
V = np.zeros((N, M))
W = np.zeros((N, M))
dt = 0.1
dx = 0.5
dy = dx
t = 0
N * dx / 2, M * dy / 4
snapshots = []
for _ in range(30):
U, V, W, t = n_steps(U, V, W, dt, dx, dy, g, t, f, nsteps=50)
snapshots.append(U.copy())
%%output holomap='scrubber'
hv.HoloMap({ind: hv.Image(U.T).opts(cmap='seismic', tools=['hover'], colorbar=True, width=700, height=600).redim.range(z=(-1e-1, 1e-1)) for ind, U in enumerate(snapshots)})