DSPS: Differentiable stellar population synthesis

Models of stellar population synthesis (SPS) are the fundamental tool that relates the physical properties of a galaxy to its spectral energy distribution (SED). In this paper, we present DSPS: a python package for stellar population synthesis. All of the functionality in DSPS is implemented natively in the JAX library for automatic differentiation, and so our predictions for galaxy photometry are fully differentiable, and directly inherit the performance benefits of JAX, including portability onto GPUs. DSPS also implements several novel features, such as i) a flexible empirical model for stellar metallicity that incorporates correlations with stellar age, and ii) support for the diffstar model that provides a physically-motivated connection between the star formation history of a galaxy (SFH) and the mass assembly of its underlying dark matter halo. We detail a set of theoretical techniques for using autodiff to calculate gradients of predictions for galaxy SEDs with respect to SPS parameters that control a range of physical effects, including SFH, stellar metallicity, nebular emission, and dust attenuation. When forward modeling the colors of a synthetic galaxy population, we find that DSPS can provide a factor of 5 speedup over standard SPS codes on a CPU, and a factor of 300-400 on a modern GPU. When coupled with gradient-based techniques for optimization and inference, DSPS makes it practical to conduct expansive likelihood analyses of simulation-based models of the galaxy--halo connection that fully forward model galaxy spectra and photometry.
Comment: 13 pages, 2 appendices. Version accepted for publication in MNRAS