jf1uids 0.2.9

differentiable (magneto)hydrodynamics for astrophysics in JAX

jf1uids - differentiable (magneto)hydrodynamics for astrophysics in JAX

[!WARNING]
This project is still an alpha version (but generally functional), updates might break examples. We are aiming for the first stable release in Fall 2025. The documentation might not reflect the current state of the code.

Features

• 1D, 2D and 3D hydrodynamics and magnetohydrodynamics simulations
• provably divergence free magnetic fields and provably positivity preserving hydrodynamics based on the approach of Pang and Wu (2024)
• the basic Lax-Friedrichs, HLL and HLLC Riemann solvers as well as the HLLC-LM (Fleischmann et al., 2020) and HYBRID-HLLC & AM-HLLC (Hu et al., 2025) (sequels to HLLC-LM) variants
• novel (possibly) conservative self gravity scheme, stable at strong discontinuities
• spherically symmetric simulations such that mass and energy are conserved based on the scheme of Crittenden and Balachandar (2018)
• backwards and forwards differentiable with adaptive timestepping
• easily extensible, all code is open source

Installation

jf1uids can be installed via pip

pip install jf1uids

Note that if JAX is not yet installed, only the CPU version of JAX will be installed as a dependency. For a GPU-compatible installation of JAX, please refer to the JAX installation guide.

Notebooks for Getting Started

• hydrodynamics
  • 1d shock tube
  • 1d spherical check of conservational properties
  • 2d Kelvin-Helmholtz instability
• magnetohydrodynamics
  • 2d Orszag-Tang vortex
• self-gravity
  • 3d simulation of Evrard's collapse
• stellar wind
  • 1d stellar wind with gradient showcase
  • 1d stellar wind with parameter optimization
  • 3d stellar wind

Showcase

Orszag-Tang Vortex	3D Collapse

Gradients Through Stellar Wind

Novel (Possibly) Conservative Self Gravity Scheme, Stable at Strong Discontinuities

Wind Parameter Optimization

Citable References

We have currently only published on the previous 1d-only version of this code.

Storcks, L., & Buck, T. (2024). Differentiable Conservative Radially Symmetric Fluid Simulations and Stellar Winds--jf1uids. arXiv preprint arXiv:2410.23093.