Heat Equation solvers

This notebook can be found in $UWHPSC/homeworks/project/Heat_Equation.ipynb

The "heat equation" or "diffusion equation" is a partial differential equation that in the simplest case takes the form $u_t(x,t) = u_{xx}(x,t)$ in one space dimension $x$ and time $t$. You could add a source term to this equation and then the steady state to obtain a boundary value problem, as solved in Homework 4. In this notebook we explore the behavior of the time dependent problem with no source term.

This equation is generally solved on some fixed domain in $x$ (e.g. in this notebook we will assume we are solving it for $0 < x < \pi$) and for time $t>t_0$. In general we must also specify initial conditions (the temperature $u(x,t_0)$) and boundary conditions $u(0,t)$ at the left edge and $u(\pi,t)$ at the right edge of the domain.

To keep thing simple, we will also assume that the boundary conditions are constant for all time, of the form $u(0,t) = 0$ and $u(\pi,t) = 0$.

This models, for example, the temperature distribution in a one-dimensional metal rod of length $\pi$, where some initial temperature is specified along the length of the rod and each end of the rod is kept fixed at 0 degrees (e.g. by placing a block of ice at each end that absorbs heat from the rod).

For any initial conditions, if we wait long enough to some time $T$, we expect the temperature to reach a constant state $u(x,T) = 0$ for all $x$. This is a steady state with $u_t = u_{xx} = 0$ and both boundary conditions satisfied.


In [ ]:
%pylab inline

Import some things we'll use below...


In [ ]:
from scipy import sparse  # to define sparse matrices
from scipy.sparse.linalg import spsolve   # to solve sparse systems
from JSAnimation import IPython_display
from matplotlib import animation

The next module is provided in this local directory:


In [ ]:
import JSAnimation_frametools as J  
plotdir = '_plots'  # to store png files for each figure
J.make_plotdir(plotdir, clobber=True)  # ok to clobber if it already exists

A function to plot the true solution at various times and make frames for an animation:


In [ ]:
def make_frames(u_true):
    x = linspace(0., pi, 200)
    tfinal = 0.25
    num_frames = 25
    for nstep in range(num_frames + 1):
        t = (float(nstep)/num_frames) * tfinal
        plot(x, u_true(x,t), 'b')
        title("u(x,t) at t = %6.4f" % t)
        axis([0,pi,-1.2,1.2])
        # Save this frame:
        J.save_frame(nstep, plotdir,verbose=False)
        clf()

Some exact solutions -- Fourier modes

For any integer $k$, the function $\sin(kx)$ satisfies the boundary conditions at $0$ and $\pi$. The second derivative of this function is $-k^2 \sin(kx)$ and the function $u(x,t) = e^{-k^2t}\sin(kx)$ satisfies the heat equation $u_t = u_{xx}$ with initial conditions $u(x,0) = \sin(kx)$.

Note that the larger $k$ is, the more rapidly $u(x,t)$ decays to zero. Below are a few examples...


In [ ]:
k = 1
u = lambda x,t: exp(-k**2 * t)*sin(k*x)
make_frames(u)
J.make_anim(plotdir,figsize=(8,4))

In [ ]:
k = 2
u = lambda x,t: exp(-k**2 * t)*sin(k*x)
make_frames(u)
J.make_anim(plotdir,figsize=(8,4))

In [ ]:
k = 8
u = lambda x,t: exp(-k**2 * t)*sin(k*x)
make_frames(u)
J.make_anim(plotdir,figsize=(8,4))

Note that any linear combination of Fourier modes is also a solution, for example...


In [ ]:
k1 = 1
k2 = 8
u = lambda x,t: exp(-k1**2 * t)*sin(k1*x) + 0.2 * exp(-k2**2 * t)*sin(k2*x)
make_frames(u)
J.make_anim(plotdir,figsize=(8,4))

Fourier sine series

A general linear combination of all possible sine functions with integer values of $k$ is called a Fourier sine series,

$$ f(x) = \sum_{k=1}^\infty c_k \sin(kx) $$

For a given function $f(x)$ there's a simple formula to determine its Fourier coefficients,

$$ c_k = \frac 2 \pi \int_0^\pi f(x) \sin(kx) \,dx $$

