In this notebook, I explore simulating the 1-dimensional time-dependent Schrödinger equation (TDSE). Assuming $\hbar = 1$, the equation is
$$ \frac{\partial}{\partial t} \Psi(x, t) = -i \hat{H} \Psi(x, t) $$where $\hat{H}$ is the Hamiltonian operator) and $\Psi(x, t)$ is the wave function of the quantum system. If the Hamiltonian is independent of time, then the equation can be written in terms of a time evolution operator,
$$\Psi(x, t)=e^{-i \hat{H} t} \, \Psi(x, 0)$$For a single non-relativistic spinless particle of mass, $m = 1$, in a time-invariant potential, $V(x)$, the Hamiltonian is given by
$$\hat{H} = -\frac{1}{2} \nabla^2 + V(x) $$where $\nabla^2 = \frac{\partial^2}{\partial{x}^2}$ is the Laplace operator.
To numerically simulate a quantum system using the TDSE, we work on a discrete grid of $N$ equally-spaced points (with spacing, $\delta x$). To do this, we have to find a way to discretize the various components of the TDSE. The wave function at a point in time, $\Psi(t)$, and potential, $V$, can be discretized as vectors of length $N$ where each element contains the value of the wave function or the potential at the corresponding point on this grid.
A discrete Laplace operator can be given as convolution with the finite-difference kernel, $\begin{bmatrix}1 & -2 & 1\end{bmatrix}$, which results in a symmetric tridiagonal Toeplitz matrix with the finite-difference coefficients along the diagonals. The discrete Laplace operator is then given by,
$$\mathbf{\nabla^2} = \frac{1}{\delta{x}^2} \begin{bmatrix}-2 & 1 \\ 1 & -2 & 1\\ & & \ddots \\ & & 1 & -2 & 1 \\ & & & 1 & -2\end{bmatrix}$$where $\delta{x}$ is the spacing between points. If we have a potential, $\mathbf{V} = \begin{bmatrix} V_0 & V_1 & \dots & V_N\end{bmatrix}$, then the full Hamiltonian operator in matrix form is given by,
$$\mathbf{\hat{H}} = -\frac{1}{2\Delta{x}^2} \begin{bmatrix}-2 & 1 \\ 1 & -2 & 1\\ & & \ddots \\ & & 1 & -2 & 1 \\ & & & 1 & -2\end{bmatrix} + \begin{bmatrix}V_0\\ & V_1\\ & & \ddots \\ & & & V_{N-1} \\ & & & & V_{N}\end{bmatrix}$$Note: Because of the way I define the finite-differences at the boundaries, the wave function will reflect completely at the boundaries. The edges of the simulation are effectively infinite potential barriers.
import numpy as np
import scipy.sparse
def hamiltonian(N, dx, V=None):
"""Returns Hamiltonian using finite differences.
Args:
N (int): Number of grid points.
dx (float): Grid spacing.
V (array-like): Potential. Must have shape (N,).
Default is a zero potential everywhere.
Returns:
Hamiltonian as a sparse matrix with shape (N, N).
"""
L = scipy.sparse.diags([1, -2, 1], offsets=[-1, 0, 1], shape=(N, N))
H = -L / (2 * dx**2)
if V is not None:
H += scipy.sparse.spdiags(V, 0, N, N)
return H.tocsc()
With the Hamiltonian, we can now get the time evolution operator $\mathbf{U} = e^{-i \mathbf{\hat{H}} \delta t}$ via matrix exponentiation, where $\delta{t}$ is the size of a time step. Given some wave function at some time $t$ as a vector, $\mathbf{\Psi}(t) = \begin{bmatrix} \Psi_0 & \Psi_1 & \dots & \Psi_N\end{bmatrix}$, the time evolution operator can be applied to get the wave function after one time step
$$\mathbf{\Psi}(t + \delta{t}) = \mathbf{U} \, \mathbf{\Psi}(t)$$You can iteratively multiply the wave function with $\mathbf{U}$ to advanced time steps. So, the wave function after $m$ time steps is,
$$\mathbf{\Psi}(t + m \delta{t}) = \mathbf{U}^m \mathbf{\Psi}(t)$$Since the Hamiltonian is a Hermitian matrix, $\mathbf{U}$ will necessarily be a unitary matrix. This is important as matrix-vector multiplication with a unitary matrix preserves norms and, therefore, the total probability is preserved.
import scipy.linalg
def time_evolution_operator(H, dt):
"""Time evolution operator given a Hamiltonian and time step."""
U = scipy.linalg.expm(-1j * H * dt).toarray()
U[(U.real**2 + U.imag**2) < 1E-10] = 0
return scipy.sparse.csc_matrix(U)
def simulate(psi, H, dt):
"""Generates wavefunction and time at the next time step."""
U = time_evolution_operator(H, dt)
t = 0
while True:
yield psi, t * dt
psi = U @ psi
t += 1
The wave function, $\Psi(x, t)$, is the complex-valued probability amplitude in position space for the quantum system - in our case, a single particle. The square modulus of the wave function, $|\Psi(x,t)|^2$ is the probability density of the particle's position. By the normalization condition,
$$\int_{-\infty}^\infty \, |\Psi(x,t)|^2dx = 1$$which basically says that there is a $100\%$ probability that the particle must be somewhere.
def probability_density(psi):
"""Position-space probability density."""
return psi.real**2 + psi.imag**2
We use a localized Gaussian wave packet centered at $x_0$ with average initial momentum, $p_0$, and Gaussian width, $\sigma_0$. The wave function is given by
$$\psi(x) = \left(\frac{1}{2 \pi {\sigma_0}^2}\right)^{\frac{1}{4}} \, \exp{\left(-\frac{(x - x_0)^2}{4 {\sigma_0}^2} + i p_0 x\right)} $$As a free particle, the Gaussian wave packet's center moves by $p_0 t$ over time $t$ (since we set the mass to $1$ earlier, the momentum is equal to the velocity). The uncertainties in position and momentum are given by $\Delta{x} = \sigma_0$ and $\Delta{p} = (2 \sigma_0)^{-1}$, respectively. The Gaussian wave packet is a minimum uncertainty wave packet with $\Delta{x} \Delta{p} = 1/2$. Consistent with the motion of a free particle, $\Delta{p}$ is constant in time.
The energy of the particle is $E = p_0^2 / 2$.
def gaussian_wavepacket(x, x0, sigma0, p0):
"""Gaussian wavepacket at x0 +/- sigma0, with average momentum, p0."""
A = (2 * np.pi * sigma0**2)**(-0.25)
return A * np.exp(1j*p0*x - ((x - x0)/(2 * sigma0))**2)
Let's try out a few standard potential functions to test the simulator!
A Gaussian wave packet centered at $x_0 = 0$ with initial momentum of $p_0 = 1$ and Gaussian width of $\sigma_0 = 5$. In the absence of a potential, this will behave like a free particle.
N = 1800
x, dx = np.linspace(-25, 155, N, endpoint=False, retstep=True)
psi0 = gaussian_wavepacket(x, x0=0.0, sigma0=5.0, p0=1.0)
H = hamiltonian(N, dx)
sim = simulate(psi0, H, dt=1.0)