Simulation of the Schrödinger equation¶

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.

Discretizing the Hamiltonian¶

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()

Running the simulation¶

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

Wave function¶

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.

Gaussian wave packet¶

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$.

Potential Function¶

Let's try out a few standard potential functions to test the simulator!

Free Particle ($V = 0$)¶

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.

The dashed red line represents the expectation value in position, computed by

$$\operatorname{E}[x] = \int_{-\infty}^{+\infty} \, x \, |\Psi(x)|^2 \, dx$$

Since the Gaussian wave packet is symmetric, this expected value should coincide with the peak of the probability density. The average momentum is $p_0 = 1$, and thus, the position of the peak is $x_\mathrm{peak} = p_0 t$.

The wave packet also spreads out over time. This is a consequence of having a non-zero uncertainty in momentum. Since we don't know exactly how fast the particle is moving, the uncertainty in position, $\Delta{x}$, grows over time due to dispersion. The wave packet moves as $x(t) = p_0 t + x_0$. By propagation of uncertainty, we get

$$\Delta{x}(t) = \sqrt{\Delta{p}^2 t^2 + \Delta{x}(0)^2}$$

At $t = 100$, the position-uncertainty has grown to $\Delta{x}(100) \approx 11.18$ - over twice the uncertainty we started with.

Quantum Harmonic Oscillator ($V = \frac{1}{2} \omega^2 x^2$)¶

Let's try a harmonic oscillator with $\omega = 1/32$ with an initial Gaussian wave packet centered at $x = -32$ with average initial momentum of $p_0 = 0$. The particle should oscillate.

As expected, the particle's position oscillates with a period of $T = 2 \pi / \omega \approx 201$. So, does the momentum and position uncertainty. Why this happens is left as an exercise to the reader (as the answer isn't particularly exciting).

Quantum Tunnelling¶

Let's place a rectangular potential barrier of height $V_0$ and width $a$ in the path of a Gaussian wave packet with energy $E = {p_0}^2 / 2 < V_0$. It should tunnel through (unlike a classical particle) with probability

$$T = \frac{1}{1 + \frac{{V_0}^2 \sinh^2{\left(a \sqrt{2 V_0 - 2 E}\right)}}{4 E (V_0 - E)}}$$

def rectangular_potential_barrier(x, V0, a):
    """Rectangular potential barrier of height V0 and width a."""
    return np.where((0 <= x) & (x < a), V0, 0.0)


def transmission_probability(E, V0, a):
    """Transmission probability of through a rectangular potential barrier."""
    k = (V0 * np.sinh(a * np.sqrt(2 * (V0 - E))))**2 / (4 * E * (V0 - E))
    return 1 / (1 + k)

We set parameters: $T = 0.2$, $a = 1.25$, and $E/V_0 = 3/4$.

from scipy.optimize import root_scalar

N = 1280
x, dx = np.linspace(-80, +80, N, endpoint=False, retstep=True)

T, r = 0.20, 3/4
k1 = root_scalar(lambda a: transmission_probability(0.5*r, 0.5, a) - T,
                 bracket=(0.0, 10.0)).root

a = 1.25
V0 = ((k1 / a)**2) / 2
E = r * V0

x0, sigma0, p0 = -48.0, 3.0, np.sqrt(2*E)
psi0 = gaussian_wavepacket(x, x0=x0, sigma0=sigma0, p0=p0)

V = rectangular_potential_barrier(x, V0, a)
H = hamiltonian(N, dx, V=V)
sim = simulate(psi0, H, dt=0.2)

The red region shows the rectangular potential barrier. As we expect based on the transmission coefficient, $20\%$ of the wave packet is transmitted through the barrier and the rest is reflected back.

Let's compare the transmitted wave packet ($x > a$) to a free particle wave packet in the absence of a barrier (after the same amount of time has passed after the initial state).

In the plot below, the dashed lines indicate the expected value of position for the respective wave packets. The transmitted wave packet has a lower height as it only represents $20\%$ of the entire wave function.

A wave packet tunneling through a barrier seems to come out ahead of a free particle in absence of a barrier! This occurs because the potential barrier effectively filters out lower energies - the higher energy components preferentially tunnel through successfully (since the transmission coefficient increases with the energy of the particle). These higher energy components naturally have higher momentum on average.

We can confirm this by looking at the probability density of the momentum-space wave function (we get this by taking the Fourier transform of the position-space wave function) of the transmitted wave packet compared to the free particle wave packet.

p = np.fft.fftshift(np.fft.fftfreq(N, dx)) * 2 * np.pi

def momentum_probability_density(psi, dx):
    """Momentum-space probability density of a given wave function."""
    norm = dx / np.sqrt(2 * np.pi)
    mom = np.fft.fftshift(np.fft.fft(psi, axis=-1))
    return probability_density(mom * norm)

p_psi_free = momentum_probability_density(psi_free, dx)
p_psi_barr = momentum_probability_density(psi_barr, dx)

The black dashed line represents the average momentum $p_0$ of the free particle Gaussian wave packet. The momentum-space probability density for the transmitted wave packet skews towards higher momentum (and thus, higher energies). As a result of this, the transmitted wave packet moves faster, on average, compared to the free particle wave packet!

Locality vs. time evolution¶

The time evolution operator, $\mathbf{U}$, can be used with as large or small a time step ($\delta{t}$) as required. With a small enough time step, a point in the wave function should only be able to affect change in its immediate neighbors. Let's call this "locality". With increasingly larger time steps, the time evolution operator should have reduced locality. At extremely large time steps, the time evolution operator should have an almost global effect. We can see this by plotting the squared modulus of the time evolution operator, $|\mathbf{U}|^2$, with different time steps.

At small time steps, all the power in $|\mathbf{U}|^2$ is concentrated around the main diagonal. As the time step increases, we see the further off-diagonal elements being more relevant. At $\delta{t} = 300$, the time evolution operator is essentially global.

One interesting thing to note here is that for a given $\delta{t}$, the time evolution operator cannot affect change beyond a certain distance. This is because our choice of sample spacing, $\delta{x}$, leads to a maximum momentum that can be simulated (this is the Nyquist frequency). In practice, the average momentum of our Gaussian wave packet must be much less than due to the limitations of the finite difference method in approximating the second derivative.