The Crank-Nicolson method rewrites a discrete time linear PDE as a matrix multiplication . Nonlinear PDE’s pose some additional problems, but are solvable as well this way (by linearizing every timestep). A major advantage here is that going steps into the future is just , and calculating a matrix power is polynomial time. The method is in general very stable.

For an assignment I had to construct the Crank-Nicolson matrix for a simple linear 1 dimension PDE, which had to be derived by hand. That’s a bit labourus, so I made a method to derive it automatically. As I came across it while sifting through some older files, and could still not find that method elsewhere online, I thought I might as well write it up here.

Derivation (1D)

This is the standard Crank-Nicolson expression for a linear, 1 dimensional PDE. The time stepping matrix would then be .

Since is linear, we can expand it into the set of its impulse responses . These are the vector outputs by applying it to a shifted unit impulse . We can then rewrite in terms of its impulse responses:

The Crank-Nicolson equation can be rewritten by substituting for this definition with the goal of reformulating it as a matrix equation. Focusing on a single element and moving the terms not resulting from into the sum one obtains

Equating this with the definition of the matrix-vector dot product
gives the elements of the Crank-Nicolson matrices and :

Implementation

Here is a simple diffusion-advection example in Javascript. Click to reset and apply new values.

Advection(1st derivative):
Diffusion(2nd derivative):

The implementation in Python is a straightforward translation of the equations above. Note that int(i == k) is the Kronecker delta .

import numpy as np
import numpy.linalg as LA

def cn_lin_mat(f, n, dt):
    d = np.eye(n)                       # shifted impulses {delta^k}
    h = [f(di) for di in d]             # shifted impulse responses {h^k}
    A = np.zeros((n, n))
    B = np.zeros((n, n))
    for i in range(n):
        for k in range(1, n+1):
            k = k % n
            A[i, k] = int(i == k) - dt/2*h[k][i]
            B[i, k] = int(i == k) + dt/2*h[k][i]
    Ainv = LA.inv(A)
    C = Ainv @ B
    return C

Higher Dimensions

This generalizes trivially to higher dimensions. You just have to rewrite f so that it works on 1 dimensional data, by reshaping it at appriopiate times. So, for a 2d linear PDE over a field with width W and height H it would result in the following code. You’d probably want to use sparse arrays for and , since they have elements, of which most are 0.

def f1d(phi):
    phi = phi.reshape((H, W))
    phi_next = f(phi)
    return phi_next.flatten()

C = cn_lin_mat(f1d, W*H, dt)
phi1 = (C @ phi0.flatten()).reshape((H, W))