Solving the heat equation¶
AMath 586, Spring Quarter 2019 at the University of Washington. For other notebooks, see Index.ipynb or the Index of all notebooks on Github.
Sample program to solve the heat equation with the Crank-Nicolson method.
We solve the heat equation $u_t = \kappa u_{xx}$ on the interval $0\leq x \leq 1$ with Dirichlet boundary conditions $u(0,t) = g_0(t)$ and $u(1,t) = g_1(t)$.
To test accuracy, we use cases where an exact solution to the heat equation for all $x$ is known. This utrue function is used to set initial conditions. It is also used in each time step to set boundary values on whatever finite interval we consider.
%pylab inline
from matplotlib import animation
from IPython.display import HTML
def make_animation(hs_input, hs_output, nplot=1):
"""
Plot every `nplot` frames of the solution and turn into
an animation.
"""
xfine = linspace(hs_input.ax,hs_input.bx,1001)
fig, ax = plt.subplots()
ax.set_xlim((hs_input.ax,hs_input.bx))
#ax.set_ylim((-0.2, 1.2))
ax.set_ylim((-1.2, 1.2))
line1, = ax.plot([], [], '+-', color='b', lw=2, label='computed')
line2, = ax.plot([], [], color='r', lw=1, label='true')
ax.legend()
title1 = ax.set_title('')
def init():
line1.set_data(hs_output.x_computed, hs_output.u_computed[:,0])
line2.set_data(xfine, hs_input.utrue(xfine, hs_input.t0))
title1.set_text('time t = %8.4f' % hs_input.t0)
return (line1,line2,title1)
def animate(n):
line1.set_data(hs_output.x_computed, hs_output.u_computed[:,n])
line2.set_data(xfine, hs_input.utrue(xfine, hs_output.t[n]))
title1.set_text('time t = %8.4f' % hs_output.t[n])
return (line1,line2,title1)
frames = range(0, len(hs_output.t), nplot) # which frames to plot
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=frames,
interval=200,
blit=True)
close('all') # so one last frame plot doesn't remain
return anim
class HeatSolutionInput(object):
def __init__(self):
# inputs:
self.t0 = 0.
self.tfinal = 1.
self.ax = 0.
self.bx = 1.
self.mx = 39.
self.utrue = None
self.kappa = 0.02
self.nsteps = 10
class HeatSolutionOutput(object):
def __init__(self):
# outputs:
self.h = None
self.dt = None
self.t = None
self.x_computed = None
self.u_computed = None
self.errors = None
Forward Euler time stepping¶
def heat_FE(heat_solution_input):
"""
Solve u_t = kappa * u_{xx} on [ax,bx] with Dirichlet boundary conditions,
using centered differences in space and the Forward Euler method for time stepping,
with m interior points, taking nsteps time steps.
Input:
`heat_solution_input` should be on object of class `HeatSolutionInput`
specifying inputs.
Output:
an object of class `HeatSolutionOutput` with the solution and other info.
This routine can be embedded in a loop on m to test the accuracy.
Note: the vector x defined below is of length m+2 and includes both boundary points.
The vector uint is of length m and is only the interior points that we solve for,
by solving an m by m linear system each time step.
The vector u is of length m+2 and obtained by extending uint with the boundary values,
so that plotting (x,u) includes the boundary values.
"""
from scipy import sparse
from scipy.sparse.linalg import spsolve
# unpack the inputs for brevity:
ax = heat_solution_input.ax
bx = heat_solution_input.bx
kappa = heat_solution_input.kappa
m = heat_solution_input.mx
utrue = heat_solution_input.utrue
t0 = heat_solution_input.t0
tfinal = heat_solution_input.tfinal
nsteps = heat_solution_input.nsteps
h = (bx-ax)/float(m+1) # h = delta x
x = linspace(ax,bx,m+2) # note x(1)=0 and x(m+2)=1
# u(1)=g0 and u(m+2)=g1 are known from BC's
dt = tfinal / float(nsteps)
# initial conditions:
u0 = utrue(x,t0)
# initialize u and plot:
tn = 0
u = u0
t = empty((nsteps+1,), dtype=float)
errors = empty((nsteps+1,), dtype=float)
u_computed = empty((m+2,nsteps+1), dtype=float)
t[0] = tn
errors[0] = 0.
u_computed[:,0] = u0
# main time-stepping loop:
for n in range(1,nsteps+1):
tnp = tn + dt # = t_{n+1}
# indices of interior points as in integer numpy array:
jint = array(range(1,m+1), dtype=int)
# Then the numerical method can be written without a loop
# or matrix-vector multiply:
u[jint] = u[jint] + kappa * dt/h**2 * (u[jint-1] - 2*u[jint] + u[jint+1])
# evaluate true solution to get new boundary values at tnp:
g0np = utrue(ax,tnp)
g1np = utrue(bx,tnp)
# augment with boundary values:
u[0] = g0np
u[-1] = g1np
error = abs(u-utrue(x,tnp)).max() # max norm
t[n] = tnp
u_computed[:,n] = u
errors[n] = error
tn = tnp # for next time step
heat_solution_output = HeatSolutionOutput() # create object for output
heat_solution_output.dt = dt
heat_solution_output.h = h
heat_solution_output.t = t
heat_solution_output.x_computed = x
heat_solution_output.u_computed = u_computed
heat_solution_output.errors = errors
return heat_solution_output
A smooth solution¶
We first use the decaying Gaussian $$ u(x,t) = \frac{1}{\sqrt{4\beta\kappa t + 1}} \exp\left(\frac{-(x-x_0)^2}{4\kappa t + 1/\beta}\right). $$ The initial data and boundary conditions are obtained by evaluating this function at $t=0$ or at $x=0$ or $x=1$. In particular, the initial conditions are simply $$ u(x,0) = \eta(x) = \exp(-\beta(x-x_0)^2). $$
beta = 150
x0 = 0.4
kappa = 0.02
utrue_gaussian = lambda x,t: exp(-(x-0.4)**2 / (4*kappa*t + 1./beta)) \
/ sqrt(4*beta*kappa*t+1.)
Recall that the forward Euler time stepping on the heat equation is only stable if the time step satisfies $k \leq 0.5h^2$. However, for smooth solutions with very small components of the high wave number Fourier modes, it can take a long time for the instability to appear even if we take much larger $k$. Here's an example. Note that it is the highest wave number (the saw-tooth mode) that grows fastest and hence appears first...
t0 = 0.
tfinal = 4.
ax = 0.
bx = 1.
mx = 39
h = 1./((mx+1))
dt_stab = 0.5*h**2 / kappa
nsteps_stab = int(floor(tfinal-t0)/dt_stab) + 1
print('For stability, need to take at least %i time steps' % nsteps_stab)
heat_solution_input = HeatSolutionInput()
heat_solution_input.t0 = t0
heat_solution_input.tfinal = tfinal
heat_solution_input.ax = ax
heat_solution_input.bx = bx
heat_solution_input.mx = mx
heat_solution_input.utrue = utrue_gaussian
heat_solution_input.kappa = kappa
heat_solution_input.nsteps = 240
heat_solution_output = heat_FE(heat_solution_input)
error_tfinal = heat_solution_output.errors[-1] # last element
print('Using %i time steps' % heat_solution_input.nsteps)
print('Max-norm Error at t = %6.4f is %12.8f' % (heat_solution_input.tfinal, error_tfinal))
# make an animation of the results, plotting every 10th frame:
anim = make_animation(heat_solution_input, heat_solution_output, nplot=10)
HTML(anim.to_jshtml()) # or use the line below...
#HTML(anim.to_html5_video())
Discontinous initial data¶
The instability is observed much more quickly if the initial data contains more high wave numbers, e.g. if it is discontinuous.
Consider the exact solution $$ u(x,t) = \text{erf}\left(x/\sqrt{4\kappa t}\right) $$ where erf is the error function defined as the integral of the Gaussian,
$$ \text{erf}(z) = \frac{2}{\pi} \int_0^z \exp(-t^2)\, dt. $$
See e.g. https://en.wikipedia.org/wiki/Error_function.
As $t \rightarrow 0$, this approaches the discontinous function jumping from $-1$ for $x<0$ to $+1$ for $x>0$.
The error function is implemented in the scipy.special library of special functions.
kappa = 0.02
def utrue_erf(x,t):
from scipy.special import erf
if t==0:
return where(x>0, 1., -1.)
else:
return erf(x/sqrt(4*kappa*t))
t0 = 0.
tfinal = 2.
ax = -1.
bx = 1.
mx = 40
h = (bx-ax)/((mx+1))
dt_stab = 0.5*h**2 / kappa
nsteps_stab = int(floor(tfinal-t0)/dt_stab) + 1
print('For stability, need to take at least %i time steps' % nsteps_stab)
heat_solution_input = HeatSolutionInput()
heat_solution_input.t0 = t0
heat_solution_input.tfinal = tfinal
heat_solution_input.ax = ax
heat_solution_input.bx = bx
heat_solution_input.mx = mx
heat_solution_input.utrue = utrue_erf
heat_solution_input.kappa = kappa
heat_solution_input.nsteps = 32
heat_solution_output = heat_FE(heat_solution_input)
error_tfinal = heat_solution_output.errors[-1] # last element
print('Using %i time steps' % heat_solution_input.nsteps)
print('Max-norm Error at t = %6.4f is %12.8f' % (heat_solution_input.tfinal, error_tfinal))
anim = make_animation(heat_solution_input, heat_solution_output, nplot=1)
HTML(anim.to_jshtml())
Crank-Nicolson method¶
This method uses the same centered difference spatial discretization with the Trapezoidal method for time stepping. That method is A-stable so this method is stable for any size time step (though not necessarily accurate).
Implementing this method requires solving a tridiagonal linear system in each time step, which we do using the sparse matrix routines from scipy.sparse.linalg.
def heat_CN(heat_solution_input):
"""
Solve u_t = kappa * u_{xx} on [ax,bx] with Dirichlet boundary conditions,
using the Crank-Nicolson method with m interior points, taking nsteps
time steps.
Input:
`heat_solution_input` should be on object of class `HeatSolutionInput`
specifying inputs.
Output:
an object of class `HeatSolutionOutput` with the solution and other info.
Note: the vector x defined below is of length m+2 and includes both boundary points.
The vector uint is of length m and is only the interior points that we solve for,
by solving an m by m linear system each time step.
The vector u is of length m+2 and obtained by extending uint with the boundary values,
so that plotting (x,u) includes the boundary values.
"""
from scipy import sparse
from scipy.sparse.linalg import spsolve
# unpack the inputs for brevity:
ax = heat_solution_input.ax
bx = heat_solution_input.bx
kappa = heat_solution_input.kappa
m = heat_solution_input.mx
utrue = heat_solution_input.utrue
t0 = heat_solution_input.t0
tfinal = heat_solution_input.tfinal
nsteps = heat_solution_input.nsteps
h = (bx-ax)/float(m+1) # h = delta x
x = linspace(ax,bx,m+2) # note x(1)=0 and x(m+2)=1
# u(1)=g0 and u(m+2)=g1 are known from BC's
dt = tfinal / float(nsteps)
# initial conditions:
u0 = utrue(x,t0)
# Each time step we solve MOL system U' = AU + g using the Trapezoidal method
# set up matrices:
r = 0.5 * kappa* dt/(h**2)
em = ones(m)
em1 = ones(m-1)
A = sparse.diags([em1, -2*em, em1], [-1, 0, 1], shape=(m,m))
A1 = sparse.eye(m) - r * A
A2 = sparse.eye(m) + r * A
# initialize u and plot:
tn = 0
u = u0
t = empty((nsteps+1,), dtype=float)
errors = empty((nsteps+1,), dtype=float)
u_computed = empty((m+2,nsteps+1), dtype=float)
t[0] = tn
errors[0] = 0.
u_computed[:,0] = u0
# main time-stepping loop:
for n in range(1,nsteps+1):
tnp = tn + dt # = t_{n+1}
# boundary values u(0,t) and u(1,t) at times tn and tnp:
# boundary values are already set at time tn in array u:
g0n = u[0]
g1n = u[m+1]
# evaluate true solution to get new boundary values at tnp:
g0np = utrue(ax,tnp)
g1np = utrue(bx,tnp)
# compute right hand side for linear system:
uint = u[1:m+1] # interior points (unknowns)
rhs = A2.dot(uint) # sparse matrix-vector product A2 * uint
# fix-up right hand side using BC's (i.e. add vector g to A2*uint)
rhs[0] = rhs[0] + r*(g0n + g0np)
rhs[m-1] = rhs[m-1] + r*(g1n + g1np)
# solve linear system:
uint = spsolve(A1,rhs) # sparse solver
# augment with boundary values:
u = hstack([g0np, uint, g1np])
error = abs(u-utrue(x,tnp)).max() # max norm
t[n] = tnp
u_computed[:,n] = u
errors[n] = error
tn = tnp # for next time step
heat_solution_output = HeatSolutionOutput() # create object for output
heat_solution_output.dt = dt
heat_solution_output.h = h
heat_solution_output.t = t
heat_solution_output.x_computed = x
heat_solution_output.u_computed = u_computed
heat_solution_output.errors = errors
return heat_solution_output
Test this with k = h:¶
With this method we can get a fine solution with only 40 steps (on a grid with 39 interior points). We only go out to time 1 but it would stay stable forever...
heat_solution_input = HeatSolutionInput()
heat_solution_input.t0 = 0.
heat_solution_input.tfinal = 1.
heat_solution_input.ax = 0.
heat_solution_input.bx = 1.
heat_solution_input.mx = 39
heat_solution_input.utrue = utrue_gaussian
heat_solution_input.kappa = kappa
heat_solution_input.nsteps = 40
heat_solution_output = heat_CN(heat_solution_input)
error_tfinal = heat_solution_output.errors[-1] # last element
print('dt = %6.4f' % heat_solution_output.dt)
print('Max-norm Error at t = %6.4f is %12.8f' % (heat_solution_input.tfinal, error_tfinal))
anim = make_animation(heat_solution_input, heat_solution_output)
HTML(anim.to_jshtml())
We can also plot how the max-norm error evolves with time:
plot(heat_solution_output.t,heat_solution_output.errors)
xlabel('time')
ylabel('max-norm error');
Test for second-order accuracy¶
If dt and h are both reduced by 2, the error should go down by a factor 4 (for sufficiently small values).
Here we loop over a range of dt and h values, with dt = h in each solve.
nsteps_vals = [20,40,80,160,320] # values to test
E = empty(len(nsteps_vals))
# print table header:
print(" h dt error ratio estimated order")
for j,nsteps in enumerate(nsteps_vals):
heat_solution_input.nsteps = nsteps
heat_solution_input.mx = nsteps - 1
heat_solution_output = heat_CN(heat_solution_input)
E[j] = heat_solution_output.errors[-1] # last element
h = heat_solution_output.h
dt = heat_solution_output.dt
if j>0:
ratio = E[j-1] / E[j]
else:
ratio = nan
p = log(ratio)/log(2)
print("%8.6f %8.6f %12.8f %4.2f %4.2f" % (h, dt, E[j], ratio, p))
loglog(nsteps_vals, E, '-o')
title('Log-log plot of errors')
xlabel('nsteps')
ylabel('error')
Observe oscillations if dt is too large¶
We know that Crank-Nicolson is stable for any time step, but the amplification factor approaches $-1$ as $k\lambda \rightarrow \infty$, so we expect high wavenumber modes to oscillate in time if we take the time step too large. This can be observed with the Gaussian initial data used here.
heat_solution_input.mx = 39
heat_solution_input.nsteps = 2
heat_solution_output = heat_CN(heat_solution_input)
error_tfinal = heat_solution_output.errors[-1] # last element
print('h = %6.4f, dt = %6.4f' % (heat_solution_output.h, heat_solution_output.dt))
print('Max-norm Error at t = %6.4f is %12.8f' % (heat_solution_input.tfinal, error_tfinal))
anim = make_animation(heat_solution_input, heat_solution_output)
HTML(anim.to_jshtml()) # or use the line below...
#HTML(anim.to_html5_video())
Discontinous data¶
With a sufficiently small time step, Crank-Nicolson behaves well on the problem with discontinous data. Note that we use an even number of grid points m = 40 so that they are symmetric about $x=0$. Try m=39 and see how the asymmetry gives a larger error!
heat_solution_input = HeatSolutionInput()
heat_solution_input.t0 = 0.
heat_solution_input.tfinal = 1.5
heat_solution_input.ax = -1.
heat_solution_input.bx = 1.
heat_solution_input.mx = 40
heat_solution_input.utrue = utrue_erf
heat_solution_input.kappa = kappa
heat_solution_input.nsteps = 40
heat_solution_output = heat_CN(heat_solution_input)
error_tfinal = heat_solution_output.errors[-1] # last element
print('h = %6.4f, dt = %6.4f' % (heat_solution_output.h, heat_solution_output.dt))
print('Max-norm Error at t = %6.4f is %12.8f' % (heat_solution_input.tfinal, error_tfinal))
anim = make_animation(heat_solution_input, heat_solution_output)
HTML(anim.to_jshtml()) # or use the line below...
#HTML(anim.to_html5_video())
The issue with oscillations is more apparent with this discontinuous initial data. Taking a much larger time step on the same grid gives the results below. Note that the Crank-Nicolson method remains stable, but the saw-tooth mode is apparent near the interface if we try to step over the rapid transient behavior in this stiff problem.
heat_solution_input.nsteps = 3
heat_solution_output = heat_CN(heat_solution_input)
error_tfinal = heat_solution_output.errors[-1] # last element
print('h = %6.4f, dt = %6.4f' % (heat_solution_output.h, heat_solution_output.dt))
print('Max-norm Error at t = %6.4f is %12.8f' % (heat_solution_input.tfinal, error_tfinal))
anim = make_animation(heat_solution_input, heat_solution_output)
HTML(anim.to_jshtml()) # or use the line below...
#HTML(anim.to_html5_video())
An L-stable method like TR-BDF2 would do better in this case.