astronomix 0.6.2

differentiable (magneto)hydrodynamics for astrophysics in JAX

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.

Features

• 1D, 2D and 3D hydrodynamics and magnetohydrodynamics simulations scaling to multiple GPUs
• a 5th 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) (the WENO scheme is also available standalone for hydrodynamics)
• isothermal hydrodynamics and magnetohydrodynamics are also supported (currently only in the finite difference scheme)
• 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
• 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

Contents

• Installation
• Hello World! Your first astronomix simulation
• Notebooks for Getting Started
• Showcase
• Scaling tests
• Documentation
• Methodology
• Limitations
• Citing astronomix

Installation

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.

Hello World! Your first astronomix simulation

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 .

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

Performance

Methods paper incoming :)

Frequently Asked Questions (FAQ)

How to store intermediate information during the simulation?

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, ...)

Documentation

See here.

Citing astronomix

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.