Last Updated on June 13, 2025 by Sushanta Barman
The imaginary time evolution method is a powerful computational tool for finding ground states in quantum mechanics, especially for systems where analytical solutions are challenging.
In this article, we demonstrate how to use imaginary time evolution to obtain the ground state of a 2D quantum harmonic oscillator.
We start with a trial wave function composed of multiple Gaussian packets, numerically evolve it, and compare the results with the known analytical solution.
Physical Model and Method
2D Quantum Harmonic Oscillator
A 2D quantum harmonic oscillator has the potential
\[V(x, y) = \frac{1}{2}m\omega^2 (x^2 + y^2)\]
where \(m\) is the mass and \(\omega\) is the angular frequency. The system is fundamental in quantum mechanics and serves as a basis for understanding more complex quantum systems.
Imaginary Time Evolution
The imaginary time evolution technique involves propagating the wave function in imaginary time (\(\tau = it\)), effectively damping higher energy components and filtering out the ground state. The evolution equation becomes:
\[\frac{\partial \psi}{\partial \tau} = -\frac{1}{\hbar} \hat{H} \psi \]
where \(\hat{H}\) is the Hamiltonian operator.
Read related article:
Imaginary Time Evolution Method for Finding the Ground State
Numerical Implementation
Spatial Grid and Potential
We discretize the 2D spatial domain using a high-resolution grid (\(N = 200\)), covering \(x, y \in [-5, 5]\). The harmonic oscillator potential is applied at each grid point.
Trial Wave Function
Instead of a single Gaussian, we use a superposition of four Gaussians centered at \((\pm d, 0)\) and \((0, \pm d)\) to create a trial wave function that is deliberately different from the ground state. This tests the robustness of imaginary time evolution, which will still converge to the ground state regardless of the initial guess.
The initial trial wave function is given by:
\[\psi_{\text{init}}(x, y) = \exp\left(-\frac{(x – d)^2 + y^2}{2\sigma^2}\right) + \exp\left(-\frac{(x + d)^2 + y^2}{2\sigma^2}\right) + \\ \exp\left(-\frac{x^2 + (y – d)^2}{2\sigma^2}\right) + \exp\left(-\frac{x^2 + (y + d)^2}{2\sigma^2}\right),\]
where, \(d\) is the displacement from the origin and \(\sigma\) is the width of each Gaussian. After summing, the wave function is normalized to ensure the total probability is unity.
Laplacian and Time Evolution
A finite-difference scheme is used to approximate the 2D Laplacian. At each step, the wave function is updated according to the discretized imaginary time evolution equation, followed by normalization to maintain physical accuracy.
Python Code for Animation
Based on the numerical procedure described above, the following Python code demonstrates how to obtain the ground state of a two-dimensional quantum harmonic oscillator using the imaginary time evolution method. This code numerically evolves an initial trial wave function toward the ground state and visualizes the results through an animated 3D surface plot.
The required packages for running the code are summarized in Table I.
| Package | Purpose |
|---|---|
| numpy | Numerical array computations |
| matplotlib | Plotting and animation |
| mpl_toolkits.mplot3d | 3D plotting (part of matplotlib) |
| ffmpeg (external tool) | Export animation as video (.mp4) |
Here is the Python code.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
# Parameters
hbar = 1.0
m = 1.0
omega = 1.0
# 1. Spatial grid and potential
N = 200 # High resolution for smoothness
x_min, x_max = -5.0, 5.0
y_min, y_max = x_min, x_max
x = np.linspace(x_min, x_max, N)
y = np.linspace(y_min, y_max, N)
dx = x[1] - x[0]
X, Y = np.meshgrid(x, y, indexing='ij')
V = 0.5 * m * omega**2 * (X**2 + Y**2)
# Proper normalization function
def normalize(psi, dx):
norm = np.sqrt(np.sum(np.abs(psi)**2) * dx * dx)
return psi / norm
# 2. Initial wavefunction: double Gaussian
sigma = 1.0 # width of each Gaussian
d = 2 # separation from origin
gauss1 = np.exp(-((X - d)**2 + Y**2) / (2 * sigma**2))
gauss2 = np.exp(-((X + d)**2 + Y**2) / (2 * sigma**2))
gauss3 = np.exp(-(X**2 + (Y - d)**2) / (2 * sigma**2))
gauss4 = np.exp(-(X**2 + (Y + d)**2) / (2 * sigma**2))
psi_init = gauss1 + gauss2 + gauss3 + gauss4
psi_init = normalize(psi_init, dx)
# 3. Laplacian operator in 2D (finite difference)
def laplacian(psi, dx):
lap = np.zeros_like(psi)
lap[1:-1,1:-1] = (
psi[2:,1:-1] + psi[:-2,1:-1] +
psi[1:-1,2:] + psi[1:-1,:-2] -
4*psi[1:-1,1:-1]
) / dx**2
return lap
# 4. Analytical ground state (for z-limit)
psi_analytical = (m * omega / (np.pi * hbar))**0.5 * np.exp(-m * omega * (X**2 + Y**2) / (2 * hbar))
psi_analytical = normalize(psi_analytical, dx)
# 5. Imaginary-time evolution parameters
dtau = 0.001
total_steps = 6000
frames = 300
steps_per_frame = total_steps // frames
wave_evolution = []
energies = []
psi_temp = psi_init.copy()
for _ in range(frames):
for _ in range(steps_per_frame):
lap = laplacian(psi_temp, dx)
H_psi = -(hbar**2) / (2 * m) * lap + V * psi_temp
psi_temp -= dtau / hbar * H_psi
psi_temp = normalize(psi_temp, dx)
wave_evolution.append(psi_temp.copy())
lap = laplacian(psi_temp, dx)
H_psi = -(hbar**2) / (2 * m) * lap + V * psi_temp
E = np.sum(np.conj(psi_temp) * H_psi) * dx * dx
energies.append(np.real(E))
# 6. 3D surface animation
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')
fig.subplots_adjust(top=0.85)
ax.set_xlim(x_min, x_max)
ax.set_ylim(y_min, y_max)
ax.set_zlim(0, np.max(np.abs(psi_analytical)))
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel(r'$|\psi(x,y)|$')
# initial surface
surf = [ax.plot_surface(
X, Y, np.abs(wave_evolution[0]),
cmap='viridis', linewidth=0, antialiased=False, edgecolor='none'
)]
def update_surface(i):
ax.clear()
Z = np.abs(wave_evolution[i])
surf[0] = ax.plot_surface(
X, Y, Z,
cmap='viridis', linewidth=0, antialiased=False, edgecolor='none'
)
ax.set_xlim(x_min, x_max)
ax.set_ylim(y_min, y_max)
ax.set_zlim(0, np.max(np.abs(psi_analytical)))
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel(r'$|\psi(x,y)|$')
iteration = (i + 1) * steps_per_frame
fig.suptitle(f"Iteration: {iteration}, Energy: {energies[i]:.6f}",
fontsize=14, y=0.92)
return surf
ani = animation.FuncAnimation(
fig, update_surface, frames=frames, interval=20, blit=False
)
plt.tight_layout()
plt.show()
ani.save("2d_double_gaussian_surface.mp4", writer="ffmpeg", fps=30)
# Print the ground state energy
print("Ground state energy (numerical):", energies[-1])
print("Ground state energy (analytical):", 1.0)
Results: Ground State Convergence and Animation
- The wave function is iteratively evolved in the above Python code for 6,000 steps, saving snapshots for visualization.
- The ground state is approached as the system evolves, confirmed by comparing with the analytical ground state solution:
\[ \psi_0(x, y) = \left(\frac{m\omega}{\pi\hbar}\right)^{1/2} \exp\left(-\frac{m\omega}{2\hbar}(x^2 + y^2)\right)\] - The energy of the system also converges to the ground state energy, which is given by the analytical formula for a 2D quantum harmonic oscillator:
\[E_{0,0} = \hbar \omega (n_x + n_y + 1)\]
For the ground state, both quantum numbers are zero (\(n_x = 0, n_y = 0\)), so the ground state energy is:
\[E_{0,0} = \hbar \omega (0 + 0 + 1) = \hbar \omega\]
3D Surface Animation
A smooth 3D animation shown in Video 1 visualizes the wave function \(|\psi(x, y)|\) evolving from the initial multi-Gaussian structure to the ground state distribution.
The animation demonstrates how high-energy features are washed out, leaving the characteristic bell-shaped ground state.
Energy Convergence
At each frame, the total energy is computed numerically and shown to converge to the analytical ground state value:
- Numerical ground state energy (from code output): close to 1.0 (using natural units: \(\hbar =1\) and \(\omega=1\))
- Analytical ground state energy: \(E_0 = 1.0\) (using natural units: \(\hbar =1\) and \(\omega=1\))
Read related article:
Imaginary Time Evolution vs Real Time Evolution in Quantum Mechanics
Conclusion
This simulation illustrates the efficiency of imaginary time evolution for finding the ground state of quantum systems.
The approach works even with complex initial guesses, such as a sum of Gaussians, and the results match theoretical expectations. This method is highly relevant for research in quantum mechanics, computational physics, and quantum chemistry.
I am a senior research scholar in the Department of Physics at IIT Kanpur. My work focuses on ion-beam optics and matter-wave phenomena. I am also interested in emerging matter-wave technologies.
