jf1uids - differentiable (magneto)hydrodynamics for astrophysics in JAX

Features

• 1D, 2D and 3D hydrodynamics and magnetohydrodynamics simulations
• provably divergence free magnetic fields based on the approach of Pang and Wu (2024) (different from their approach not yet provably positivity preserving)
• 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.