Boo-AI — Master Artificial Intelligence by Building from Scratch

Learning Objectives

By the end of this section, you will be able to:

Explain why the vast majority of Schrödinger problems must be solved numerically, and why "exact" hydrogen is the rare exception.
Discretise the kinetic-energy operator $-\tfrac{1}{2}\,\partial_x^2$ by a 3-point finite-difference stencil and write the Hamiltonian as a tridiagonal matrix.
Convert a continuous eigenvalue problem $\hat H \psi = E \psi$ into a matrix eigenproblem you can solve with one library call.
Apply the shooting method and Numerov's higher-order integrator to find bound states one at a time.
Evolve a wave packet in time using the split-operator (FFT) method, and use it to watch tunnelling happen.
Estimate the convergence rate of each method and pick the right tool for a given physics problem.
Implement a finite-difference solver from scratch in NumPy and a batched GPU version in PyTorch.

Why We Need Numerical Methods

"The exactly solvable problems are a measure-zero subset of all the problems you actually want to solve." — every computational physicist, ever

The Schrödinger equation looks simple. It is a single linear PDE: $\hat H \psi(x) = E \psi(x)$ . Yet of all the atoms, molecules, semiconductors, and quantum dots in the universe, only a handful can be solved with pencil and paper. The list is essentially: free particle, infinite square well, finite square well (transcendentally), the harmonic oscillator, the hydrogen atom, the delta-function potential, and a few exotic toys (Pöschl–Teller, Morse, ...). Everything else — including the helium atom with its measly two electrons — needs a computer.

The intuition in one sentence

Discrete numerical methods replace the differential operator

\partial_x^2

with a matrix, and replace the wavefunction

\psi(x)

with a vector of values on a grid. The continuous eigenvalue problem becomes a finite-dimensional linear algebra problem — and computers are very, very good at linear algebra.

This is more than a workaround. The numerical viewpoint is more general than the analytical one: you can change the potential to anything you want, change the boundary conditions, add time dependence, couple multiple particles — the code structure stays nearly identical. Once you understand the recipe, the variety of physics you can simulate explodes.

The three problems we will solve

Problem	What we want	Best tool
Bound states of a 1D potential	Energies E_n and stationary states ψ_n(x)	Finite-difference + dense eigensolve
A single bound state at a known approximate energy	Refine E to high precision; localised eigenfunction	Shooting method (Numerov)
How a wave packet evolves in time	ψ(x, t) for t > 0 from a given ψ(x, 0)	Split-operator method (FFT)

The Recipe: Turn an Operator Into a Matrix

Take a moment to look at the time-independent Schrödinger equation in one dimension:

\displaystyle -\frac{\hbar^2}{2m}\,\frac{d^2\psi}{dx^2} + V(x)\,\psi(x) = E\,\psi(x).

For the rest of this section we use the standard atomic units $\hbar = m = 1$ , which absorbs all the physical constants into the choice of length and energy scales. The equation simplifies to:

\displaystyle -\tfrac{1}{2}\,\psi''(x) + V(x)\,\psi(x) = E\,\psi(x).

Think of the left-hand side as an operator $\hat H$ acting on the function ψ:

\hat H = -\tfrac{1}{2}\,\partial_x^2 + V(x).

Question: if ψ is a function on a continuous line, how do we make a computer understand it?

Answer: pretend you only care about ψ at a finite list of points $x_1, x_2, \ldots, x_N$ . The wavefunction is now a column vector $\boldsymbol{\psi} = (\psi_1, \ldots, \psi_N)^T$ of N numbers, and $\hat H$ becomes an $N \times N$ matrix $\mathbf H$ . The eigenvalue problem becomes:

\mathbf H\, \boldsymbol\psi = E\, \boldsymbol\psi.

That is just a matrix eigenvalue problem. We already have one-line library calls for it: numpy.linalg.eigh, torch.linalg.eigh, scipy.sparse.linalg.eigsh. The only question left is: what does the matrix $\mathbf H$ look like?

Finite Differences for the Second Derivative

The Schrödinger equation has a single hard piece: the second derivative $\psi''(x)$ . Everything else (V, E, ψ itself) is pointwise. So if we can write $\psi''$ on a grid, we are done.

Take three neighbouring points $x_{i-1}, x_i, x_{i+1}$ separated by spacing $dx$ . Taylor-expand ψ at the two outer points around $x_i$ :

\psi(x_{i+1}) = \psi_i + \psi_i'\,dx + \tfrac{1}{2}\psi_i'' dx^2 + \tfrac{1}{6}\psi_i'''\,dx^3 + \tfrac{1}{24}\psi_i^{(4)}\,dx^4 + \cdots

\psi(x_{i-1}) = \psi_i - \psi_i'\,dx + \tfrac{1}{2}\psi_i''\,dx^2 - \tfrac{1}{6}\psi_i'''\,dx^3 + \tfrac{1}{24}\psi_i^{(4)}\,dx^4 - \cdots

