Numerical solution of PDE:s, Part 7: 2D Schrödinger equation

Haven’t been posting for a while, but here’s something new… Earlier I showed how to solve the 1D Schrödinger equation numerically in different situations. Now I’m going to show how to calculate the evolution of a 2D wavepacket in a potential energy field that has been constructed to mimic the classical “two-slit experiment” which shows how the mechanics of low-mass particles like electrons can exhibit interference similar to the mechanics of classical waves (sound, light, water surface, and so on).

A 2D Schrödinger equation for a single particle in a time-independent background potential V(x,y) is

Where the particle mass has been set to 1 and the Planck’s constant to $2\pi$.

To solve this numerically, we need the Crank-Nicolson method, as was the case when solving the 1D problem. More specifically, the linear system to be solved is

with

where the wavefunction now has two position indices and one time index, and the potential energy has only two position indices.

To form a model of the two-slit experiment, we choose a domain 0 < x < 6; 0 < y < 6 and make a potential energy function defined by

IF (x < 2.2 OR x > 3.8 OR (x > 2.7 AND x < 3.3)) THEN IF (3.7 < y < 4) THEN V(x,y) = 30

IF (x < 0.5 OR x > 5.5 OR y < 0.5 OR y > 5.5) THEN V(x,y) = 30

Otherwise V(x,y) = 0.

which corresponds to having hard walls surrounding the domain and a barrier with two holes around the line y = 3.85

For an initial condition, we choose a Gaussian wavepacket that has a nonzero expectation value of the momentum in y-direction:

An R-Code that solves this problem for a time interval 0 < t < 1 is

library(graphics) #load the graphics library needed for plotting

lx <- 6.0 #length of the computational domain (both x and y)
lt <- 1.0 #length of the simulation time interval
nx <- 47 #number of discrete lattice points in x and y directions
nt <- 60 #number of timesteps
dx <- lx/nx #length of one discrete lattice cell, same in x and y directions
dt <- lt/nt #length of timestep

V = matrix(nrow=nx,ncol=nx) #potential energies at discrete xy points

for(j in c(1:nx)) { #construct the potential field with a double-slit barrier
for(k in c(1:nx)) {
V[j,k] = 0+0i
if((j*dx<2.2)||(j*dx>3.8)||((j*dx>2.7) && (j*dx<3.3))) {
if((k*dx>3.7) && (k*dx<4.0)) {
V[j,k] = 30+0i #No significant density is going to go through these barriers
}
}
if((j*dx>5.5) || (j*dx<0.5) || (k*dx>5.5) || (k*dx<0.5)) {
V[j,k] = 30+0i
}
}
}

kappa1 = (1i)*dt/(2*dx*dx) #an element needed for the matrices
kappa2 <- c(1:(nx*nx)) #another element

for(j in c(1:nx)) {
for(k in c(1:nx)) {
kappa2[(j-1)*nx+k] <- kappa1*2*dx*dx*V[j,k]
}
}

psi = c(1:(nx*nx)) #array for the wave function values

for(j in c(1:nx)) {
for(k in c(1:nx)) {
psi[(j-1)*nx+k] = exp(-2*(k*dx-1.7)*(k*dx-1.7)-2*(j*dx-3)*(j*dx-3)+(2i)*k*dx) #Gaussian initial wavefunction
if((j*dx > 5.5)||(j*dx < 0.5)||(k*dx > 5.5)||(k*dx < 0.5)) {
psi[(j-1)*nx+k] = as.complex(0)
}
}
}

xaxis <- c(1:nx)*dx #the x values corresponding to the discrete lattice points
yaxis <- c(1:nx)*dx #the y values

A = matrix(nrow=nx*nx,ncol=nx*nx) #matrix for forward time evolution
B = matrix(nrow=nx*nx,ncol=nx*nx) #matrix for backward time evolution
P = matrix(nrow=nx,ncol=nx) #matrix for the solution after time stepping
IP = matrix(nrow=4*nx,ncol=4*nx) #matrix for the higher resolution image obtained by interpolation

