$\texttt{synax}$: A Differentiable and GPU-accelerated Synchrotron Simulation Package
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2024
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866909550667366400 |
|---|---|
| author | Diao, Kangning Li, Zack Grumitt, Richard D. P. Mao, Yi |
| author_facet | Diao, Kangning Li, Zack Grumitt, Richard D. P. Mao, Yi |
| contents | We introduce synax, a novel library for automatically differentiable simulation of Galactic synchrotron emission. Built on the JAX framework, synax leverages JAX's capabilities, including batch acceleration, just-in-time compilation, and hardware-specific optimizations (CPU, GPU, TPU). Crucially, synax uses JAX's automatic differentiation (AD) mechanism, enabling precise computation of derivatives with respect to any model parameters. This feature facilitates powerful inference algorithms, such as Hamiltonian Monte Carlo (HMC) and gradient-based optimization, which enables inference over models that would otherwise be computationally prohibitive. In its initial release, synax supports synchrotron intensity and polarization calculations down to GHz frequencies, alongside several models of the Galactic magnetic field (GMF), cosmic ray (CR) spectra, and thermal electron density fields. We demonstrate the transformative potential of AD for tasks involving full posterior inference using gradient-based techniques or Maximum Likelihood Estimation (MLE) optimization. Notably, we show that GPU acceleration brings a twenty-fold enhancement in efficiency, while HMC achieves a two-fold improvement over standard random walk Metropolis-Hastings (RWMH) when performing inference over a four-parameter test model. HMC still works on a more complex, 16-parameter model while RWMH fails to converge. Additionally, we showcase the application of synax in optimizing the GMF based on the Haslam 408 MHz map, achieving residuals with a standard deviation below 1 K. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2410_01136 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | $\texttt{synax}$: A Differentiable and GPU-accelerated Synchrotron Simulation Package Diao, Kangning Li, Zack Grumitt, Richard D. P. Mao, Yi Instrumentation and Methods for Astrophysics Astrophysics of Galaxies We introduce synax, a novel library for automatically differentiable simulation of Galactic synchrotron emission. Built on the JAX framework, synax leverages JAX's capabilities, including batch acceleration, just-in-time compilation, and hardware-specific optimizations (CPU, GPU, TPU). Crucially, synax uses JAX's automatic differentiation (AD) mechanism, enabling precise computation of derivatives with respect to any model parameters. This feature facilitates powerful inference algorithms, such as Hamiltonian Monte Carlo (HMC) and gradient-based optimization, which enables inference over models that would otherwise be computationally prohibitive. In its initial release, synax supports synchrotron intensity and polarization calculations down to GHz frequencies, alongside several models of the Galactic magnetic field (GMF), cosmic ray (CR) spectra, and thermal electron density fields. We demonstrate the transformative potential of AD for tasks involving full posterior inference using gradient-based techniques or Maximum Likelihood Estimation (MLE) optimization. Notably, we show that GPU acceleration brings a twenty-fold enhancement in efficiency, while HMC achieves a two-fold improvement over standard random walk Metropolis-Hastings (RWMH) when performing inference over a four-parameter test model. HMC still works on a more complex, 16-parameter model while RWMH fails to converge. Additionally, we showcase the application of synax in optimizing the GMF based on the Haslam 408 MHz map, achieving residuals with a standard deviation below 1 K. |
| title | $\texttt{synax}$: A Differentiable and GPU-accelerated Synchrotron Simulation Package |
| topic | Instrumentation and Methods for Astrophysics Astrophysics of Galaxies |
| url | https://arxiv.org/abs/2410.01136 |