Add them. The odd-order terms cancel:

\psi(x_{i+1}) + \psi(x_{i-1}) = 2\,\psi_i + \psi_i''\,dx^2 + \tfrac{1}{12}\psi_i^{(4)}\,dx^4 + O(dx^6).

Solve for $\psi_i''$ :

\displaystyle \psi_i'' \;\approx\; \frac{\psi_{i+1} - 2\psi_i + \psi_{i-1}}{dx^2} \;+\; O(dx^2).

This is the famous 3-point central difference stencil. It says: to find the curvature at point i, compare the average of your neighbours to your own value, and divide by $dx^2$ . The error is $O(dx^2)$ — second-order accurate. Cut $dx$ in half, and the error drops by a factor of four.

Geometrically the stencil is asking: how much does ψ deviate from the straight line joining its two neighbours? A flat function has zero deviation (ψ″ = 0); a strongly curved one has a lot.

Now plug this into the kinetic-energy operator $-\tfrac{1}{2}\,\partial_x^2$ :

\displaystyle (\hat T \psi)_i \;\approx\; \frac{1}{dx^2}\,\psi_i \;-\; \frac{1}{2\,dx^2}\,\bigl(\psi_{i+1} + \psi_{i-1}\bigr).

Read off the matrix entries: on row i, the diagonal entry is $+1/dx^2$ , the two off-diagonal entries are $-1/(2\,dx^2)$ , and the rest of the row is zero.

The Finite-Difference Eigenproblem

Add the potential and the entire Hamiltonian appears in front of us:

\mathbf H = \begin{pmatrix} \frac{1}{dx^2}+V_1 & -\frac{1}{2 dx^2} & & & \\ -\frac{1}{2 dx^2} & \frac{1}{dx^2}+V_2 & -\frac{1}{2 dx^2} & & \\ & -\frac{1}{2 dx^2} & \frac{1}{dx^2}+V_3 & \ddots & \\ & & \ddots & \ddots & -\frac{1}{2 dx^2} \\ & & & -\frac{1}{2 dx^2} & \frac{1}{dx^2}+V_N \end{pmatrix}

This is a real, symmetric, tridiagonal $N \times N$ matrix. Its eigenvalues are the approximate bound-state energies; its eigenvectors are the approximate bound-state wavefunctions sampled at the grid points.

Boundary conditions are baked in implicitly. By excluding the endpoints from the unknown vector we forced

\psi_0 = \psi_{N+1} = 0

— Dirichlet conditions on the box

[x_\text{min}, x_\text{max}]

. For unbounded potentials (like the oscillator) we just pick a domain wide enough that the true ψ is already negligible at the walls.

Let's look at three rows of $\mathbf H$ for $N = 5$ , $dx = 0.2$ , and a harmonic potential $V(x) = \tfrac{1}{2} x^2$ on $[-0.6, 0.6]$ , so $x = -0.4, -0.2, 0, 0.2, 0.4$ . Then $1/dx^2 = 25$ and $1/(2\,dx^2) = 12.5$ :

i	x_i	V(x_i)	H_{i,i-1}	H_{i,i}	H_{i,i+1}
1	−0.4	0.080	—	25.080	−12.5
2	−0.2	0.020	−12.5	25.020	−12.5
3	0.0	0.000	−12.5	25.000	−12.5
4	0.2	0.020	−12.5	25.020	−12.5
5	0.4	0.080	−12.5	25.080	—

This 5×5 matrix is enough that you could diagonalise it by hand with characteristic polynomials. With a computer we push N up to a few hundred or a few thousand and the answers converge to machine precision.

Interactive: Build the Hamiltonian Yourself

Pick a potential. Move the slider for the shape parameter. Drag N to change the grid resolution. The viz builds $\mathbf H$ , diagonalises it in your browser, and shows the lowest six eigenfunctions stacked on their energy levels — exactly the figure you would get from the Python solver below. The right-hand table compares the numerical energies to the analytic ones (when available) so you can see the $O(dx^2)$ convergence as you slide N up.

Loading finite-difference solver…

Things to try

On the harmonic potential, start with N = 20 and watch the error in E₀ shrink as you move toward N = 200. Doubling N really does cut the error by ~4×.
Switch to double well and decrease the parameter a. The two lowest eigenvalues collapse toward each other — you are seeing the tunnel-doublet form in real time.
On the finite well, count the bound states (energies below zero). Make the well shallower (smaller V₀) and bound states disappear from the top of the well one by one — exactly the transcendental-equation story from §29.4.

The Shooting Method

The finite-difference matrix approach finds all bound states at once. That is wasteful if you only care about one — say, the ground state to ten digits of accuracy. The shooting method is the surgical alternative.

The idea is delightfully physical. Treat the time-independent Schrödinger equation as an initial-value problem in x:

