# 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[0], 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)[0]

# 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

It is working, but slowly. Your timestep is 125 times smaller than your grid separation. Given that all the other constants are unity, your cell-crossing time is 125 timesteps. It takes a long time to see the changes.

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.

• Thanks Chris, you were exactly right. I've got a followup Q: I've always been taught that dt << dx, i.e. about two orders of magnitude. Is this not true, even for the nonlinear Schroedinger equation (propagated in real time as opposed to imaginary)? Also, is there a particular reason why only every other figure looks right? – hadsed Mar 31 '13 at 18:51
• Stability is a complicated subject. You can get started with Von Neumann and the CFL condition. Numerical recipes also has a good intro. In short, there is no rule that fits all situations. It depends on the equation, and the diffusion equation (evolved forward) is very good at smoothing things (though dispersion might become an issue perhaps?). – user3224 Apr 1 '13 at 0:14
• The every other figure thing is tied to the imaginary time stuff. I'm not entirely sure about the details (this is a method I've never dealt with), but it makes sense considering you've thrown away half the wavefunction by evolving it as a purely real entity. – user3224 Apr 1 '13 at 0:16
• i thought i should mention that i was able to get rid of the strange pulsing behavior of the snapshots at each timestep, but i'm not sure why. i believe it was because of normalization (but i can't remember if i was doing that wrong, it was a while ago). anyway, here is the relevant [fixed] code, if anyone is looking at this thread in the future: github.com/hadsed/PySchro/blob/… – hadsed Aug 24 '13 at 5:35