# Partial Differential Equations

Exercise 9.8: The Schrödinger equation and the Crank–Nicolson method

Perhaps the most important partial differential equation, at least for physicists, is the Schrödinger equation. This exercise uses the Crank–Nicolson method to solve the full time-dependent Schrödinger equation and hence develop a picture of how a wavefunction evolves over time. The following exercise, Exercise 9.9, solves the same problem again, but using the spectral method.

We will look at the Schrödinger equation in one dimension. The techniques for calculating solutions in two or three dimensions are basically the same as for one dimension, but the calculations take much longer on the computer, so in the interests of speed we’ll stick with one dimension. In one dimension the Schrödinger equation for a particle of mass M with no potential energy reads

$-{\hbar^2\over{2M}} {\partial^2\psi\over{\partial x^2}} = i \hbar {\partial\psi\over{\partial t}}$

For simplicity, let’s put our particle in a box with impenetrable walls, so that we only have to solve the equation in a finite-sized space. The box forces the wavefunction ψ to be zero at the walls, which we will put at x = 0 and x = L.

Replacing the second derivative in the Schrödinger equation with a finite difference and applying Euler’s method, we get the FTCS equation

$\psi(x,t+h) = \psi(x,t) + h {i \hbar \over{2ma^2}} \big[ \psi(x+a,t) + \psi(x-a,t) - 2 \psi(x,t) \big]$

where a is the spacing of the spatial grid points and h is the size of the time-step. (Be careful not to confuse the time-step h with Planck’s constant $\hbar$.) Performing a similar step in reverse, we get the implicit equation

$\psi(x,t) = \psi(x,t+h) - h {i \hbar \over {2ma^2}} \big[ \psi(x+a,t+h) + \psi(x-a,t+h) -2\psi(x,t+h) \big].$

And taking the average of these two, we get the Crank–Nicolson equation for the Schrödinger equation:

$\psi(x,t+h) - h {i\hbar\over 4ma^2} \big[ \psi(x+a,t+h) + \psi(x-a,t+h) -2 \psi(x,t+h) \big]$
$= \psi(x,t) - h {i\hbar\over 4ma^2} \big[ \psi(x+a,t) + \psi(x-a,t) -2 (x,t) \big]$

This gives us a set of simultaneous equations, one for each grid point.

The boundary conditions on our problem tell us that ψ = 0 at x = 0 and x = L for all t. In between these points we have grid points at a, 2a, 3a, and so forth. Let us arrange the values of ψ at these interior points into a vector

$\boldsymbol{\psi}(t) = \begin{pmatrix} \psi(a,t)\\ \psi(2a,t)\\ \psi(3a,t)\\ \vdots \\ \end{pmatrix}.$

Then the Crank–Nicolson equations can be written in the form Aψ(t + h) = Bψ(t),
where the matrices A and B are both symmetric and tridiagonal:

$\boldsymbol{A} = \begin{pmatrix} a_1 & a_2\\ a_2 & a_1 & a_2\\ & a_2 & a_1 & a_2\\ & & & \ddots \\ \end{pmatrix}, \boldsymbol{B} = \begin{pmatrix} b_1 & b_2\\ b_2 & b_1 & b_2\\ & b_2 & b_1 & b_2\\ & & & \ddots \\ \end{pmatrix}$

with

$a_1 = 1 + h {i \hbar \over{2ma^2}}, \, \, \, \, \, \, a_2 = -h {i \hbar \over{4ma^2}}, \, \, \, \, \, b_1 = 1 - h {i \hbar \over{2ma^2}}, \, \, \, \, \, \, b_2 = h {i \hbar \over{4ma^2}}.$

(Note the different signs and the factors of 2 and 4 in the denominators.)

The equation Aψ(t + h) = Bψ(t) has precisely the form Ax = v of the simultaneous equation problems we studied in Chapter 6 and can be solved using the same methods. Specifically, since the matrix A is tridiagonal in this case, we can use the fast tridiagonal version of Gaussian elimination that we looked at in Section 6.1.6.

Consider an electron (mass M = 9.109 × 10−31 kg) in a box of length L = 10−8 m. Suppose that at time t = 0 the wavefunction of the electron has the form

$\psi(x,0) = \exp[-{(x-x_0)^2 \over {2\sigma^2}}] e^{i \kappa x},$

where

$x_0 = {L\over2}, \, \, \, \, \sigma = 1 \times 10^{-10} \,m, \, \, \, \, \kappa = 5 \times 10^{10} \, m^{-1},$

and ψ = 0 on the walls at x = 0 and x = L. (This expression for ψ(x, 0) is not normalized— there should really be an overall multiplying coefficient to make sure that the probability density for the electron integrates to unity. It’s safe to drop the constant, however, because the Schrödinger equation is linear, so the constant cancels out on both sides of the equation and plays no part in the solution.)

1. Write a program to perform a single step of the Crank–Nicolson method for this electron, calculating the vector ψ(t) of values of the wavefunction, given the initial wavefunction above and using N = 1000 spatial slices with a = L/N. Your program will have to perform the following steps. First, given the vector ψ(0) at t = 0, you will have to multiply by the matrix B to get a vector v = Bψ. Because of the tridiagonal form of B, this is fairly simple. The ith component of v is given by

vi = b1ψi + b2i+1 + ψi-1).

You will also have to choose a value for the time-step h. A reasonable choice is h = 10−18 s.

