Quantum Mechanics with Python
Solving the 1D Time Independent Schrödinger Equation Numerically
It seems that no one, even the experts, really understands quantum mechanics. No really! Richard Feynman, the famous Nobel laureate in physics and one of the most brilliant scientists to have walked the earth once quoted “I think I can safely say that nobody really understands quantum mechanics”. Despite this lack of understanding, quantum mechanics has also turned out to be perhaps the most successful scientific theory ever produced by mankind.
Quantum mechanics is the basis of chemistry, which is the basis of biology and therefore life. Quantum mechanics explains why gold is golden and not silvery like the other metals, why the sun shines, how solar cells turn light into electricity, and how nuclear bombs work. Without quantum mechanics, the world would have turned out very differently.
One reason why I chose to pursue physics was because I simply had to know more about quantum mechanics. This was to me the proverbial unbearable itch that simply had to be scratched! Today I want to share a small bit of what I know, and hope that you become more interested in this fascinating field too!
The Time Independent Schrödinger Equation
The Schrödinger equation is to quantum mechanics what Newton’s laws of motion are to kinematics. The Schrödinger equation describes how sub-atomic particles behave in a quantum mechanical system. In the general case, the Schrödinger equation contains both time and spatial derivatives which are not trivial to solve even with computational methods.
Fortunately it can be assumed that the temporal and spatial portions of the solution can be obtained through the separation of variables, and we can focus on solving the portion involving spatial derivatives only. This portion of the equation is the time independent Schrödinger equation, which we will explore today.
Eigenvalues and Eigenvectors
The time independent Schrödinger equation is an eigenvalue problem. This means that it can be cast in matrix form as: Hψ = Eψ, where H is the Hamiltonian matrix (the Hamiltonian is essentially the sum of a particle’s kinetic and potential energies), ψ is the wave function vector and E is the energy eigenvalue. For the less mathematically inclined, this relation simply means that multiplying the matrix H to the vector ψ gives the same result as multiplying the scalar value E to the vector ψ. Physically this means that operating the Hamiltonian which is the sum of kinetic and potential energies returns the total energy of the particle!
If we can determine the form of the matrix H, we can then numerically solve the eigenvalue equation to get the both quantum mechanical wave function ψ, as well as the corresponding total particle energy E.
As can be seen in the equation at the top of this post, the matrix H is the sum of a second order derivative (physically this is the kinetic energy) and the potential energy V. For simplicity’s sake, we will assume that the Planck constant ħ and particle mass m are both equal to 1. Also, we will solve the Schrödinger equation numerically in one dimension only.
Discretizing the 1D Schrödinger Equation
The first step is to implement the spatial derivative numerically. By using second order central finite differences in one dimension, the wave function ψ is differentiated numerically as: -1/(2dx²)(ψ[n-1]–2ψ[n] + ψ[n+1]). The result is the discretised 1 dimensional Schrödinger equation shown below. Note that we have set ħ = m = 1 for simplicity’s sake, and that dx is the step size of the spatial grid used.
Therefore if we scale the values of V and E by a factor of 2, the Hamiltonian matrix H takes the form of a tri-diagonal matrix with 2/dx² in the main diagonal, and -1/dx² in the first diagonals above and below the main diagonal. We then add the potential energy V to complete the Hamiltonian matrix H.
Note that setting up the Hamiltonian matrix in such a manner results in Dirichlet boundary conditions being intrinsically applied to the numerical system. This means that the wave function ψ disappears at the boundaries of our system. This should work for most situations in general, but might break down for situations which require other forms of boundary conditions such as periodic boundary conditions.
Solving the Eigenvalue Problem Numerically
Instead of using standard linear algebra libraries, we use the sparse matrix linear algebra library. This is because the Hamiltonian matrix H is composed mainly of 0s (i.e. the matrix is sparse), and using the sparse matrix libraries will help to speed up computation for extremely large systems. Also using the sparse matrix eigenvalue calculators allow us to control the number of solutions to find, rather than calculating every single possible solution which will take a lot of time and resources.
After constructing H, we calculate its eigenvectors and eigenvalues using [evl, evt] = sla.eigs(H, k = neigs, which = 'SM'). We choose the number of solutions to obtain by setting the value of k = neigs, and also choose to obtain the smallest magnitude solutions with which = 'SM'. The eigenvectors are then normalized.
The Python code used to construct H and calculate its eigenvectors and eigenvalues is presented below. Note that the function schrodinger1D has the following arguments: xmin = minimum x grid value, xmax = maximum x grid value, Nx = number of grid points, Vfun = potential energy function, neigs = number of eigenvalues to solve for, and findpsi = flag to tell the function to return both the eigenvalues and eigenvectors if flag = True.
import numpy as np
import matplotlib.pyplot as plt
from scipy import sparse
from scipy.sparse import linalg as sla
def schrodinger1D(xmin, xmax, Nx, Vfun, params,
neigs=20, findpsi=False): x = np.linspace(xmin, xmax, Nx) # x axis grid
dx = x[1] - x[0] # x axis step size # Obtain the potential function values:
V = Vfun(x, params) # create the Hamiltonian Operator matrix:
H = sparse.eye(Nx, Nx, format='lil') * 2 # implement the numerical derivative
for i in range(Nx - 1):
H[i, i + 1] = -1
H[i + 1, i] = -1
H = H / (dx ** 2) # Add in the potential energy V
for i in range(Nx):
H[i, i] = H[i, i] + V[i] # convert to csc sparse matrix format:
H = H.tocsc() # obtain neigs solutions from the sparse matrix:
[evl, evt] = sla.eigs(H, k=neigs, which='SM')
for i in range(neigs):
# normalize the eigenvectors:
evt[:, i] = evt[:, i] / np.sqrt(
np.trapz(np.conj(
evt[:,i])*evt[:,i],x))
# eigen values MUST be real:
evl = np.real(evl) if findpsi == False:
return evl
else:
return evl, evt, x
The Quantum Harmonic Oscillator
Now that we have the functionality to solve the 1 dimensional Schrödinger equation numerically, all we need to do is to give it something to solve by specifying the potential energy Vfun! For the first example we will explore the quantum simple harmonic oscillator.
The quantum simple harmonic oscillator has a potential energy function given by V = kx² where k is the “spring constant”. To simplify matters we set k = 1, so we can set V = x². We then pass this potential energy function to the solver above to solve for the eigenvalues and eigenvectors. The theoretical energy eigenvalues of the quantum harmonic oscillator are given by E = ħω(n + 1/2). In our case we set ħ = ω = 1, and remember that we scaled the energies by a factor of 2 in our numerical differentiation scheme, so the adjusted theoretical energy eigenvalues take the form E = 2n + 1.
All of this is implemented in sho_wavefunctions_plot below.
def sho_wavefunctions_plot(xmin = -10, xmax = 10, Nx = 500,
neigs = 20, params = [1]):
def Vfun(x, params):
V = params[0] * x**2
return V
eval_wavefunctions(xmin, xmax, Nx, Vfun,
params, neigs, True)The function eval_wavefunctions calls schrodinger1D to solve the given eigenvalue problem, calculates the probability of the particle existing at some particular grid point from the wave functions, and plots the probabilities. Also as a sanity check, we print the eigenvalues sorted in ascending magnitude!
def eval_wavefunctions(xmin, xmax, Nx, Vfun, params, neigs,
findpsi = True):
# call the 1D Schrodinger solver:
H = schrodinger1D(xmin, xmax, Nx, Vfun, params, neigs, findpsi)
evl = H[0] # energy eigenvalues
indices = np.argsort(evl)
print("Energy eigenvalues:")
for i,j in enumerate(evl[indices]):
print("{}: {:.2f}".format(i+1,j))
evt = H[1] # eigenvectors
x = H[2] # x grid
i = 0
plt.figure(figsize=(8,8))
while i < neigs:
n = indices[i]
# obtain probabilities from wave functions
y = np.real(np.conj(evt[:, n]) * evt[:, n])
plt.subplot(neigs, 1, i+1)
plt.plot(x, y)
plt.axis('off')
i = i + 1
plt.show()Running the function sho_wavefunctions_plot plots the first 20 probabilities for the quantum harmonic oscillator shown below in order of increasing energy states! We see that for the lowest energy level at the top of the plot, the particle is most likely to be found at the centre of the potential well where the potential energy is the lowest. As energy increases, the particle is able to move up to regions of higher potential energy on both sides of the centre of the well, hence the probability of it being found at other locations increases!
Also, the numerically calculated energy eigenvalues for the 20 discrete energy levels n turn out to match the theoretical values given by E = 2n + 1 very closely especially for the smaller values of n, so we know our 1 dimensional Schrödinger equation solver has worked! Note that n starts from 0 and not 1!
sho_wavefunctions_plot()
Energy eigenvalues:
1: 1.00
2: 3.00
3: 5.00
4: 7.00
5: 9.00
6: 10.99
7: 12.99
8: 14.99
9: 16.99
10: 18.98
11: 20.98
12: 22.97
13: 24.97
14: 26.96
15: 28.96
16: 30.95
17: 32.95
18: 34.94
19: 36.93
20: 38.92Closing Remarks
This concludes how to solve the time independent Schrödinger equation numerically for one dimension. I hope you gained some insight on how physicists study quantum mechanics using computational methods!
Unfortunately, one dimensional systems are usually not too realistic, and we need to get to at least two dimensions in order to simulate many physical systems realistically. Can you think of a way on how to improve the Python functions to solve the Schrödinger equation in two or three dimensions? Thank you for reading!
GitHub
The full 1D Schrödinger equation solver code can be found on my Github repository. The Jupyter notebook version of the contents can also be found on my GitHub repository.
References
[1] Bransden and Joachain (2000). Quantum Mechanics (2nd Edition), Prentice Hall.
