GPU Lattice-Boltzmann Solver: von Kármán Vortex Street

# Large-scale fluid motion from Boltzmann's kinetic equation — GPU Lattice-Boltzmann The **lattice-Boltzmann method (LBM)** is a discrete velocity solver for the Boltzmann kinetic equation. In the hydrodynamic (large-scale, low-Knudsen) limit, its moments provably converge to the incompressible Navier–Stokes equations — the very *kinetic → fluid* passage recognized by the **2026 Fields Medal work of Yu Deng (with Zaher Hani and Xiao Ma) on Hilbert's sixth problem**. Here we watch that emergence directly: kinetic particle distributions on a **D2Q9** lattice, evolved by a single-relaxation **BGK collision**, spontaneously produce macroscopic turbulence — a **von Kármán vortex street** behind a cylinder. **Everything is pure `torch` on the GPU** — no CFD library, no extra installs: - **Streaming** = a tensor `roll` along lattice directions. - **Collision** = elementwise BGK relaxation toward local equilibrium. **Deliverables:** a smooth vorticity animation (blue-white-red, inline + downloadable GIF) and a physics check — the measured **Strouhal number** from an FFT of the wake velocity vs. the textbook `St ≈ 0.2`.

import torch, time, math
import numpy as np
import matplotlib.pyplot as plt

# --- Hard requirement: real GPU, never fall back to CPU ---
assert torch.cuda.is_available(), "CUDA GPU not available — refusing to run on CPU."
device = torch.device("cuda")
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
print("CUDA available :", torch.cuda.is_available())
print("Device         :", torch.cuda.get_device_name(0))
print("torch version  :", torch.__version__)

CUDA available : True
Device         : NVIDIA A100-SXM4-40GB
torch version  : 2.7.1+cu128

# ============================================================
# D2Q9 lattice constants + simulation parameters (all on GPU)
# Array layout: f has shape (9, nx, ny); axis0 = lattice direction,
# axis1 = streamwise x, axis2 = transverse y.
# ============================================================

# 9 discrete velocities (Latt ordering) and BGK weights
v = torch.tensor([[ 1, 1],[ 1, 0],[ 1,-1],
                  [ 0, 1],[ 0, 0],[ 0,-1],
                  [-1, 1],[-1, 0],[-1,-1]], device=device, dtype=torch.long)
w = torch.tensor([1/36,1/9,1/36, 1/9,4/9,1/9, 1/36,1/9,1/36], device=device, dtype=torch.float32)
opp = torch.tensor([8,7,6,5,4,3,2,1,0], device=device)   # opposite direction indices

col_R = torch.tensor([0,1,2], device=device)   # vx = +1 (right-moving)
col_0 = torch.tensor([3,4,5], device=device)   # vx =  0
col_L = torch.tensor([6,7,8], device=device)   # vx = -1 (left-moving)

# ---- Domain / physics ----
…

nu = 0.01200   tau = 0.5360   omega = 1.8657  (stable: 0<omega<2)
lattice = 720 x 200   cylinder D = 40   Re = 200

from tqdm.auto import tqdm

maxIter      = 45000        # total LBM steps
record_start = 8000         # begin capturing animation frames after wake develops
record_every = 175          # frame stride  -> ~210 frames
n_frames_est = (maxIter - record_start)//record_every + 1

# probe point in the wake (centreline, 3 diameters downstream) for Strouhal FFT
px, py = cx0 + 3*D, cy0
probe_uy = torch.zeros(maxIter, device=device)   # GPU-resident, no per-step sync

frames = []          # list of (ny, nx) vorticity arrays (CPU numpy) for the animation
frame_steps = []

def vorticity(u):
    ux, uy = u[0], u[1]
    duy_dx = (torch.roll(uy,-1,0) - torch.roll(uy,1,0)) * 0.5
    dux_dy = (torch.roll(ux,-1,1) - torch.roll(ux,1,1)) * 0.5
…

Done: 45000 steps in 161.8 s (278 steps/s)
Captured 212 vorticity frames (every 175 steps).

from matplotlib import animation
from matplotlib.animation import PillowWriter
from IPython.display import Image, display
import matplotlib as mpl

frames_arr = np.stack(frames)                     # (n_frames, ny, nx)
# robust symmetric color scale for the diverging map
vmax = np.nanpercentile(np.abs(frames_arr), 99.0)
vmin = -vmax
print(f"{len(frames_arr)} frames, shape {frames_arr.shape[1:]},  color scale +/-{vmax:.4f}")

cmap = mpl.colormaps['bwr'].copy()                # blue -> white -> red diverging
cmap.set_bad('#dddddd')                           # cylinder / walls (NaN) shown grey

fig, ax = plt.subplots(figsize=(12, 3.6), dpi=100)
im = ax.imshow(frames_arr[0], origin='lower', cmap=cmap, vmin=vmin, vmax=vmax,
               interpolation='bilinear', aspect='equal')
circ = plt.Circle((cx0, cy0), r, fill=False, color='k', lw=1.2)   # cylinder outline
…

212 frames, shape (200, 720),  color scale +/-0.0186
Saved /home/user/vortex_street.gif  (4.4 MB, 212 frames @ 24 fps)

# ---- Strouhal number from the wake transverse-velocity signal ----
# Use the fully-developed shedding window (after the initial transient).
transient = 12000
sig = probe_uy_cpu[transient:].astype(np.float64)
sig = sig - sig.mean()                     # remove DC
n   = sig.size
dt  = 1.0                                   # 1 lattice time unit per step

freqs = np.fft.rfftfreq(n, d=dt)            # cycles per step
amp   = np.abs(np.fft.rfft(sig))
kpeak = 1 + int(np.argmax(amp[1:]))         # skip DC bin

# sub-bin parabolic interpolation around the peak for a finer frequency
a, b, c = amp[kpeak-1], amp[kpeak], amp[kpeak+1]
delta   = 0.5*(a - c)/(a - 2*b + c)         # in [-0.5, 0.5]
f_shed  = (kpeak + delta)/(n*dt)            # refined shedding frequency [1/step]

St_measured = f_shed * D / U
…

probe @ (x=300, y=100), samples analysed = 33000
shedding frequency f = 3.734653e-04  cycles/step   (period ~ 2,678 steps)
====================================================
  measured Strouhal  St = f*D/U = 0.2490
  textbook  (Re~200)    St ~ 0.2000
  relative difference   = 24.5 %
====================================================

print(f"RUN SUMMARY  |  lattice {nx}x{ny} (D2Q9)  |  Re = {Re:.0f}  |  "
      f"{maxIter:,} steps  |  wall-clock {wall_clock:.1f} s ({wall_clock/60:.2f} min) on {torch.cuda.get_device_name(0)}  |  "
      f"measured St = {St_measured:.3f} vs textbook ~0.20")