Second you will have to solve the linear system Ax = v for x, which gives you the new value of ψ. You could do this using a standard linear equation solver like the function solve in numpy.linalg, but since the matrix A is tridiagonal a better approach would be to use the fast solver for banded matrices given in Appendix E, which can be imported from the file banded.py (which you can find in the on-line resources). Note that although the wavefunction of a particle in principle has a complex value, in this case the wave- function is always real—all the coefficients in the equations above are real numbers so if, as here, the wavefunction starts off real, then it remains real. Thus you do not need to use a complex array to represent the vector ψ. A real one will do the job.

Third, once you have the code in place to perform a single step of the calculation, extend your program to perform repeated steps and hence solve for ψ at a sequence of times a separation h apart. Note that the matrix A is independent of time, so it doesn’t change from one step to another. You can set up the matrix just once and then keep on reusing it for every step.

2. Extend your program to make an animation of the solution by displaying the real part of the wavefunction at each time-step. You can use the function rate from the package visual to ensure a smooth frame-rate for your animation—see Section 3.5 on page 117.There are various ways you could do the animation. A simple one would be to just place a small sphere at each grid point with vertical position representing the value of the real part of the wavefunction. A more sophisticated approach would be to use the curve object from the visual package—see the on-line documentation at www.vpython.org for details. Depending on what coordinates you use for measuring x, you may need to scale the values of the wavefunction by an additional constant to make them a reasonable size on the screen. (If you measure your x position in meters then a scale factor of about 10−9 works well for the wavefunction.)
3. Run your animation for a while and describe what you see. Write a few sentences explaining in physics terms what is going on in the system.

Exercise 9.9: The Schrödinger equation and the spectral method

This exercise uses the spectral method to solve the time-dependent Schrödinger equation

$-{\hbar^2\over{2M}} {\partial^2\psi\over{\partial x^2}} = i \hbar {\partial\psi\over{\partial t}}$

for the same system as in Exercise 9.8, a single particle in one dimension in a box of length L with impenetrable walls. The wavefunction in such a box necessarily goes to zero on the walls and hence one possible (unnormalized) solution of the equation is

$\psi_k(x,t) = \sin \big( {\pi k x\over L} \big)\,e^{i Et/\hbar},$

where the energy E can be found by substituting into the Schrödinger equation, giving

$E = {\pi^2\hbar^2k^2\over{2ML^2}}.$

As with the vibrating string of Section 9.3.4, we can write a full solution as a linear combination

of such individual solutions, which on the grid points xn = nL/N takes the value

$\psi(x_n,t) = {1\over N} \sum_{k=1}^{N-1} b_k \sin ( {\pi k n\over N} )\> \exp ( i{\pi^2\hbar k^2\over{2ML^2}} t),$

where the bk are some set of (possibly complex) coefficients that specify the exact shape of the wavefunction and the leading factor of 1/N is optional but convenient.

Since the Schrödinger equation (unlike the wave equation) is first order in time, we need only a single initial condition on the value of ψ(x, t) to specify the coefficients bk, although, since the coefficients are in general complex, we will need to calculate both real and imaginary parts of each coefficient.

As in Exercise 9.8 we consider an electron (mass M = 9.109 × 10−31 kg) in a box of length L = 10−8 m. At time t = 0 the wavefunction of the electron has the form

$\psi(x,0) = \exp [ -{(x-x_0)^2\over{2\sigma^2}} ] e^{i \kappa x},$

where

$x_0 = {L\over2}, \, \, \, \, \sigma = 1 \times 10^{-10} \,m, \, \, \, \, \kappa = 5 \times 10^{10} \, m^{-1},$

and ψ=0 on the walls at x=0 and x=L.

1. Write a program to calculate the values of the coefficients bk, which for convenience can be broken down into their real and imaginary parts as bk = αk + iηk. Divide the box into N = 1000 slices and create two arrays containing the real and imaginary parts of ψ(xn, 0) at each grid point. Perform discrete sine transforms on each array separately and hence calculate the values of the αk and ηk for all k=1…N−1.To perform the discrete sine transforms, you can use the fast transform function dst from the package dcst, which you can find in the on-line resources in the file named dcst.py. A copy of the code for the package can also be found in Appendix E. The function takes an array of N real numbers and returns the discrete sine transform as another array of N numbers.(Note that the first element of the input array should in principle always be zero for a sine transform, but if it is not the dst function will simply pretend that it is. Similarly the first element of the returned array is always zero, since the k = 0 coefficient of a sine transform is always zero. So in effect, the sine transform really only takes N − 1 real numbers and transforms them into another N − 1 real numbers. In some implementations of the discrete sine transform, therefore, though not the one in the package dsct used here, the first element of each array is simply omitted, since it’s always zero anyway, and the arrays are only N − 1 elements long.)
2. b)  Putting bk = αk + iηk in the solution above and taking the real part we get

$\Re \psi(x_n,t) = {1\over N} \sum_{k=1}^{N-1} [ \alpha_k \cos ({\pi^2 \hbar k^2 \over {2ML^2}} t ) - \eta_k \sin ({\pi^2 \hbar k^2 \over {2ML^2}} t )] \sin ( {\pi k n\over N} )$

for the real part of the wavefunction. This is an inverse sine transform with coefficients equal to the quantities in the square brackets. Extend your program to calculate the real part of the wavefunction ψ(x, t) at an arbitrary time t using this formula and the inverse discrete sine transform function idst, also from the package dcst. Test your program by making a graph of the wavefunction at time t = 10−16 s.

3. c)  Extend your program further to make an animation of the wavefunction over time, similar to that described in part (b) of Exercise 9.8 above. A suitable time interval for each frame of the animation is about 10−18 s.
4. d)  Run your animation for a while and describe what you see. Write a few sentences explaining in physics terms what is going on in the system.