AeroJAX

A real-time, JAX-native CFD framework for interactive flow research, control, and inverse design.

What makes AeroJAX different?

Most CFD software is batch‑oriented: simulations are configured, run to completion, and analysed afterwards. Changing anything means starting over.

AeroJAX is an interactive CFD framework in which solver parameters, boundary conditions, and immersed geometries can be modified during runtime - without restarting the simulation.

• Draw an obstacle with your mouse - the solver injects it on the next timestep using Brinkman penalisation.
• Drag a NACA airfoil across the domain with a slider and watch the wake evolve in real time.
• Pressure solvers (FFT, Conjugate Gradient, Multigrid, or Neural Operator) can be swapped mid-simulation without restarting.
• Toggle LES models (Smagorinsky / dynamic Smagorinsky) on the fly.
• Inject dye at any point and watch it advect.
• Record video or export frames with one click.
• Inspect and replay solver internals step-by-step using the Trace Viewer (no black-box timesteps).

AeroJAX is built on JAX, making each solver step end-to-end differentiable. You can run gradient‑based inverse design (optimise an airfoil shape to minimise drag) without writing a separate adjoint solver.

The framework is CPU‑optimised for real-time performance. Typical performance reaches ~297 FPS at 512×96 on a laptop CPU. No GPU is required for smaller grids.

Quick start

git clone https://github.com/arriemeijer-creator/AeroJAX
cd AeroJAX
pip install -r requirements.txt
python main.py

Requirements

Python 3.11+

Optional NVIDIA Hardware Acceleration

python -m pip install --upgrade "jax[cuda12]" #For CUDA 12
python -m pip install --upgrade "jax[cuda13]" #For CUDA 13

The GUI

First run tips

• Start with 512×96 grid for real‑time performance.
• Try the default initialized von Kármán with NACA 0012 airfoil at 10° AoA.
• Enable "Adaptive dt" to see the PID controller in action.

What you can control

The UI is organised into collapsible panels. Everything below can be changed mid‑simulation unless noted.

Floating Quick-Control Bar

A semi-transparent overlay provides instant access to the most frequent interactions without needing to navigate the sidebar:

• Playback: Instant Start (▶), Pause (⏸), and Reset (↺).
• Layer Toggles: Hot-swap visibility for Velocity, Divergence, Vorticity, Pressure, and Dye layers.
• Diagnostics Toggles: One-click activation of Error Metrics, Airfoil Metrics, and Profiling Overlays.
• Live Colormaps: Change colormaps for V, ω, and P on the fly to highlight different physical features.
• Precision Dye Injection: Dual-slider (X, Y) control for real-time scalar tracer injection.

Simulation control (top bar)

• Start / Pause / Reset
• Inverse Design - experimental module for adjoint‑based shape optimisation
• Thermal - placeholder for future heat transfer (Boussinesq approximation buoyancy flow is already conceptually implemented).
• Theme toggle - light / dark mode

Grid & solver

Flow parameters

Solver parameters

Boundary conditions

Obstacles (all live)

• Type: Cylinder (single), NACA airfoil (4‑/5‑digit), Cow (arbitrary complex geometry demo), Three‑cylinder array.
• NACA: Choose from selected 4‑digit and 5‑digit series. Set chord length, angle of attack (AoA) with slider/spinbox.
• Cylinder: Radius (live preview as you input).
• Cylinder array: Diameter, spacing between centres.
• Position: X and Y sliders - drag the obstacle across the domain live.
• Draw custom obstacle: Click and draw a shape with your mouse in the PyGame drawing window. The SDF generator injects it immediately.

Time stepping

• dt (fixed) - apply live.
• Adaptive dt - PID controller based on divergence error with CFL monitoring (experimental).

Visualisation

Dye injection

• X / Y position (spinbox or sliders)
• Amount (0‑100%)
• Inject dye - adds a scalar tracer that advects with the flow.

Diagnostics & metrics

Wake-based signal tracking (real-time physics extraction)

AeroJAX now supports signal-based flow diagnostics as an alternative to explicit structure tracking.

Instead of identifying and tracking vortices, the solver extracts flow physics directly from field signals:

• Sample vertical velocity downstream of bluff bodies
• Detect wake center via zero-crossing in the velocity field
• Construct a time series of wake oscillation
• Apply FFT to extract dominant frequency → Strouhal number

This approach is significantly more stable than vortex identification and tracking, especially in transitional and noisy regimes.

It reflects a broader shift in AeroJAX:

• Not tracking structures explicitly
• Tracking the signals they generate

This enables robust, real-time extraction of flow characteristics without relying on fragile feature detection.

This approach is particularly well-suited for real-time and differentiable workflows.

Solver Trace Viewer (step-by-step numerical inspection)

AeroJAX now includes a fully decoupled solver trace viewer for inspecting the numerical pipeline at each timestep.

This is not a visualization layer. It reconstructs and exposes the solver itself.

What it enables

• Step-by-step traversal of simulation timesteps

• Full decomposition of the solver pipeline:

  1. Advection-Diffusion
  2. Divergence computation
  3. Pressure solve
  4. Velocity correction
  5. State update

• Inspection of intermediate reconstructed fields (u*, divergence, etc.)

• Subdomain-level numerical inspection (cell-by-cell values)

• Live metrics (CFL, Reynolds number, extrema)

• LaTeX-rendered equations for each step

Architecture

The viewer is strictly decoupled from the solver:

• Solver writes snapshots during runtime
• Trace viewer reads snapshots only
• No runtime coupling, no performance impact

This enables:

• Deterministic replay
• Deep debugging of numerical behaviour
• Post-run analysis without rerunning simulations

Why this matters

Most CFD tools treat the solver as a black box.

AeroJAX exposes the solver as an inspectable pipeline.

This aligns with the broader goal:

• Making CFD not just interactive
• But interrogatable at the numerical level

Neural Operator training (experimental)

• Generate dataset: Run simulation for N steps, save chosen fields (u, v, p, mask, divergence) to .npz.
• Select operator: Choose from any .py file in neural_operators/.
• Architecture: Linear, NonLinear, Advanced (FNO‑like).
• Train: Set epochs, learning rate, batch size, cancel anytime. Progress bar.
• Load trained model: Replace the pressure solver with a neural operator directly inside the simulation loop.

Export & recording

• Export Frame - save current visualisation as PNG.
• Record - toggle video recording (saves to disk).
• Save State - not yet implemented.

Numerical methods & Solver details

• Navier‑Stokes (incompressible) finite‑difference solver on collocated or MAC grid.
• Lattice Boltzmann (D2Q9 / D2Q7) with BGK or MRT collision (low-Mach regime).
• Advection: RK3, RK2, multiple specialised schemes.
• Pressure projection: FFT (periodic), Conjugate Gradient, Multigrid (V‑cycle).
• Turbulence: Smagorinsky & dynamic Smagorinsky LES.
• Immersed boundaries: Brinkman penalisation with smooth Heaviside mask.
• Differentiability: Every step is JAX‑native. Use jax.grad to optimise shapes and controls. This enables gradient-based inverse design and future differentiable control workflows.
• Performance: Decoupled solver / render / metrics threads with shared‑memory zero‑copy buffers. CPU‑optimised - ~297 FPS (512×96) on a laptop CPU.
• Live parameter updates: Redux‑style state management. The solver checks for changes each timestep and only re-compiles when necessary.

Performance benchmarks (CPU)

Benchmarks include full solver stepping and rendering.

Limitations

• 2D only - designed for rapid prototyping and neural operator research.
• Force coefficients (CL, CD) are trend-accurate but not quantitatively reliable - Brinkman penalisation prevents accurate pressure integration.
• Moderate Reynolds numbers - best for Re < 10,000 (laminar to early turbulent).
• CPU‑focused - GPU execution is supported, but the current architecture is tuned for workstations without dedicated GPUs.
• LBM operates in the low-Mach regime; accuracy degrades if this constraint is violated.

License

LGPL v3.0 - you can use it in proprietary software as long as you release modifications to the library itself.

Author

Arno Meijer - Mechanical Engineer | CFD-ML Systems Developer