Simulate Quantum Mechanics with Python

Solve the Schrodinger equation and animate quantum wave functions.

Introduction

30 Second Summary

Every piece of technology you use relies on quantum mechanics. The transistors in your phone, the lasers in fiber optic cables, and the LEDs in your screen all depend on particles being confined to quantized energy levels.

In this project, you will numerically solve the Schrodinger equation for a particle trapped in a one-dimensional box. You'll compute quantized energy levels, visualize wave functions, and animate a quantum superposition state evolving in time.

What You'll Build

A Python simulation that solves the 1D Schrodinger equation using the finite-difference method, producing verified energy levels, wave function plots, and a time-evolution animation.

By the end of this project, you'll have:

A working numerical solver that constructs the Hamiltonian matrix and extracts eigenstates using SciPy.
Visualizations of wave functions, probability densities, and an energy level diagram that confirm the n-squared energy scaling.
An animated simulation showing a quantum superposition state oscillating inside the potential well.
Secret Mission: Replace the infinite well with a finite potential well and observe quantum tunneling.

Are there any prerequisites?

Basic Python knowledge (functions, arrays) and familiarity with high-school physics concepts like the Bohr model and de Broglie waves. Some linear algebra basics (matrices, eigenvalues) will help, but we explain everything as we go.

Set Up the Grid and Physical Parameters

Your tools are installed and your project file is ready. Now it's time to define the quantum system you'll simulate: a particle trapped in a one-dimensional box.

In quantum mechanics, the "particle in a box" is solved by discretizing continuous space into a grid of points. You'll set up this grid along with the physical constants that govern the system, using natural units to keep the math clean.

In this step, get ready to:

Define physical constants using natural units.
Create a discretized spatial grid with boundary conditions.
Verify the setup runs and prints correct grid parameters.

Set up the physical parameters and spatial grid

Before writing any physics, you need to import the libraries that handle numerical arrays, eigenvalue solving, and plotting.

Add the following imports at the top of quantum_box.py:

import numpy as np
from scipy.linalg import eigh
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

What do these imports do?

numpy provides arrays and math functions for building the grid and matrices.
eigh from SciPy solves eigenvalue problems for symmetric matrices, which is exactly what the Schrodinger equation becomes once discretized.
matplotlib.pyplot handles plotting wave functions and energy levels.
FuncAnimation creates the animated visualization of quantum time evolution.

You won't see a change yet. The next chunk defines the actual physics.

Now you'll define the physical constants and discretize the box into a grid of points. In natural units, we set hbar = 1 and m = 1 to eliminate unit clutter.

Add the following code below your imports in quantum_box.py:

# =============================================================
# STEP 1: Setup - Physical parameters and spatial grid
# =============================================================

# Natural units: hbar = 1, m = 1
hbar = 1.0
m = 1.0
L = 1.0  # Box length

# Discretization
N = 200  # Number of interior grid points
dx = L / (N + 1)
x_inner = np.linspace(dx, L - dx, N)  # Interior points only (boundary = 0)

print(f"Grid parameters: N = {N}, dx = {dx:.6f}")
print(f"Domain: x from {x_inner[0]:.4f} to {x_inner[-1]:.4f}")

What does this code do?

The constants hbar, m, and L define the quantum system. Setting hbar and m to 1 means the analytical energy formula simplifies to E_n = n^2 * pi^2 / 2, making verification straightforward.

The grid splits the box into N = 200 interior points. The spacing dx is computed as L / (N + 1) because the box has N + 2 total points (including the two boundaries), but the wave function is always zero at the walls.

The np.linspace(dx, L - dx, N) call creates only the interior points. The boundary conditions (psi = 0 at x = 0 and x = L) are enforced by excluding those endpoints.

Why use natural units?

In real-world quantum mechanics, hbar is a tiny number (1.055 x 10^-34 J*s) and particle masses are similarly small. Carrying these around creates floating-point headaches and obscures the physics.

Natural units let you focus on the mathematical structure of quantum mechanics. The results translate directly to physical units whenever needed.

Verify the setup by running the script

Save quantum_box.py.
Run the script in your terminal:

python quantum_box.py

You should see output like this:

Grid parameters: N = 200, dx = 0.004975
Domain: x from 0.0050 to 0.9950

What should I see?

The grid spacing dx is about 0.005, meaning you have 200 evenly spaced sample points between the box walls. The domain starts just inside x = 0 and ends just inside x = 1, confirming the boundary points are excluded.

Seeing an import error?

If you see ModuleNotFoundError, confirm your libraries installed correctly by running pip install numpy scipy matplotlib again.
Make sure you're running the script with the same Python interpreter where you installed the packages. Try python3 quantum_box.py if python doesn't work.

Help me debug my import error in quantum_box.py

✔️ Awesome, I've got everything!

Great. Double check you've saved your file before moving on.

ⓧ I'd like to double check the full code

import numpy as np
from scipy.linalg import eigh
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# =============================================================
# STEP 1: Setup - Physical parameters and spatial grid
# =============================================================

# Natural units: hbar = 1, m = 1
hbar = 1.0
m = 1.0
L = 1.0  # Box length

# Discretization
N = 200  # Number of interior grid points
dx = L / (N + 1)
x_inner = np.linspace(dx, L - dx, N)  # Interior points only (boundary = 0)

print(f"Grid parameters: N = {N}, dx = {dx:.6f}")
print(f"Domain: x from {x_inner[0]:.4f} to {x_inner[-1]:.4f}")

