Solving 1-D Schrodinger Equation in Python

Published in

GWLab

4 min readMay 26, 2021

As a student majoring in Physics, Quantum Mechanics is an important subject to learn. Solving the time-dependent Schrodinger Equation, thereby seeing the time-evolution of wave-function numerically, can be an important experience to achieve a good understanding of Quantum Dynamics. In this article, I’ll show you how to use python to generate a short animation about a simple-harmonic-oscillator, a wavepacket moving back and forth in a quadratic potential well:

A wavepacket moving back and forth in a quadratic potential well

Discretization (Finite Difference)

This is the one-dimensional time-dependent Schrodinger Equation:

To solve it numerically, we first discretize it in the spatial direction, i.e., using the finite-difference scheme:

The corresponding Laplace operator can be expressed in terms of a matrix D2:

, where the dx is the spacing between discretized spatial grid points.

Since the matrix D2 is a sparse banded matrix, we used the function scipy.sparse.diags() to generate the sparse version of it, which gives us better performance.

dx    = 0.02                       # spatial separation
x     = np.arange(0, 10, dx)       # spatial grid points# Laplace Operator (Finite Difference)
D2 = scipy.sparse.diags([1, -2, 1], 
                        [-1, 0, 1],
                        shape=(x.size, x.size)) / dx**2

By using this D2 operator, we may write down the so-called Right-Hand-Side (RHS) of the Schrodinger Equation, psi_t(t,ψ), in the discretized form:

# RHS of Schrodinger Equation
hbar = 1
def psi_t(t, psi):
    return -1j * (- 0.5 * hbar / m * D2.dot(psi) + V / hbar * psi)

Time-evolution (Runge–Kutta method)

The function psi_t(t,ψ) is the RHS of our system that we will feed into the Initial-Value-Problem solver, scipy.integrate.solve_ivp():

dt = 0.005  # time interval for snapshots
t0 = 0.0    # initial time
tf = 1.0    # final time
t_eval = np.arange(t0, tf, dt)  # recorded time shots# Solve the Initial Value Problem
sol = scipy.integrate.solve_ivp(psi_t, 
                                t_span = [t0, tf],
                                y0 = psi0, 
                                t_eval = t_eval,
                                method="RK23")

Conservation of Probability

The time evolution driven by the Schrodinger equation guarantees the conservation of probability. This can be verified by printing out the total probability in each snapshot, which is a good sanity check:

# Print Total Probability (Should = 1)
for i, t in enumerate(sol.t):
    print(np.sum(np.abs(sol.y[:,i])**2) * dx)

Plotting

The simulation result can be plotted by:

for i, t in enumerate(sol.t):
    plt.plot(x, np.abs(sol.y[:,i])**2)

Animation

To generate an animation like the one shown at the beginning of this article, we can use the matplotlib.animation package:

import matplotlib.pyplot as plt
from matplotlib import animationfig = plt.figure()
ax1 = plt.subplot(1,1,1)
ax1.set_xlim(0, 10)
ax1.set_ylim(0, 6)
title = ax1.set_title('')
line1, = ax1.plot([], [], "k--")
line2, = ax1.plot([], [])def init():
    line1.set_data(x, V * 0.01)
    return line1,def animate(i):
    line2.set_data(x, np.abs(sol.y[:,i])**2)
    title.set_text('Time = {0:1.3f}'.format(sol.t[i]))
    return line1,anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=len(sol.t), interval=50,                 
                               blit=True)# Save the animation into a short video
anim.save('sho.mp4', fps=15, 
          extra_args=['-vcodec', 'libx264'], dpi=600)

Complete source (Jupyter notebook)

You may find the complete Jupyter notebook in this Github repository: https://github.com/c0rychu/SchrodingerEq_1D_tutorial
If you find it useful, please give me a ⭐️ on github~

Homework?!

Could you try to modify the potential in the code to generate an animation about a wavepacket scattered by a step potential like this?

The solution is here :)

Discussion about the initial condition

Given that the Schrodinger equation is first order in time, we only need to specify ψ(x, t=0) as our initial condition. This is very different from the classical wave equation, which is second order in time and we have to specify both ψ(x, t=0) and its first time-derivative ∂ψ(x, t=0)/∂t as our initial condition.

C++ Version

If you would like to write the whole thing in C++, here is an example code:
https://github.com/c0rychu/SchrodingerEq_1D_tutorial/tree/master/step_potential_cpp
Especially, we wrote the RK4 time-evolution solver part by ourselves while the matrix operations are tackled by the Eigen library.