pythonnumpymatplotlibfinite-difference
2D finite difference scheme of reaction diffusion equation
I want to visualize the solution of the following partial differential equation:
u_t=u_xx+u_yy+f(u)
At different times, such as u(t=1) and u(t=3). I use a finite difference scheme and the following Python code:
import numpy as np
import matplotlib.pyplot as plt
# Parameters
D = 0.1 # Diffusion coefficient
L = 1.0 # Length of domain
T = 3.0 # Total time
Nx = 100 # Number of grid points in x-direction
Ny = 100 # Number of grid points in y-direction
Nt = 3000 # Number of time steps
dx = L / (Nx - 1) # Grid spacing in x-direction
dy = L / (Ny - 1) # Grid spacing in y-direction
dt = T / Nt # Time step size
# Initial condition
def initial_condition(x, y):
return np.exp(-((x - 0.5)**2 + (y - 0.5)**2) / 0.1)
# Reaction term
def reaction_term(u):
return u * (1 - u)
# Set boundary condition (u = 0 at boundary)
def apply_boundary_condition(u):
u[0, :] = 0 # Bottom boundary
u[-1, :] = 0 # Top boundary
u[:, 0] = 0 # Left boundary
u[:, -1] = 0 # Right boundary
# Create grid
x = np.linspace(0, L, Nx)
y = np.linspace(0, L, Ny)
X, Y = np.meshgrid(x, y)
# Initialize solution array
u = np.zeros((Nx, Ny))
# Set initial condition
u = initial_condition(X, Y)
# Time points to visualize
times_to_visualize = [1, 3] # Time points at which to visualize solution
# Iterate over time
for n in range(Nt + 1):
# Compute spatial derivatives using central differences
u_xx = (np.roll(u, -1, axis=0) - 2*u + np.roll(u, 1, axis=0)) / dx**2
u_yy = (np.roll(u, -1, axis=1) - 2*u + np.roll(u, 1, axis=1)) / dy**2
# Compute reaction term
f_u = reaction_term(u)
# Update solution using forward Euler method
u += dt * (D * (u_xx + u_yy) + f_u)
# Apply boundary condition
apply_boundary_condition(u)
# Check if current time is in times_to_visualize
if n * dt in times_to_visualize:
# Plot solution
plt.figure()
plt.contourf(X, Y, u, cmap='coolwarm')
plt.colorbar(label='u')
plt.xlabel('x')
plt.ylabel('y')
plt.title('2D Reaction-Diffusion Equation at t = {:.2f}'.format(n * dt))
plt.show()
My aim is to have something like this:
But my code doesn't give me such a result. Instead I get the following:
I am unsure where my code fails. I am not someone who has a lot of experience with numerical methods, so I was wondering what are some things I could do to correct this result?
Solution
The main issue was numerical instability. You can resolve this by increasing the temporal resolution (smaller dt), or decreasing the spatial resolution (larger dx and dy). I found that, for your combination of constants, increasing the temporal resolution 5-fold resolved the divergence problem. Also, I think the axis= arguments for u_xx and u_yy were the wrong way round.
I've included the more accurate second-order central differences (np.gradient()) for your reference, though they're not used. I've also included plotting that I used for debugging, as well as the surface plots you were aiming for.
Reproducible example
import numpy as np
import matplotlib.pyplot as plt
# Parameters
D = 0.1 # Diffusion coefficient
L = 1.0 # Length of domain
T = 3.0 # Total time
Nx = 100 # Number of grid points in x-direction
Ny = 100 # Number of grid points in y-direction
Nt = 5 * 3000 # Number of time steps
dx = L / (Nx - 1) # Grid spacing in x-direction
dy = L / (Ny - 1) # Grid spacing in y-direction
dt = T / (Nt - 1) # Time step size
# Initial condition
def initial_condition(x, y):
return np.exp(-((x - 0.5)**2 + (y - 0.5)**2) / 0.009)
# Reaction term
def reaction_term(u):
return u * (1 - u)
# Set boundary condition (u = 0 at boundary)
def apply_boundary_condition(u):
u[0, :] = 0 # top boundary
u[-1, :] = 0 # bottom boundary
u[:, 0] = 0 # Left boundary
u[:, -1] = 0 # Right boundary
# Create grid
x = np.linspace(0, L, Nx)
y = np.linspace(0, L, Ny)
X, Y = np.meshgrid(x, y)
# Set initial condition
u = initial_condition(X, Y)
plt.close('all')
plot_every_n = int(0.5 / dt) #plot every 0.5 seconds
n_to_visualise_3d = [int(t / dt) for t in [0, 0.1, 3]] #viz. in 3D at 0s, 0.1s, 3s
# Iterate over time
for n in range(Nt):
u_start = u.copy() #for visualisation
# Apply boundary condition
apply_boundary_condition(u)
# Compute spatial derivatives using central differences
u_xx = (np.roll(u, -1, axis=1) - 2*u + np.roll(u, 1, axis=1)) / dx**2
u_yy = (np.roll(u, -1, axis=0) - 2*u + np.roll(u, 1, axis=0)) / dy**2
#Could use second-order central differences
# u_xx = np.gradient(np.gradient(u, dx, axis=1, edge_order=1), dx, axis=1, edge_order=2)
# u_yy = np.gradient(np.gradient(u, dy, axis=0, edge_order=1), dy, axis=0, edge_order=2)
# Compute reaction term
f_u = reaction_term(u)
# Update solution using forward Euler method
delta_u = dt * (D * (u_xx + u_yy) + f_u)
u = u + delta_u
#
# 3D plotting of solution
#
if n in n_to_visualise_3d:
ax = plt.figure(figsize=(6, 6)).add_subplot(projection='3d')
ax.view_init(azim=45, elev=25)
ax.set_box_aspect(aspect=None, zoom=0.8)
ax.set_title(f'u(t={n * dt:.1f}s)', weight='bold')
ax.set(xlabel='x', ylabel='y', zlabel='u(t)')
#A regular surface plot
res = 100 #use either rcount/ccount or rstride/cstride to control surface res
im = ax.plot_surface(X, Y, u, rcount=res, ccount=res, cmap='magma', antialiased=False)
ax.contour(X, Y, u, zdir='z', offset=u.max() * 1.2, cmap='magma', alpha=0.25, linewidths=1)
ax.contour(X, Y, u, zdir='x', offset=x.min(), cmap='magma', alpha=0.25, linewidths=2)
ax.contour(X, Y, u, zdir='y', offset=y.min(), cmap='magma', alpha=0.25, linewidths=2)
#Alternatively, a coloured wireframe
# from matplotlib.colors import Normalize
# norm = Normalize(u.min(), u.max())
# im = ax.plot_surface(
# X, Y, u,
# facecolors=plt.cm.magma(norm(u)), facecolor=[0] * 4,
# rcount=20, ccount=20, cmap='magma'
# )
#Colorbar
ax.figure.colorbar(
im, shrink=0.5, pad=0.001,
label='u(t)', orientation='horizontal'
)
#
#Detailed plotting at start, end, and every 500
#
if not ( n==0 or n==Nt-1 or (n+1)%plot_every_n==0 ):
continue
u_limits = dict(vmin=u.min(), vmax=u.max())
deriv_limits = dict( vmin=min(u_xx.min(), u_yy.min()), vmax=max(u_xx.max(), u_yy.max()) )
variable_titles = [
r'$u_{start}$',
r'$\partial^2 u / \partial x^2$',
r'$\partial^2 u / \partial y^2$',
'f(u)',
r'$\Delta u$ x 1000',
r'$u_{end} = u_{start} + \Delta u$',
]
f, axs = plt.subplots(ncols=6, figsize=(15, 3.6), layout='tight', sharex=True, sharey=True)
f.suptitle(f'after {n + 1} steps | t={n * dt:4.2f}s', weight='bold')
for ax, name, data in zip(
axs, variable_titles, [u_start, u_xx, u_yy, f_u, delta_u * 1000, u]
):
#Use a separate scale for derivatives and non-derivatives
if 'partial^2 u' in name:
cmap_limits = deriv_limits
cmap = 'Greys'
else:
cmap_limits = u_limits
cmap = 'magma'
im = ax.imshow(
data, aspect='auto', interpolation='none',
extent=[x.min(), x.max(), y.min(), y.max()],
cmap=cmap, **cmap_limits
)
f.colorbar(im, ax=ax, orientation='horizontal', location='top', pad=0.12)
ax.set_title(name)
if ax == axs[0]:
ax.set(xlabel='x', ylabel='y')
else:
ax.axis('off')