Astronomix

screenshot of Astronomix

differentiable (magneto)hydrodynamics for astrophysics in JAX

Overview

Astronomix is an innovative code designed specifically for differentiable hydrodynamics and magnetohydrodynamics, written in JAX. Its primary focus is on astrophysical applications, making it a vital tool for researchers and developers engaging in high-performance simulations. The project is currently active, having reached a stable and user-friendly state while continuing to evolve with ongoing development. The differentiability of astronomix highlights its potential for gradient-based inverse modeling and sampling, allowing scientists to enhance their simulations significantly.

This software combines flexibility, ease of use, and cutting-edge techniques to provide both beginner and advanced users an effective platform for their astrophysics projects. It's particularly exciting for those keen on fast method development and looking to leverage GPU scaling capabilities.

Features

  • 1D, 2D, and 3D Simulations: Supports scalable hydrodynamics and magnetohydrodynamics simulations across multiple GPUs, enhancing computational efficiency for complex astrophysical models.

  • Advanced MHD Schemes: Utilizes a high-order finite difference WENO MHD scheme compliant with contemporary research, ensuring high accuracy and reliability in simulations.

  • Robust Riemann Solvers: Integrates various Riemann solvers, including HLLC and its advanced variants, providing a comprehensive toolset for tackling hydrodynamic problems and enhancing solution stability.

  • Self-Gravity Scheme: Features a novel self-gravity scheme with improved stability during strong discontinuities, improving the range of scenarios that can be accurately modeled.

  • Spherically Symmetric Simulations: Implements a scheme that conserves mass and energy, allowing for accurate representations of spherical astrophysical phenomena.

  • Differentiable and Adaptive: Fully differentiable with adaptive timestepping capabilities, making it suitable for complex gradient-based optimization work.

  • Additional Modules: Includes modules for turbulent driving, stellar wind modeling, and radiative cooling, broadening its application potential in astrophysical research.

  • Open Source and Extensible: The entire codebase is open source, allowing for customization and contributions from the community, fostering innovation and collaboration.