How leo1200/astronomix works
Astronomix is a differentiable JAX hydrodynamics and magnetohydrodynamics simulation package. A typical run configures a simulation, registers its state variables, builds an initial condition, and advances the state through the time-integration driver using a selected finite-difference or finite-volume solver. Major product branches include physics modules, multi-device execution, checkpointing and snapshots, and analysis. The README documents the user workflow and solver capabilities; the source index confirms several internal references, but many subsystem implementations were not sampled, so those links are limited to documented relationships.
Main entry points: time_integration.py Time integration · _finite_difference Finite difference · _finite_volume Finite volume · _pallas_helpers.py Pallas utilities
Components 19
Run setup
- Configuration
astronomix/option_classes/simulation_config.py - Runtime parameters
astronomix/option_classes/simulation_params.py - Variable registry
astronomix/variable_registry/registered_variables.py - Geometry helpers
astronomix/data_classes/simulation_helper_data.py - Initial state
astronomix/initial_condition_generation/construct_primitive_state.py
Numerical solvers
- Finite difference
astronomix/_finite_difference - WENO kernels
astronomix/_finite_difference/_interface_fluxes/_weno_pallas.py - Finite volume
astronomix/_finite_volume - Riemann solvers
astronomix/_finite_volume/_riemann_solver/_athena_riemann_solvers.py
Physics extensions
- Physics modules
astronomix/_modules - Fluid equations
astronomix/_fluid_equations/_equations.py
Execution and output
- Time integration
astronomix/time_stepping/time_integration.py - Loop driver
astronomix/time_stepping/_time_loop.py - Distributed execution
astronomix/parallel/distributed.py - Pallas utilities
astronomix/_pallas_helpers.py - Snapshots and checkpoints
astronomix/_snapshotting - Restart support
astronomix/setup_helpers/restart.py
Analysis
- Analysis helpers
astronomix/analysis_helpers
Other components
- Simulation user
Connections 19
- Simulation user → Configuration: configures
- Simulation user → Runtime parameters: sets parameters
- Simulation user → Variable registry: requests registry
- Simulation user → Geometry helpers: requests geometry
- Simulation user → Initial state: builds state
- Simulation user → Time integration: runs simulation
- Time integration → Finite difference: advances with FD
astronomix/time_stepping/time_integration.py - Time integration → Finite volume: advances with FV
astronomix/time_stepping/time_integration.py - Finite difference → WENO kernels: uses WENO
astronomix/_finite_difference/_interface_fluxes/_weno.py - Finite volume → Riemann solvers: uses Riemann solvers
astronomix/_finite_volume/_timestep_estimation/_timestep_estimator.py - Time integration → Loop driver: drives loop
astronomix/time_stepping/time_integration.py - Time integration → Physics modules: applies modules
astronomix/time_stepping/time_integration.py - Time integration → Snapshots and checkpoints: records diagnostics
astronomix/time_stepping/time_integration.py - Time integration → Distributed execution: supports sharding
- Finite difference → Pallas utilities: uses GPU utilities
astronomix/_finite_difference/_interface_fluxes/_weno_pallas.py - Finite volume → Pallas utilities: uses GPU utilities
- Time integration → Fluid equations: evolves fluid state
- Simulation user → Restart support: restarts run
- Simulation user → Analysis helpers: analyzes results