\psi''(x) = 2\,\bigl(V(x) - E\bigr)\,\psi(x).

Choose a trial energy E. Start at the left wall with $\psi(x_\text{min}) = 0$ and $\psi'(x_\text{min}) = \epsilon$ (any small number — the equation is linear, so the overall scale is irrelevant). Integrate ψ rightward step by step using your favourite ODE method. Look at the value at the right wall.

If $\psi(x_\text{max}) \approx 0$ , E is an eigenvalue: you found a bound state.
If $\psi(x_\text{max})$ blows up positive, E is too low.
If $\psi(x_\text{max})$ blows up negative, E is too high.

The sign change of $\psi(x_\text{max})$ as you scan E across a true eigenvalue is exactly what makes this method work. Wrap a one-dimensional root finder (bisection, Brent) around it and you get bound-state energies to arbitrary precision.

Why does ψ blow up at wrong energies? The Schrödinger ODE has two independent solutions in the classically forbidden region (where E < V): one decaying like

e^{-\kappa x}

, one growing like

e^{+\kappa x}

. Generic initial data is a mixture of both. Only at very special E does the growing piece have coefficient zero — and at those E the leftover decaying solution dies at the right wall instead of exploding. That is what we are searching for.

Numerov: A Smarter Integrator

For an equation of the form $\psi'' = f(x)\,\psi$ there is a beautiful 4th-order method that costs the same as plain finite differences: Numerov's method. Define $f_i = 2 (V_i - E)$ . The recurrence is

\displaystyle \psi_{i+1}\,\bigl(1 - \tfrac{h^2}{12}\,f_{i+1}\bigr) = 2\,\psi_i\,\bigl(1 + \tfrac{5 h^2}{12}\,f_i\bigr) - \psi_{i-1}\,\bigl(1 - \tfrac{h^2}{12}\,f_{i-1}\bigr),

with $h = dx$ . The local truncation error is $O(h^6)$ ; the global error is $O(h^4)$ — two orders better than the simple 3-point stencil, for the same number of arithmetic operations per step.

Where does the magic come from? Numerov absorbs the leading $\psi^{(4)}$ term in the Taylor expansion by cleverly distributing the function values over three points. Read this as: the second-derivative formula has a built-in correction that exploits the particular structure $\psi'' = f \psi$ — it would not work for an arbitrary right-hand side, but our Schrödinger ODE is precisely of this form.

Interactive: Shoot the Harmonic Oscillator

Set the trial energy with the slider. The orange curve is the result of integrating ψ from the left wall using Numerov. Whenever E is far from an eigenvalue, the wavefunction explodes near the right wall. Walk E slowly toward 0.5, 1.5, 2.5, 3.5, … and watch the curve collapse back to zero — those are the bound states $E_n = n + \tfrac{1}{2}$ .

Loading shooting-method demo…

Why is the slider so sensitive?

Near a true eigenvalue the tail

\psi(x_\text{max})

changes sign through zero. Off by 0.01 in E and the tail can be 10⁻¹; off by 10⁻⁴ in E and the tail is 10⁻³; that exponential sensitivity is exactly what makes the shooting method precise — and exactly why you wrap a root finder around it instead of dragging a slider in production.

Time Evolution: The Split-Operator Method

The time-dependent Schrödinger equation

i\,\partial_t \psi = \hat H\,\psi

has the formal solution $\psi(t + dt) = e^{-i\,\hat H\,dt}\,\psi(t)$ . The matrix exponential of an N×N matrix is expensive — but the kinetic and potential parts of $\hat H$ have a beautiful property: each is diagonal in its natural basis.

The potential operator $\hat V = V(x)$ is diagonal in position space.
The kinetic operator $\hat T = \tfrac{1}{2}\hat p^2$ is diagonal in momentum space, with eigenvalue $k^2/2$ .

And the Fast Fourier Transform shuffles between the two bases in $O(N \log N)$ time. That suggests:

\displaystyle \psi(t + dt) \approx e^{-i V(x)\,dt/2}\;\mathcal F^{-1}\,\Bigl[ e^{-i\,k^2/2\,dt}\;\mathcal F\bigl[\,e^{-i V(x)\,dt/2}\,\psi(t)\,\bigr]\Bigr].

This is the Strang splitting form of the split-operator method. The error per step is $O(dt^3)$ ; the global error is $O(dt^2)$ . The scheme is unitary by construction (every multiplied factor is a complex phase of magnitude 1), so it conserves the total probability $\int |\psi|^2\,dx$ to floating-point precision indefinitely. That is enormous — naive Euler schemes on the Schrödinger equation are catastrophically wrong because they slowly leak probability and eventually make ψ explode.

The intuition: each step is three rotations

Multiplying ψ by

e^{-i V dt/2}

in position space is a per-point phase rotation that depends on the local potential. Multiplying by

