astronomix - differentiable mhd for astrophysics in JAX

astronomix (formerly jf1uids) is a differentiable hydrodynamics and magnetohydrodynamics code written in JAX with a focus on astrophysical applications. astronomix is easy to use, well-suited for fast method development, scales to multiple GPUs and its differentiability opens the door for gradient-based inverse modeling and sampling as well as surrogate / solver-in-the-loop training.

• 1D, 2D and 3D hydrodynamics (HD) and magnetohydrodynamics (MHD) simulations scaling to multiple GPUs and nodes
• fifth order WENO-JS and WENO-Z finite difference schemes (FD) for both HD and MHD simulations (adiabtatic and isothermal)
• for the fifth order finite difference scheme a constrained transport MHD approach following HOW-MHD by Seo & Ryu 2023
• for the fifth order finite difference scheme a positivity-preserving reconstruction (based on Zhang and Shu, 2012 and Wu, 2018) that keeps density and pressure positive for C_cfl <= 0.75 without floors, additionally slight dissipation for the electric field can be activated for stability in extreme flows (where the cell-centered WENO dissipation does not translate into enough damping of the interface electric field at the Nyquist scale)
• for the finite difference scheme a novel semi-discretely energy conserving self-gravity scheme
• a re-implementation of AthenaPK's default second-order finite volume (FV) scheme (VL2 predictor-corrector, PLM, HLLD / HLLE / HLLC, GLM divergence cleaning, first-order flux correction), agreeing with AthenaPK to round-off (validation), further supported Riemann solvers include Lax-Friedrichs, HLL, HLLC-LM (Fleischmann et al., 2020), HYBRID-HLLC & AM-HLLC (Hu et al., 2025)
• backwards and forwards differentiable with adaptive timestepping
• physics modules: turbulent driving (white-in-time and Ornstein-Uhlenbeck, AthenaK / AthenaPK driving spectra), radiative cooling with an implicit solver, thermal conduction, viscosity and resistivity, passive scalars (e.g. composition tracers), dual energy, stellar winds
• spherically symmetric simulations such that mass and energy are conserved based on the scheme of Crittenden and Balachandar (2018)
• Pallas (Triton) GPU kernels for the finite difference and the VL2 finite volume schemes, multi-GPU via shard_map halo exchange
• easily extensible, all code is open source

• Features
• Installation
• Hello World! Your first astronomix simulation
• Examples for Getting Started
• Showcase
• Performance
• Frequently Asked Questions (FAQ)
  • How to store intermediate information during the simulation?
  • How to fit larger simulations into GPU memory?
  • What if my simulation crashes?
• Documentation
• Citing astronomix

astronomix can be installed via pip

pip install astronomix

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.

Below is a minimal example of a 1D hydrodynamics shock tube simulation using astronomix.

import jax.numpy as jnp
from astronomix import (
    SimulationConfig, SimulationParams,
    get_helper_data, finalize_config,
    get_registered_variables, construct_primitive_state,
    time_integration
)

# the SimulationConfig holds static 
# configuration parameters
config = SimulationConfig(
    box_size = 1.0,
    num_cells = 101,
    progress_bar = True
)

# the SimulationParams can be changed
# without causing re-compilation
params = SimulationParams(
    t_end = 0.2,
)

# the variable registry allows for the principled
# access of simulation variables from the state array
registered_variables = get_registered_variables(config)

# next we set up the initial state using the helper data
helper_data = get_helper_data(config)
shock_pos = 0.5
r = helper_data.geometric_centers
rho = jnp.where(r < shock_pos, 1.0, 0.125)
u = jnp.zeros_like(r)
p = jnp.where(r < shock_pos, 1.0, 0.1)

# get initial state
initial_state = construct_primitive_state(
    config = config,
    registered_variables = registered_variables,
    density = rho,
    velocity_x = u,
    gas_pressure = p,
)

# finalize and check the config
config = finalize_config(config, initial_state.shape)

# now we run the simulation
final_state = time_integration(initial_state, config, params, registered_variables)

# the final_state holds the final primitive state, the 
# variables can be accessed via the registered_variables
rho_final = final_state[registered_variables.density_index]
u_final = final_state[registered_variables.velocity_index]
p_final = final_state[registered_variables.pressure_index]

You've just run your first astronomix simulation! You can continue with the notebooks below and we have also prepared a more advanced use-case (stellar wind in driven MHD tubulence) which you can .

Runnable scripts live under examples/scripts/, notebooks under examples/notebooks/.

• forward simulations
  • 1D Sod shock tube, a minimal notebook
  • 2D Kelvin-Helmholtz instability (notebook)
  • 2D Orszag-Tang vortex (notebook)
  • 3D MHD jet (notebook)
  • 3D driven MHD turbulence and the small-scale dynamo (astronomix against AthenaPK)
  • 3D self-gravitating collapse (notebook)
• differentiability
  • Field-level inference (notebook)
  • KHI eigenmode initialization
  • Solver-in-the-loop correction network
• output options
  • On-the-fly movie via callback
  • Orbax checkpointing and restart
  • In-memory snapshot diagnostics
• multi-GPU
  • Sharded multi-GPU run
  • Multi-node run
• validation
  • AthenaPK's VL2 scheme in astronomix
  • Positivity-preserving WENO

More involved, physics-module showcase scripts — turbulent radiative mixing layer, stellar wind, cooling — live in examples/gallery/.

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.

3D MHD jet propagating into a magnetized medium.

Novel semi-discretely energy conserving self-gravity scheme.

Methods paper incoming :)

There are two main options for storing intermediate states:

• activate return_snapshots in the SimulationConfig, set either num_snapshots in the SimulationConfig for equidistant snapshots or snapshot_timepoints in the SimulationParams for custom snapshot timepoints, then activate specific variables to be stored via snapshot_settings in the SimulationConfig (e.g. the full states, the total energy, ..., see here) - these snapshots are stored on the GPU directly, which avoids host-device transfer and is e.g. very useful for losses over intermediate states
• activate snapshot_callable in the SimulationConfig and provide a callable to the time_integration function which is called at the same timepoints as the snapshots in the first option, this function can then offload specific data (possibly after processing) to the host via jax.debug.callback and save it appropriately (e.g. directly make plots, save to disk in the preferred format, ...)

Make sure that you only have the initial state in GPU memory when you start the simulation, e.g. have a function which constructs the initial state (otherwise e.g. the intermediate helper data you used to construct the state might still be in GPU memory). To further save storage, you can donate the initial state to the time integration (activate donate_state in the SimulationConfig), which allows JAX to reuse the same memory for the state throughout the simulation (but you can also no longer access the initial state after the simulation has started).

First of all check if the initial conditions are valid. For the finite difference scheme, turn on the positivity-preserving WENO reconstruction (weno_positivity_preserving=True in the SimulationConfig), which keeps density and pressure positive for C_cfl <= 0.75. The PositivityConfig in the SimulationConfig additionally offers a per-step density / pressure floor.

See here.

If you use astronomix in your research, please cite via

@misc{storcks_astronomix_2025,
  doi = {10.5281/ZENODO.17782162},
  url = {https://zenodo.org/doi/10.5281/zenodo.17782162},
  author = {Storcks, Leonard},
  title = {astronomix - differentiable MHD in JAX},
  publisher = {Zenodo},
  year = {2025},
  copyright = {MIT License}
}

There is also a workshop paper on an earlier stage of the project:

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