# Schroedinger/Diffusion equation with Crank-Nicolson in Python/SciPy

I tried to make the question as detailed as possible. I have an extremely simple solver written for the Schroedinger equation but with imaginary time, which transforms it basically into the diffusion equation (with a potential term). The method is pretty well documented on this page, and I basically followed the steps almost exactly. Here's the link for that:

http://jkwiens.com/2010/01/02/finite-difference-heat-equation-using-numpy/

And the corresponding Python code:

http://jkwiens.com/heat-equation-using-finite-difference/

Here is my code:

import scipy as sp
from scipy import integrate, sparse, linalg
import scipy.sparse.linalg
import pylab as pl

nx = 8000
dx = 0.0025
dt = 0.00002
niter = 20
nonlin = 0.0
gridx = sp.zeros(nx)
igridx = sp.array(range(nx))
psi = sp.zeros(nx)
pot = sp.zeros(nx)
depth = 0.01

# Set up grid, potential, and initial state
gridx = dx*(igridx - nx/2)
pot = depth*gridx*gridx
psi = sp.pi**(-1/4)*sp.exp(-0.5*gridx*gridx)

# Normalize Psi
#psi /= sp.integrate.simps(psi*psi, dx=dx)

# Plot parameters
xlimit = [gridx, gridx[-1]]
ylimit = [0, 2*psi[nx/2]]

# Set up diagonal coefficients
Asup = sp.empty(nx)
Asub = sp.empty(nx)
bdiag = sp.empty(nx)
bsup = sp.empty(nx)
bsub = sp.empty(nx)
Asup.fill(dt/(2*dx**2))
Asub.fill(dt/(2*dx**2))
bdiag.fill(1 + dt/dx**2)
bsup.fill(-dt/(2*dx**2))
bsub.fill(-dt/(2*dx**2))

# Construct tridiagonal matrix
A = sp.sparse.spdiags([Adiag, Asup, Asub], [0, 1, -1], nx, nx)
b = sp.sparse.spdiags([bdiag, bsup, bsub], [0, 1, -1], nx, nx)

# Loop through time
for t in range(0, niter) :
# Calculate effect of potential and nonlinearity
psi *= sp.exp(-dt*(pot + nonlin*psi*psi))

# Calculate spacial derivatives
psi = sp.sparse.linalg.bicg(A, b*psi)

# Normalize Psi
psi /= sp.integrate.simps(psi*psi, dx=dx)

# Output figures
pl.plot(gridx, psi)
pl.plot(gridx, psi*psi)
pl.plot(gridx, pot)
pl.xlim(xlimit)
pl.ylim(ylimit)
pl.savefig('outputla/fig' + str(t))
pl.clf()


The nonlin variable is just an extra term for nonlinearity (I'll be solving the nonlinear Schroedinger equation later), but if I set it to zero then it is irrelevant.

My problem now is that the spatial operator doesn't do anything (it's supposed to diffuse rather quickly), and eventually the entire thing blows up (typical instability.. lines get jagged and eventually the plot just looks crazy).

I've checked everything I can think of, and I have absolutely no clue why this is happening. I have changed linear solvers, dense matrices instead of sparse, the coefficients for A and b are correct (I think.. they are different from the blog link because Schro. eq. has an i term, so some negative-positive sign swappings happen.. this is mentioned in the paper referenced below).

Also, here is a paper that describes the calculations (but I don't use the recursion relation, just an operator): http://arxiv.org/abs/0904.3131

• Not that this is your problem, but innocent changes to your code could quickly lead to some unexpected behavior due to integer division. Unless of course you're doing this in 3.0, in which case ignore this. – user3224 Mar 31 '13 at 5:26
• I am using 3.0, but thank you for pointing that out as I do want to support 2.7 as well. – hadsed Mar 31 '13 at 14:44

If you change dt to 0.001, for instance, there should just be a noticeable difference between figure 1 and figure 19. Of course, only alternating frames can be compared, a bizarre artifact of trying to de-complexify the Schrödinger equation.
Note though that you are going forward in "timesteps" of $-\mathrm{i} \Delta t$ (note the sign!). Compare the tridiagonal matrices you have with the matrix for spatial differencing, $$+\frac{\partial^2}{\partial x^2} \to A \equiv \begin{pmatrix} -2 & 1 \\ 1 & -2 & 1 \\ & 1 & -2 \\ & & & \ddots \end{pmatrix}.$$ Applying this and the typical Crank-Nicolson method to $$\mathrm{i} \frac{\partial}{\partial t} \psi = \left(-\frac{\partial^2}{\partial x^2} + V\right) \psi$$ results in $$\left(I - \frac{\mathrm{i}\Delta t}{2(\Delta x)^2} A\right) \psi_{n+1} = \left(I + \frac{\mathrm{i}\Delta t}{2(\Delta x)^2} A\right) \psi_n.$$
Either change the sign of dt in all the entries in matrices A and b, or for a quick fix put a negative sign in your definition of dt (while making it larger - no need to take such small timesteps). Then you actually do see numerically stable diffusion. As written, your code is running the diffusion equation backward, so blowing up is not unexpected.