but it involves integrals that may not be be computable in closed form. However, if you do know the Fourier sine series for the initial data $u_0(x) = u(x,t)$, then the solution to the heat equation can be written down in closed form in terms of a (generally infinite) series.

Numerical methods

One way to approximate the solution to this equation with more arbitrary initial data is to express the data as a Fourier series and use what we found above to determine the solution. For some initial data functions this can be done analytically by taking the "Fourier transform". For more complicated functions one can use the "discrete Fourier transform" based on discrete samples of the data to determine an approximate solution. (The famous Fast Fourier Transform (FFT) algorithm will compute the $n$ Fourier coefficients based on $n$ discrete data values in ${\cal O}(n\log n)$ time.) This sort of numerical method is called a "spectral method" and for smooth initial data can give a very accurate solution.

But here we will instead consider a "finite difference method", using the same approximation to the second derivative $u_{xx}$ used in Homework 4,

$$ u_{xx}(x_i, t_N) \approx \frac{U_{i-1}^N - 2U_i^N + U_{i+1}^N}{\Delta x^2}. $$

This has to be used with some time stepping method: Given an approximate solution at all grid points at time $t_N$, we want to step forward to time $t_{N+1}$.

Explicit numerical method

A simple thing to try is the Forward Euler method, which we saw in lecture for an ODE $u'(t) = F(u(t))$ takes the form

$$ U^{N+1} = U^N + \Delta t F(U^N) $$

When applied to the heat equation, this gives the method

$$ U_i^{N+1} = U_i^N + \frac{\Delta t}{\Delta x^2} (U_{i-1}^N - 2U_i^N + U_{i+1}^N). $$

If the interval $0 < x < \pi$ is discretized with $x_i = i\Delta x$ for $i=0,1,\ldots,n+1$, with $\Delta x = \pi / (n+1)$, then the equation above is used to update $U_i^N$ at all interior points ($i=1,2,\ldots,n$) while $U_0^N$ and $U_{n+1}^N$ are zero for all $N$.

This is called an explicit method since we have an explicit formula to update each point based on past data. We do not have to solve any linear system.


In [ ]:
def solve_heat_explicit(x, u0, t0, tfinal, nsteps):
    
    dt = (tfinal - t0) / float(nsteps)
    n = len(x) - 2  # number of interior points
    dx = x[1]-x[0]
    dtdx2 = dt/dx**2
    print "dt / dx**2 = ", dtdx2
    if dtdx2 > 0.5:
        print "*** Warning, the explicit method is not stable for dt / dx**2 > 0.5"
        
    u = u0.copy()  # to leave u0 alone below
    
    for nstep in range(nsteps):
        # define uxx of length n approximating second derivative at interior points:
        uxx = (u[0:-2] - 2*u[1:-1] + u[2:]) / dx**2   
        u[1:-1] = u[1:-1] + dt*uxx
        
    return u

Test it out...


In [ ]:
k = 2.
u_true = lambda x,t: exp(-k**2 * t) * sin(k*x)

t0 = 0.
tfinal = 0.2
nsteps = 20
n = 20
x = linspace(0, pi, n+2)
u0 = u_true(x, t0)
plot(x,u0,'go-',label='initial u0')

u = solve_heat_explicit(x, u0, t0, tfinal, nsteps)

utrue_tfinal = u_true(x,tfinal)
error = abs(u - utrue_tfinal).max()
print "Maximum error at final time is %8.4e" % error

plot(x,u,'bo-',label='computed')
plot(x,utrue_tfinal,'r',label='exact')
legend()
ax = axis([0, pi, -1.2, 1.2])

Explicit methods are simple to implement, but do not always work so well. In particular, it can be shown that the method above is only stable if the time step is small enough that $\Delta t / \Delta x^2 \leq 1/2$.

This is easiest to observe if we take initial data that is not a single Fourier mode, e.g. $u(x,0) = x(\pi - x)$. Here are two experiments with different choices of $\Delta t$:


In [ ]:
t0 = 0.
tfinal = 0.2
n = 20
x = linspace(0, pi, n+2)
u0 = x * (pi - x)