for(j in c(1:(nx*nx))) { #Set the values for time evolution matrix elements
for(k in c(1:(nx*nx))) {
A[j,k]=0
B[j,k]=0
if(j==k) {
A[j,k] = 1 + 4*kappa1 + kappa2[j]
B[j,k] = 1 – 4*kappa1 – kappa2[j]
}
if((k==j+1) || (k==j-1)) {
A[j,k] = -kappa1
B[j,k] = kappa1
}
if((k==j+nx)||(k==j-nx)) {
A[j,k] = -kappa1
B[j,k] = kappa1
}
}
}

for(k in c(1:nt)) { #main time stepping loop

sol <- solve(A,B%*%psi) #solve the system of equations

for(l in c(1:(nx*nx))) {
psi[l] <- sol[l]
}

for(l in c(1:nx)) {
for(m in c(1:nx)) {
P[l,m] = abs(psi[(l-1)*nx + m])*abs(psi[(l-1)*nx + m]) #square of the absolute value of wave function
if(abs(V[l,m]) > 5) P[l,m] = 2
}
}
for(l in c(1:(nx-1))) {
for(m in c(1:(nx-1))) { #make a bitmap with 4 times more pixels, using linear interpolation
IP[4*l-3,4*m-3] = P[l,m]
IP[4*l-2,4*m-3] = P[l,m]+0.25*(P[l+1,m]-P[l,m])
IP[4*l-1,4*m-3] = P[l,m]+0.5*(P[l+1,m]-P[l,m])
IP[4*l,4*m-3] = P[l,m]+0.75*(P[l+1,m]-P[l,m])
}
}

for(l in c(1:(4*nx))) {
for(m in c(1:(nx-1))) {
IP[l,4*m-2] = IP[l,4*m-3]+0.25*(IP[l,4*m+1]-IP[l,4*m-3])
IP[l,4*m-1] = IP[l,4*m-3]+0.5*(IP[l,4*m+1]-IP[l,4*m-3])
IP[l,4*m] = IP[l,4*m-3]+0.75*(IP[l,4*m+1]-IP[l,4*m-3])
}
}

jpeg(file = paste(“plot_abs_”,k,”.jpg”,sep=””)) #save the image
image(IP, zlim = c(0,0.15))
dev.off()

}

The code produces a sequence of image files, where the probability density is plotted with colors, as an output. Some representative images from this sequence (converted to grayscale) is shown below:

A video of the time evolution is shown below:

The treshold for maximum white color has been chosen to be quite low, to make the small amount of probability density that crosses the barrier visible.

The discrete grid of points has been made quite coarse here to keep the computation time reasonable, and the resolution has been increased artificially by using linear interpolation between the discrete points.

So, now we’ve seen how to solve the motion of 2D wavepackets moving around obstacles. In the next numerical methods post, I’ll go through the numerical solution of a nonlinear PDE.

Numerical solution of PDE:s, Part 4: Schrödinger equation

In the earlier posts, I showed how to numerically solve a 1D or 2D diffusion or heat conduction problem using either explicit or implicit finite differencing. In the 1D example, the relevant equation for diffusion was

and an important property of the solution was the conservation of mass,

i.e. the integral of the concentration field over whole domain stays constant.

Next, I will show how to integrate the 1D time-dependent Schrödinger equation, which in a nondimensional form where we set $\hbar = 1$ and $m = 1$ reads:

Here $i$ is the imaginary unit and $V(x)$ is the potential energy as a function of $x$. The solutions of this equation must obey a conservation law similar to the mass conservation in the diffusion equation, the conservation of norm:

where the quantity $|\Psi (x,t)|$ is the modulus of the complex-valued function $\Psi (x,t)$ . This property of the solution is also called unitarity of the time evolution.

Apart from the TDSE, another way to represent the time development of this system is to find the normalized solutions $\psi_0 (x)$, $\psi_1 (x)$, $\psi_2 (x) \dots$ of the time-independent Schrödinger equation

and write the initial state $\Psi (x,0)$ as a linear combination of those basis functions:

This is possible because the solutions of the time-independent equation form a basis for the set of acceptable wave functions $\psi (x)$. Then, every term in that eigenfunction expansion is multiplied by a time dependent phase factor $\exp(-iE_n t)$:

The numbers $E_n$ are the eigenvalues corresponding to the solutions $\psi_n (x)$ and the function $\psi_0 (x)$ is called the ground state corresponding to potential $V(x)$, while the functions $\psi_1 (x)$ is the first excited state and so on.

The Schrödinger equation can’t be discretized by using either the explicit or implicit method that we used when solving the diffusion equation. The method is either numerically unstable or doesn’t conserve the normalization of the wave function (or both) if you try to do that. The correct way to discretize the Schrödinger equation is to replace the wave function with a discrete equivalent

and the potential energy function $V(x)$ with $V_{i;j}$ (or $V_i$ in the case of time-independent potential), and write an equation that basically tells that propagating the state $\Psi_{i;j}$ forward by half a time step gives the same result as propagating the state $\Psi_{i;j+1}$ backwards by half a time step:

Here we have

and

This kind of discretization is called the Crank-Nicolson method. As boundary conditions, we usually set that at the boundaries of the computational domain the wavefunction stays at value zero: $\Psi (0,t) = \Psi (L,t) = 0$ for any value of $t$. In the diffusion problem, this kind of a BC corresponded to infinite sinks at the boundaries, that annihilated anything that diffused through them. In the Schrödinger equation problem, which is a complex diffusion equation, the equivalent condition makes the boundaries impenetrable walls that deflect elastically anything that collides with them.

An R-Code that calculates the time evolution of a Gaussian initial wavefunction

in an area of zero potential:

for a domain 0 < x < 6, a lattice spacing $\Delta x = 0.05$, time interval 0 < t < 2 and time step $\Delta t = 0.01$, is given below:

library(graphics) #load the graphics library needed for plotting

lx <- 6.0 #length of the computational domain
lt <- 2.0 #length of the simulation time interval
nx <- 120 #number of discrete lattice points
nt <- 200 #number of timesteps
dx <- lx/nx #length of one discrete lattice cell
dt <- lt/nt #length of timestep

V = c(1:nx) #potential energies at discrete points

for(j in c(1:nx)) {
V[j] = 0 #zero potential
}

kappa1 = (1i)*dt/(2*dx*dx) #an element needed for the matrices
kappa2 <- c(1:nx) #another element

for(j in c(1:nx)) {
kappa2[j] <- as.complex(kappa1*2*dx*dx*V[j])
}

psi = as.complex(c(1:nx)) #array for the wave function values

for(j in c(1:nx)) {
psi[j] = as.complex(exp(-2*(j*dx-3)*(j*dx-3))) #Gaussian initial wavefunction
}

xaxis <- c(1:nx)*dx #the x values corresponding to the discrete lattice points

A = matrix(nrow=nx,ncol=nx) #matrix for forward time evolution
B = matrix(nrow=nx,ncol=nx) #matrix for backward time evolution

for(j in c(1:nx)) {
for(k in c(1:nx)) {
A[j,k]=0
B[j,k]=0
if(j==k) {
A[j,k] = 1 + 2*kappa1 + kappa2[j]
B[j,k] = 1 - 2*kappa1 - kappa2[j]
}
if((j==k+1) || (j==k-1)) {
A[j,k] = -kappa1
B[j,k] = kappa1
}
}
}

for (k in c(1:nt)) { #main time stepping loop

sol <- solve(A,B%*%psi) #solve the system of equations

for (l in c(1:nx)) {
psi[l] <- sol[l]
}

if(k %% 3 == 1) { #make plots of |psi(x)|^2 on every third timestep
jpeg(file = paste("plot_",k,".jpg",sep=""))
plot(xaxis,abs(psi)^2,xlab="position (x)", ylab="Abs(Psi)^2",ylim=c(0,2))
title(paste("|psi(x,t)|^2 at t =",k*dt))
lines(xaxis,abs(psi)^2)
dev.off()
}
}

The output files are plots of the absolute squares of the wavefunction, and a few of them are shown below.

In the next simulation, I set the domain and discrete step sizes the same as above, but the initial state is:

Which is a Gaussian wave packet that has a nonzero momentum in the positive x-direction. This is done by changing the line

for(j in c(1:nx)) {
psi[j] = as.complex(exp(-2*(j*dx-3)*(j*dx-3))) #Gaussian initial wavefunction
}+(1i)*j*dx

into

for(j in c(1:nx)) {
psi[j] = as.complex(exp(-2*(j*dx-3)*(j*dx-3)+(1i)*j*dx)) #Gaussian initial wavefunction
}

The plots of $|\Psi (x,t)|^2$ for several values of $t$ are shown below

and there you can see how the wave packet collides with the right boundary of the domain and bounces back.

In the last simulation, I will set the potential function to be

which is a harmonic oscillator potential, and with the nondimensional mass $m =1$ and Planck constant $\hbar = 1$ the ground state $\psi _0 (x)$ of this system is

If I’d set the initial state to be $\Psi (x,0) = \psi_0 (x)$, or any other solution of the time-independent SE, the modulus of the wavefunction would not change at all. To get something interesting to happen, I instead set an initial state that is a displaced version of the ground state:

The solution can be obtained with the code shown below:

library(graphics) #load the graphics library needed for plotting

lx <- 6.0 #length of the computational domain
lt <- 3.0 #length of the simulation time interval
nx <- 360 #number of discrete lattice points
nt <- 300 #number of timesteps
dx <- lx/nx #length of one discrete lattice cell
dt <- lt/nt #length of timestep

V = c(1:nx) #potential energies at discrete points

for(j in c(1:nx)) {
V[j] = as.complex(2*(j*dx-3)*(j*dx-3)) #Harmonic oscillator potential with k=4
}

kappa1 = (1i)*dt/(2*dx*dx) #an element needed for the matrices
kappa2 <- c(1:nx) #another element

for(j in c(1:nx)) {
kappa2[j] <- as.complex(kappa1*2*dx*dx*V[j])
}

psi = as.complex(c(1:nx)) #array for the wave function values

for(j in c(1:nx)) {
psi[j] = as.complex(exp(-(j*dx-2)*(j*dx-2))) #Gaussian initial wavefunction, displaced from equilibrium
}

xaxis <- c(1:nx)*dx #the x values corresponding to the discrete lattice points

A = matrix(nrow=nx,ncol=nx) #matrix for forward time evolution
B = matrix(nrow=nx,ncol=nx) #matrix for backward time evolution

for(j in c(1:nx)) {
for(k in c(1:nx)) {
A[j,k]=0
B[j,k]=0
if(j==k) {
A[j,k] = 1 + 2*kappa1 + kappa2[j]
B[j,k] = 1 - 2*kappa1 - kappa2[j]
}
if((j==k+1) || (j==k-1)) {
A[j,k] = -kappa1
B[j,k] = kappa1
}
}
}

for (k in c(1:nt)) { #main time stepping loop

sol <- solve(A,B%*%psi) #solve the system of equations

for (l in c(1:nx)) {
psi[l] <- sol[l]
}

if(k %% 3 == 1) { #make plots of Abs(psi(x))^2 on every third timestep
jpeg(file = paste("plot_",k,".jpg",sep=""))
plot(xaxis,abs(psi)^2, xlab="position (x)", ylab="Abs(Psi)^2",ylim=c(0,2))
title(paste("|psi(x,t)|^2 at t =",k*dt))
lines(xaxis,abs(psi)^2)
lines(xaxis,V)
dev.off()
}
}

and the solution at different values of $t$ look like this (images and video):

Here the shape of the Hookean potential energy is plotted in the same images. So, here you see how the center of the Gaussian wavefunction oscillates around the point $x = 3$, just like a classical mechanical harmonic oscillator does when set free from a position that is displaced from equilibrium.

By changing the code that produces the output images, we can also get a sequence of plots of the imaginary part of the wavefunction:

if(k %% 3 == 1) { #make plots of Im(psi(x)) on every third timestep
jpeg(file = paste("plot_",k,".jpg",sep=""))
plot(xaxis,Im(psi), xlab="position (x)", ylab="Im(Psi)",ylim=c(-1.5,1.5))
title(paste("Im(psi(x,t)) at t =",k*dt))
lines(xaxis,Im(psi))
lines(xaxis,V)
dev.off()
}

and the resulting plots look like this: