import numpy as np
from matplotlib import rc
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import axes3d
import seaborn as sns
%matplotlib inline
sns.set()
sns.set_style("whitegrid")
rc('animation', html='html5')
plt.rcParams['figure.figsize'] = 8, 6
def explicit_step(last_state, D, dt, dx):
new_state = np.copy(last_state)
alpha = D * dt / (dx ** 2)
filt = np.array([alpha, 1 - 2 * alpha, alpha])
new_state[1:-1] = np.convolve(last_state, filt, mode='valid')
return new_state
LATTICE = 500
TOTAL_TIME = 200
D = 1
dx = 1
state = np.zeros(LATTICE)
state[LATTICE // 2] = 1
dt = 1
states = [np.copy(state)]
for i in range(TOTAL_TIME):
state = explicit_step(state, D, dt, dx)
states.append(np.copy(state))
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_title("DeltaT = %.1f" % dt)
ax.set_xlim(0, LATTICE)
ax.grid()
line = ax.plot([], [])
x = np.arange(0,LATTICE)
plt.close()
def update_anim(i):
line[0].set_data(x, states[i])
ax.set_ylim(np.min(states[i]), max(1, np.max(states[i])))
return line,
ani = animation.FuncAnimation(fig, update_anim, range(TOTAL_TIME), blit=False, interval=50, repeat=True)
ani
state = np.zeros(LATTICE)
state[LATTICE // 2] = 1
dt = 0.5
states = [np.copy(state)]
for i in range(TOTAL_TIME):
state = explicit_step(state, D, dt, dx)
states.append(np.copy(state))
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_title("DeltaT = %.1f" % dt)
ax.set_ylim(-0.5, 1)
ax.set_xlim(0, LATTICE)
ax.grid()
line = ax.plot([], [])
x = np.arange(0,LATTICE)
plt.close()
def update_anim(i):
line[0].set_data(x, states[i])
return line,
ani = animation.FuncAnimation(fig, update_anim, range(TOTAL_TIME), blit=False, interval=50, repeat=True)
ani
state = np.zeros(LATTICE)
state[LATTICE // 2] = 1
dt = 0.1
states = [np.copy(state)]
for i in range(TOTAL_TIME):
state = explicit_step(state, D, dt, dx)
states.append(np.copy(state))
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_ylim(-0.5, 1)
ax.set_title("DeltaT = %.1f" % dt)
ax.set_xlim(0, LATTICE)
ax.grid()
line = ax.plot([], [])
x = np.arange(0,LATTICE)
plt.close()
def update_anim(i):
line[0].set_data(x, states[i])
return line,
ani = animation.FuncAnimation(fig, update_anim, range(TOTAL_TIME), blit=False, interval=50, repeat=True)
ani
According to the Courant condition, the required time step size (for $\Delta x = 1$) is:
$$ \frac{2 D \Delta t}{(\Delta x)^2} \leq 2 \implies \Delta t \leq 0.5$$
def implicit_step(y, D, dt, dx):
alpha = D * dt / (dx ** 2)
# init a, b, c
a = np.array([-alpha for i in range(LATTICE)])
a[0] = 0
a[-1] = 0
b = np.array([1 + 2 * alpha for i in range(LATTICE)])
b[0] = 1
b[-1] = 1
c = np.array([-alpha for i in range(LATTICE)])
c[0] = 0
c[-1] = 0
# forward pass
ctag = np.zeros_like(c)
ctag[0] = c[0] / b[0]
ytag = np.zeros_like(y)
ytag[0] = y[0] / b[0]
for i in range(1, LATTICE):
ctag[i] = c[i] / (b[i] - a[i] * ctag[i - 1])
ytag[i] = (y[i] - a[i] * ytag[i - 1]) / (b[i] - a[i] * ctag[i - 1])
# backward pass
x = np.zeros_like(y)
x[-1] = ytag[-1]
for i in range(2, LATTICE + 1):
x[-i] = ytag[-i] - ctag[-i] * x[-i+1]
return x
state = np.zeros(LATTICE)
state[LATTICE // 2] = 1
dt = 1
states = [np.copy(state)]
for i in range(TOTAL_TIME):
state = implicit_step(state, D, dt, dx)
states.append(np.copy(state))
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_title("DeltaT = %.1f" % dt)
ax.set_ylim(-0.5, 1)
ax.set_xlim(0, LATTICE)
ax.grid()
line = ax.plot([], [])
x = np.arange(0,LATTICE)
plt.close()
def update_anim(i):
line[0].set_data(x, states[i])
return line,
ani = animation.FuncAnimation(fig, update_anim, range(TOTAL_TIME), blit=False, interval=50, repeat=True)
ani
state = np.zeros(LATTICE)
state[LATTICE // 2] = 1
dt = 0.5
state_ex = state
states_ex = [np.copy(state)]
state_im = state
states_im = [np.copy(state)]
for i in range(TOTAL_TIME):
state_ex = explicit_step(state_ex, D, dt, dx)
states_ex.append(np.copy(state_ex + 0.5))
state_im = implicit_step(state_im, D, dt, dx)
states_im.append(np.copy(state_im))
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_title("DeltaT = %.1f" % dt)
ax.set_ylim(-0.5, 1.5)
ax.set_xlim(0, LATTICE)
ax.yaxis.set_ticklabels([])
ax.grid()
x = np.arange(0,LATTICE)
line_ex = ax.plot(x, state)
line_im = ax.plot(x, state)
ax.legend(['Explicit', 'Implicit'])
plt.close()
def update_anim(i):
line_ex[0].set_data(x, states_ex[i])
line_im[0].set_data(x, states_im[i])
return line_ex, line_im,
ani = animation.FuncAnimation(fig, update_anim, range(TOTAL_TIME), blit=False, interval=50, repeat=True)
ani
state = np.zeros(LATTICE)
state[LATTICE // 2] = 1
dt = 10
states = [np.copy(state)]
for i in range(TOTAL_TIME):
state = implicit_step(state, D, dt, dx)
states.append(np.copy(state))
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_title("DeltaT = %.1f" % dt)
ax.set_ylim(-0.1, 0.5)
ax.set_xlim(0, LATTICE)
ax.grid()
line = ax.plot([], [])
x = np.arange(0,LATTICE)
plt.close()
def update_anim(i):
line[0].set_data(x, states[i])
return line,
ani = animation.FuncAnimation(fig, update_anim, range(TOTAL_TIME), blit=False, interval=50, repeat=True)
ani
Even stable for $\Delta t = 10$!
In a similar notation to 1D, let's define
$$\alpha = \frac{D \Delta t}{2 (\Delta x)^2}$$
And now we can re-write the first equation (half-step) of (8.29) as:
$$ -\alpha u^{n+0.5}_{j-1, k} + (1 + 2 \alpha) u^{n+0.5}_{j, k} -\alpha u^{n+0.5}_{j+1, k} = \alpha u^{n}_{j, k - 1} + (1 - 2 \alpha) u^{n}_{j, k} + \alpha u^{n}_{j, k + 1}$$
otice that if we fix $j$, we got a nice Tridiagnoal system:
$$ a_i x_{i-1} + b_i x_i + c_i x_{i+1} = d_i $$
(Just that $d_i$ is a bit more "ugly" than it was in the previous question...) So we solve just life before, forward sweeps and backward passes.
Similarly, we can express the second half step as:
$$ -\alpha u^{n+1}_{j, k - 1} + (1 + 2 \alpha) u^{n+1}_{j, k} -\alpha u^{n+1}_{j, k + 1} = \alpha u^{n+0.5}_{j - 1, k} + (1 - 2 \alpha) u^{n + 0.5}_{j, k} + \alpha u^{n+0.5}_{j + 1, k}$$
Which we solve in a symmetrical fashion (same shtick, switch j and k)
Instead of coding the same function twice, I wrote one function which completes a half step. Then for the second half step - we input the transposed state matrix and transpose back the output.
Useful links:
LATTICE_DIM = 48
TOTAL_TIME = 500
D = 1
dx = 1
def ADI_half_step(state, D, dt, dx):
alpha = D * dt / (2 * (dx ** 2))
new_state = np.zeros_like(state)
# init a, b, c
a = np.array([-alpha for i in range(LATTICE_DIM)])
a[0] = 0
a[-1] = 0
b = np.array([1 + 2 * alpha for i in range(LATTICE_DIM)])
b[0] = 1
b[-1] = 1
c = np.array([-alpha for i in range(LATTICE_DIM)])
c[0] = 0
c[-1] = 0
# we fix j, and solve the tridiaognal system
for j in range(LATTICE_DIM):
# init d
d = np.zeros_like(a)
d[0] = state[j, 0]
d[-1] = state[j, -1]
for i in range(1, LATTICE_DIM - 1):
d[i] = alpha * state[j, i - 1] + (1 - 2 * alpha) * state[j, i] + alpha * state[j, i + 1]
# forward pass
ctag = np.zeros_like(c)
ctag[0] = c[0] / b[0]
dtag = np.zeros_like(d)
dtag[0] = d[0] / b[0]
for i in range(1, LATTICE_DIM):
ctag[i] = c[i] / (b[i] - a[i] * ctag[i - 1])
dtag[i] = (d[i] - a[i] * dtag[i - 1]) / (b[i] - a[i] * ctag[i - 1])
# backward pass
x = np.zeros_like(d)
d[-1] = dtag[-1]
for i in range(2, LATTICE_DIM + 1):
x[-i] = dtag[-i] - ctag[-i] * x[-i+1]
new_state[j, :] = x
return new_state
def ADI_step(state, D, dt, dx):
# first half step, fix j (the 1st dim)
state = ADI_half_step(state, D, dt, dx)
# second half step, fix k (the 2nd dim) by transposing the state and transposing back the result
state = np.transpose(ADI_half_step(np.transpose(state), D, dt, dx))
return state
3D Animation in matplotlib: https://matplotlib.org/2.0.0/examples/mplot3d/wire3d_animation_demo.html
state = np.zeros((LATTICE_DIM, LATTICE_DIM))
def add_seed(state):
start_pos = np.random.randint(1, LATTICE_DIM - 1, 2)
state[start_pos[0], start_pos[1]] = LATTICE_DIM / 4
return state
dt = 1
states = [np.copy(state)]
for i in range(TOTAL_TIME):
# Every X frame, add a random
if i % 6 == 0:
state = add_seed(state)
state = ADI_step(state, D, dt, dx)
states.append(np.copy(state))
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
ax.yaxis.set_ticklabels([])
ax.xaxis.set_ticklabels([])
plt.axis('off')
plt.close()
# # Make the X, Y meshgrid.
xs = np.linspace(-1, 1, LATTICE_DIM)
ys = np.linspace(-1, 1, LATTICE_DIM)
X, Y = np.meshgrid(xs, ys)
# Set the z axis limits so they aren't recalculated each frame.
ax.set_zlim(-0.5, 1)
wframe = None
def update_anim(i):
global wframe
# If a line collection is already remove it before drawing.
if wframe:
ax.collections.remove(wframe)
# Plot the new wireframe and pause briefly before continuing.
wframe = ax.plot_wireframe(X, Y, states[i], rstride=1, cstride=1)
ani = animation.FuncAnimation(fig, update_anim, range(TOTAL_TIME), blit=False, interval=100, repeat=True)
ani
We are going to solve using equation (8.69) for a Laplace equation, so $\rho = 0$
def SOR_step(u, alpha):
new_u = (1 - alpha) * u
new_u[0, :] = 0
new_u[:, 0] = 0
new_u[-1, :] = 1
new_u[:, -1] = -1
for j in range(1, len(u) - 1):
for k in range(1, len(u) - 1):
new_u[j, k] += (alpha / 4) * u[j + 1, k]
new_u[j, k] += (alpha / 4) * new_u[j - 1, k]
new_u[j, k] += (alpha / 4) * u[j, k + 1]
new_u[j, k] += (alpha / 4) * new_u[j, k - 1]
return new_u
lattice_size = 50
alpha = 1
T = 300
u = np.zeros((lattice_size, lattice_size))
us = []
for i in range(T):
u = SOR_step(u, alpha)
us.append(np.copy(u))
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
fig.suptitle("SOR Laplace - Lattice Size: %dx%d, alpha: %.1f" % (lattice_size, lattice_size, alpha))
plt.close()
# # Make the X, Y meshgrid.
xs = np.linspace(-1, 1, len(u))
ys = np.linspace(-1, 1, len(u))
X, Y = np.meshgrid(xs, ys)
# Set the z axis limits so they aren't recalculated each frame.
ax.set_zlim(-1, 1)
wframe = None
def update_anim(i):
global wframe
# If a line collection is already remove it before drawing.
if wframe:
ax.collections.remove(wframe)
# Plot the new wireframe and pause briefly before continuing.
wframe = ax.plot_wireframe(X, Y, us[i], rstride=1, cstride=1)
ani = animation.FuncAnimation(fig, update_anim, range(len(us)), blit=False, interval=50, repeat=True)
ani
lattice_size = 50
alpha = 1.9
T = 300
u = np.zeros((lattice_size, lattice_size))
us = []
for i in range(T):
u = SOR_step(u, alpha)
us.append(np.copy(u))
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
fig.suptitle("SOR Laplace - Lattice Size: %dx%d, alpha: %.1f" % (lattice_size, lattice_size, alpha))
plt.close()
# # Make the X, Y meshgrid.
xs = np.linspace(-1, 1, len(u))
ys = np.linspace(-1, 1, len(u))
X, Y = np.meshgrid(xs, ys)
# Set the z axis limits so they aren't recalculated each frame.
# ax.set_zlim(-1, 1)
wframe = None
def update_anim(i):
global wframe
# If a line collection is already remove it before drawing.
if wframe:
ax.collections.remove(wframe)
# Plot the new wireframe and pause briefly before continuing.
wframe = ax.plot_wireframe(X, Y, us[i], rstride=1, cstride=1)
ani = animation.FuncAnimation(fig, update_anim, range(len(us)), blit=False, interval=50, repeat=True)
ani
lattice_size = 50
alpha = 0.5
T = 300
u = np.zeros((lattice_size, lattice_size))
us = []
for i in range(T):
u = SOR_step(u, alpha)
us.append(np.copy(u))
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
fig.suptitle("SOR Laplace - Lattice Size: %dx%d, alpha: %.1f" % (lattice_size, lattice_size, alpha))
plt.close()
# # Make the X, Y meshgrid.
xs = np.linspace(-1, 1, len(u))
ys = np.linspace(-1, 1, len(u))
X, Y = np.meshgrid(xs, ys)
# Set the z axis limits so they aren't recalculated each frame.
ax.set_zlim(-1, 1)
wframe = None
def update_anim(i):
global wframe
# If a line collection is already remove it before drawing.
if wframe:
ax.collections.remove(wframe)
# Plot the new wireframe and pause briefly before continuing.
wframe = ax.plot_wireframe(X, Y, us[i], rstride=1, cstride=1)
ani = animation.FuncAnimation(fig, update_anim, range(len(us)), blit=False, interval=50, repeat=True)
ani
Let's try to see what's the largest $\alpha$ for which a 50x50 lattice still converges... just by playing around with diffeent values, 1.3 seems to work fine.
lattice_size = 50
alpha = 1.99
T = 300
u = np.zeros((lattice_size, lattice_size))
us = []
for i in range(T):
u = SOR_step(u, alpha)
us.append(np.copy(u))
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
fig.suptitle("SOR Laplace - Lattice Size: %dx%d, alpha: %.1f" % (lattice_size, lattice_size, alpha))
plt.close()
# # Make the X, Y meshgrid.
xs = np.linspace(-1, 1, len(u))
ys = np.linspace(-1, 1, len(u))
X, Y = np.meshgrid(xs, ys)
# Set the z axis limits so they aren't recalculated each frame.
ax.set_zlim(-1, 1)
wframe = None
def update_anim(i):
global wframe
# If a line collection is already remove it before drawing.
if wframe:
ax.collections.remove(wframe)
# Plot the new wireframe and pause briefly before continuing.
wframe = ax.plot_wireframe(X, Y, us[i], rstride=1, cstride=1)
ani = animation.FuncAnimation(fig, update_anim, range(len(us)), blit=False, interval=50, repeat=True)
ani