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

Question

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
Adiag = sp.empty(nx)
Asup = sp.empty(nx)
Asub = sp.empty(nx)
bdiag = sp.empty(nx)
bsup = sp.empty(nx)
bsub = sp.empty(nx)
Adiag.fill(1 - dt/dx**2)
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

Answer 1

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.