Programming the finite difference method using Python

Lately I found myself needing to solve the 1D spherical diffusion equation using the Python programming language. To make sure that I can remember how to do this in the far future (because I will forget), this post goes over a few examples of how it can be done.

Diffusion in a sphere happens all the time, mostly when chemical reactions are involved and a reactant or a product has to make its way to or from the reaction site. In my case, the application is lithium-ion batteries where lithium diffuses into and out-of a particle of the active material.

First of all, some good web resources that already exist:

The last link is mostly what I based my code on.

The math:
The spherical diffusion equation is:
Equation: $\frac{\partial C}{\partial t} = \frac{D}{r^2} \frac{\partial}{\partial r}\left(r^2 \frac{\partial C}{\partial r}\right)$

But, to make things a bit more simple, I'm going to:

  1. Use dimensionless numbers: $\theta=\frac{C}{C_0}$, $x=\frac{r}{R_p}$ (where $R_p$ is the radius of the particle)
  2. Use the expanded form of the spherical diffusion equation to make it easier to solve.

All of these transformations lead to:

Equation: $\frac{\partial \theta}{\partial t} = K \frac{\partial^2 \theta}{\partial x^2} + \frac{2K}{x} \frac{\partial \theta}{\partial x}$
where $K=\frac{D}{R_p^2}$ and the domain falls within x=0 (the center of the particle) and x=1 ($=R_p$)
Note: In this simple example, we'll just set $\frac{2K}{x}=2K$

The example:
The following example doesn't have any physical significance, but it demonstrates how to implement two different types of boundary conditions.

We'll solve the problem and assume the following boundary conditions:

  1. Boundary conditions: at x=0 $\theta = 1$, at x=1 $\frac{d \theta}{dx}=\delta$
  2. $\theta(0)=1$
  3. Parameters: $K=1$, $\delta=-1$

The finite difference discretization:
The system of equations that we want to solve is: $A u = b$ where $A$ is an NxN matrix of coefficients, $u$ is a vector containing $\theta_i$ for each node, i, and $b$ is a vector of size N containing the source terms and some other contributions from the PDE.

For this problem, the relevant finite-difference discretizations (in space) are:

  1. $\frac{d^2 \theta}{dx^2} = \frac{1}{\Delta x^2} (\theta_{i-1} - 2 \theta_i + \theta_{i+1})$
  2. $\frac{d \theta}{dx} = \frac{1}{2 \Delta x} (\theta_{i+1} - \theta_{i-1})$

The matrices that we want to make are then (assuming a small N=5 node grid including boundary points $\theta_1$ and $\theta_5$):

$\frac{d \theta}{d t} = \frac{K}{\Delta x^2} A_1 u + \frac{2 K}{x 2 \Delta x} A_2 u + F$

$A_1 = \left[ {\begin{array}{cc} 0 & 0 & 0 & 0 & 0 \\ 1 & -2 & 1 & 0 & 0 \\ 0 & 1 & -2 & 1 & 0 \\ 0 & 0 & 1 & -2 & 1 \\ 0 & 0 & 0 & 2 & -2 \\ \end{array} } \right]$ $A_2 = \left[ {\begin{array}{cc} 0 & 0 & 0 & 0 & 0 \\ -1 & 0 & 1 & 0 & 0 \\ 0 & -1 & 0 & 1 & 0 \\ 0 & 0 & -1 & 0 & 1 \\ 0 & 0 & 0 & 0 & 0 \\ \end{array} } \right]$ $u = \left[ {\begin{array}{cc} \theta_1 \\ \theta_2 \\ \theta_3 \\ \theta_4 \\ \theta_5 \\ \end{array} } \right]$ $F = \left[ {\begin{array}{cc} 0 \\ 0 \\ 0 \\ 0 \\ 2 K \delta \left(\frac{1}{\Delta x} + 1\right) \\ \end{array} } \right]$

The inner elements of the $A$ matrices (for nodes i=2..N-1) can be derived based on the discretization above, but the elements corresponding to N=1 and N=5 are different because here we have the boundary conditions worked in. Also, the vector $F$ arises because of the gradient boundary condition at N=5.

Boundary conditions:
This problem has a constant value boundary condition at x=0 and a gradient boundary condition at x=1.

1. To handle a constant value boundary condition such as the one given in this problem at x=0, it's quite easy:
Just set $\theta_1=1$ in the $u$ matrix and in the $A$ matrix, set the equivalent position (ie. A[1][1]) to 1 and all others in that row equal to 0. Note: We actually set A[1][1]=0 in the matrix above, this is because the 1 comes out of rht eidentity matrix that we use below.

2. To handle a gradient boundary condition such as the one given in this problem at x=1, we need to use an imaginary grid point, in this case, $\theta_{i+1}=\theta_6$. So, at the boundary x=1, which is equivalent to node i=5:
$\frac{d \theta}{dx} = \delta = \frac{1}{2 \Delta x} (\theta_{i+1} - \theta_{i-1})$ therefore $\theta_{6}=\theta_{4}+2\delta \Delta x$
Now, we can substitute this "imaginary grid point", $\theta_{6}$ into the full equation:

