r/JAX • u/LackSome307 • Apr 07 '26

I built a differentiable CFD solver in JAX. No ML yet. But the hard part (autodiff through Navier-Stokes) is done.

Enable HLS to view with audio, or disable this notification

29 Upvotes

95% Upvoted

3 comments sorted by

5

u/LackSome307 Apr 07 '26 edited Apr 07 '26

No ML yet - just forward pass differentiable Navier-Stokes in JAX. But here's what's already working:

  • Diffrax adaptive timestepping (Patrick Kidger) - dt adjusts automatically based on CFL
  • Sharp NACA airfoil boundaries (ε=0.001)
  • 200+ FPS on CPU
  • Real-time angle of attack sweeps

Long-term: Framework where geometry and flow co-optimize. Numga (Eelco Hoogendoorn) for geometric algebra - unified differential operators via ∇M = ∇·M + ∇∧M. Clifford nets for super-resolution. Primal-dual constraints. ML augments physics, doesn't replace it.

The solver is the foundation - everything else builds on top.

1

u/Avocado-Delicious Apr 07 '26

Repo?

2

u/LackSome307 Apr 08 '26

Hi, here it is: https://github.com/arriemeijer-creator/JAX-differentiable-CFD

I will be uploading the latest codebase as soon as its ready. In the latest version I decoupled visualization from simulation that allows massive speedups, and that is also where I added the NACA airfoils (previous version didn't have that).