Your spatial grid is ready and your physics constants are defined. Next up, you'll build the Hamiltonian matrix that encodes the particle's kinetic energy as a finite-difference operator.

Build the Hamiltonian Matrix

Your spatial grid is ready. Now you need something that acts on that grid to describe how the particle behaves.

In quantum mechanics, the Hamiltonian is the operator that encodes the total energy of the system. For a particle in a box with no internal potential, the Hamiltonian is purely kinetic energy. You will translate the continuous second derivative from the Schrodinger equation into a matrix that a computer can solve.

In this step, get ready to:

Construct the kinetic energy operator using finite differences.
Assemble the full Hamiltonian matrix H.
Verify the matrix is symmetric and inspect its structure.

Construct the Hamiltonian matrix using finite differences

The Schrodinger equation contains a second derivative of the wave function. On a discrete grid, you approximate that derivative using the finite difference method.

The three-point stencil replaces the continuous second derivative with a pattern: each grid point references itself and its two neighbours. This produces a tridiagonal matrix with -2 on the main diagonal and +1 on the off-diagonals.

In quantum_box.py, add the following code below the existing Step 1 code:

# =============================================================
# STEP 2: Build the Hamiltonian Matrix
# =============================================================

# Kinetic energy coefficient: -hbar^2 / (2 * m * dx^2)
coeff = -hbar**2 / (2 * m * dx**2)

# Tridiagonal matrix for the second derivative
main_diag = -2 * np.ones(N)
off_diag = np.ones(N - 1)

# Assemble the Hamiltonian: H = T (kinetic energy only, since V=0 inside the box)
H = coeff * (np.diag(main_diag) + np.diag(off_diag, k=1) + np.diag(off_diag, k=-1))

What does this code do?

coeff is the scaling factor from the kinetic energy term of the Schrodinger equation. It combines the physical constants (hbar, mass) with your grid spacing.

main_diag is a vector of -2s (one per grid point) and off_diag is a vector of +1s (one fewer, connecting neighbours). Together they form the finite-difference stencil for the second derivative.

np.diag() builds a full matrix from a diagonal vector. The k=1 and k=-1 arguments place the off-diagonals one row above and below the main diagonal. For the infinite well, the potential V(x) = 0 inside the box, so H equals the kinetic energy matrix alone.

Why a tridiagonal matrix?

The finite-difference approximation of the second derivative only involves three neighbouring points: the point itself and one on each side. This "locality" produces a matrix where only the main diagonal and its immediate neighbours are non-zero. Tridiagonal matrices are fast to solve and appear throughout computational physics.

You won't see output from this code yet. The next chunk adds the verification prints.

Verify the Hamiltonian is symmetric

A symmetric matrix (H equals its own transpose) guarantees real eigenvalues. This is physically required because energy levels must be real numbers.

Add these print statements below the H = coeff * (...) line to verify the matrix structure:

print(f"\nHamiltonian shape: {H.shape}")
print(f"Hamiltonian is symmetric: {np.allclose(H, H.T)}")
print(f"\nH[0:5, 0:5] (top-left corner):")
print(np.round(H[0:5, 0:5], 2))

What does this verification do?

H.shape confirms the matrix is (200, 200), matching your N interior grid points.

np.allclose(H, H.T) checks that H equals its transpose within floating-point tolerance. This must return True for the eigenvalue solver in the next step to work correctly.

The 5x5 corner printout lets you visually confirm the tridiagonal pattern: large negative values on the diagonal, positive values on the immediate off-diagonals, and zeros everywhere else.

Save quantum_box.py.
Run the script in your terminal by running:

python quantum_box.py

You should see the grid parameters from Step 1, followed by new output showing Hamiltonian shape: (200, 200), Hamiltonian is symmetric: True, and a 5x5 matrix corner with the scaled -2 / +1 tridiagonal pattern.

Seeing False for the symmetry check?

Make sure you used N - 1 for the off-diagonal length (not N). The off-diagonals have one fewer element than the main diagonal.

Check that both k=1 and k=-1 are included in the np.diag() calls. Missing one side breaks the symmetry.

Help me debug my Hamiltonian matrix construction

✔️ Awesome, I've got everything!

Great. Double check you have saved your file before moving on.

ⓧ I'd like to double check the full code

import numpy as np
from scipy.linalg import eigh
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# =============================================================
# STEP 1: Setup - Physical parameters and spatial grid
# =============================================================

hbar = 1.0
m = 1.0
L = 1.0  # Box length

N = 200
dx = L / (N + 1)
x_inner = np.linspace(dx, L - dx, N)

print(f"Grid parameters: N = {N}, dx = {dx:.6f}")
print(f"Domain: x from {x_inner[0]:.4f} to {x_inner[-1]:.4f}")

# =============================================================
# STEP 2: Build the Hamiltonian Matrix
# =============================================================

coeff = -hbar**2 / (2 * m * dx**2)

main_diag = -2 * np.ones(N)
off_diag = np.ones(N - 1)

H = coeff * (np.diag(main_diag) + np.diag(off_diag, k=1) + np.diag(off_diag, k=-1))

print(f"\nHamiltonian shape: {H.shape}")
print(f"Hamiltonian is symmetric: {np.allclose(H, H.T)}")
print(f"\nH[0:5, 0:5] (top-left corner):")
print(np.round(H[0:5, 0:5], 2))

