image of me

Schrödinger's Equation with Python and Numpy


The Schrödinger Equation

We start out with the equation which tells us how our state will change in time (tψ\frac{\partial}{\partial t}\ket{\psi}), given its energy (H^ψ\hat H \ket{\psi}):

itψ=H^ψ\begin{equation*} i\hbar\frac{\partial}{\partial t}\ket{\psi} = \hat H \ket{\psi} \end{equation*}

For a single particle system of mass mm, we can use the Hamiltonian

H^=T^+V^=22m2x2+V^\begin{align*} \hat H &= \hat T + \hat V \\ &= -\frac{\hbar^{2}}{2m}\frac{\partial^{2}}{\partial x^{2}} + \hat V \end{align*}

Cool. How do we put this on the computer?


We want to be able to write an equation:

ψt+1=M^ψt\begin{equation*} \ket{\psi_{t+1}} = \hat M \ket{\psi_{t}} \end{equation*}

Where M^\hat M is some matrix which tells us: if our state looks like ψt\ket{\psi_{t}} at time tt what it'll look like at time t+1t+1. We get the state at time t+1t+1 by applying M^\hat M to our state at time tt.

First Derivative

Lets start by splitting up the partial derivative with respect to time (t\frac{\partial}{\partial t}): Since a derivative is the difference in a function over two closely spaced points, divided by that difference, we can do:

tψt=limΔt0ψt+ΔtψtΔt\begin{equation*} \frac{\partial}{\partial t} \ket{\psi_{t}} = \lim_{\Delta t \rightarrow 0} \frac{\ket{\psi_{t+\Delta t}} - \ket{\psi_{t}}}{\Delta t} \end{equation*}

One interesting thing to note here is the choice of going from tt+Δtt\rightarrow t+\Delta t to get the value of the derivative at tt. We should hopefully be able to choose any two points in {(tαΔt,t+(1α)Δt)α(0,1)}\left\{\left( t-\alpha\Delta t, t+(1-\alpha)\Delta t \right) \mid \forall \alpha \in \left(0, 1\right)\right\}, and this will give the same answer as Δt0\Delta t \rightarrow 0, but you can get some nasty functions that don't share this behaviour (eg. the Weierstrass function). Make sure you're dealing with a smooth one :^).

Second Derivative

We'll deal with the second derivative in xx in a similar way to the first derivative in tt. First, work out the first derivative at the points xΔx2x-\frac{\Delta x}{2} and x+Δx2x+\frac{\Delta x}{2}, then calculate the derivative between these.

2x2ψx=limΔx0xψx+Δx2xψxΔx2Δx=limΔx0ψx+ΔxψxΔxψxψxΔxΔxΔx=limΔx0ψx+Δx+ψxΔx2ψxΔx2\begin{align*} \frac{\partial^{2}}{\partial x^{2}}\ket{\psi_{x}} &= \lim_{\Delta x \rightarrow 0} \frac{\frac{\partial}{\partial x}\ket{\psi_{x+\frac{\Delta x}{2}}} - \frac{\partial}{\partial x}\ket{\psi_{x-\frac{\Delta x}{2}}}}{\Delta x} \\ &= \lim_{\Delta x \rightarrow 0} \frac{\frac{\ket{\psi_{x+\Delta x}} - \ket{\psi_{x}}}{\Delta x} - \frac{\ket{\psi_x} - \ket{\psi_{x-\Delta x}}}{\Delta x}}{\Delta x} \\ &= \lim_{\Delta x \rightarrow 0} \frac{\ket{\psi_{x+\Delta x}} + \ket{\psi_{x-\Delta x}} - 2\ket{\psi_{x}}}{\Delta x^2} \end{align*}

Getting an equation out of it

We've got discretised versions of the two derivative operators, so we're close to just plugging them into the Schrödinger equation but first those pesky limits need to go.

Removing the limits

