# leo1200/astronomix architecture

> Architecture diagram of the GitHub repository leo1200/astronomix, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-11-13). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/leo1200/astronomix
- Repository: https://github.com/leo1200/astronomix

## Overview

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: [`astronomix/time_stepping/time_integration.py`](https://github.com/leo1200/astronomix/blob/main/astronomix/time_stepping/time_integration.py) (Time integration), [`astronomix/_finite_difference`](https://github.com/leo1200/astronomix/tree/main/astronomix/_finite_difference) (Finite difference), [`astronomix/_finite_volume`](https://github.com/leo1200/astronomix/tree/main/astronomix/_finite_volume) (Finite volume), [`astronomix/_pallas_helpers.py`](https://github.com/leo1200/astronomix/blob/main/astronomix/_pallas_helpers.py) (Pallas utilities)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_setup["Run setup"]
  node_config["Configuration"]
  node_params["Runtime parameters"]
  node_registry["Variable registry"]
  node_helpers["Geometry helpers"]
  node_initial["Initial state"]
end

subgraph group_solvers["Numerical solvers"]
  node_fd_solver["Finite difference"]
  node_weno["WENO kernels<br/>[_weno_pallas.py]"]
  node_fv_solver["Finite volume"]
  node_riemann["Riemann solvers"]
end

subgraph group_physics["Physics extensions"]
  node_physics_modules["Physics modules"]
  node_fluid["Fluid equations<br/>[_equations.py]"]
end

subgraph group_runtime["Execution and output"]
  node_integrator["Time integration"]
  node_time_loop["Loop driver<br/>[_time_loop.py]"]
  node_parallel["Distributed execution<br/>[distributed.py]"]
  node_pallas["Pallas utilities<br/>[_pallas_helpers.py]"]
  node_snapshots["Snapshots and checkpoints"]
  node_restart["Restart support<br/>[restart.py]"]
end

subgraph group_analysis["Analysis"]
  node_analysis["Analysis helpers"]
end

node_user(("Simulation user"))

node_user -->|"configures"| node_config
node_user -->|"sets parameters"| node_params
node_user -->|"requests registry"| node_registry
node_user -->|"requests geometry"| node_helpers
node_user -->|"builds state"| node_initial
node_user -->|"runs simulation"| node_integrator
node_integrator -->|"advances with FD"| node_fd_solver
node_integrator -->|"advances with FV"| node_fv_solver
node_fd_solver -.->|"uses WENO"| node_weno
node_fv_solver -.->|"uses Riemann solvers"| node_riemann
node_integrator -->|"drives loop"| node_time_loop
node_integrator -->|"applies modules"| node_physics_modules
node_integrator -->|"records diagnostics"| node_snapshots
node_integrator -->|"supports sharding"| node_parallel
node_fd_solver -.->|"uses GPU utilities"| node_pallas
node_fv_solver -.->|"uses GPU utilities"| node_pallas
node_integrator -.->|"evolves fluid state"| node_fluid
node_user -->|"restarts run"| node_restart
node_user -.->|"analyzes results"| node_analysis

click node_config "https://github.com/leo1200/astronomix/blob/main/astronomix/option_classes/simulation_config.py"
click node_params "https://github.com/leo1200/astronomix/blob/main/astronomix/option_classes/simulation_params.py"
click node_registry "https://github.com/leo1200/astronomix/blob/main/astronomix/variable_registry/registered_variables.py"
click node_helpers "https://github.com/leo1200/astronomix/blob/main/astronomix/data_classes/simulation_helper_data.py"
click node_initial "https://github.com/leo1200/astronomix/blob/main/astronomix/initial_condition_generation/construct_primitive_state.py"
click node_integrator "https://github.com/leo1200/astronomix/blob/main/astronomix/time_stepping/time_integration.py"
click node_time_loop "https://github.com/leo1200/astronomix/blob/main/astronomix/time_stepping/_time_loop.py"
click node_fd_solver "https://github.com/leo1200/astronomix/tree/main/astronomix/_finite_difference"
click node_weno "https://github.com/leo1200/astronomix/blob/main/astronomix/_finite_difference/_interface_fluxes/_weno_pallas.py"
click node_fv_solver "https://github.com/leo1200/astronomix/tree/main/astronomix/_finite_volume"
click node_riemann "https://github.com/leo1200/astronomix/blob/main/astronomix/_finite_volume/_riemann_solver/_athena_riemann_solvers.py"
click node_physics_modules "https://github.com/leo1200/astronomix/tree/main/astronomix/_modules"
click node_fluid "https://github.com/leo1200/astronomix/blob/main/astronomix/_fluid_equations/_equations.py"
click node_parallel "https://github.com/leo1200/astronomix/blob/main/astronomix/parallel/distributed.py"
click node_pallas "https://github.com/leo1200/astronomix/blob/main/astronomix/_pallas_helpers.py"
click node_snapshots "https://github.com/leo1200/astronomix/tree/main/astronomix/_snapshotting"
click node_restart "https://github.com/leo1200/astronomix/blob/main/astronomix/setup_helpers/restart.py"
click node_analysis "https://github.com/leo1200/astronomix/tree/main/astronomix/analysis_helpers"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_config,node_params,node_registry,node_helpers,node_initial,node_user toneBlue
class node_fd_solver,node_weno,node_fv_solver,node_riemann toneAmber
class node_physics_modules,node_fluid toneMint
class node_integrator,node_time_loop,node_parallel,node_pallas,node_snapshots,node_restart toneRose
class node_analysis toneIndigo
```

## Components

### Run setup

- [Configuration](https://github.com/leo1200/astronomix/blob/main/astronomix/option_classes/simulation_config.py): `astronomix/option_classes/simulation_config.py`
- [Runtime parameters](https://github.com/leo1200/astronomix/blob/main/astronomix/option_classes/simulation_params.py): `astronomix/option_classes/simulation_params.py`
- [Variable registry](https://github.com/leo1200/astronomix/blob/main/astronomix/variable_registry/registered_variables.py): `astronomix/variable_registry/registered_variables.py`
- [Geometry helpers](https://github.com/leo1200/astronomix/blob/main/astronomix/data_classes/simulation_helper_data.py): `astronomix/data_classes/simulation_helper_data.py`
- [Initial state](https://github.com/leo1200/astronomix/blob/main/astronomix/initial_condition_generation/construct_primitive_state.py): `astronomix/initial_condition_generation/construct_primitive_state.py`

### Numerical solvers

- [Finite difference](https://github.com/leo1200/astronomix/tree/main/astronomix/_finite_difference): `astronomix/_finite_difference`
- [WENO kernels](https://github.com/leo1200/astronomix/blob/main/astronomix/_finite_difference/_interface_fluxes/_weno_pallas.py): `astronomix/_finite_difference/_interface_fluxes/_weno_pallas.py`
- [Finite volume](https://github.com/leo1200/astronomix/tree/main/astronomix/_finite_volume): `astronomix/_finite_volume`
- [Riemann solvers](https://github.com/leo1200/astronomix/blob/main/astronomix/_finite_volume/_riemann_solver/_athena_riemann_solvers.py): `astronomix/_finite_volume/_riemann_solver/_athena_riemann_solvers.py`

### Physics extensions

- [Physics modules](https://github.com/leo1200/astronomix/tree/main/astronomix/_modules): `astronomix/_modules`
- [Fluid equations](https://github.com/leo1200/astronomix/blob/main/astronomix/_fluid_equations/_equations.py): `astronomix/_fluid_equations/_equations.py`

### Execution and output

- [Time integration](https://github.com/leo1200/astronomix/blob/main/astronomix/time_stepping/time_integration.py): `astronomix/time_stepping/time_integration.py`
- [Loop driver](https://github.com/leo1200/astronomix/blob/main/astronomix/time_stepping/_time_loop.py): `astronomix/time_stepping/_time_loop.py`
- [Distributed execution](https://github.com/leo1200/astronomix/blob/main/astronomix/parallel/distributed.py): `astronomix/parallel/distributed.py`
- [Pallas utilities](https://github.com/leo1200/astronomix/blob/main/astronomix/_pallas_helpers.py): `astronomix/_pallas_helpers.py`
- [Snapshots and checkpoints](https://github.com/leo1200/astronomix/tree/main/astronomix/_snapshotting): `astronomix/_snapshotting`
- [Restart support](https://github.com/leo1200/astronomix/blob/main/astronomix/setup_helpers/restart.py): `astronomix/setup_helpers/restart.py`

### Analysis

- [Analysis helpers](https://github.com/leo1200/astronomix/tree/main/astronomix/analysis_helpers): `astronomix/analysis_helpers`

### Other components

- **Simulation user**

## Connections

- 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 (evidence: `astronomix/time_stepping/time_integration.py`)
- Time integration → Finite volume: advances with FV (evidence: `astronomix/time_stepping/time_integration.py`)
- Finite difference → WENO kernels: uses WENO (evidence: `astronomix/_finite_difference/_interface_fluxes/_weno.py`)
- Finite volume → Riemann solvers: uses Riemann solvers (evidence: `astronomix/_finite_volume/_timestep_estimation/_timestep_estimator.py`)
- Time integration → Loop driver: drives loop (evidence: `astronomix/time_stepping/time_integration.py`)
- Time integration → Physics modules: applies modules (evidence: `astronomix/time_stepping/time_integration.py`)
- Time integration → Snapshots and checkpoints: records diagnostics (evidence: `astronomix/time_stepping/time_integration.py`)
- Time integration → Distributed execution: supports sharding
- Finite difference → Pallas utilities: uses GPU utilities (evidence: `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