Your Hamiltonian matrix is assembled and verified. Next, you will solve the eigenvalue problem to extract the quantized energy levels and wave functions from this matrix.

Solve for Eigenstates and Verify

In the previous step, you assembled the Hamiltonian matrix that encodes the physics of your particle in a box. Now it's time to extract the quantum states from it.

Solving the eigenvalue equation Hψ = Eψ gives you the allowed energy levels and corresponding wave functions. You'll use SciPy's eigenvalue solver to diagonalize H, then verify the numerical results against known analytical solutions.

In this step, get ready to:

Solve the eigenvalue problem to get quantized energy levels and wave functions.
Verify that each wave function is properly normalized.
Confirm numerical energies match the analytical formula to within 1%.

Solve the eigenvalue problem

The Hamiltonian matrix H is a real symmetric matrix. Diagonalizing it yields eigenvalues (the allowed energies) and eigenvectors (the wave functions at each grid point).

Add the following code to the bottom of quantum_box.py to solve the eigenvalue problem and extract the first 5 states:

# =============================================================
# STEP 3: Solve for Eigenstates and Verify
# =============================================================

# Solve the eigenvalue problem
energies, wavefuncs = eigh(H)

# Extract first 5 states
num_states = 5
E = energies[:num_states]
psi = wavefuncs[:, :num_states]

What does this code do?

eigh(H) solves the eigenvalue equation for a real symmetric (Hermitian) matrix. It returns all 200 eigenvalues in ascending order and their corresponding eigenvectors as columns of a matrix.

You only need the lowest 5 energy states for this project, so E = energies[:num_states] slices the first 5 eigenvalues and psi = wavefuncs[:, :num_states] grabs the first 5 eigenvector columns.

Save quantum_box.py.
You won't see new output from this chunk yet. The next section verifies these results by checking normalization.

Verify normalization of wave functions

A properly normalized wave function satisfies the condition that the integral of |ψ|² over all space equals 1. This means the total probability of finding the particle somewhere in the box is 100%.

Add this normalization check below the code you just wrote:

# Verify normalization
print(f"\nNormalization check (should be ~1.0):")
for n in range(num_states):
    norm = np.sum(np.abs(psi[:, n])**2) * dx
    print(f"  |psi_{n+1}|^2 integral = {norm:.6f}")

What does this code do?

For each of the 5 wave functions, this loop computes the discrete integral of |ψ|² by summing the squared magnitudes at every grid point and multiplying by the spacing dx.

If eigh returned properly normalized eigenvectors (which it does for real symmetric matrices), each integral should print a value very close to 1.0.

Save quantum_box.py and run the script:

python quantum_box.py

Scroll to the bottom of the output. You should see all 5 normalization integrals printed as values very close to 1.000000.

Seeing values far from 1.0?

Make sure you're multiplying by dx at the end of the sum. Without the grid spacing factor, the discrete sum doesn't approximate the continuous integral correctly.

help me debug my normalization check

Compare with analytical solutions

The infinite potential well has a known analytical solution. The energy of the nth state is E_n = n² π² ℏ² / (2mL²).

Add the comparison table below your normalization check:

# Compare with analytical energies: E_n = n^2 * pi^2 * hbar^2 / (2 * m * L^2)
print(f"\nEnergy comparison (numerical vs analytical):")
print(f"{'n':<4}{'Numerical':<15}{'Analytical':<15}{'Error %':<10}")
for n in range(num_states):
    E_analytical = (n + 1)**2 * np.pi**2 * hbar**2 / (2 * m * L**2)
    error = abs(E[n] - E_analytical) / E_analytical * 100
    print(f"{n+1:<4}{E[n]:<15.6f}{E_analytical:<15.6f}{error:<10.4f}")

What does this code do?

This loop computes the analytical energy for each quantum number n = 1 through 5 using the exact formula. It then prints both values side-by-side along with the percentage error.

With 200 grid points, you should see errors well below 1%. The energies scale as 1:4:9:16:25, confirming the n² relationship.

Save quantum_box.py and run it again:

python quantum_box.py

At the bottom of the output, you should see a table with 5 rows. The Numerical and Analytical columns should nearly match, with Error % values below 0.1 for all states.

Errors larger than 1%?

Double check that your coeff variable from Step 2 uses -hbar**2 / (2 * m * dx**2). A sign error or missing factor of 2 in the denominator will shift all energy levels.

Also confirm that N = 200 in your grid setup. Fewer grid points reduce accuracy significantly for higher states.

help me fix my energy level errors

✔️ Awesome, I've got everything!

Great. Double check you've saved quantum_box.py before moving on.

ⓧ I'd like to double check the full code

import numpy as np
from scipy.linalg import eigh
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# =============================================================
# STEP 1: Setup - Physical parameters and spatial grid
# =============================================================

hbar = 1.0
m = 1.0
L = 1.0

N = 200
dx = L / (N + 1)
x_inner = np.linspace(dx, L - dx, N)

print(f"Grid parameters: N = {N}, dx = {dx:.6f}")
print(f"Domain: x from {x_inner[0]:.4f} to {x_inner[-1]:.4f}")

# =============================================================
# STEP 2: Build the Hamiltonian Matrix
# =============================================================

coeff = -hbar**2 / (2 * m * dx**2)

main_diag = -2 * np.ones(N)
off_diag = np.ones(N - 1)

H = coeff * (np.diag(main_diag) + np.diag(off_diag, k=1) + np.diag(off_diag, k=-1))