Since we're on a computer with finite memory there's no way we can approach limα0\lim_{\alpha\rightarrow 0}, so we need to just choose a small enough value for our differential distances Δα\Delta \alpha. This size can be tuned when we actually come to run the simulations. We want a value small enough that if we decrease it the results of the simulation don't change appreciably, but large enough that the simulations don't take forever or run out of space.


After substitution for the derivatives, the Schrödinger equation becomes (now explicitly looking at the components of ψ=dtdxψ(x,t)x,t\ket{\psi}=\int \mathrm{d}t \mathrm{d}x \psi(x, t)\ket{x, t}):

iψ(x,t+Δt)ψ(x,t)Δt=22mψ(x+Δx,t)+ψ(xΔx,t)2ψ(x,t)Δx2+V(x,t)ψ(x,t)ψ(x,t+Δt)=ψ(x,t)+iΔt22mψ(x+Δx,t)+ψ(xΔx,t)2ψ(x,t)Δx2iΔtV(x,t)ψ(x,t)\begin{align*} i\hbar \frac{\psi(x, t+\Delta t) - \psi(x, t)}{\Delta t} &= -\frac{\hbar^{2}}{2m}\frac{\psi(x+\Delta x, t) + \psi(x - \Delta x, t) - 2\psi(x, t)}{\Delta x^{2}}\\ &+ V(x, t)\psi(x, t) \\ \psi(x, t+\Delta t) &= \psi(x, t) \\ &+ \frac{i\Delta t}{\hbar}\frac{\hbar^{2}}{2m}\frac{\psi(x+\Delta x, t) + \psi(x - \Delta x, t) - 2\psi(x, t)}{\Delta x^{2}} \\&- \frac{i\Delta t}{\hbar}V(x, t)\psi(x, t) \end{align*}

Letting α=Δt2mΔx2\alpha = \frac{\hbar \Delta t}{2m \Delta x^{2}}, we can get it into a more recognisable matrix equation form:

ψ(x,t+Δt)=iαψ(xΔx,t)+(1+2iαiΔtV(x,t))ψ(x,t)+iαψ(x+Δx,t)=(iα(1+2iαiΔtV(x,t))iα)(ψ(xΔx,t)ψ(x,t)ψ(x+Δx,t))\begin{align*} \psi(x, t+\Delta t) &= i\alpha \psi(x-\Delta x, t)\\ &+\left(1 + 2i\alpha -\frac{i\Delta t}{\hbar}V(x, t)\right)\psi(x, t) + i\alpha \psi(x+\Delta x, t) \\ &= \Big( \begin{matrix} i\alpha & (1 + 2i\alpha -\frac{i\Delta t}{\hbar}V(x, t)) & i\alpha \end{matrix} \Big) \begin{pmatrix} \psi(x-\Delta x, t) \\ \psi(x, t) \\ \psi(x + \Delta x, t) \end{pmatrix} \end{align*}

We can see that the state ψ(x,t+Δt)\psi(x, t + \Delta t) at xx in the next timestep depends on the state of that point in the previous timestep, and the two points at xΔxx-\Delta x and x+Δxx + \Delta x. If we discretise the state, so that instead of having continuous values over a continuous basis of positions, it becomes a sum over points separated by Δx\Delta x:

dxψ(x,t)xxψxtx\begin{equation*} \int \mathrm{d}x \psi(x, t) \ket{x} \rightarrow \sum_{x} \psi_{xt} \ket{x} \end{equation*}