$\frac{d \theta}{d t} = \frac{K}{\Delta x^2} (\theta_{i-1} - 2 \theta_i + \theta_{i+1}) + \frac{2 K}{x 2 \Delta x} (\theta_{i+1} - \theta_{i-1})$

$\frac{d \theta}{d t} = \frac{K}{\Delta x^2} (\theta_{4} - 2 \theta_5 + \theta_{4}+2\delta \Delta x) + \frac{2 K}{x 2 \Delta x} (\theta_{4}+2\delta \Delta x - \theta_{4})$

which results in the $A_1$ entry of [0 0 0 2 -2], the $A_2$ entry of [0 0 0 0 0] and an extra term $2 K \delta \left(\frac{1}{\Delta x} + 1\right)$ which is placed into the $F$ vector.

Discretization in time:
The last part before we start to program this up is the time discretization. In this example, we'll use the Crank-Nicolson method for the time discretiztion:

$\frac{d \theta}{dt} = \frac{ \theta^{k+1} - \theta^k}{\Delta t} = \frac{1}{2} F(\theta^{k+1}) + \frac{1}{2} F(\theta^{k})$
where $k$ is the current time step for which we know the solution of $\theta$, $k+1$ is the future timestep at time $t=t+\Delta t$ for which we want to solve and $F(\theta)$ is the function that we're evaluating: $\frac{K}{\Delta x^2} A_1 u + \frac{2 K}{x 2 \Delta x} A_2 u + F$

A little bit of rearrangement and we get:
$(I - 0.5 \Delta t \times A_1 - 0.5 \Delta t \times A_2) u^{k+1} = (I + 0.5 \Delta t \times A_1 + 0.5 \Delta t \times A_2) u^{k} + \Delta t F$

or, in other words, for the system $A u = b$:

$A = (I - 0.5 \Delta t \times A_1 - 0.5 \Delta t \times A_2)$
$u = u^{k+1}$
$b = (I + 0.5 \Delta t \times A_1 + 0.5 \Delta t \times A_2) u^{k} + \Delta t F$

The solution:
This can all be solved quite easily using Python. IMHO, programming in Python is like cutting through butter using a ceramic knife. It's fantastic. And the numpy and scipy packages make it perfect for scientific computing.

The basics of the program:

  1. Create the $A_1$, $A_2$, $u^0$ and $F$, $I$ matrices
  2. Use numpy/scipy to invert the matrix
  3. Profit

The python code:
Download here.

#!/usr/bin/python
# By Ben Kenney - http://www.benk.ca
# 1D Time dependent spherical diffusion equation
# dC/dt = K div(grad(C)) + 2K/x grad(c) on x=[0,1]
# BC: @x=0 C=1, @x=1 dC/dC=-1
# Time discretization using Crank Nicolson scheme
# Isn't python beautiful?

import scipy
import scipy.sparse as sparse
import scipy.sparse.linalg
import numpy

N = 50
dx = 1/(N-1.0)
delta = -1.0
K = 1.0

#grid points
x = numpy.linspace(0,1,N)

#create time steps
k = 0.5/100
TFinal = 1
NumOfTimeSteps = int(TFinal/k)

#initial solution
u = numpy.transpose([numpy.ones(N)*1.0])

#source term
F = numpy.transpose([numpy.zeros(N)])
F[-1]=2.0*K*delta*(1.0/dx+1.0)
print F

#create matrices with boundary conditions
A1=numpy.zeros([N])
A2=numpy.zeros([N])
A1[0]=0.0 # constant value boundary
for i in range(1,N-1):
        array = numpy.zeros([N])
        array[i-1:i-1+3] = [1,-2,1]
        A1=numpy.vstack([A1,array])
        array = numpy.zeros([N])
        #array[i-1:i-1+3] = [-1.0/x[i],0.0,1.0/x[i]] #x is the grid spacing
        array[i-1:i-1+3] = [-1.0,0.0,1.0]
        A2=numpy.vstack([A2,array])
array = numpy.zeros([N])
array[-2:]=[2,-2] #gradient boundary condition
A1=numpy.vstack([A1,array])
print A1
A1=A1*K/dx/dx
A1=scipy.sparse.csr_matrix(A1)
array = numpy.zeros([N])
A2=numpy.vstack([A2,array])
print A2
A2=A2*2.0*K/2.0/dx #note: grid factor, 1/x, is built into A2 matrix already
A2=scipy.sparse.csr_matrix(A2)

data = []

#identity matrix
I = sparse.identity(N)

print("Time step = %g \t Time = %g"%(0, 0))
print(u)
for i in range(NumOfTimeSteps):
        A = (I - k/2.0*A1 - k/2.0*A2)
        b = (I + k/2.0*A1 + k/2.0*A2)*u+k*F
        u = numpy.transpose(numpy.mat(sparse.linalg.spsolve(A, b)))
        print("Time step = %g \t Time = %g"%(i+1, k*(i+1)))
        data.append(u)
print u[:,-1]

A comparison between this code and Comsol at time=1s: