texttt{synax}$: A Differentiable and GPU-accelerated Synchrotron Simulation Package
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). Cruciall...
Gespeichert in:
Hauptverfasser: | , , , |
---|---|
Format: | Artikel |
Sprache: | eng |
Schlagworte: | |
Online-Zugang: | Volltext bestellen |
Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
Zusammenfassung: | 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. |
---|---|
DOI: | 10.48550/arxiv.2410.01136 |