# stable:
plot(x,u0,'go-')
nsteps = 20
u = solve_heat_explicit(x, u0, t0, tfinal, nsteps)
plot(x,u,'bo-')
ax = axis([0, pi, -4, 4])

In [ ]:
# unstable:
plot(x,u0,'go-')
nsteps = 10
u = solve_heat_explicit(x, u0, t0, tfinal, nsteps)
plot(x,u,'bo-')
ax = axis([0, pi, -4, 4])

Implicit method

The explicit method requires $\Delta t \leq \frac 1 2 \Delta x^2$, which means we must take very small time steps if the spatial grid is finely resolved.

It is often much more efficient and equally accurate to instead use an implicit method, which requires solving a linear system at every time step but allows much larger time steps.

One standard method for the heat equation is the Crank-Nicolson method, which is based on the ODE solver known as the Trapezoid rule for the ODE $u'(t) = F(t)$:

$$ U^{N+1} = U^N + \frac{\Delta t}{2} (F(U^N) + F(U^{N+1}). $$

When applied to the heat equation, this gives:

$$ U_i^{N+1} = U_i^N + \frac{\Delta t}{2\Delta x^2} [(U_{i-1}^N - 2U_i^N + U_{i+1}^N) + (U_{i-1}^{N+1} - 2U_i^{N+1} + U_{i+1}^{N+1})]. $$

This can be rearranged to give a tridiagonal linear system of equations to solve in each time step, as implemented in the next function. The matrix equation takes the form

$$ AU^{N+1} = BU^N $$

where $U^N$ is the vector with $n$ components $U_1^N, \ldots, U_n^N$ at the interior points and

$$ A = I - \frac{\Delta t}{2} D_2, \qquad B = I + \frac{\Delta t}{2} D_2 $$

and $D_2$ is the tridiagonal matrix approximating the second derivative operator. The right hand side would be have to be modified if the boundary conditions were values other than $U_0=U_n=0$, but for this simple case they drop out.

For $n=3$ the matrix $D_2$ is $$ D_2 = \frac{1}{\Delta x^2} \left[\begin{array}{rrr} -2&1&0\\ 1&-2&1 \\ 0&1&-2 \end{array}\right] $$


In [ ]:
def solve_heat_implicit(x, u0, t0, tfinal, nsteps):
    
    dt = (tfinal - t0) / float(nsteps)
    n = len(x) - 2  # number of interior points
    dx = x[1]-x[0]
    dtdx2 = dt/dx**2
    print "dt / dx**2 = ", dtdx2
        
    u = u0.copy()  # leave u0 alone below
    
    # Form the matrix:
    d1 = ones(n)
    d0 = -2 * ones(n)
    D2 = sparse.spdiags([d1,d0,d1], [-1,0,1],n,n,format='csc') / dx**2
    I = sparse.eye(n)
    
    A = (I - 0.5*dt*D2)
    B = (I + 0.5*dt*D2)
    
    for nstep in range(nsteps):
        rhs = B*u[1:-1]
        u[1:-1] = spsolve(A,rhs)
        
    return u

Try it out...


In [ ]:
k = 2.
u_true = lambda x,t: exp(-k**2 * t) * sin(k*x)

t0 = 0.
tfinal = 0.2
nsteps =20
n = 20
x = linspace(0, pi, n+2)
u0 = u_true(x, t0)
plot(x,u0,'go-',label='initial u0')

u = solve_heat_implicit(x, u0, t0, tfinal, nsteps)

utrue_tfinal = u_true(x,tfinal)
error = abs(u - utrue_tfinal).max()
print "Maximum error at final time is %8.4e" % error

plot(x,u,'bo-',label='computed')
plot(x,utrue_tfinal,'r',label='exact')
legend()
ax = axis([0, pi, -1.2, 1.2])

The advantage of a numerical method is that we can approximate the solution for any smooth initial data, e.g.


In [ ]:
n = 100
x = linspace(0, pi, n+2)
u0 = x*(pi-x)*sin(x**2)*exp(0.1*x)  # some function of x as initial data
plot(x,u0,'go-')

t0 = 0.
tfinal = 0.2
nsteps = 100

u = solve_heat_implicit(x, u0, t0, tfinal, nsteps)
plot(x,u,'bo-')
ax = axis([0, pi, -4, 4])