Finally, we can write out the full matrix equation (letting Vxt=1+2iαiΔtVxtV'_{xt}=1+2i\alpha-\frac{i\Delta t}{\hbar}V_{xt}):

(ψx(t+Δt))=(iαV(xΔx)tiα00iαVxtiα)(ψ(x2Δx)tψ(xΔx)tψxtψ(x+Δx)tψ(x+2Δx)t)=M^tψt\begin{align*} \begin{pmatrix} \vdots \\ \psi_{x(t+\Delta t)} \\ \vdots \end{pmatrix} &= \begin{pmatrix} &&\dots \\ \dots & i\alpha & V'_{(x-\Delta x)t} & i\alpha & 0 &\dots \\ \dots & 0&i\alpha & V'_{xt} & i\alpha & \dots \\ \dots & & \dots \end{pmatrix} \begin{pmatrix} \vdots \\ \psi_{(x-2\Delta x)t} \\ \psi_{(x-\Delta x)t} \\ \psi_{xt} \\ \psi_{(x + \Delta x)t} \\ \psi_{(x + 2\Delta x)t} \\ \vdots \end{pmatrix} \\ &= \hat M_t \ket{\psi_t} \end{align*}

Our evolution matrix, which takes us from the state at time tt to time t+Δtt+\Delta t is tri-diagonal. The diagonal depends on the potential and the off-diagonals are all constant.

Crank-Nicolson Method

Our current equation uses the forward Euler method which has some stability problems. If we instead average the value of the Hamiltonian at times tt and t+Δtt+\Delta t we can get a much more stable equation.

ψt+Δt=12(M^tψt+M^t+Δtψt+Δt)\begin{align*} \ket{\psi_{t+\Delta t}} = \frac{1}{2}\left(\hat M_{t}\ket{\psi_{t}} + \hat M_{t+\Delta t}\ket{\psi_{t+\Delta t}}\right) \end{align*}

Rearranging, we've got another equation for ψt+Δt\ket{\psi_{t+\Delta t}}:

(I12M^t+Δt)ψt+Δt=12M^tψtψt+Δt=12(I12M^t+Δt)1M^tψt=C^ψt\begin{align*} \left(I-\frac{1}{2}\hat M_{t+\Delta t}\right)\ket{\psi_{t+\Delta t}} &= \frac{1}{2}\hat M_{t}\ket{\psi_{t}} \\ \ket{\psi_{t+\Delta t}} &= \frac{1}{2}\left(I-\frac{1}{2}\hat M_{t+\Delta t}\right)^{-1}\hat M_{t}\ket{\psi_{t}} \\ &= \hat C \ket{\psi_{t}} \end{align*}

Where C^\hat C is the Crank-Nicolson matrix.

1-Dimensional Simulation

We've got our equation! Now's time for some code. Our first function will generate the Crank Nicolson matrix, given arguments of: VxtV_{xt}, Vx(t+Δt)V_{x(t+\Delta t)} (the potential over positions at that those specific times); Δx\Delta x and Δt\Delta t - the spacing between our coordinates; and mm, the mass of the particle.

def crank_nicolson(Δx, Δt, mass, vt1, vt2=None):
    """Generate the Crank-Nicolson matrix that will do our evolution
    vt2 allows us to set a different potential at the next timestep.
    if vt2 is None: vt2 = vt1
    ħ = 1.0545718e-34 # kgm^2/s
    α = ħ*Δt/(2*mass*Δx**2)
    Mt = np.zeros((len(vt1), len(vt1)), dtype=np.complex128)
    off_diag = np.full(len(vt1) - 1, 1j*α, dtype=np.complex128)
    Mt += np.diag(off_diag, 1)
    Mt += np.diag(off_diag, -1)
    Mtdt = np.copy(Mt)
    Mt += np.diag(1 + 2*1j*α - 1j*Δt/ħ*vt1)
    Mtdt += np.diag(1 + 2*1j*α - 1j*Δt/ħ*vt2)
    return 0.5 * np.linalg.inv(np.eye(len(vt1), dtype=np.complex128) - 0.5 * Mtdt) @ Mt

We should also make a function which gets the probability distribution from a state:

def probability(state):
    prob = np.abs(state)**2
    return prob / prob.sum()

And a function which creates a gaussian enveloped state with a given spatial frequency (kk):

def gaussian(k, σ, x, xs):
    state = np.exp(-(xs-x)**2/(2*σ**2)) * np.exp(1j * k * xs)
    state = state / np.sqrt((np.abs(state)**2).sum())
    return state

Finally, a function which plots a state on a given axis:

def plot_state(xs, state, axis):
    axis.plot(xs, state.real, color="red", label=r"$\mathrm{Re}(\psi)$")
    axis.plot(xs, state.imag, color="blue", label=r"$\mathrm{Im}(\psi)$")
    axis.plot(xs, probability(state), color="black", label=r"$\mathrm{P}(\psi)$")

This function plots the real and imaginary parts of the state, and its probability. Lets test it out on a gaussian:

xs = np.linspace(0, 1, 300)
state = gaussian(100, 0.1, 0.5, xs)
fig, axis = plt.subplots()
plot_state(xs, state, axis)
fig.savefig(filename, dpi=200, bbox_inches="tight")


Lets test out the Crank-Nicolson method with an eigenstate of the infinite square well:

xs = np.linspace(0, 1e-32, 100)
state = np.sin(3*xs*np.pi/(1e-32))
state = state / np.sqrt((np.abs(state)**2).sum())
max_val = np.abs(np.max(state)) * 1.2
steps = 31
fig, axs = plt.subplots(5, 10)
C = crank_nicolson(np.abs(xs[1]-xs[0]), 1e-37, 1, np.zeros_like(xs))
C = np.linalg.matrix_power(C, steps)
for i, ax in enumerate(axs.flatten()):
    ax.set_ylim(-max_val, max_val)
    plot_state(xs, state, ax)
    state = C@state

fig.savefig(filename, dpi=200, bbox_inches="tight")


We can see that the particle stays in the eigenstate, and only the phase of the state changes over the evolution.

If we instead start with a gaussian state, we can see a much more complicated evolution, with the particle reflecting off the walls.

xs = np.linspace(0, 1e-32, 100)
state = gaussian(1e33, 4e-34, xs[50], xs)
max_val = np.abs(np.max(state)) * 1.2
steps = 1000
fig, axs = plt.subplots(5, 10)
C = crank_nicolson(np.abs(xs[1]-xs[0]), 1e-37, 1, np.zeros_like(xs))
C = np.linalg.matrix_power(C, steps)
for i, ax in enumerate(axs.flatten()):
    ax.set_ylim(-max_val, max_val)
    plot_state(xs, state, ax)
    state = C@state

fig.savefig(filename, dpi=200, bbox_inches="tight")


We can see the gaussian state disperses into the infinite square well - it's not an eigenstate.

Confined particle

Lets check out what happens if you confine a particle inbetween two finite potential walls.

xs = np.linspace(0, 1e-32, 500)
state = gaussian(5e34, 4e-35, xs[len(xs)//2], xs)
barrier_width = 1
barrier_positions = [150, 350]
potential = np.zeros_like(xs)
potential_value = 36.0
for pos in barrier_positions: potential[pos:pos+barrier_width] = potential_value
fig, axs = plt.subplots(6, 1, figsize=(5, 7))
steps = 2000
max_val = np.abs(np.max(state))
C = crank_nicolson(np.abs(xs[1]-xs[0]), 1e-37, 1, potential)
C = np.linalg.matrix_power(C, steps)
for i, ax in enumerate(axs.flatten()):
    for pos in barrier_positions:
        ax.axvline(xs[pos], color=(0, 0, 0, 0.2))
    ax.set_ylim(-max_val, max_val)
    plot_state(xs, state, ax)
    state = C@state

fig.savefig(filename, dpi=200, bbox_inches="tight")


The particle bounces between the walls, releasing a small amount of the wavefunction at each barrier.

What Else?

At the moment when the particle hits the boundaries of the space it reflects, but a small change to the Crank-Nicolson matrix could let us simulate cyclic boundary conditions. Physically, this would be like connecting both ends of the simulation space, so the particle's travelling around on a ring. It's also not too big a change to do a 2-dimensional simulation (which may be the subject of a further post??? 🐵)