print(f"\nHamiltonian shape: {H.shape}")
print(f"Hamiltonian is symmetric: {np.allclose(H, H.T)}")
print(f"\nH[0:5, 0:5] (top-left corner):")
print(np.round(H[0:5, 0:5], 2))

# =============================================================
# STEP 3: Solve for Eigenstates and Verify
# =============================================================

energies, wavefuncs = eigh(H)

num_states = 5
E = energies[:num_states]
psi = wavefuncs[:, :num_states]

print(f"\nNormalization check (should be ~1.0):")
for n in range(num_states):
    norm = np.sum(np.abs(psi[:, n])**2) * dx
    print(f"  |psi_{n+1}|^2 integral = {norm:.6f}")

print(f"\nEnergy comparison (numerical vs analytical):")
print(f"{'n':<4}{'Numerical':<15}{'Analytical':<15}{'Error %':<10}")
for n in range(num_states):
    E_analytical = (n + 1)**2 * np.pi**2 * hbar**2 / (2 * m * L**2)
    error = abs(E[n] - E_analytical) / E_analytical * 100
    print(f"{n+1:<4}{E[n]:<15.6f}{E_analytical:<15.6f}{error:<10.4f}")

Your solver is producing verified quantum energy levels and normalized wave functions. Next up, you'll bring these results to life with plots of the wave functions, probability densities, and an energy level diagram.

Visualize Wave Functions and Probability Densities

In the previous step, you confirmed that your numerical solver produces quantized energy levels matching the analytical formula. Numbers in a table are convincing, but the real power of computational physics is turning abstract mathematics into pictures you can see and interpret.

In this step, you will plot the wave functions and probability densities for the first four quantum states. You will also create an energy level diagram showing the characteristic n-squared spacing.

In this step, get ready to:

Plot wave functions for the first four eigenstates.
Plot probability densities and count nodes.
Create an energy level diagram showing quantized spacing.

Plot the first four wave functions

The eigenvectors from eigh only contain values at the 200 interior grid points. The boundary conditions require the wave function to be zero at x=0 and x=L, so you need to add those zeros before plotting.

In quantum_box.py, add the following code at the bottom of the file to create the full spatial array, a helper function, and the wave function plots:

# =============================================================
# STEP 4: Visualize Wave Functions and Probability Densities
# =============================================================

# Full x array including boundaries
x_full = np.linspace(0, L, N + 2)

# Add boundary zeros to wave functions
def add_boundaries(psi_inner):
    return np.concatenate(([0], psi_inner, [0]))

# Plot wave functions
fig1, axes1 = plt.subplots(2, 2, figsize=(10, 8))
fig1.suptitle("Wave Functions psi_n(x)")
for i, ax in enumerate(axes1.flat):
    psi_full = add_boundaries(psi[:, i])
    ax.plot(x_full, psi_full, linewidth=2)
    ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.5)
    ax.set_title(f"n = {i+1}, E = {E[i]:.2f}")
    ax.set_xlabel("x")
    ax.set_ylabel("psi(x)")
    ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("wavefunctions.png", dpi=100)
plt.show()

What does this code do?

x_full creates an array of 202 points (200 interior + 2 boundaries) spanning the full domain from 0 to L. This gives you the complete x-axis for your plots.

add_boundaries() prepends and appends a zero to any interior wave function array. This enforces the physical boundary condition that the particle cannot exist at the box walls.

The plotting loop iterates over a 2x2 grid of subplots, drawing each of the first four eigenstates. Each subplot also includes a dashed zero line to help you visually count nodes (zero crossings).

Save quantum_box.py and run it:

python quantum_box.py

A figure window opens showing four wave functions. Close the figure window to let the script continue (it will stop at plt.show() until you do).
Verify that psi_1 has 0 nodes (no zero crossings), psi_2 has 1 node, psi_3 has 2 nodes, and psi_4 has 3 nodes.

Why n-1 nodes?

A fundamental result in quantum mechanics is that the nth eigenstate of any 1D potential has exactly n-1 nodes. This is a consequence of the Sturm-Liouville theorem. Your plots should confirm this pattern visually.

Plot window not appearing?

Some environments (remote servers, WSL) do not have a display backend configured. If no window appears, check that you are running the script on your local machine where Matplotlib can open GUI windows.

Even if the window does not appear, the file wavefunctions.png is saved to your project folder. You can open it directly to verify the output.

Help me debug my Matplotlib display issue.

Plot probability densities

The wave function itself can be positive or negative, but |psi(x)|^2 gives the probability density. This is what you would actually measure in an experiment.

Add the following code below the wave function plot to create probability density plots:

# Plot probability densities
fig2, axes2 = plt.subplots(2, 2, figsize=(10, 8))
fig2.suptitle("Probability Densities |psi_n(x)|^2")
for i, ax in enumerate(axes2.flat):
    psi_full = add_boundaries(psi[:, i])
    prob_density = np.abs(psi_full)**2
    ax.plot(x_full, prob_density, linewidth=2, color='darkorange')
    ax.fill_between(x_full, prob_density, alpha=0.3, color='orange')
    ax.set_title(f"n = {i+1}, nodes = {i}")
    ax.set_xlabel("x")
    ax.set_ylabel("|psi(x)|^2")
    ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("probability_densities.png", dpi=100)
plt.show()

What does this code do?

