jf1uids 0.4.0

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 mid-december 2025. The documentation might not reflect the current state of the code.

[!WARNING]
Version 0.4.0 removed the helper_data argument from the time_integration function. This breaks existing code, just remove the argument to fix it.

Features

• 1D, 2D and 3D hydrodynamics and magnetohydrodynamics simulations
• a high-order finite difference constrained transport WENO MHD scheme following HOW-MHD by Seo & Ryu 2023 as well as the provably divergence free and provably positivity preserving finite volume approach of Pang and Wu (2024)
• for finite volume simulations 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, with improved stability at strong discontinuities (currently only available for the finite volume solver)
• 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
• turbulent driving, simple stellar wind, simple radiative cooling modules
• 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
  • 3D MHD blast with the 5th order FD scheme
• 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

Magnetohydrodynamics simulation with driven turbulence at a resolution of 512³ cells in a fifth order CT MHD scheme run on 4 H200 GPUs.

Magnetohydrodynamics simulation with driven turbulence and stellar wind at a resolution of 512³ cells in a fifth order CT MHD scheme run on 4 H200 GPUs.

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.