// physics · numerical-physics·

Jan. 2026·10 min read

Solving efficiently the 2D Schrödinger equation using a split-step FFT method.

Split-Step Fourier for the Schrödinger Equation

Introduction

In the life of a young physicist, quantum mechanics, and especially the Schrödinger equation, often triggers a fascinating, almost masochistic curiosity during our teenage years, one from which we never truly recover...

During my studies, I had the opportunity to implement the Crank–Nicolson algorithm in 1D, which aim to solve the Schrodinger equation with finite differences and semi imlpicit method (ref). 1D is cool, but not cool enough, unfortunately I was limited by my computational resources when trying to extend it to 2D.

More recently, however, I discovered the Split-Step Fourier Method, which saves a lot of resources and is surprisingly easy to implement. In just a few lines of code, we can visualize the famous double-slit experiment, and the younger Joseph (me), who once struggled to visualize QM, would probably be very proud !

Problem

If you are still interested by this article, I assume you already are familiar with quantum mechanics and the Schrödinger equation we aim to solve:

i \hbar \frac{ \partial }{\partial t} \, |\psi(t)\rangle = H(t) |\psi(t)\rangle

Where, as usual, the Hamiltonian is the sum of kinetic and potential energy:

\begin{aligned} H &= T + V\\ T &= \frac{P^2}{2m} = - \frac{\hbar^2}{2m} \Delta\\ P &= -i \hbar \nabla \end{aligned}

In the continuum, the kinetic energy is a differential operator. A naive approach would be to solve this PDE using finite differences, but this is computationally expensive in 2D and requires careful treatment to preserve unitarity (e.g., Crank–Nicolson).

The Split-Step Fourier method provides an efficient alternative by working with the time evolution operator, which comes from the exact solution:

\begin{aligned} |\psi(t)\rangle &= U(t, t') |\psi(t')\rangle\\ U(t, t') &= \mathcal{T} \exp\left(-\frac{i}{\hbar} \int_{t'}^t H(\tau)\, d\tau\right) \end{aligned}

where $\mathcal{T}$ denotes time-ordering. For small time steps $\Delta t$ , or when the potential varies slowly, we can approximate (even exact for non-varying potential):

U(t+\Delta t, t) \approx \exp\left(-\frac{i \Delta t}{\hbar} H(t)\right)

The Key Idea: Splitting the Exponential

Directly exponentiating $H = T + V$ is challenging because $T$ and $V$ are not simultaneously diagonal. However:

$V$ is diagonal in position space.
$T$ is diagonal in momentum space.

Mathematically:

\begin{aligned} \langle x | f(X) | \psi \rangle &= f(x) \psi(x)\\ \langle p | f(P) | \psi \rangle &= f(p) \tilde{\psi}(p) \end{aligned}

where $\tilde{\psi}$ is the Fourier transform of $\psi$ .
This means we can efficiently compute exponentials in the basis where the operator is diagonal. Computationally, the FFT will allows us to switch between these bases at low cost, this is the heart of the high efficiency of our algorithm.

Exponential Magic: Trotter and Strang Splitting

Using the Baker–Campbell–Hausdorff formula, we can approximate $e^{\epsilon(A+B)}$ by splitting:

First-order (Lie–Trotter)

e^{\epsilon(A+B)} = e^{\epsilon A} e^{\epsilon B} + \mathcal{O}(\epsilon^2)

Second-order symmetric (Strang)

e^{\epsilon(A+B)} = e^{\epsilon A/2} e^{\epsilon B} e^{\epsilon A/2} + \mathcal{O}(\epsilon^3)

Higher-order generalizations exist (Suzuki–Trotter), e.g.

Applying Strang splitting to the time evolution operator:

U(t+\Delta t, t) \approx e^{-i \Delta t V / (2\hbar)} \, e^{-i \Delta t T / \hbar} \, e^{-i \Delta t V / (2\hbar)}

Because $V$ and $T$ are diagonal in respectively position and momentum bases:

\psi(t+\Delta t) = e^{-i \Delta t V/2\hbar} * \mathcal{F}^{-1} \Big( e^{-i \Delta t T/\hbar} * \mathcal{F} ( e^{-i \Delta t V/2\hbar} * \psi (t) ) \Big)

where $*$ denotes pointwise multiplication and $\mathcal{F}$ / $\mathcal{F}^{-1}$ are spatial Fourier transform / inverse.

Discretization

To implement on a computer:

Discretize space: a lattice with $N_x \times N_y$ points and spacing $\Delta x$ .
Discretize time: $N_t$ steps of size $\Delta t$ .
Represent the wavefunction as a finite 2D complex array and compute the FFT efficiently.

The FFT reduces the cost of basis transformation from $O(N^2)$ to $O(N \log N)$ , which is a huge speedup in 2D or 3D simulations.

Summary of the Algorithm

Discretize the spatial lattice $(x, y)$ and time.
Apply half-step potential phase shift in position space (diagonal).
Free propagation (kinetic step):
- FFT to momentum space
- Apply full-step kinetic phase shift (diagonal)
- IFFT back to position space
Reapply half-step potential phase shift in position space.

Python Implementation

Initialize position lattice and potential:

x = np.arange(0, dx*Nx, dx)
y = np.arange(0, dy*Ny, dx)
X, Y = np.meshgrid(x, y, indexing='ij')
V = potential(X, Y)

Initialize momentum lattice and kinetic energy:

kx = 2*np.pi*np.fft.fftfreq(Nx,dx)
ky = 2*np.pi*np.fft.fftfreq(Ny,dx)
Kx,Ky = np.meshgrid(kx, ky, indexing='ij')
K2 = Kx**2+Ky**2
T = hbar**2*K2/(2*m)

Evaluate evolution operators:

U_V = np.exp(-1j * V * dt/(2*hbar))
U_T = np.exp(-1j * K * dt/hbar)

Perform one time step:

psi = U_V * np.fft.ifft2( U_T * np.fft.fft2( U_V * psi) )

Easy to generalise to 3D with the numpy function:

numpy.fft.fftn

Visualization

The phase is represented using a $2\pi$ -periodic colormap, while the density represents a scaled version of the probability distribution to enhance visibility. The potential can be represented in the background.

Resolution, Stability, Error and Complexity

Spatial Resolution:

To resolve the wavefunction properly:

\Delta x < \frac{\sigma}{4}, \quad \Delta x < \frac{\pi \hbar}{p_{\max}}

where $\sigma$ is the typical width of the wavefunction and $p_{\max}$ is the maximum momentum.

Time Step / Stability:

To ensure stable evolution:

\Delta t < \frac{m \Delta x^2}{\hbar}, \quad \Delta t \ll \frac{\hbar}{V_{\max}}

The first condition controls phase rotation due to the kinetic term, while the second controls phase rotation due to the potential term.

Error:

Local error per step: $\mathcal{O}(\Delta t^3)$
Global error over total time $T=N_t\Delta t$ :

\mathcal{O}(N_t \Delta t^3) = \mathcal{O}(T \Delta t^2)

Computational Complexity:

For a lattice of size $N_x \times N_y$ and $N_t$ time steps:

\mathcal{O}(N_t \, N_x N_y \, \log(N_x N_y))

Results

To illustrate all this method we build two interesting potentials : the double-slit experiment and the proton diffusion. In the first we simply create walls using a high potential energy, and in the second case we use a $1/(|r|+\epsilon)$ potential. The initial wave packet is a simple Gaussian wave packet with a given momentum.

Double-Slit Simulation.

Proton Diffusion Simulation.

Outlooks

Go to 3D, but require a good visualisation framework. Switch to Blender ? Can be cool for a video.
Some fancy potential used to illustrate other phenomenons
2 electrons ? Does it work ? How to handle indistiguishability ?
As always, OPTIMIZATIONS !
...