np.abs(psi_full)**2 computes the probability density at every point. The fill_between call shades the area under the curve, making it visually clear where the particle is most likely to be found.

Notice that higher-n states spread the probability more evenly across the box. The ground state (n=1) peaks in the center, while n=4 has four distinct lobes.

Save quantum_box.py and run it again. Close the first figure window when it appears, then inspect the probability density figure.
Confirm that each subplot title shows the correct node count (0, 1, 2, 3 for n=1, 2, 3, 4).

Probability densities look identical to wave functions?

Make sure you are squaring the absolute value (np.abs(psi_full)**2), not just plotting psi_full again. The probability density should always be non-negative (no dips below zero).

Help me fix my probability density plot.

Create an energy level diagram

An energy level diagram shows the quantized energies as horizontal lines stacked vertically. This is the classic visualization that makes the discrete nature of quantum mechanics immediately obvious.

Add the following code below the probability density plot to draw the energy level diagram:

# Energy level diagram
fig3, ax3 = plt.subplots(figsize=(6, 8))
for n in range(num_states):
    ax3.hlines(E[n], 0.2, 0.8, linewidth=2, color='blue')
    ax3.text(0.85, E[n], f"n={n+1}, E={E[n]:.1f}", va='center')
ax3.set_xlim(0, 1.5)
ax3.set_ylim(0, E[num_states - 1] * 1.2)
ax3.set_ylabel("Energy")
ax3.set_title("Quantized Energy Levels")
ax3.set_xticks([])
ax3.grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.savefig("energy_levels.png", dpi=100)
plt.show()

What does this code do?

ax3.hlines() draws a horizontal line at each energy value. The x-axis is purely cosmetic here (the lines span from 0.2 to 0.8) since the only meaningful axis is the vertical energy axis.

Look at the spacing between levels. The gap between E_1 and E_2 is smaller than the gap between E_4 and E_5. This is the quadratic scaling (E proportional to n^2) that distinguishes quantum confinement from, say, a harmonic oscillator where levels are equally spaced.

Save quantum_box.py and run it one more time. Close the first two figures, then inspect the energy level diagram.
Confirm you see five horizontal blue lines with increasing vertical gaps between them. The annotations should read approximately E=4.9, 19.7, 44.4, 79.0, and 123.3.

Energy levels all bunched at the bottom?

If the lines are not visibly spread out, check that your y-axis limit uses E[num_states - 1] * 1.2 so the axis scales to fit all five levels with some headroom.

Help me debug my energy level diagram.

✔️ Awesome, I've got everything!

Great. Double check you have saved quantum_box.py and that all three figures (wavefunctions.png, probability_densities.png, and energy_levels.png) are saved in your project folder.

ⓧ I'd like to double check the full code

import numpy as np
from scipy.linalg import eigh
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# =============================================================
# STEP 1: Setup - Physical parameters and spatial grid
# =============================================================

hbar = 1.0
m = 1.0
L = 1.0

N = 200
dx = L / (N + 1)
x_inner = np.linspace(dx, L - dx, N)

print(f"Grid parameters: N = {N}, dx = {dx:.6f}")
print(f"Domain: x from {x_inner[0]:.4f} to {x_inner[-1]:.4f}")

# =============================================================
# STEP 2: Build the Hamiltonian Matrix
# =============================================================

coeff = -hbar**2 / (2 * m * dx**2)

main_diag = -2 * np.ones(N)
off_diag = np.ones(N - 1)

H = coeff * (np.diag(main_diag) + np.diag(off_diag, k=1) + np.diag(off_diag, k=-1))

print(f"\nHamiltonian shape: {H.shape}")
print(f"Hamiltonian is symmetric: {np.allclose(H, H.T)}")
print(f"\nH[0:5, 0:5] (top-left corner):")
print(np.round(H[0:5, 0:5], 2))

# =============================================================
# STEP 3: Solve for Eigenstates and Verify
# =============================================================

energies, wavefuncs = eigh(H)

num_states = 5
E = energies[:num_states]
psi = wavefuncs[:, :num_states]

print(f"\nNormalization check (should be ~1.0):")
for n in range(num_states):
    norm = np.sum(np.abs(psi[:, n])**2) * dx
    print(f"  |psi_{n+1}|^2 integral = {norm:.6f}")

print(f"\nEnergy comparison (numerical vs analytical):")
print(f"{'n':<4}{'Numerical':<15}{'Analytical':<15}{'Error %':<10}")
for n in range(num_states):
    E_analytical = (n + 1)**2 * np.pi**2 * hbar**2 / (2 * m * L**2)
    error = abs(E[n] - E_analytical) / E_analytical * 100
    print(f"{n+1:<4}{E[n]:<15.6f}{E_analytical:<15.6f}{error:<10.4f}")

# =============================================================
# STEP 4: Visualize Wave Functions and Probability Densities
# =============================================================

# Full x array including boundaries
x_full = np.linspace(0, L, N + 2)

# Add boundary zeros to wave functions
def add_boundaries(psi_inner):
    return np.concatenate(([0], psi_inner, [0]))

# Plot wave functions
fig1, axes1 = plt.subplots(2, 2, figsize=(10, 8))
fig1.suptitle("Wave Functions psi_n(x)")
for i, ax in enumerate(axes1.flat):
    psi_full = add_boundaries(psi[:, i])
    ax.plot(x_full, psi_full, linewidth=2)
    ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.5)
    ax.set_title(f"n = {i+1}, E = {E[i]:.2f}")
    ax.set_xlabel("x")
    ax.set_ylabel("psi(x)")
    ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("wavefunctions.png", dpi=100)