e^{-i k^2 dt/2}

in momentum space is a per-mode phase rotation that depends on how fast that mode wiggles. The FFT shuffles between the two views. Three rotations, no amplitude change — that is why the method conserves probability exactly.

Interactive: Watch a Wave Packet Hit a Barrier

A Gaussian wave packet is launched from the left with mean momentum $k_0$ . The energy is approximately $E \approx k_0^2/2$ . Move the barrier height above E and watch most of the packet bounce back — but a thin transmitted tail leaks through the wall. That tail is quantum tunnelling, visible because we did the time evolution properly.

Loading split-operator demo…

Things to try

Free potential: a Gaussian packet just translates and slowly spreads — the dispersion of a quantum free particle.
Barrier with $V_0 > k_0^2/2$ : classical particles reflect with probability 1, but the green |ψ|² shows a small transmitted bump on the far side.
Harmonic: an off-centre Gaussian oscillates back and forth in the trap forever — that is a coherent state, the classical-like motion you met in §29.5.

Convergence, Error, and Choosing a Method

Every numerical method has a budget: dx (or dt), N, polynomial order, boundary placement. The art is matching the method to the physics.

Method	Accuracy	Cost	Best for
3-pt finite difference + dense eigh	O(dx²)	O(N³) to diagonalise	Bound states of arbitrary V(x), small N (≤ 2000)
5-pt or 7-pt FD + sparse eigsh	O(dx⁴) or O(dx⁶)	O(k · N) per Lanczos step	Large N, only a few eigenstates needed
Numerov + shooting + Brent root-find	O(dx⁴) per integration	~ N · (# bisections) per eigenvalue	High-precision single eigenvalue
Split-operator (Strang) + FFT	O(dt²)	O(N log N) per step	Time evolution; preserves unitarity
Spectral (Hermite / Chebyshev)	Exponential in N	O(N²) per eigensolve	Smooth potentials, very high accuracy

For most non-pathological 1D bound-state problems, the recipe is: start with finite differences. Once it works, refine N until the answer stops moving. If you need higher accuracy or larger N, upgrade to a higher-order stencil or a spectral basis. If you need time evolution, reach for split-operator.

Worked Example: Particle in a Box by Hand

Let's walk through one full finite-difference calculation by hand, using the simplest non-trivial potential: the infinite square well of width $L = \pi$ , with V = 0 inside and ψ = 0 at both walls. Analytic energies: $E_n = n^2 \pi^2 / (2 L^2) = n^2/2$ for $L = \pi$ , i.e. 0.5, 2.0, 4.5, 8.0, ….

📐 Click to expand — hand calculation with N = 4

Use four interior grid points $x_1, x_2, x_3, x_4$ on $[0, \pi]$ . The spacing is $dx = \pi/(N+1) = \pi/5 \approx 0.6283$ , so $1/dx^2 \approx 2.533$ and $-1/(2\,dx^2) \approx -1.267$ .

The Hamiltonian (V = 0 inside) is the 4×4 matrix:

H ≈ [  2.533  -1.267   0       0     ]
    [ -1.267   2.533  -1.267   0     ]
    [  0      -1.267   2.533  -1.267 ]
    [  0       0      -1.267   2.533 ]

For a tridiagonal Toeplitz matrix with diagonal $d$ and off-diagonal $e$ there is a closed-form spectrum:

\lambda_k = d + 2 e \cos\!\Bigl(\tfrac{k\pi}{N+1}\Bigr),\quad k = 1, \ldots, N.

With $d = 2.533$ , $e = -1.267$ , $N+1 = 5$ :

k	cos(kπ/5)	λ_k = 2.533 + 2(−1.267)·cos	analytic E = k²/2	error
1	0.8090	2.533 − 2.050 = 0.483	0.500	0.017
2	0.3090	2.533 − 0.783 = 1.750	2.000	0.250
3	−0.3090	2.533 + 0.783 = 3.316	4.500	1.184
4	−0.8090	2.533 + 2.050 = 4.583	8.000	3.417

With only 4 grid points the lowest eigenvalue is already within 3% of the truth (0.483 vs 0.500). The higher modes are awful because their wavelength is comparable to dx — a wave with three or four nodes simply cannot live faithfully on a 4-point grid. This is aliasing, exactly the same phenomenon as in digital signal processing.

Repeat with N = 40: dx = π/41 ≈ 0.0766, 1/dx² ≈ 170.5, −1/(2 dx²) ≈ −85.2. Plug into the same closed-form spectrum with k = 1, 2, 3, 4 and you get:

k	λ_k (N = 40)	analytic	error
1	0.49975	0.50000	2.5 × 10⁻⁴
2	1.99602	2.00000	4.0 × 10⁻³
3	4.47980	4.50000	2.0 × 10⁻²
4	7.93722	8.00000	6.3 × 10⁻²

From N = 4 to N = 40 the error on $E_1$ dropped from 1.7 × 10⁻² to 2.5 × 10⁻⁴ — a factor of ≈ 68, very close to the (N₂/N₁)² = 100 you would predict from the $O(dx^2)$ error scaling (the discrepancy is because dx scales as 1/(N+1), not 1/N). The convergence is real, predictable, and exactly what the interactive viz above lets you watch live.

Python: Tiny Finite-Difference Solver

Here is a clean, complete NumPy implementation of the recipe. It is about 30 lines of code and reproduces every plot in §§29.4–29.5 to four-digit precision. Read it once, then read the per-line annotations for every subtlety.

🐍python

Explanation(12)

Code(75)

1What this script does

We are going to turn the Schrödinger differential equation into a finite-dimensional matrix and let numpy find its eigenvalues. That single move — replace an operator with a matrix on a grid — is the workhorse of every numerical PDE method in physics.

13build_hamiltonian: the grid

We sample x at N evenly spaced interior points. By choosing dx = (xmax − xmin)/(N+1) and starting at xmin + dx, the two outer ghost points x₀ and x_{N+1} sit exactly on the boundary where we want ψ = 0. The wavefunction at the boundary is not a free variable — it is forced to be zero — so we leave it out of the matrix entirely.

EXECUTION STATE

x_min, x_max = domain endpoints (Dirichlet ψ = 0 there)

N = number of free grid points

dx = (x_max − x_min) / (N + 1)

17Kinetic energy as a tridiagonal matrix

The 3-point central finite-difference approximation of ψ″(x) is (ψ_{i+1} − 2 ψ_i + ψ_{i−1}) / dx². Multiplying by −½ gives the discrete kinetic-energy operator. On the matrix that means: 1/dx² on the main diagonal, −1/(2 dx²) on both first off-diagonals, zeros elsewhere. T is symmetric and banded — perfect for fast eigensolvers.

LOOP TRACE · 2 iterations

main diag

T_ii = + 1/dx²

neighbours

T_{i, i±1} = − 1/(2 dx²)

24Potential energy is diagonal

In the position basis, V(x) acts by multiplication: (V ψ)(x_i) = V(x_i) ψ(x_i). The matrix representation is just a diagonal matrix whose entries are V evaluated at each grid point. There is no off-diagonal coupling from V.

27Assemble H = T + V

Two symmetric matrices added together stay symmetric. The result is still tridiagonal: only the diagonal got modified. We return x, dx (we need them later for normalisation and plotting), and H.

30solve: diagonalise H

np.linalg.eigh is for symmetric/Hermitian matrices. It is faster and more numerically stable than the general np.linalg.eig because it can use Householder tridiagonalisation followed by a QR/divide-and-conquer step. Eigenvalues come back already sorted ascending, which is exactly the order we want for bound states.

34Slice the lowest k_want states

We only care about the ground state and a handful of excited states. eigvecs has eigenvectors as columns, so psis = eigvecs[:, :k_want].T gives rows. Each row is a discrete sample of ψ_n(x) at the N grid points.

38Renormalise discrete eigenvectors

numpy returns eigenvectors with Euclidean norm 1, i.e. Σ_i |ψ_i|² = 1. But we want ∫ |ψ(x)|² dx ≈ Σ_i |ψ_i|² dx = 1. So we divide each ψ by √(Σ|ψ|² dx). This rescaling is critical when you compare to analytic formulas — otherwise the maximum amplitude depends on N.

47Harmonic oscillator potential

V(x) = ½ x² with ℏ = m = ω = 1. Exact spectrum E_n = n + ½ for n = 0, 1, 2, …. This is the gold standard test: every numerical method on Earth is first tried on the harmonic oscillator because the answers are known to a thousand digits.

50Run the solver

We sweep N = 200 grid points over [−6, 6]. The eigenfunctions decay like e^{−x²/2}, so by x = ±6 they are well below 10⁻⁸ — the Dirichlet boundary at ±6 is harmless. The whole eigensolve takes a few milliseconds.

EXECUTION STATE

x = shape (200,), spacing dx ≈ 0.0299

H = shape (200, 200), tridiagonal symmetric

E = lowest 4 eigenvalues

54Compare numerical and exact

Loop n = 0..3 and print the FD energy, the analytic energy, and the absolute error. With N = 200 you should see errors around 10⁻⁴ for E₀ and 10⁻³ for E₃. The error grows for higher states because their wavefunctions have more oscillations per dx and the second-derivative formula has trouble keeping up.

LOOP TRACE · 4 iterations

n=0

E_num = 0.49994

E_exact = 0.50000

error = ≈ 6 × 10⁻⁵

n=1

E_num = 1.49975

E_exact = 1.50000

error = ≈ 2.5 × 10⁻⁴

n=2

E_num = 2.49938

E_exact = 2.50000

error = ≈ 6 × 10⁻⁴

n=3

E_num = 3.49883

E_exact = 3.50000

error = ≈ 1.2 × 10⁻³

60Plot stacked eigenfunctions

A clean way to visualise eigenstates of a 1D Schrödinger problem is to draw V(x) once, then draw each ψ_n shifted vertically to sit at its energy level E_n. The horizontal axis is x, the vertical axis is energy, and the wavefunction is an artistic device showing the wave shape on top of its level. This is the same picture you saw in §29.4 and §29.5 — and it is what the interactive viz above is showing you live.

63 lines without explanation

1import numpy as np
2import matplotlib.pyplot as plt
3
4# We work in dimensionless units where hbar = m = 1.
5# The time-independent Schrodinger equation on an interval [x_min, x_max] is:
6#
7#     - (1/2) psi''(x)  +  V(x) psi(x)  =  E psi(x)
8#
9# We discretise psi at N interior grid points x_1, ..., x_N with spacing
10#   dx = (x_max - x_min) / (N + 1).
11# Dirichlet boundary conditions psi(x_0) = psi(x_{N+1}) = 0 are enforced
12# implicitly by leaving those points out of the unknown vector.
13
14
15def build_hamiltonian(x_min, x_max, N, V_func):
16    """Return the (N, N) finite-difference Hamiltonian matrix and the grid."""
17    dx = (x_max - x_min) / (N + 1)
18    x  = np.linspace(x_min + dx, x_max - dx, N)        # interior grid
19
20    # Kinetic energy: T = -(1/2) d^2/dx^2 with a 3-point stencil
21    #   T_{ii}     =  1 / dx^2
22    #   T_{i,i+1}  = -1 / (2 dx^2)
23    main_diag = np.full(N,  1.0 / dx**2)
24    off_diag  = np.full(N - 1, -0.5 / dx**2)
25    T = (np.diag(main_diag)
26         + np.diag(off_diag, k=1)
27         + np.diag(off_diag, k=-1))
28
29    # Potential energy is diagonal in position basis
30    V = np.diag(V_func(x))
31
32    H = T + V
33    return x, dx, H
34
35
36def solve(H, dx, k_want=6):
37    """Diagonalise H and return the k_want lowest eigenpairs."""
38    # H is symmetric => use eigh (returns sorted ascending).
39    eigvals, eigvecs = np.linalg.eigh(H)
40
41    energies = eigvals[:k_want]
42    psis     = eigvecs[:, :k_want].T            # rows are eigenvectors
43
44    # Renormalise so that integral |psi|^2 dx = 1.
45    for i in range(k_want):
46        norm = np.sqrt(np.sum(psis[i]**2) * dx)
47        psis[i] /= norm
48
49    return energies, psis
50
51
52# ----- Harmonic oscillator: V(x) = (1/2) x^2 -----------------------
53V_harmonic = lambda x: 0.5 * x**2
54
55x, dx, H = build_hamiltonian(-6.0, 6.0, N=200, V_func=V_harmonic)
56E, psi   = solve(H, dx, k_want=4)
57
58# Exact answers are E_n = n + 1/2 for n = 0, 1, 2, ...
59print("level | numeric | exact   | error")
60for n in range(4):
61    print(f"  {n}   | {E[n]:7.5f} | {n + 0.5:6.5f} | {abs(E[n] - (n + 0.5)):.2e}")
62
63# Plot the first four eigenfunctions stacked on their energy levels.
64fig, ax = plt.subplots(figsize=(8, 5))
65ax.plot(x, V_harmonic(x), color="gray", lw=1, label="V(x)")
66for n in range(4):
67    ax.plot(x, psi[n] * 0.6 + E[n], lw=1.8, label=f"n = {n}")
68    ax.axhline(E[n], color="black", lw=0.4, alpha=0.4)
69ax.set_xlabel("x")
70ax.set_ylabel("E,  psi (stacked at E_n)")
71ax.set_title("Harmonic oscillator — finite-difference eigenstates")
72ax.legend()
73plt.tight_layout()
74plt.savefig("harmonic_fd.png", dpi=120)
75print("Saved harmonic_fd.png")

PyTorch: Batched Eigensolves on the GPU

Real physics problems — varying a parameter, sweeping a barrier height, training a neural surrogate — often need hundreds of Schrödinger solves. NumPy does them one at a time; PyTorch does them all in parallel on a GPU. The code structure is essentially the same. What changes is the leading batch dimension and the device.

🐍python

Explanation(10)

Code(71)

1Why PyTorch for a numerical PDE?

PyTorch is far more than a deep-learning library. Its tensors are dense CPU/GPU arrays, its linalg module is fully batched, and torch.linalg.eigh handles symmetric eigenproblems in a single call. For parameter sweeps over many potentials we get true parallelism on a GPU — and we can later feed the results into a neural network (e.g. as a physics-informed loss).

9Pick a device

Setting device = 'cuda' if available lets the entire pipeline run on the GPU without further code changes. A 1000×1000 symmetric eigensolve takes ~1s on CPU but ~50 ms on a modern GPU; the batched version scales further still.

14fd_hamiltonian_batch: shared T, per-batch V

Kinetic energy is identical for every member of the batch — it only depends on the grid. So we build T once and broadcast it to shape (B, N, N) via expand. Then we add each batch member's diagonal V on top. Crucially, expand does not copy memory; .clone() is what materialises a writable copy so we can mutate the diagonals.

EXECUTION STATE

B = batch size (e.g. 8 different wells)

N = grid size

T = (N, N) shared kinetic operator

H = (B, N, N) per-batch Hamiltonian

22Fancy-indexed diagonal update

H[:, idx, idx] is the standard PyTorch trick for accessing the main diagonal of every batch matrix at once. The shape of that view is (B, N), so adding V_batch (also (B, N)) is a single broadcast operation. No Python loop over the batch dimension.

26solve_batch: batched eigh

torch.linalg.eigh is batched: pass a (B, N, N) symmetric input and you get (B, N) eigenvalues + (B, N, N) eigenvectors back, each batch element diagonalised independently in parallel. We then slice the lowest k_want and transpose so the per-state eigenvector becomes the last axis.

32Per-batch, per-state renormalisation

Each ψ_{b, k} has its own √(Σ ψ² dx) normalisation constant. We compute the (B, k) tensor of norms, then divide with broadcasting along the grid axis. After this line every state in every batch satisfies ∫|ψ|² dx = 1 to floating-point precision.

38Parameter sweep: 8 double wells

We build V_b(x) = 0.1 (x² − a_b²)² for a_b ∈ [0.8, 2.2]. Small a means two wells close together — the barrier in the middle is short and the ground state spreads across both wells (think a single fused well). Large a means the wells are far apart — the ground state localises into one minimum at a time and the lowest two eigenvalues become almost degenerate (the famous double-well tunnel splitting).

LOOP TRACE · 3 iterations

a = 0.80

E₀ = 0.286 (well-mixed)

E₁−E₀ = ≈ 0.60 (large)

a = 1.50

E₀ = 0.122

E₁−E₀ = ≈ 0.16

a = 2.20

E₀ = 0.025 (deeply localised)

E₁−E₀ = ≈ 4 × 10⁻³ (near-degenerate)

43Diagonalise the whole batch at once

On a GPU this single call does 8 independent eigensolves in parallel. The wall-clock time is roughly the same as one solve. Scaling this to 1000+ potentials is how you train surrogate ML models on quantum chemistry data without writing a CUDA kernel by hand.

45Print eigenvalues across the family

Watch the splitting E₁ − E₀ shrink as a grows. For a ≳ 2 the gap drops below 10⁻² and you have what chemists call a tunnel doublet: two states whose symmetric/antisymmetric combination corresponds to 'electron in left well' or 'electron in right well'.

51Plot |ψ₀|² for each parameter

A 2×4 grid of panels shows the ground state's probability density change continuously from a single peak (small a) into two well-separated peaks (large a). The plot is the visual story of how tunnel coupling decays exponentially with barrier width — a single equation, eight pictures, and you have built intuition that no closed-form formula gives.

61 lines without explanation

1import torch
2import matplotlib.pyplot as plt
3
4# Same physics, but now:
5#  - we use torch instead of numpy
6#  - we batch many potentials at once and diagonalise them in parallel
7#  - we move to GPU when one is available (10-100x speedup for N ~ 1000)
8
9
10device = "cuda" if torch.cuda.is_available() else "cpu"
11print("device:", device)
12
13
14def fd_hamiltonian_batch(x, V_batch):
15    """
16    Build a (B, N, N) batch of finite-difference Hamiltonians.
17    x:        (N,)           the shared grid
18    V_batch:  (B, N)          one potential per batch element
19    """
20    B, N = V_batch.shape
21    dx = (x[1] - x[0])
22    main = torch.full((N,), 1.0 / dx**2, device=device, dtype=torch.float64)
23    off  = torch.full((N - 1,), -0.5 / dx**2, device=device, dtype=torch.float64)
24    T = torch.diag(main) + torch.diag(off, 1) + torch.diag(off, -1)   # (N, N)
25
26    H = T.unsqueeze(0).expand(B, N, N).clone()                        # (B, N, N)
27    idx = torch.arange(N, device=device)
28    H[:, idx, idx] += V_batch                                         # add diagonal V
29    return H, dx
30
31
32def solve_batch(H, dx, k_want=4):
33    """Symmetric eigendecomposition of a batch of (N, N) matrices."""
34    # torch.linalg.eigh batches over leading dimensions.
35    eigvals, eigvecs = torch.linalg.eigh(H)
36    E   = eigvals[:, :k_want]                  # (B, k)
37    psi = eigvecs[:, :, :k_want].transpose(-1, -2)   # (B, k, N)
38    # Renormalise per-batch, per-state
39    norm = torch.sqrt((psi ** 2).sum(dim=-1) * dx)   # (B, k)
40    psi = psi / norm.unsqueeze(-1)
41    return E, psi
42
43
44# ----- Build a family of double-well potentials parameterised by a -----
45N = 400
46x = torch.linspace(-3.5, 3.5, N, device=device, dtype=torch.float64)
47
48a_vals = torch.linspace(0.8, 2.2, 8, device=device, dtype=torch.float64)  # 8 wells
49V_batch = 0.1 * (x.unsqueeze(0)**2 - a_vals.unsqueeze(1)**2)**2           # (8, N)
50
51H, dx = fd_hamiltonian_batch(x, V_batch)
52E, psi = solve_batch(H, dx, k_want=4)
53
54print("Lowest 4 eigenvalues for each well separation a:")
55for i, a in enumerate(a_vals):
56    print(f"  a = {a.item():.2f}:  E = "
57          + ", ".join(f"{e.item():.4f}" for e in E[i]))
58
59# Plot the ground-state probability density as a function of well separation:
60# tightly coupled wells (small a) -> single bump in the middle.
61# well-separated wells (large a) -> two bumps, one in each well.
62fig, axes = plt.subplots(2, 4, figsize=(14, 6), sharex=True)
63for i, ax in enumerate(axes.flat):
64    a = a_vals[i].item()
65    ax.plot(x.cpu(), V_batch[i].cpu(), color="gray", lw=1)
66    ax.fill_between(x.cpu(), 0, (psi[i, 0]**2).cpu(), alpha=0.5, color="tab:green")
67    ax.set_title(f"a = {a:.2f},  E0 = {E[i, 0].item():.3f}")
68    ax.set_ylim(-1, 6)
69plt.tight_layout()
70plt.savefig("double_well_batch.png", dpi=120)
71print("Saved double_well_batch.png")

From physics to ML

The output of the batched solver — a tensor of shape (B, N) for the ground-state densities — is exactly the kind of dataset you would feed into a neural network that learns the map from potential to ground-state density. That is the core idea behind machine-learning density-functional theory: you let a small CNN do in a millisecond what the eigensolver does in a minute, after training on a few thousand pairs.

Where This Is Used

Field	What is solved numerically	Why exact solutions fail
Quantum chemistry	Electronic structure of molecules (Hartree–Fock, DFT, post-HF)	Many-electron Coulomb correlation has no closed form
Solid-state physics	Band structure, defect states, quantum dots	Periodic potentials are smooth but not analytically tractable
Cold-atom experiments	Bose-Einstein condensates in shaped optical traps	Trap potentials are designed in software, not given by formulas
Nuclear physics	Bound nucleon orbitals in nuclei	Nuclear potential is empirical; no analytic eigenfunctions
Photonics & quantum optics	Modes of curved waveguides and microresonators	Helmholtz / paraxial equations are Schrödinger-like with arbitrary V
Machine learning	Generating training data for neural force fields and surrogates	Need thousands of solves at different geometries — only numerics scale

Every commercial quantum-chemistry package — Gaussian, ORCA, NWChem, Psi4, VASP — at its core does some flavour of the recipe in this section: discretise the Schrödinger operator, build a (huge, sparse, often non-Hermitian) matrix, and diagonalise. The methods get fancier (Gaussian basis sets, plane waves, multi-reference perturbation theory), but the spirit is the same: turn an operator into a matrix and let linear algebra do the rest.

Summary

The Schrödinger equation is solvable in closed form only in a handful of exceptional cases. Every interesting real-world problem requires numerics.
The universal recipe is: sample ψ on a grid → approximate the second derivative by finite differences → build a sparse symmetric matrix H → diagonalise. Eigenvalues are energies, eigenvectors are wavefunctions.
The 3-point central stencil is $O(dx^2)$ accurate. Doubling N divides the error by ≈ 4.
The shooting method turns the eigenproblem into a 1D root-finding problem on the boundary residue $\psi(x_\text{max}, E)$ . Pair it with Numerov for $O(dx^4)$ accuracy at the same cost as the basic FD method.
For time evolution, use the split-operator method: alternate phase rotations in position and momentum space, glued together by FFTs. The scheme is unitary and conserves probability exactly.
PyTorch's batched eigensolver lets you do hundreds of Schrödinger problems in parallel on a GPU — the bridge from computational physics to quantum machine learning.

The deepest lesson of this section is not any specific algorithm. It is the realisation that continuous operators are matrices in disguise. Once you internalise that, you stop being afraid of differential equations: every PDE in physics, from Maxwell's to Navier–Stokes, surrenders to the same recipe.