A differentiable magnetohydrodynamics simulation toolkit that lets you model stellar winds, magnetic reconnection, and cosmic turbulence with GPU acceleration and gradient-based optimization.

The cosmos is alive with flowing plasma—from the turbulent atmospheres of stars to the magnetic fields threading galaxy clusters. Understanding these magnetohydrodynamic (MHD) flows has traditionally required computationally intensive simulations that could take weeks to run and couldn’t easily incorporate observational constraints or uncertainty quantification.

Astronomix revolutionizes this landscape by bringing differentiable programming to astrophysical fluid dynamics. Built on JAX, it implements high-order finite difference schemes and advanced Riemann solvers (including the cutting-edge HLLC variants) that scale seamlessly across multiple GPUs. The toolkit handles everything from 1D shock tubes to 3D magnetospheric dynamics, with built-in modules for turbulent driving, stellar winds, and radiative cooling. What sets it apart is full differentiability—you can compute gradients through entire simulations, enabling gradient-based parameter estimation and machine learning integration.

This opens unprecedented possibilities for inverse modeling of astrophysical systems. Researchers can now fit simulation parameters directly to observational data, train neural network surrogates for rapid parameter exploration, or embed physics-based simulations within larger machine learning pipelines. Whether you’re modeling solar flare dynamics, interstellar medium turbulence, or accretion disk physics, astronomix provides the computational foundation to bridge theory, simulation, and observation in ways previously impossible.

⭐ Stars: 49
💻 Language: Jupyter Notebook
🔗 Repository: leo1200/astronomix