plt.show()

# Plot probability densities
fig2, axes2 = plt.subplots(2, 2, figsize=(10, 8))
fig2.suptitle("Probability Densities |psi_n(x)|^2")
for i, ax in enumerate(axes2.flat):
    psi_full = add_boundaries(psi[:, i])
    prob_density = np.abs(psi_full)**2
    ax.plot(x_full, prob_density, linewidth=2, color='darkorange')
    ax.fill_between(x_full, prob_density, alpha=0.3, color='orange')
    ax.set_title(f"n = {i+1}, nodes = {i}")
    ax.set_xlabel("x")
    ax.set_ylabel("|psi(x)|^2")
    ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("probability_densities.png", dpi=100)
plt.show()

# Energy level diagram
fig3, ax3 = plt.subplots(figsize=(6, 8))
for n in range(num_states):
    ax3.hlines(E[n], 0.2, 0.8, linewidth=2, color='blue')
    ax3.text(0.85, E[n], f"n={n+1}, E={E[n]:.1f}", va='center')
ax3.set_xlim(0, 1.5)
ax3.set_ylim(0, E[num_states - 1] * 1.2)
ax3.set_ylabel("Energy")
ax3.set_title("Quantized Energy Levels")
ax3.set_xticks([])
ax3.grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.savefig("energy_levels.png", dpi=100)
plt.show()

You now have a complete visual picture of the quantum particle in a box. Next up, you will bring these static states to life by animating a superposition state that sloshes back and forth inside the well.

Animate Time Evolution of a Superposition

Your static plots reveal the spatial structure of quantum states. But quantum mechanics is fundamentally about dynamics.

A particle in a superposition of energy eigenstates is not stationary. Its probability density changes with time, sloshing back and forth inside the box. In this step, you will animate that time evolution so you can see quantum dynamics in action.

In this step, get ready to:

Create a superposition of the first two eigenstates.
Implement the time evolution formula for the probability density.
Animate the evolving probability density with Matplotlib FuncAnimation.

Create a superposition state and define time parameters

Each eigenstate you solved for in Step 3 is a stationary state. If a particle is purely in state psi_1, its probability density never changes. The interesting behaviour happens when you combine two or more eigenstates into a superposition.

In quantum_box.py, scroll to the bottom of your file (after the energy level diagram code).
Add the Step 5 section header and superposition setup by copying this code:

# =============================================================
# STEP 5: Animate Time Evolution of a Superposition
# =============================================================

# Create superposition of first two states: (1/sqrt(2))(psi_1 + psi_2)
c1 = 1 / np.sqrt(2)
c2 = 1 / np.sqrt(2)

# Time parameters
omega_1 = E[0] / hbar
omega_2 = E[1] / hbar
T_period = 2 * np.pi / (omega_2 - omega_1)  # Period of oscillation
num_frames = 100
t_values = np.linspace(0, 2 * T_period, num_frames)

# Precompute probability density at each time step
psi1_full = add_boundaries(psi[:, 0])
psi2_full = add_boundaries(psi[:, 1])

What does this code do?

You are building the initial state Psi(x, 0) = (1/sqrt(2)) * (psi_1 + psi_2). The equal coefficients c1 and c2 mean the particle has a 50/50 chance of being measured in either energy state.

The angular frequencies omega_1 and omega_2 determine how fast each component rotates in the complex plane. The difference between them sets the oscillation period T_period of the probability density.

Why does a superposition evolve in time?

Each eigenstate picks up a phase factor exp(-i * E_n * t / hbar) that rotates at a rate proportional to its energy. When you add two eigenstates with different energies, their phases drift apart over time. This interference pattern is what makes the probability density move.

Implement the time evolution function

You need a function that computes the probability density |Psi(x, t)|^2 at any given time. This function will be called once per animation frame.

Add the compute_prob_density function directly below the psi2_full line you just wrote:

def compute_prob_density(t):
    """Compute |Psi(x,t)|^2 for the superposition state."""
    Psi_t = (c1 * psi1_full * np.exp(-1j * omega_1 * t) +
             c2 * psi2_full * np.exp(-1j * omega_2 * t))
    return np.abs(Psi_t)**2

What does this function do?

At each time t, the function multiplies each eigenstate by its complex phase factor exp(-i * omega * t), sums them, and takes the absolute value squared. The result is a real-valued probability density that changes shape as the phases drift.

Save quantum_box.py.

Seeing a syntax error?

Make sure the second line of the function body (starting with c2 * psi2_full) is indented to align with c1 * psi1_full inside the parentheses. Python requires consistent indentation for continuation lines.

Help me fix a syntax error in my compute_prob_density function.

Animate with Matplotlib FuncAnimation

Now you will set up the Matplotlib animation. FuncAnimation works by calling an update function on every frame, redrawing the plot data each time.

Add the animation setup below the compute_prob_density function:

# Set up animation
fig4, ax4 = plt.subplots(figsize=(8, 5))
line, = ax4.plot(x_full, compute_prob_density(0), linewidth=2, color='purple')
ax4.set_xlim(0, L)
ax4.set_ylim(0, 5)
ax4.set_xlabel("x")
ax4.set_ylabel("|Psi(x, t)|^2")
ax4.set_title("Time Evolution of Superposition State")
ax4.grid(True, alpha=0.3)
time_text = ax4.text(0.02, 0.95, '', transform=ax4.transAxes)

def update(frame):
    t = t_values[frame]
    prob = compute_prob_density(t)
    line.set_ydata(prob)
    time_text.set_text(f"t = {t:.3f}")
    return line, time_text

ani = FuncAnimation(fig4, update, frames=num_frames, interval=50, blit=True)
plt.show()

What does this code do?

The figure is created with the initial probability density at t=0 drawn as a purple line.
The update function is called once per frame. It recomputes the probability density at the current time and updates the line's y-data.
FuncAnimation ties everything together: it loops through 100 frames with 50ms between each, creating the illusion of smooth motion.
The time_text label in the top-left corner shows the current simulation time so you can track where you are in the oscillation cycle.

Save quantum_box.py.
Run the script by executing this command in your terminal:

python quantum_box.py

Close each static plot window as it appears (the wave function, probability density, and energy level plots from earlier steps).
After the static plots, the animation window opens. You should see a purple probability density curve oscillating back and forth inside the box.

What should I see?

The probability density "sloshes" from one side of the box to the other, like a wave bouncing between walls. The time label in the top-left corner advances as the animation plays.

The oscillation period is T = 2*pi*hbar / (E_2 - E_1). The animation runs for two full periods so you can see the repeating pattern. The total area under the curve stays constant at 1.0, meaning probability is conserved at all times.

Animation window is blank or freezes?

Some systems need an interactive Matplotlib backend to display animations. If the window is unresponsive, try closing all plot windows and running the script again. Make sure no other Matplotlib windows are still open from a previous run.

Also check that blit=True is included in the FuncAnimation call. Without it, the animation may render too slowly to appear smooth.

Help me debug my FuncAnimation that is not displaying correctly.

✔️ Awesome, I've got everything!

Great. Double check you have saved your file.

ⓧ I'd like to double check the full code

Here is the complete quantum_box.py file including all steps up to this point:

import numpy as np
from scipy.linalg import eigh
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# =============================================================
# STEP 1: Setup - Physical parameters and spatial grid
# =============================================================

# Natural units: hbar = 1, m = 1
hbar = 1.0
m = 1.0
L = 1.0  # Box length

# Discretization
N = 200  # Number of interior grid points
dx = L / (N + 1)
x_inner = np.linspace(dx, L - dx, N)  # Interior points only (boundary = 0)

print(f"Grid parameters: N = {N}, dx = {dx:.6f}")
print(f"Domain: x from {x_inner[0]:.4f} to {x_inner[-1]:.4f}")

# =============================================================
# STEP 2: Build the Hamiltonian Matrix
# =============================================================

# Kinetic energy coefficient
coeff = -hbar**2 / (2 * m * dx**2)

# Tridiagonal matrix for the second derivative
main_diag = -2 * np.ones(N)
off_diag = np.ones(N - 1)

H = coeff * (np.diag(main_diag) + np.diag(off_diag, k=1) + np.diag(off_diag, k=-1))

print(f"\nHamiltonian shape: {H.shape}")
print(f"Hamiltonian is symmetric: {np.allclose(H, H.T)}")
print(f"\nH[0:5, 0:5] (top-left corner):")
print(np.round(H[0:5, 0:5], 2))

# =============================================================
# STEP 3: Solve for Eigenstates and Verify
# =============================================================

# Solve the eigenvalue problem
energies, wavefuncs = eigh(H)

# Extract first 5 states
num_states = 5
E = energies[:num_states]
psi = wavefuncs[:, :num_states]

# Verify normalization
print(f"\nNormalization check (should be ~1.0):")
for n in range(num_states):
    norm = np.sum(np.abs(psi[:, n])**2) * dx
    print(f"  |psi_{n+1}|^2 integral = {norm:.6f}")

# Compare with analytical energies: E_n = n^2 * pi^2 * hbar^2 / (2 * m * L^2)
print(f"\nEnergy comparison (numerical vs analytical):")
print(f"{'n':<4}{'Numerical':<15}{'Analytical':<15}{'Error %':<10}")
for n in range(num_states):
    E_analytical = (n + 1)**2 * np.pi**2 * hbar**2 / (2 * m * L**2)
    error = abs(E[n] - E_analytical) / E_analytical * 100
    print(f"{n+1:<4}{E[n]:<15.6f}{E_analytical:<15.6f}{error:<10.4f}")

# =============================================================
# STEP 4: Visualize Wave Functions and Probability Densities
# =============================================================

# Full x array including boundaries
x_full = np.linspace(0, L, N + 2)

# Add boundary zeros to wave functions
def add_boundaries(psi_inner):
    return np.concatenate(([0], psi_inner, [0]))

# Plot wave functions
fig1, axes1 = plt.subplots(2, 2, figsize=(10, 8))
fig1.suptitle("Wave Functions psi_n(x)")
for i, ax in enumerate(axes1.flat):
    psi_full = add_boundaries(psi[:, i])
    ax.plot(x_full, psi_full, linewidth=2)
    ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.5)
    ax.set_title(f"n = {i+1}, E = {E[i]:.2f}")
    ax.set_xlabel("x")
    ax.set_ylabel("psi(x)")
    ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("wavefunctions.png", dpi=100)
plt.show()

