synax: A Differentiable and GPU-accelerated Synchrotron Simulation Package

Speaker: Kangning Diao (UC Berkeley) In Person and zoom 

Zoom info: https://stanford.zoom.us/j/98604058568

Understanding Galactic synchrotron emission is critical for both radio cosmology and Galactic physics. However, extracting valuable information from these signals requires advanced inference techniques. To address this challenge, we present synax (https://synax.readthedocs.io), a differentiable forward model for Galactic synchrotron emission. Implemented in JAX, synax takes advantage of automatic differentiation and GPU acceleration, enabling gradient-based methods such as Hamiltonian Monte Carlo (HMC) and gradient descent. Compared to conventional sampling approaches like random-walk Metropolis-Hastings (RWMH), synax with HMC achieves a 40-fold speedup on a simple four-parameter model, reaching an effective sample size of 0.016 per second, while RWMH fails to converge. Additionally, we demonstrate the utility of synax by optimizing the Galactic magnetic field model against the Haslam 408 MHz map, achieving residuals with a standard deviation below 1 K.