neural-fluid-engine

Can a compact U-Net learn the next step of a turbulent flow?

A two-part experiment: a pseudo-spectral solver generates 2D Kolmogorov-flow turbulence, and a PyTorch U-Net learns to predict the next vorticity field from the current one — optionally conditioned on the physical parameters themselves.

The Kolmogorov flow leaving its laminar state: velocity and vorticity early on, through transition, into sustained chaos.

The setup

Kolmogorov flow — 2D incompressible Navier–Stokes on a periodic $[0, 2\pi]^2$ domain, driven by a sinusoidal body force $A \sin(n y)$ — is one of the simplest systems that transitions to genuine 2D turbulence. That makes it an ideal benchmark for a question at the heart of ML-for-physics: how much of the solver can a neural network replace?

The pipeline here goes end to end:

1. Simulate — solver/kolmogorov_2d.py integrates the flow with the Dedalus 3 spectral framework (default 128², adaptive timestep), writing HDF5 snapshots, PNG frames, and an optional MP4.
2. Package — data/ turns raw runs into an ML dataset: a JSONL index of runs, train/val/test splits, and normalization statistics (make_dataset.py, splits.py, stats.py, with on-demand HDF5 loading in dataset.py).
3. Learn — models/unet2d.py is a compact U-Net (GroupNorm, optional residual head $\hat\omega_{t+1} = \omega_t + f(\omega_t)$). With --use_params the fields $(\nu, A)$ are broadcast as constant input channels, so one network can learn across physical regimes.
4. Train & evaluate — train/train_unet_nextstep.py (AdamW, mixed precision) and eval/eval_onestep.py (one-step MSE/RMSE on held-out trajectories).

Running it

# 1. Generate a turbulent run (writes HDF5 + frames under data/raw/)
python solver/kolmogorov_2d.py --outdir data/raw/run0 --T 50 --make_video

# 2. Build the dataset index, splits, and normalization stats
python scripts/prepare_ml_dataset.py --tag kolmogorov_128 \
    --index data/index_kolmogorov_128.jsonl --out data/

# 3. Train the next-step U-Net
python train/train_unet_nextstep.py --tag kolmogorov_128 \
    --index_path data/index_kolmogorov_128.jsonl \
    --splits_path data/splits_kolmogorov_128.json \
    --stats_path data/stats_kolmogorov_128.json \
    --batch_size 32 --epochs 50 --amp

# 4. Evaluate one-step prediction on the test split
python eval/eval_onestep.py --checkpoint checkpoints/unet_kolmogorov_128/best.pt \
    --index_path data/index_kolmogorov_128.jsonl \
    --stats_path data/stats_kolmogorov_128.json

Full pipeline options (parameter conditioning, prediction horizon, strides, debug runs) are documented in README_ml.md.

Generated data (HDF5 snapshots, frames, videos, checkpoints) stays out of the repository — everything above regenerates it from scratch.