# Plot probability densities
fig2, axes2 = plt.subplots(2, 2, figsize=(10, 8))
fig2.suptitle("Probability Densities |psi_n(x)|^2")
for i, ax in enumerate(axes2.flat):
    psi_full = add_boundaries(psi[:, i])
    prob_density = np.abs(psi_full)**2
    ax.plot(x_full, prob_density, linewidth=2, color='darkorange')
    ax.fill_between(x_full, prob_density, alpha=0.3, color='orange')
    ax.set_title(f"n = {i+1}, nodes = {i}")
    ax.set_xlabel("x")
    ax.set_ylabel("|psi(x)|^2")
    ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("probability_densities.png", dpi=100)
plt.show()

# Energy level diagram
fig3, ax3 = plt.subplots(figsize=(6, 8))
for n in range(num_states):
    ax3.hlines(E[n], 0.2, 0.8, linewidth=2, color='blue')
    ax3.text(0.85, E[n], f"n={n+1}, E={E[n]:.1f}", va='center')
ax3.set_xlim(0, 1.5)
ax3.set_ylim(0, E[num_states - 1] * 1.2)
ax3.set_ylabel("Energy")
ax3.set_title("Quantized Energy Levels")
ax3.set_xticks([])
ax3.grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.savefig("energy_levels.png", dpi=100)
plt.show()

# =============================================================
# STEP 5: Animate Time Evolution of a Superposition
# =============================================================

# Create superposition of first two states: (1/sqrt(2))(psi_1 + psi_2)
c1 = 1 / np.sqrt(2)
c2 = 1 / np.sqrt(2)

# Time parameters
omega_1 = E[0] / hbar
omega_2 = E[1] / hbar
T_period = 2 * np.pi / (omega_2 - omega_1)  # Period of oscillation
num_frames = 100
t_values = np.linspace(0, 2 * T_period, num_frames)

# Precompute probability density at each time step
psi1_full = add_boundaries(psi[:, 0])
psi2_full = add_boundaries(psi[:, 1])

def compute_prob_density(t):
    """Compute |Psi(x,t)|^2 for the superposition state."""
    Psi_t = (c1 * psi1_full * np.exp(-1j * omega_1 * t) +
             c2 * psi2_full * np.exp(-1j * omega_2 * t))
    return np.abs(Psi_t)**2

# Set up animation
fig4, ax4 = plt.subplots(figsize=(8, 5))
line, = ax4.plot(x_full, compute_prob_density(0), linewidth=2, color='purple')
ax4.set_xlim(0, L)
ax4.set_ylim(0, 5)
ax4.set_xlabel("x")
ax4.set_ylabel("|Psi(x, t)|^2")
ax4.set_title("Time Evolution of Superposition State")
ax4.grid(True, alpha=0.3)
time_text = ax4.text(0.02, 0.95, '', transform=ax4.transAxes)

def update(frame):
    t = t_values[frame]
    prob = compute_prob_density(t)
    line.set_ydata(prob)
    time_text.set_text(f"t = {t:.3f}")
    return line, time_text

ani = FuncAnimation(fig4, update, frames=num_frames, interval=50, blit=True)
plt.show()

Your quantum simulation is complete. You have gone from a static eigenvalue problem all the way to a living, breathing animation of quantum dynamics. Next up: the Secret Mission challenges you to break the walls of the infinite well and observe quantum tunneling.

Secret mission

Observe Quantum Tunneling in a Finite Well

In the infinite well, the particle is perfectly trapped. The wave function is forced to zero at the walls, with no chance of escape. But real quantum systems are not so rigid. Can you modify your simulation to let the particle leak through the barrier?

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your files for future experiments, or delete them entirely.

Resources you used:

Python script: quantum_box.py
Generated plot images: wavefunctions.png, probability_densities.png, energy_levels.png, tunneling.png
Python packages: NumPy, SciPy, Matplotlib.

Keep everything running

No action needed. Choose this if you plan to keep experimenting with the simulation, tweaking parameters, or extending the code further.

Your script and plot files stay on your Desktop.
NumPy, SciPy, and Matplotlib remain installed and available for other Python projects.

Pause - I'll come back to this later

There are no running processes or servers to shut down. Your files will stay on disk until you return.

No action required. Since this is a standalone script (not a running service), there is nothing to pause.
Your quantum_box.py and all generated images will be waiting on your Desktop when you come back.

Delete - I don't want to use this again

Remove all project files and uninstall the Python packages you installed for this project.

Delete the project files by running this command:

cd ~/Desktop
rm quantum_box.py wavefunctions.png probability_densities.png energy_levels.png tunneling.png

Uninstall the Python packages by running:

pip uninstall numpy scipy matplotlib -y

Keep the packages if you use Python often

NumPy, SciPy, and Matplotlib are foundational Python libraries used across data science, machine learning, and scientific computing. If you plan to do more Python work, consider keeping them installed.

That's a Wrap!

Nice work! You just built a quantum mechanics simulation from scratch, turning the mathematical formalism of the Schrodinger equation into working Python code that produces real, verifiable physics.

You've learned how to:

Build a numerical solver for the 1D Schrodinger equation using the finite-difference method, and verify that quantized energy levels emerge naturally from boundary conditions.
Visualize wave functions and probability densities to confirm quantum properties like nodes, symmetry, and normalization.
Animate a superposition state evolving in time to see probability density oscillating back and forth inside the potential well.
Secret Mission: Observed quantum tunneling by switching to a finite potential well and seeing wave functions leak into classically forbidden regions.

Ready to quiz yourself?