synax: A Differentiable and GPU-accelerated Synchrotron Simulation Package

Authors:
Open Access
Publisher website
Google Scholar
Abstract
We introduce synax (https://github.com/dkn16/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 analytical derivatives with respect to any model parameters. This 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 spectra, and thermal electron density fields. When running synax on the CPU we obtain identical performance to hammurabi, a state-of-the-art synchrotron simulation package, while on the GPU synax brings a 20-fold enhancement in efficiency. We further demonstrate the potential of AD in enabling full posterior inference using gradient-based inference algorithms. Using synax with HMC to perform inference over a four-parameter test model, we attain a twofold improvement compared to standard random walk Metropolis–Hastings (RWMH). When applied to a more complex 16-parameter model, HMC is still able to obtain accurate posterior expectations, while RWMH fails to converge. We also showcase the application of synax to optimizing the GMF based on the Haslam 408 MHz map, achieving residuals with a standard deviation below 1 K.
Funding Information
  • DOE ∣ Office of Science (DE- AC02-05CH11231)

This publication has 39 references indexed in Scilit: