Numerical Quantum Mechanics: The Time-dependent Schrödinger Equation (III)

Back in 2016, I wrote two blog posts (I and II) concerning how to solve the time-dependent Schrödinger equation (TDSE) numerically. Although I gave recipes for computing the spatial derivative that occurs in the TDSE, and for integrating the TDSE itself, I omitted a rather important aspect of the computation: boundary conditions (BC’s).

In this post, I will rectify that situation, and fully describe how to implement two frequently used BC’s, namely, periodic and reflective. I have created a Python script that implements the equations I shared in the previous post and the BC’s described here. To keep things focused, I will stay in one dimension, i.e., a particle on a line.

Periodic boundary conditions

Periodic BC’s connect one side of the computational domain to another side, creating a closed domain with no ‘edges’ (much like the surface of the earth). In the case of a particle in a 1D universe, it may help to think of that universe as a ring. Our domain is discretized using gridpoints; periodic BC’s then imply that

,

because the wave function must be continuous (we use indexing from to as in most programming languages).

Two things are needed for a correct computation: 1) we must choose an initial wave function that is compatible with the BC’s, and 2) we must design our second-spatial-derivative function (see post II) to take the periodicity of the space into account.

Note that we used the central-difference method for computing the second derivative. This means that, at the indices and , special cases arise, where one of the two neighboring indices points to a non-existent array entry. To solve this problem in a way that is consistent with periodic boundary conditions, we simply identify the `left neighbor’ of index – which would be index – with index , and the `right neighbor’ of index with index .

To show how this might be implemented in code, please examine the following Python fragment, which implements the second spatial derivative for a complex-valued vector psi using periodic BC’s:

def second_deriv(psi):
    secondderiv = np.zeros(Npoints, dtype=complex)
    for i in range(1, Npoints - 1):
        secondderiv[i] = (psi[i+1] + psi[i-1] - 2. * psi[i]) / (dx * dx)
    secondderiv[0]         = (psi[1] + psi[Npoints-1] - 2. * psi[0]) / (dx * dx)  
    secondderiv[Npoints-1] = (psi[0] + psi[Npoints-2] - 2. * psi[Npoints-1]) / (dx * dx)   
    return secondderiv

And finally, here is an animation that shows periodic BC’s in action (the Python script used to generate this animation is included at the end of this post). The video lasts just over half a period for this state. Note that the initial wave function is symmetric, so that the BC’s are respected. I track the wave function’s norm to get an indication of the accumulated error of the integration.

Time evolution of an initial Gaussian wave function with periodic BC’s.

Reflecting boundary conditions

These BC’s are also known as ‘hard wall’ or Dirichlet BC’s. Implementing them renders the edges of the domain into impenetrable barriers. This situation corresponds to the well-known `infinite square well’ potential (so we are sneaking in a potential using BC’s, rather than explicitly specifying !) Reflecting BC’s imply that

,

i.e., the wave function vanishes at the domain edges. To implement these boundary conditions, we must once again choose a compatible wave function (the Gaussian will again suffice for a reasonable result, even though it is not quite zero at the domain boundaries). Now, since we have `forced’ boundary conditions where and are constant in time, we only update the indices through . Due to the nature of the implementation of our RK4 integration algorithm (see Python script below), it is convenient to accomplish this by setting the first and last indices of our second derivative to zero:

def second_deriv(psi):
    secondderiv = np.zeros(Npoints, dtype=complex)
    for i in range(1, Npoints - 1):
        secondderiv[i] = (psi[i+1] + psi[i-1] - 2. * psi[i]) / (dx * dx)
    secondderiv[0]         = 0.
    secondderiv[Npoints-1] = 0.  
    return secondderiv

And once again, here is an animation that shows the BC’s in action, this time lasting just over one period for this state:

Time evolution of an initial Gaussian wave function with hard-wall (Dirichlet) BC’s.

Final comments

We did not consider potentials. When specifying a potential in periodic BC’s, please ensure the potential is itself periodic. We only considered a 1D domain, but extending these BC’s from 1D to ND should be straightforward (contact me if you have issues!) Finally, another type of BC’s that could be of interest are lossy or absorbing BC’s. These are much more complicated to implement. For a freely available, wonderfully pessimistic book about this subject, please see Chapter 6 of Trefethen.

Python script

This script should produce an animated plot. To save your animation, please install ffmpeg and uncomment lines 12, 99, and 100.

# This Python script integrates the time-dependent Schrodinger equation 
# using a 4th-order Runge-Kutta algorithm and plots an animation of the
# result.
# To save the animation, install FFMPEG and uncomment lines 12, 99, and 100.
# Author: Thomas Bronzwaer
# Date: 10 Aug. 2020

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import cmath
#plt.rcParams['animation.ffmpeg_path'] = '/usr/local/bin/ffmpeg'

# NUMERICAL METHODS
###################

def get_norm(psi):
    norm = 0.
    for i in range(Npoints):
        norm += dx * abs(psi[i])**2
    return norm

# Return a complex vector containing the second spatial derivative of psi.
def second_deriv(psi):
    secondderiv = np.zeros(Npoints, dtype=complex)
    for i in range(1, Npoints - 1):
        secondderiv[i] = (psi[i+1] + psi[i-1] - 2. * psi[i]) / (dx * dx)
    secondderiv[0]         = (psi[1] + psi[Npoints-1] - 2. * psi[0]) / (dx * dx)  
    secondderiv[Npoints-1] = (psi[0] + psi[Npoints-2] - 2. * psi[Npoints-1]) / (dx * dx)   
    return secondderiv

# Update psi using a single RK4 step with timestep dt.
def update_psi_rk4(psi, dt):
    k1 = -(1. / (2.j)) * second_deriv(psi)
    k2 = -(1. / (2.j)) * second_deriv(psi + 0.5 * dt * k1)
    k3 = -(1. / (2.j)) * second_deriv(psi + 0.5 * dt * k2)
    k4 = -(1. / (2.j)) * second_deriv(psi + dt * k3)
    
    psinew = psi + dt * 1. / 6. * (k1 + 2. * k2 + 2. * k3 + k4)
    return psinew

# SIMULATION PARAMETERS
#######################

L         = 1. # Domain length is 2L. Specified in units of a_0.
Npoints   = 201    
sigma     = 1./12.
x         = np.linspace(-L, L, Npoints)
dx        = x[1]-x[0]
time_unit = 2.4188843265857e-17
timestep  = 0.0001
psi       = np.zeros(Npoints, dtype=complex)

# Initialize and normalize psi.
for i in range(Npoints):
    psi[i] = np.exp(-x[i]**2 / 2 / sigma / sigma)  
norm = get_norm(psi)
psi = psi/np.sqrt(norm)
   
# Set up figure.
fig, ax = plt.subplots()
line, = ax.plot(x, psi.real)
plt.ylim(-.0,abs(np.max(psi))**2 * 1.05)
plt.xlim(-1.,1.)
plt.xlabel(r'$x$ [$a_0$]')
plt.ylabel(r'$\left| \psi \right|^2$')
textheight = abs(np.max(psi))**2
steptext = ax.text(-0.98, textheight * 1, 'Integration step: 0')
timetext = ax.text(-0.98, textheight * 0.95, 'Elapsed time [s]: 0')
normtext = ax.text(-0.98, textheight * 0.9, 'Norm: 0.99999999999')
plt.title(r'Time evolution of Gaussian $\psi_0$, periodic boundary conditions')

rk4_steps_per_frame = 4

# ANIMATION FUNCTIONS
#####################

def animate(i):
    global psi
    for q in range(rk4_steps_per_frame):
        psinew = update_psi_rk4(psi, timestep)
        psi = psinew
    currentnorm = get_norm(psi)
    line.set_ydata(abs(psi)**2)  # update the data
    steptext.set_text('Integration step: %s'%(i * rk4_steps_per_frame))
    timetext.set_text('Elapsed time [s]: %.3e'%(i * rk4_steps_per_frame * timestep * time_unit))
    normtext.set_text('Norm: %.15s'%(currentnorm))
    return line,

def init():
    return line,

# RUN ANIMATION AND SAVE OR SHOW
################################

ani = animation.FuncAnimation(fig, animate, np.arange(1, 1620), init_func=init,
                              interval=25, save_count=1620)

#FFwriter=animation.FFMpegWriter(fps=60, extra_args=['-vcodec', 'libx264'])
#ani.save('/Users/thomasbronzwaer/Desktop/Dirichlet_BC_.mp4', writer = FFwriter)

plt.show()

Advertisement