DSPS: Differentiable stellar population synthesis

Abstract

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, 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.

As Simple as Possible but No Simpler: Optimizing the Performance of Neural Net Emulators for Galaxy SED Fitting

“…jl (Revels et al 2016) in Julia. One such implementation is DSPS (Hearin et al 2023), which can generate predicted SEDs with roughly the same execution time as ANN emulators. However, DSPS is currently only capable of generating SEDs using a precomputed SSP spectral grid, and bringing it to feature-parity with FSPS will likely result in either a slower execution time (in the case of using stellar spectral libraries and isochrones as the base) or an extremely high memory requirement (in the case of using high-fidelity precomputed SSP grids).…”

PopSED: Population-level Inference for Galaxy Properties from Broadband Photometry with Neural Density Estimation

“…To showcase the capabilities of POPSED on such population-level analysis, we focus on constraining the SFMS for galaxies at z < 0.1. With the samples drawn from the inferred population distribution in hand, we select samples with z < 0.1 and calculate the average SFR within the past 0.1 Gyr and the surviving stellar mass using the fitting function in Hearin et al (2023). 10 Here we choose the surviving stellar mass rather than the total formed mass to better compare with literature results.…”

“…Using traditional SED fitting methods, analyzing 10 5 galaxies will take up to 2 × 10 6 CPU hr. Even with the development of accelerated SED fitting (e.g., Alsing et al 2020;Hearin et al 2023;Khullar et al 2022;Wang et al 2023), an analysis of 10 5 galaxies will still take up to ∼10 3 GPU hr. POPSED is able to recover the posterior of the population distribution for ∼10 5 galaxies within ∼10 GPU hr, 100 times faster than the SBIbased methods.…”

“…This problem has been partially mitigated by recent developments in accelerating SED modeling by emulating the SPS models with neural networks (Alsing et al 2020), building differentiable SPS models with a high-performance library (Hearin et al 2023), and speeding up the sampling by using amortized simulation-based inference (SBI; Khullar et al 2022;Wang et al 2023). However, some of these methods need sophisticated training and are costly to retrain when adapting to different SPS models or noise properties.…”

SBI++: Flexible, Ultra-fast Likelihood-free Inference Customized for Astronomical Applications

“…First, the generation of an SPS model can be accelerated by using a differentiable SPS code or an artificial neural network emulator. The former generates exact SPS with a boost in speed enabled by specific code libraries (e.g., Hearin et al 2023); while the latter uses a quick-to-evaluate neural network that approximates SPS (e.g., Alsing et al 2020). Second, the sampling time can be decreased by switching to a gradientbased sampler (Duane et al 1987;Hoffman & Gelman 2011), or combining the sampler with a normalizing flow (Karamanis et al 2022;Wong et al 2023).…”