DSPS: Differentiable Stellar Population Synthesis

TL;DR

DSPS is a differentiable stellar population synthesis Python package built with JAX, enabling fast, GPU-accelerated galaxy modeling and advanced gradient-based inference for galaxy properties and evolution.

Contribution

It introduces a fully differentiable SPS tool with novel features like metallicity-age correlation and dark matter halo connection, enhancing galaxy modeling capabilities.

Findings

DSPS achieves 5x speedup on CPU and 300-400x on GPU compared to standard codes.

Enables gradient-based optimization and inference for galaxy spectral modeling.

Supports complex physical effects including SFH, metallicity, nebular emission, and dust attenuation.

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