$\texttt{synax}$: A Differentiable and GPU-accelerated Synchrotron Simulation Package

TL;DR

synax is a GPU-accelerated, differentiable simulation library for Galactic synchrotron emission that enables efficient gradient-based inference and optimization, significantly improving computational performance and enabling complex model fitting.

Contribution

It introduces synax, a novel JAX-based library that provides automatic differentiation and GPU acceleration for synchrotron emission modeling, facilitating advanced inference techniques.

Findings

GPU acceleration yields twenty-fold speedup.

HMC outperforms RWMH in complex models.

Successful optimization of GMF with low residuals.

Abstract

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.