Python implementation of Crank-Nicolson scheme

Since at this point we know everything about the Crank-Nicolson scheme, it is time to get our hands dirty. In this post, the third on the series on how to numerically solve 1D parabolic partial differential equations, I want to show a Python implementation of a Crank-Nicolson scheme for solving a heat diffusion problem.

While Phython is certainly not the best choice for scientific computing, in terms of performance and optimization, it is a good language for rapid prototyping and scripting (and in many cases even for complex production-level code).

For a problem of this type, Python is more than sufficient at doing the job. For more complicated problems involving multiple dimensions, more coupled equations and many extra terms, other languages are typically preferred (Fortran, C, C++,…), often with the inclusion of parallel programming using the Message Passing Interface (MPI) paradigm.

The problem that I am going to present is the one proposed in the excellent book “Numerical Solution of Partial Differential Equations” by G.D. Smith (Oxford University Press),

$$\left\{ \begin{eqnarray} \frac{\partial u}{\partial t} \ \ & = & \ \frac{\partial^2 u}{\partial x^2}, \ \ & 0\leq x \leq 1 \nonumber \\ u(t,0) & = & u(t,1)=0 & \ \ \ \ \ \ \forall t\nonumber\\ u(0,x) & = & 2x & \mathrm{if} \ \ x\leq 0.5 \nonumber\\ u(0,x) & = & 2(1-x) & \mathrm{if} \ \ x> 0.5 \nonumber \end{eqnarray} \right.$$

Before I turn to the numerical implementation of a Crank-Nicolson scheme to solve this problem, let’s see already how the solution looks like in the video below.

As you can see, the maximum of the function $u$ occurs at $t=0$, after which $u$ keeps decreasing. This behaviour is in line with the maximum principle for parabolic equations, which essentially states that the solution of a parabolic equation takes its maximum on the boundary (intended as the time $t=0$ or space boundaries).

The reason for this can be made intuitive by comparing to the case of a metallic one-dimensional rod with the sides that are kept at some fixed temperature and with a temperature distribution that is linear and maximum at the centre, as in the initial condition above. As time progresses, the two “heat sources” (or sinks) at the sides are kept at constant low temperature. The diffusion of heat results in the rod becoming colder and colder until its temperature becomes equal to the temperature at the boundaries.

Let’s now talk about the numerical solution of the problem above. As already discussed, the numerical solution has to solve for the following matrix equation

$$\begin{equation} A\textbf{u}_{j+1} = B\textbf{u}_{j} + \textbf{b}_{j}, \end{equation}$$

where

$$\begin{equation} A = \left( \begin{matrix} 2+2r & -r & & & \\ -r & 2+2r & -r & \Huge{0} &\\ & & \ddots & & & \\ & \Huge{0} & -r & 2+2r & -r \\ & & & -r & 2+2r \\ \end{matrix} \right) , \textbf{u}_{j+1} = \left( \begin{matrix} u_{2,j+1} \\ u_{3,j+1} \\ \vdots \\ u_{N-2,j+1} \\ u_{N-1,j+1} \\ \end{matrix} \right) \\ B = \left( \begin{matrix} 2-2r & r & & & \\ r & 2-2r & r & \Huge{0} &\\ & & \ddots & & & \\ & \Huge{0} & r & 2-2r & r \\ & & & r & 2-2r \\ \end{matrix} \right), \textbf{u}_{j} = \left( \begin{matrix} u_{2,j} \\ u_{3,j} \\ \vdots \\ u_{N-2,j} \\ u_{N-1,j} \\ \end{matrix} \right), \textbf{b}_{j} = \left( \begin{matrix} r u_{1,j} \\ 0 \\ \vdots \\ 0 \\ r u_{N,j} \\ \end{matrix} \right) \end{equation}$$

and $r =\Delta t/\Delta x^2$.

The Python implementation below can be broken down into the following steps:

  1. definition of the parameters of the problem: time step, grid spacing, number of grid nodes ($\Delta t, \Delta x, N$)
  2. definition of initial and boundary conditions
  3. definition of the matrices $A$ and $B$ and of the vector of boundary conditions $b$
  4. solution of linear system of equations at each time step, using the $\texttt{linalg}$ package in $\texttt{numpy}$

The code should run for just a few seconds. To generate the m4v movie, I use the $\texttt{ffmpeg}$ library that can be downloaded here.


[download CN_example.py ]


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import numpy as np
import matplotlib.pyplot as plt
import os, sys
import matplotlib

matplotlib.rc('font', size=18)
matplotlib.rc('font', family='Arial')

#definition of numerical parameters
N = 51 #number of grid points
dt = 5.e-4 #time step
L = float(1) #size of grid
nsteps = 620 #number of time steps
dx = L/(N-1) #grid spacing
nplot = 20 #number of timesteps before plotting

r = dt/dx**2 #assuming heat diffusion coefficient == 1

#initialize matrices A, B and b array
A = np.zeros((N-2,N-2))
B = np.zeros((N-2,N-2))
b = np.zeros((N-2))
#define matrices A, B and b array
for i in range(N-2):
    if i==0:
        A[i,:] = [2+2*r if j==0 else (-r) if j==1 else 0 for j in range(N-2)]
        B[i,:] = [2-2*r if j==0 else r if j==1 else 0 for j in range(N-2)]
        b[i] = 0. #boundary condition at i=1
    elif i==N-3:
        A[i,:] = [-r if j==N-4 else 2+2*r if j==N-3 else 0 for j in range(N-2)]
        B[i,:] = [r if j==N-4 else 2-2*r if j==N-3 else 0 for j in range(N-2)]
        b[i] = 0. #boundary condition at i=N
    else:
        A[i,:] = [-r if j==i-1 or j==i+1 else 2+2*r if j==i else 0 for j in range(N-2)]
        B[i,:] = [r if j==i-1 or j==i+1 else 2-2*r if j==i else 0 for j in range(N-2)]

#initialize grid
x = np.linspace(0,1,N)
#initial condition
u = np.asarray([2*xx if xx<=0.5 else 2*(1-xx) for xx in x])
#evaluate right hand side at t=0
bb = B.dot(u[1:-1]) + b

fig = plt.figure()
plt.plot(x,u,linewidth=2)
filename = 'foo000.jpg';
fig.set_tight_layout(True);
plt.xlabel("x")
plt.ylabel("u")
plt.title("t = 0")
plt.savefig(filename)
plt.clf()

c = 0
for j in range(nsteps):
    print(j)
    #find solution inside domain
    u[1:-1] = np.linalg.solve(A,bb)
    #update right hand side
    bb = B.dot(u[1:-1]) + b
    if(j%nplot==0): #plot results every nplot timesteps
        plt.plot(x,u,linewidth=2)
        plt.ylim([0,1])
        filename = 'foo' + str(c+1).zfill(3) + '.jpg';
        plt.xlabel("x")
        plt.ylabel("u")
        plt.title("t = %2.2f"%(dt*(j+1)))
        plt.savefig(filename)
        plt.clf()
        c += 1

os.system("ffmpeg -y -i 'foo%03d.jpg' heat_equation.m4v")
os.system("rm -f *.jpg")

⁂

Today’s Mondrian random generator