††thanks<EMAIL_ADDRESS>
# DSPS: Differentiable Stellar Population Synthesis

Andrew P. Hearin1,⋆ Jonás Chaves-Montero1,2 Alex Alarcon1 Matthew R. Becker1
Andrew Benson3 1HEP Division, Argonne National Laboratory, 9700 South Cass
Avenue, Lemont, IL 60439, USA  2Donostia International Physics Centre, Paseo
Manuel de Lardizabal 4, 20018 Donostia-San Sebastian, Spain  3Carnegie
Observatories, 813 Santa Barbara Street, Pasadena, CA 91101

###### 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 20 speedup over standard SPS codes on a CPU, and a factor
of over 1000 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.

###### keywords:

Cosmology: large-scale structure of Universe; methods: N-body simulations

##  1 Introduction

Stellar population synthesis (SPS) is the prevailing framework for predicting
the spectral energy distribution (SED) of a galaxy from its fundamental
physical properties (Conroy, 2013). SPS is a mature subfield with a long
history (e.g., Tinsley, 1978; Bruzual A., 1983; Arimoto & Yoshii, 1987;
Buzzoni, 1989), and applications of SPS range from inferring the physical
properties of individual galaxies (Sawicki & Yee, 1998; Brinchmann & Ellis,
2000; Salim et al., 2007; Kriek et al., 2009; Leja et al., 2019b), to forward
modeling the galaxy distribution across cosmic time (Baugh et al., 1996;
Kauffmann et al., 1999; Somerville & Primack, 1999; Korytov et al., 2019;
Drakos et al., 2021). In reflection of the scientific breadth of this topic,
there are by now many publicly available libraries that can be used to carry
out SPS-related calculations (e.g., Bruzual & Charlot, 2003; Fioc & Rocca-
Volmerange, 1999; Le Borgne et al., 2004; Maraston, 2005; Conroy et al.,
2009a; Eldridge et al., 2017; Johnson et al., 2021; Bradley et al., 2020).

A wide range of theoretical models and techniques are considered to be
components of SPS, from the specialized interpolation algorithms used to
transform the outputs of stellar evolution codes into isochrones and stellar
evolution tracks (Morton, 2015; Dotter, 2016), to the modeling of a galaxy’s
chemical evolution history (Audouze & Tinsley, 1976; Weinberg et al., 2017),
to the computation of the observed photometry of a galaxy from its composite
stellar population (Oke & Sandage, 1968; Blanton & Roweis, 2007). These and
other related predictions of SPS models can be computationally intensive,
which can create a significant limitation on the level of physical realism
that is achievable in a practical analysis. Classical machine learning (ML)
techniques such as Gaussian Process emulation (Rasmussen & Williams, 2006;
GPy, 2012; Ambikasaran et al., 2014) or polynomial chaos expansion (Wiener,
1938; Xiu, 2010) are well-suited to address these computational challenges.
The basic idea underlying this application of ML is to carry out a pre-
processing step in which a limited number of expensive function evaluations
are used to optimize the hyper-parameters of the ML model; subsequently, one
uses the ML model to provide a fast-evaluating, “surrogate” prediction in the
performance-critical bottleneck of the analysis. Such ML “emulation” methods
have been used within the field of cosmology for years (Heitmann et al.,
2006), and are now used by many different groups to accelerate cosmological
predictions in a wide variety of contexts (e.g., Harnois-Déraps et al., 2019;
McClintock et al., 2019; Euclid Collaboration et al., 2019; Nishimichi et al.,
2019; Kokron et al., 2021; Ramachandra et al., 2021). Recently, artificial
intelligence algorithms (AI) such as a neural network (McCulloch & Pitts,
1943) have become more widely deployed as surrogate functions in cosmology and
astrophysics (e.g., Kobayashi et al., 2020; Aricò et al., 2021; Villaescusa-
Navarro et al., 2020), including applications to stellar population synthesis
(Alsing et al., 2020).

When making a computationally expensive prediction, using an AI/ML-based
surrogate function not only reduces the runtime of any particular model
evaluation, but has the additional benefit of being an exactly differentiable
function of its parameters, a highly advantageous feature in likelihood
analyses. Gradient-based optimization algorithms such as BFGS (Broyden, 1970;
Fletcher, 1970; Goldfarb, 1970; Shanno, 1970) and Adam (Kingma & Ba, 2014)
require far fewer total evaluations of the objective function relative to
other techniques. Inference algorithms such as Hamiltonian Monte Carlo (HMC,
Duane et al., 1987) offer analogous reductions in the number of likelihood
evaluations required to derive posteriors relative to conventional MCMC
methods. These performance benefits become increasingly dramatic in analyses
of models with larger numbers of free parameters. For example, contemporary
optimization algorithms leverage the differentiability of a neural network to
train its weights and biases, which can comprise millions of degrees of
freedom, and HMC is routinely used to derive posteriors on models with
hundreds of parameters or more (Hoffman & Gelman, 2014).

As an alternative to emulating the predictions of a model, the benefits of
differentiability can also be reaped by directly implementing the model in a
software library that supports automatic differentiation. Autodiff is an
algorithm for evaluating the derivative of a function defined by a computer
program, and is a distinct algorithm from numerically estimating a derivative
via finite differencing methods (see Baydin et al., 2015, and references
therein). The computational cost of finite differencing estimates scales with
the number of parameters, which is not the case for autodiff algorithms, and
so gradient computations with autodiff vastly outperform finite-difference
derivatives as the dimension of the model parameter space increases. Whereas
numerical derivatives require careful checking of the finite step size to
protect against inaccuracy (which can be especially tedious for the case of
higher-order derivatives), the autodiff algorithm ensures that the gradient
computation gives the same result as symbolic differentiation within working
precision. There are numerous publicly available libraries providing high-
performance implementations, such as TensorFlow (Abadi et al., 2015), PyTorch
(Paszke et al., 2019), and JAX (Bradbury et al., 2018); these libraries
provide a convenient Python interface to autodiff algorithms that are highly
performant on both CPUs and GPUs. Autodiff computations have now been used in
a wide range of applications in scientific computing, including cosmological
N-body simulations (Modi et al., 2021), molecular dynamics (Schoenholz &
Cubuk, 2019), and fluid dynamics (Kochkov et al., 2021).

Motivated by these developments, in this paper we present DSPS, a stellar
population synthesis code implemented in JAX. In §2 we describe how SPS
predictions for galaxy SEDs are natural to implement with autodiff, and in §3,
we illustrate autodiff-based techniques for calculating gradients of galaxy
SEDs with respect to various physical ingredients that have differentiable
implementations in DSPS. We discuss the advantages and limitations of DSPS in
§4, and conclude with a brief summary in §5.

##  2 Differentiable Predictions for Galaxy SEDs

In this section, we describe our approach for making differentiable
predictions for the SED of a galaxy based on stellar population synthesis. One
of the foundational concepts of stellar population synthesis is the Simple
Stellar Population (SSP), which is defined to be a population of stars that
formed simultaneously from a homogeneous gas cloud; by definition, all the
stars in an SSP have the same age and metallicity. If we define $L_{\rm
SS}(\lambda|M_{\star}^{i},\tau_{\rm age},Z)$ in units of $L_{\odot}{\rm
Hz}^{-1}$ to be the luminosity per unit frequency emitted at time $\tau_{\rm
age}$ after the birth of a single star with initial mass $M_{\star}^{i},$ and
initial metallicity $Z,$ then we have

(1) $\displaystyle L_{\rm SSP}(\lambda|\tau_{\rm age},Z)\equiv\int{\rm
d}M_{\star}^{i}\Phi_{\rm IMF}(M_{\star}^{i})L_{\rm
SS}(\lambda|M_{\star}^{i},\tau_{\rm age},Z),$

where $\Phi_{\rm IMF}$ is the number of stars per unit mass in the SSP, i.e.,
the initial mass function (IMF); Eq. 1 is normalized per unit of total stellar
mass formed.111For the sake of brevity, in this section and throughout the
paper we assume that the initial mass, $M_{\star}^{i},$ and total metallicity,
$Z,$ uniquely determine the SED emitted by a star at a time $\tau_{\rm age}$
after it is born, when in fact numerous other variables are known to play a
significant role, e.g., the rotation speed and chemical abundance pattern. We
furthermore neglect treatment of the effect of stellar binaries on the light
emitted by an SSP. See §4.2 for discussion of how DSPS could be extended to
take these effects into account in a differentiable fashion.

The quantity $L_{\rm SS}(\lambda|M_{\star}^{i},\tau_{\rm age},Z)$ is one of
the fundamental outputs of a stellar evolution code such as MESA (Paxton et
al., 2011, 2013, 2015). Due to the computational expense of stellar evolution
calculations, the outputs of stellar evolution codes must be computed in
advance, stored to disk, and ingested as input data in calculations of the
light emitted by a stellar population; libraries of stellar spectra such as
MESA Isochrones and Stellar Tracks (MIST, Choi et al., 2016) provide such
tabulations in a convenient and publicly available form. Conceptually,
calculating $L_{\rm SSP}(\lambda|\tau_{\rm age},Z)$ via Equation 1 is simply a
matter of interpolating $L_{\rm SS}(\lambda|M_{\star}^{i},\tau_{\rm age},Z)$
between the control points supplied by the stellar spectral library, and
computing an IMF-weighted sum. However, the $\tau_{\rm age}$-dependence of
$L_{\rm SS}(\lambda)$ has notoriously sharp features that are highly sensitive
to $M_{\star}^{i},$ particularly at late stages of stellar evolution, and so
specialized interpolation techniques are required in order to carry out the
integration in Equation 1 (see Dotter, 2016, and references therein). This is
one of the chief technical reasons that drives many SPS codes to make a fixed
choice for the IMF, and to predict galaxy SEDs as weighted sums of a discrete
collection of $L_{\rm SSP}(\lambda|\tau_{\rm age},Z)$ that are tabulated on a
grid in advance. This is the approach taken by DSPS (although see §4.2 for
further discussion).

Throughout the remainder of this section, we will suppose that the SED of a
galaxy is described by just a single SSP, $L_{\rm SSP}(\lambda),$ and in §2.1
we will demonstrate how to calculate SED-derived quantities such as
photometric flux and emission line strength in a differentiable fashion. Of
course, the SED of a physically realistic galaxy is better described by a
composite stellar population that is comprised of many SSPs, but we will
postpone our treatment of how $L_{\rm CSP}(\lambda)$ can be computed
differentiably until §3. Here and throughout the paper, we will use the
variable $\tau_{\rm age}$ to denote the length of time that has passed since
the birth of a star or stellar population. The age of the universe is the
physically natural time variable used to describe the star formation history
of a galaxy, and so we will use the variable $t$ to refer to the age of the
universe at the time of some particular event. The distinction between
$\tau_{\rm age}$ and $t$ will be more important in §3.1 when we turn attention
to the calculation of $L_{\rm CSP}(\lambda).$

###  2.1 Predicting Photometry and Emission Lines

Galaxy samples targeted by cosmological surveys are commonly defined in terms
of observable quantities that derive from the full SED, such as the
photometric flux observed through a broad-band filter, or the strength of a
particular emission line. We now describe differentiable techniques for
predicting these two quantities in turn.

If the rest-frame SED of a galaxy is described in units of luminosity per unit
frequency by $L(\lambda),$ and if we use $T_{\rm Q}$ to denote the filter
transmission curve defined as the probability of photon transmission, then
$M_{\rm Q},$ the absolute magnitude of the galaxy observed in the rest-frame
through the filter, is given by

$\displaystyle M_{\rm Q}$ $\displaystyle=$
$\displaystyle-2.5\log_{10}\left(\frac{\int\frac{{\rm
d}\lambda}{\lambda}T_{\rm Q}(\lambda)L(\lambda)}{AB_{0}\int\frac{{\rm
d}\lambda}{\lambda}T_{\rm Q}(\lambda)}\right),$ (2)

where $AB_{0}=1.13492\times 10^{-13}L_{\odot}{\rm Hz^{-1}}$ is the the AB flux
zero-point for a source observed at 10 pc. The calculation of $M_{\rm Q}$ in
Equation 2 is straightforward to implement in JAX because trapezoidal
integration is included in the library, and the quantities $T_{\rm
Q}(\lambda)$ and $L(\lambda)$ are pre-computed data as described above. As we
will see in §3, the SED of a composite stellar population is calculable as a
weighted sum of a finite collection of $L_{\rm SSP}(\lambda|\tau_{\rm
age},Z),$ and so using JAX to calculate the gradients of the rest-frame flux
of a galaxy is simply a matter of summing the derivatives of the weights, and
propagating the result through Equation 2.

In order to predict the apparent magnitude of the galaxy in the observer-
frame, $m_{\rm Q},$ we must take into account both the redshift of the SED, as
well as the dimming of the source across cosmological distances. For the
former, we need only apply redshift relation, $\lambda_{\rm obs}=(1+z_{\rm
obs})\lambda_{\rm rest},$ before carrying out the integral in the numerator of
Eq. 2; this is straightforward to compute in JAX with simple interpolation;
for notational convenience, let us denote the result by $M^{\prime}_{\rm Q}.$
We then have

$m_{\rm Q}=M^{\prime}_{\rm Q}+D_{\rm mod}(z_{\rm obs})-2.5\log_{10}(1+z_{\rm
obs}),$

where $D_{\rm mod}(z)$ is the distance modulus function that depends on
cosmology. The DSPS library implements a JAX-based calculation of $D_{\rm
mod}(z)$ for flat $w{\rm CDM}$ cosmological models. Typical SPS calculations
are carried out under a fixed cosmology, and so for more general cosmological
models, one could use jax-cosmo, or alternatively use JAX to interpolate from
a lookup table for the function $D_{\rm mod}(z)$ that can be computed in
advance with a publicly available package such as Astropy (Astropy
Collaboration et al., 2013, 2018) or Colossus (Diemer, 2018).

A second characteristic of the galaxy SED that is commonly used in the
selection of galaxy samples in cosmological surveys is the strength of a
particular emission line at wavelength, $\lambda,$ commonly quantified in
terms of the equivalent width, $W_{\lambda},$ defined as

$\displaystyle
W_{\lambda}=\int_{\lambda-\Delta\lambda}^{\lambda+\Delta\lambda}{\rm
d}\lambda^{\prime}\left[L({\lambda}^{\prime})-L_{0}({\lambda}^{\prime})\right]/L_{0}(\lambda^{\prime}),$
(3)

where $L(\lambda)$ is the spectrum of the galaxy, $\Delta\lambda$ specifies
the wavelength range bracketing the line, and $L_{0}(\lambda)$ is the
continuum spectrum evaluated at the center of the line. The summation
operation in the integral of Eq. 3 is a linear operator, and so poses no
problem for automatic differentiation. Thus the key to a differentiable
calculation of $W_{\lambda}$ is the computation of the continuum,
$L_{0}(\lambda).$

To calculate $L_{0}(\lambda),$ we fit the spectrum $L(\lambda)$ in the
neighborhood of the line with a quadratic polynomial, masking the wavelength
range containing the line. General purpose fitting algorithms can be
challenging to formulate in a differentiable fashion, but polynomial fitting
is purely a matter of linear algebra, and so can be implemented in autodiff
library without complications. Thus if the SED of a galaxy depends in a
differentiable fashion upon some model parameter, $\theta,$ so that $\partial
L(\lambda)/\partial\theta$ can be calculated with autodiff, then it is
straightforward to calculate $\partial L_{0}(\lambda)/\partial\theta,$ since
polynomial fitting is merely a sequence of linear operations, and so it is in
turn straightforward to calculate $\partial W_{\lambda}/\partial\theta$ by
differentiating through the right-hand side of Eq. 3.

Having shown how a galaxy’s full SED can be predicted differentiably, as well
as its photometry and the strength of its emission lines, in the next section
we show how autodiff can be used to compute gradients with respect to the
parameters of physical models of stellar population synthesis.

##  3 Physical Ingredients and their Gradients

In this section, we describe the basic physical ingredients currently
implemented in DSPS, as well as the autodiff-based techniques we use to
calculate gradients with respect to the parameters of these models. We present
differentiable techniques for calculating the dependence of galaxy SED upon
stellar age in §3.1, stellar metallicity in §3.2, nebular emission in §3.3,
and dust attenuation in §3.4.

###  3.1 Star Formation Histories and Stellar Ages

The light emitted by the stars in a galaxy at a time $t$ is the cumulative
contribution of all the evolved simple stellar populations that formed prior
to $t.$ Using $L_{\rm CSP}(\lambda|t)$ to denote the SED of the composite
stellar population, we have

$\displaystyle L_{\rm CSP}(\lambda|t)=\int_{0}^{t}{\rm
d}t^{\prime}\int_{0}^{Z_{\rm max}}$ $\displaystyle{\rm
d}Z\dot{M_{\star}}(t^{\prime},Z)F_{\rm surv}(\tau_{\rm age})$ (4)
$\displaystyle\times L_{\rm SSP}(\lambda|\tau_{\rm age},Z),$

where $\tau_{\rm age}\equiv t-t^{\prime},$ the star formation history of the
galaxy is $\dot{M_{\star}}(t^{\prime},Z),$ and $F_{\rm surv}(\tau_{\rm age})$
is the fraction of stellar mass formed at $t^{\prime}$ that survives until $t$
(see Appendix A).222Note that we include $F_{\rm surv}(\tau_{\rm age})$ in Eq.
4 since the definition of $L_{\rm SSP}(\lambda)$ in Eq. 1 is normalized per
unit mass of total stars formed. Throughout the remainder of §3.1, we will
work under the assumption that all the stars in a galaxy have the same
metallicity, and postpone treatment of galaxy populations with a diverse
chemical composition until §3.2. Under this assumption, Equation 4 becomes

$\displaystyle L_{\rm CSP}(\lambda|t)=\int_{0}^{t}{\rm d}t^{\prime}$
$\displaystyle\dot{M_{\star}}(t^{\prime})F_{\rm surv}(\tau_{\rm age})L_{\rm
SSP}(\lambda|\tau_{\rm age}).$

As discussed in §2, the SEDs of simple stellar populations, $L_{\rm
SSP}(\lambda|\tau_{\rm age}),$ are typically pre-computed on a discretized
grid of values of $\tau_{\rm age}^{i},$ so that the integral in Eq.
LABEL:eq:csp becomes a finite summation:

(6) $\displaystyle L_{\rm CSP}(\lambda|t)=M_{\star}^{\rm
surv}(t)\sum_{i}L_{\rm SSP}(\lambda|\tau_{\rm age}^{i})P_{\rm SFH}(\tau_{\rm
age}^{i}|t),$

where in Eq. 6, $P_{\rm SFH}(\tau_{\rm age}^{i}|t)$ is the fraction of the
stars that survive until time $t$ that have ages between the boundaries of the
$i^{\rm th}$ age bin,

(7) $\displaystyle P_{\rm SFH}(\tau_{\rm age}^{i}|t)=\frac{1}{M_{\star}^{\rm
surv}(t)}\int_{t_{\rm lo}^{i}}^{t_{\rm hi}^{i}}{\rm
d}t^{\prime}\dot{M_{\star}}(t^{\prime})F_{\rm surv}(\tau_{\rm age}),$

with $t_{\rm hi/lo}^{i}=t^{\prime}\pm\Delta\tau_{\rm age}^{\rm i}/2,$ and

$\displaystyle M_{\star}^{\rm surv}(t)\equiv\int_{0}^{t}{\rm
d}t^{\prime}\dot{M_{\star}}(t^{\prime})F_{\rm surv}(\tau_{\rm age}).$ (8)

Equation 6 shows that the influence of SFH upon the SED of a composite stellar
population is calculable via a probability-weighted sum of the spectra of
SSPs. In the computations implemented in DSPS, the finite collection of
$L_{\rm SSP}(\lambda|\tau_{\rm age}^{i})$ are loaded into memory in advance
and thereafter held fixed, and so the dependence of $L_{\rm CSP}(\lambda)$
upon star formation history is entirely contained in the SFH-dependence of
$P_{\rm SFH}(\tau_{\rm age}^{i}|t).$ Using $\theta_{\rm SFH}$ to denote the
parameter(s) encoding the behavior of the star formation history of a galaxy,
we can see that calculating $\partial L_{\rm CSP}/\partial\theta_{\rm SFH}$
poses no significant technical obstacle for an autodiff-based implementation:
the integrations in Eqs. 6-8 can be evaluated numerically as a discretized
sum, and so gradients of the SED can be achieved provided that $F_{\rm
surv}(\tau_{\rm age})$ and $\dot{M_{\star}}(t)$ admit a differentiable
implementation. The fitting function for $F_{\rm surv}(\tau_{\rm age})$
presented in Appendix A was developed for this purpose, and so calculating
derivates of $L_{\rm CSP}$ with respect to $\theta_{\rm SFH}$ rests on the
differentiability of the SFH model.

Figure 1: Stellar populations of diffstar galaxies. The top panel shows the
SFH of two different galaxies parameterized by the diffstar model, each of
which resides in a dark matter halo with $M_{\rm halo}=10^{12}{\rm M}_{\odot}$
at redshift zero. The blue curve shows the history of a galaxy that remains on
the main sequence for its entire lifetime; the red curve shows a galaxy that
experienced a quenching event that shut down its star formation at $z\approx
1.$ The middle panel illustrates the quenching function, $F_{\rm q}(t),$ used
in the diffstar model to logarithmically drop the SFR of a quenched galaxy to
some value, $q_{\rm drop},$ below the main sequence (see Eq. 3.1). The bottom
panel compares the distribution of stellar ages of these same two galaxies
observed at different redshifts as indicated in the legend.

All of the most widely-used parametric forms for SFH admit a straightforward
implementation in JAX, including flexible, piecewise-defined models in which
the SFH is calculated by interpolation between a set of control points in
time. In order to demonstrate $\partial L_{\rm CSP}/\partial\theta_{\rm SFH},$
we will use the diffstar model for SFH introduced in Alarcon et al. (2021).
The parameterization of diffstar was designed to have a close connection to
the physics of galaxy formation, and to be sufficiently flexible to
approximate the complex SFHs of galaxies predicted by contemporary simulations
in an unbiased fashion. In the basic physical picture of the diffstar model,
the accretion rate of gas, $\dot{M}_{\rm g},$ is proportional to the accretion
rate of the dark matter halo, $\dot{M}_{\rm halo};$ main sequence galaxies
transform accreted gas into stars over a gas depletion timescale, $\tau_{\rm
dep},$ with efficiency, $\epsilon_{\rm ms};$ and some galaxies experience a
quenching event that drops their SFR below the main-sequence rate by a
multiplicative factor, $F_{\rm q}.$ The key equations that define the diffstar
model are as follows:

$\displaystyle\dot{M}_{\star}(t)=\dot{M}_{\star}^{\rm ms}(t)\times F_{\rm
q}(t)$ $\displaystyle\dot{M}_{\star}^{\rm ms}(t)=\epsilon_{\rm
ms}\int_{0}^{t}{\rm d}t^{\prime}\dot{M}_{\rm g}(t^{\prime})F_{\rm
dep}(t^{\prime}|\tau_{\rm dep})$ (9) $\displaystyle\dot{M}_{\rm g}(t)=f_{\rm
b}\dot{M}_{\rm halo}(t),$

where $f_{\rm b}=\Omega_{\rm b}/\Omega_{\rm m}$ is the cosmic baryon fraction,
and $F_{\rm dep}(\tau_{\rm dep})$ controls how a fraction of the gas accreted
onto the halo is gradually transformed into stars. For $\dot{M}_{\rm
halo}(t),$ we use diffmah, a differentiable parametric model that accurately
captures the mass assembly of individual dark matter halos in both gravity-
only and hydrodynamic simulations (see Hearin et al., 2021a, for details). We
refer the reader to Alarcon et al. (2021) for a detailed description of the
diffstar model.

In the top panel of Figure 1, we show two different examples of SFHs in the
diffstar model, each pertaining to a galaxy that resides in a dark matter halo
with Milky Way mass at $z=0.$ The blue curve shows the SFH of a typical main
sequence galaxy for a halo of this mass, while the red curve shows a galaxy
that experienced a quenching event at $z\approx 1$ that shut down its specific
star formation rate to negligible levels. The behavior of the quenching
function, $F_{\rm q}(t),$ is shown in the middle panel of Figure 1, which
demonstrates the meaning of the $q_{\rm drop}$ parameter that controls the
severity of the quenching event at time $t_{\rm q};$ the quenching function is
implemented in terms of a triweight error function (see Appendix B). The
bottom panel of Fig. 1 compares the distribution of stellar ages of these same
two galaxies, with the solid and dashed curves showing results for different
times of observation as indicated in the legend. Even though the top panel of
Fig. 1 shows that the SFR of the main-sequence galaxy increases monotonically,
the specific star formation rate is a decreasing function of time, and so by
comparing dashed to solid curves of the same color in the bottom panel, we see
that the composite stellar population ages as the galaxy evolves at late
times.

In the top panel of Figure 2, we show the history of the broad-band colors of
these same two galaxies, plotted as a function of the time at which the galaxy
is observed. We calculated g-r and r-i colors in the rest-frame according to
Eq. 2, using transmission curves that mimic the filters of the Rubin
Observatory Legacy Survey of Space and Time (LSST, LSST Science Collaboration
et al., 2009). We can see in Fig. 2 that the colors of the galaxy redden with
time; this is consistent with the aging of the stellar population of the
galaxy visible in the bottom panel of Fig. 1. The bottom panel of Fig. 2 shows
the gradient of these color histories with respect to the parameter $q_{\rm
drop},$ defined as

$q_{\rm drop}\equiv\log_{10}\dot{M}_{\star}(t_{\rm
q})-\log_{10}\dot{M}_{\star}^{\rm ms}(t_{\rm q}),$

where $t_{\rm q}$ is the first moment at which the quenching function $F_{\rm
q}(t)$ reaches its lowest point (see the middle panel of Fig. 1 for a visual
demonstration).333We note that in diffstar, $F_{\rm q}(t)$ can be non-
monotonic because the parameterization of the model captures the possibility
that some quenched galaxies experience rejuvenated star formation, but here we
only demonstrate an example of permanent quenching.

When computing the gradient shown in the bottom panel of Fig. 2, for the
fiducial point in parameter space, $\theta_{\rm SFH},$ we use the star
formation history of the quenched galaxy shown with the red curves in Fig. 1,
for which $q_{\rm drop}=-2.5;$ larger values of $q_{\rm drop}$ correspond to
less extreme quenching events, and so the gradient of the galaxy color is
negative in the neighborhood of $t_{\rm q},$ since increasing $q_{\rm drop}$
corresponds to a less quenched, and thus generally bluer galaxy.

Figure 2: SFH-dependent colors and gradients. The top panel shows the history
of the rest-frame colors of these same two galaxies appearing in Figure 1,
with colors observed through different filters as indicated in the legend. The
bottom panel shows the gradient of the color of each galaxy with respect to
the parameter $q_{\rm drop}$ that describes the severity of the quenching
event. Larger values of $q_{\rm drop}$ correspond to less extreme quenching
events, and so the y-axis in the bottom panel is negative in the neighborhood
of the quenching time, $t_{\rm q},$ since increasing $q_{\rm drop}$ produces a
galaxy that is less quenched, and is thus generally bluer. We note that the
gradient in the bottom panel was not calculated based on finite differencing
methods, but was instead computed with the JAX autodiff library. See text for
details.

###  3.2 Metallicity

In the previous section, we described the techniques required to make
predictions for galaxy photometry that are differentiable with respect to star
formation history parameters, and illustrated these techniques in Figure 2
with a particular gradient of the rest-frame color history of a galaxy. When
calculating galaxy photometry in §3.1, we made the simplifying assumption that
all the stars in the galaxy have the same metallicity; from this assumption,
Equation 6 follows from the discretization of Equation LABEL:eq:csp. We will
now relax this assumption, and instead suppose that the metallicity of the
stars in a galaxy are described by a parameterized Metallicity Distribution
Function (MDF), $P_{\rm met}(Z).$

For a composite stellar population with a diversity in metallicity, the
discretization of Equation 6 becomes

$\displaystyle L_{\rm CSP}(\lambda|t)=\sum_{i,j}$ $\displaystyle L_{\rm
SSP}(\lambda|\tau_{\rm age}^{i},Z^{j})$ $\displaystyle\times P_{\rm
SSP}(\tau_{\rm age}^{i},Z^{j}|t).$

In Eq. 3.2, the quantity $P_{\rm SSP}$ encodes the joint dependence of the SSP
weights upon both the star formation history parameters, $\theta_{\rm SFH},$
and the parameters regulating chemical evolution, $\theta_{\rm Z}.$ The
generalization of Equation 6 to Equation 3.2 folds in an extra axis of
summation over a new dimension in the data block storing the SEDs of the SSPs;
by itself, this poses no technical obstacle for maintaining our autodiff-based
implementation with JAX, and so the ability to calculate gradients of $L_{\rm
CSP}$ with respect to $\theta_{\rm Z}$ rests on the differentiability of the
MDF, $\partial P_{\rm SSP}/\partial\theta_{\rm Z}.$

In principle, differentiable MDF implementations are possible even for codes
such as flexCE (Andrews et al., 2017) or VICE (Johnson & Weinberg, 2020) that
numerically solve systems of differential equations in which the influence of
$\theta_{\rm SFH}$ and $\theta_{\rm Z}$ are coupled, since JAX now supports
several algorithms for numerically integrating ODEs. We discuss this future
extension of DSPS in §4.2, but for our present purposes, we will demonstrate
the differentiable influence of metallicity on SEDs using two simpler MDFs. In
§3.2.1, we illustrate gradients of metallicity parameters for an MDF that is a
clipped Gaussian with uncorrelated scatter, and in §3.2.2 we show gradients of
an MDF that empirically captures correlations between stellar metallicity and
age.

####  3.2.1 Uncorrelated Metallicity Distribution Function

The single-metallicity assumption used in §3.1 essentially assumes the MDF is
a delta function centered on one of the particular metallicities for which
$L_{\rm SSP}(\lambda)$ has been pre-computed. As a simple generalization, in
this section we consider an MDF characterized by a clipped Gaussian
distribution with scatter that is uncorrelated with any of the other
parameters in the SPS model. Under this assumption, the two-dimensional
probability distribution $P_{\rm SSP}(\tau_{\rm age},Z|\theta_{\rm
SFH},\theta_{\rm Z})$ becomes separable into two one-dimensional PDFs, $P_{\rm
SFH}(\tau_{\rm age}|\theta_{\rm SFH})$ and $P_{\rm met}(Z|\theta_{\rm Z}).$ In
the equation below, we will rewrite Eq. 6 for this MDF, now explicitly
including the dependence upon the SPS model parameters for clarity:

$\displaystyle L_{\rm CSP}(\lambda$ $\displaystyle|t,\theta_{\rm
SFH},\theta_{\rm Z})=$ $\displaystyle M_{\star}^{\rm
surv}(t)\cdot\sum_{i,j}L_{\rm SSP}(\lambda|\tau_{\rm age}^{i},Z^{j})$
$\displaystyle\times P_{\rm SFH}(\tau_{\rm age}^{i}|t,\theta_{\rm SFH})\cdot
P_{\rm met}(Z^{j}|\theta_{\rm Z}).$

From Eq. 3.2.1, we can see that the ability to calculate exact gradients
$\partial L_{\rm CSP}/\partial\theta_{\rm Z}$ rests on the differentiability
of the MDF, $\partial P_{\rm met}/\partial\theta_{\rm Z}.$

We parameterize $P_{\rm met}$ using the triweight kernel,
$\mathcal{T}(x|\mu,\sigma),$ which is very similar to a Gaussian distribution
centered at $\mu$ with spread $\sigma$ that has been clipped to zero for
$|{x-\mu}|\geq 3\sigma$ (see Appendix B for details). In capturing the MDF
with a triweight kernel, we will use $\log_{10}(Z/{\rm Z}_{\odot})$ as our
independent variable, so that $P_{\rm met}(Z|\theta_{\rm Z})$ is essentially a
clipped log-normal distribution with uncorrelated scatter, and $\theta_{\rm
Z}=\left\\{\mu_{\rm mdf},\sigma_{\rm mdf}\right\\}.$ Figure 3 gives a simple
demonstration of the SED of a galaxy derived from this triweight-based
MDF.444Note that the SEDs we show in this section include a contribution from
both starlight as well as emission from nebular gas, and so emission lines are
visibly present in the spectra in Figure 3. However, all SPS parameters
controlling the physics of nebular emission are held fixed in this section; we
postpone discussion of gradients with respect to the parameters of nebular
emission until §3.3. For the SFH of the galaxy shown in Fig. 3, we used the
same main sequence galaxy at $z=0$ shown in the figures in §3.1; to calculate
the SED of this example galaxy, we assumed an MDF centered at $\mu_{\rm
mdf}=\log_{10}(Z/{\rm Z}_{\odot})=-0.3,$ with spread $\sigma_{\rm mdf}=0.25;$
the SSP of the smallest and largest metallicities that contribute to the MDF
are color-coded as indicated in the legend of Figure 3.

Figure 3: Galaxy SED with simple MDF. The black curve shows the SED of a main
sequence galaxy at $z=0$ derived with a simple metallicity distribution
function (MDF) centered at $\mu_{\rm mdf}=\log_{10}Z/{\rm Z}_{\odot}=-0.3$
with uncorrelated scatter of $\sigma_{\rm mdf}=0.25$ dex. The purple and
orange curves show the SED of the SSPs with the smallest and largest
metallicities that contribute to this MDF, respectively. The black curve is
calculated as the MDF-weighted sum of all SSPs with metallicity in between the
two bracketing cases.

In the top panel of Figure 4, we show the history of the broad-band colors of
this same main sequence galaxy, plotted as a function of the time at which the
galaxy is observed. The orange curve in Figure 4 lies above the purple for
most of cosmic time, so the higher-metallicity galaxy has redder rest-frame
optical colors than the lower-metallicity galaxy, with the exception of the
high-redshift behavior for the case of g-r color. This general trend is also
apparent in the bottom panel of Figure 4, which shows the gradient of these
color histories with respect to the parameter $\mu_{\rm mdf}.$

Figure 4: Metallicity-dependent colors and gradients. The top panel shows the
history of the rest-frame colors of the same main sequence galaxy shown in
Figure 3. The bottom panel shows the gradient of the color with respect to
$\mu_{\rm mdf},$ the mean metallicity of the log-normal MDF.

####  3.2.2 Age-Dependent Metallicity Distribution Function

Figure 5: Age-dependent MDF. The two shaded histograms in the top panel show
the metallicity distribution function (MDF) for stellar populations of
different ages as indicated in the legend. At fixed $\tau_{\rm age},$ a
triweight kernel is used to model the MDF as a clipped log-normal. The bottom
panel shows the triweight cumulative kernel that smoothly transitions the
central value of the MDF from a high-metallicity value for newly formed stars,
$\mu_{\rm mdf}^{\rm young},$ to a low-metallicity value for the oldest stars
in the galaxy, $\mu_{\rm mdf}^{\rm old}.$

The calculations in §3.2.1 assumed that the MDF is a clipped log-normal
distribution with uncorrelated scatter; this allowed us to simplify the
expression for $L_{\rm CSP}(\lambda),$ Eq. 3.2, by assuming that $P_{\rm
SSP}(\tau_{\rm age},Z)$ is the product of two decoupled distributions, $P_{\rm
SFH}(\tau_{\rm age})$ and $P_{\rm met}(Z).$ In this section, we will relax
this assumption with a simple empirical model for the MDF that captures
possible correlations between stellar age and metallicity. We will maintain
the use of a clipped log-normal distribution for metallicity, but allow for
$\tau_{\rm age}$-dependence of the center of the MDF:

$P_{\rm SSP}(\tau_{\rm age},Z)=P_{\rm SFH}(\tau_{\rm age})\cdot P_{\rm
met}(Z|\tau_{\rm age}).$

We will use a triweight cumulative kernel, $\mathcal{T}_{\rm
erf}(\log_{10}\tau_{\rm age}),$ to capture the dependence of $\mu_{\rm mdf}$
upon age, where $\mathcal{T}_{\rm erf}(x)$ is the integral of a triweight
kernel, as defined in Appendix B. Thus $\mathcal{T}_{\rm erf}$ smoothly
transitions the value of $\mu_{\rm mdf}$ from an young-age value, $\mu_{\rm
mdf}^{\rm young},$ to an old-age value, $\mu_{\rm mdf}^{\rm old},$ and at
fixed $\tau_{\rm age},$ the MDF has the same triweight kernel shape used in
§3.2.1. Figure 5 gives a visual demonstration of $P_{\rm met}(Z|\tau_{\rm
age}).$

For this MDF, Equation 3.2.1 becomes:

$\displaystyle L_{\rm CSP}(\lambda|t)=\sum_{i,j}$ $\displaystyle L_{\rm
SSP}(\lambda|\tau_{\rm age}^{i},Z^{j})$ $\displaystyle\times P_{\rm
met}(Z^{j}|\tau_{\rm age}^{i})\cdot P_{\rm SFH}(\tau_{\rm age}^{i}|t).$

In the top panel of Figure 6, we show how the history of the broad-band galaxy
color depends upon $\tau_{\rm age}$-dependent metallicity. For a galaxy with
the same main-sequence SFH studied above, each colored curve in the top panel
of Fig. 6 shows the history of r-i color in the rest frame for a different
combination of $\mu_{\rm mdf}^{\rm young}$ and $\mu_{\rm mdf}^{\rm old}.$ The
bottom panel of Fig. 6 shows the gradient of r-i color with respect to each
parameter.

Figure 6: Age-dependent MDF colors and gradients. The top panel shows how the
history of r-i color for stellar populations with different $\tau_{\rm
age}$-dependent metallicity as indicated in the legend. The bottom panel shows
the autodiff-based gradient of r-i color with respect to each of the two
parameters controlling the shape of MDF.

###  3.3 Nebular Emission

For the results calculated in the previous sections, the spectra of the simple
stellar populations, $L_{\rm SSP}(\lambda),$ include a contribution from both
starlight as well as emission from nebular gas, but all parameters pertaining
directly to the nebulae were held fixed, so that our SSPs could be described
by a two-dimensional grid of $\left\\{\tau_{\rm age}^{i},Z^{j}\right\\}.$ In
this section, we examine the differentiable influence on the composite SED of
nebular emission parameters, $\theta_{\rm neb},$ focusing here on the impact
of the ionization state of the nebular gas, parameterized by $u_{\rm
s}=\log_{10}U_{\rm s}.$ Adapting the differentiable techniques used in §3.2 to
describe the MDF, we capture the effect of a probability distribution of
$P_{\rm neb}(\theta_{\rm neb})$ using a triweight kernel for the distribution
of $u_{\rm s}:$

$\displaystyle L_{\rm CSP}(\lambda$ $\displaystyle|t,\theta_{\rm
SFH},\theta_{\rm Z},\theta_{\rm neb})=$ $\displaystyle M_{\star}^{\rm
surv}(t)\cdot\sum_{i,j,k}L_{\rm SSP}(\lambda|\tau_{\rm age}^{i},Z^{j},u^{k})$
$\displaystyle\times P_{\rm SFH}(\tau_{\rm age}^{i}|t,\theta_{\rm SFH})\cdot
P_{\rm met}(Z^{j}|\theta_{\rm Z})\cdot P_{\rm neb}(u^{k}|\theta_{\rm neb}).$

We use python-fsps to tabulate our collection of SSPs, $L_{\rm
SSP}(\lambda|\tau_{\rm age}^{i},Z^{j},u^{k}),$ and we use Equation 3.3
together with the techniques described in §2.1 to calculate how emission line
strength varies as a function of $u_{\rm s}.$ Figure 7 shows the results of
our calculations. For the [OIII] emission line at $\lambda=5000{\rm\AA},$ we
plot the equivalent width, $W_{\rm[OIII]},$ as a function of cosmic time for
the same star-forming and quenched galaxies defined in §3.1. The top panel of
Fig. 7 shows emission line histories for the three different values of $u_{\rm
s}$ indicated in the legend. As discussed in §3.1, the specific star formation
rate of this galaxy is a decreasing function of time, and so the strength of
[OIII] emission decreases as the galaxy evolves. We also see that the time
$t_{\rm q}$ leaves a strong signature on the emission line history of the
quenched galaxy; this is sensible, since [OIII] emission comes from the HII
regions surrounding young stars, and the SFR of the quenched galaxy is shut
down by the time $t_{\rm q}$ (see Figure 1).

The bottom panel shows the gradient of the emission line history, using
$u_{\rm s}=-2$ as our fiducial point in parameter space, with scatter in
$\sigma_{u_{\rm s}}=0.25.$ The positive sign of the gradient is sensible,
since more highly ionized gas produces stronger [OIII] emission; for the red
curve, the dropping of the gradient to zero after $t_{\rm q}$ is also
expected, since the ionization parameter $u_{\rm s}$ has no appreciable effect
on the composite SED after nearly all of star formation has been shut down.

Figure 7: Effect of nebular ionization on [OIII] emission. We show the history
of the equivalent width of the [OIII] emission line for same main-sequence and
quenched galaxies shown in Fig. 1. The top panel shows emission line histories
for the three different values of $u_{\rm s},$ as indicated in the legend; the
bottom panel shows the gradient of the emission line history with respect to
$u_{\rm s}.$

###  3.4 Dust Attenuation

Some of the starlight and nebular emission of a galaxy is obscured by dust,
and the associated reduction in flux is captured by the attenuation curve,
$A_{\lambda},$ defined as

$\displaystyle L_{\rm obs}(\lambda)\equiv F_{\rm att}(\lambda)\times L_{\rm
CSP}(\lambda),$ (14)

where $F_{\rm att}(\lambda)=10^{-0.4A_{\lambda}}.$ For the functional form of
$A_{\lambda}$ presented here, we use the same parameterization as in Salim et
al. (2018), defined as follows:

$\displaystyle A_{\lambda}=A_{\rm V}\cdot k_{\lambda}$ (15) $\displaystyle
k_{\lambda}=k_{0}(\lambda)\cdot\left(\frac{R_{\rm V,{\rm mod}}}{R_{\rm V,{\rm
Cal}}}\right)\cdot\left(\frac{\lambda}{\lambda_{0}}\right)^{\gamma}+D_{\lambda}$
$\displaystyle R_{\rm V,{\rm Cal}}/R_{\rm V,{\rm mod}}=(1+R_{\rm V,{\rm
Cal}})\cdot C^{\delta}-R_{\rm V,{\rm Cal}}$

where in Eq. 15, the quantities $\lambda_{0}=5500{\rm\AA},$ and $R_{\rm V,{\rm
Cal}}=4.05,$ and $C=4400/5500$ are constants; the quantity $k_{0}$ is defined
piecewise to be:

$\displaystyle k_{0}(\lambda)\equiv\begin{cases}k_{\lambda,{\rm
Cal}},&\lambda<0.15\micron\\\ k_{\lambda,{\rm
L+02}},&\lambda>0.15\micron\end{cases}$ (16)

with $k_{\lambda,{\rm Cal}}$ taken from Calzetti et al. (2000), and
$k_{\lambda,{\rm L+02}}$ taken from Leitherer et al. (2002); finally, the
quantity $D_{\lambda}$ is the UV attenuation bump, which we model based on a
Drude profile defined as:

$\displaystyle D_{\lambda}\equiv\frac{E_{\rm
b}(\lambda\Delta\lambda)^{2}}{\left(\lambda^{2}-\lambda_{\rm
b}^{2}\right)^{2}+(\lambda\Delta\lambda)^{2}},$ (17)

where $\Delta\lambda=350{\rm\AA}$ and $\lambda_{\rm b}=2175{\rm\AA}$ are
constants, and we take $E_{\rm b}=-1.9\delta+0.85$ as in Kriek & Conroy
(2013). See the  attenuation_kernels module in DSPS for implementation
details.

Figure 8: Dust attenuation colors and gradients. The top panel shows the SED
of the same main sequence galaxy illustrated in Figure 1; here we show this
SED after applying attenuation curves, $A_{\lambda},$ defined by different
values of $\gamma$ as indicated in the legend (see Eq. 15). Each attenuation
curve has $A_{\rm V}$ equal to unity, except for the black curve, which shows
an unattenuated SED with $A_{\rm V}=0.$ The middle panel shows the i-z color
history of the galaxy for different $\gamma,$ and the bottom panel shows the
autodiff-based gradient of color with respect to $\gamma.$

We can see from Eqns. 14-17 that the effect of $A_{\lambda}$ upon the observed
SED is calculable analytically, and so in contrast to §3.3, predicting the
influence of dust attenuation on $L_{\rm CSP}(\lambda)$ does not require
expanding the dimension of the SSP grid. The top panel of Figure 8 shows
$L_{\rm CSP}(\lambda)$ for the same main sequence galaxy studied in the
previous sections; different colored curves show SEDs with attenuation curves
defined by different values of the parameter $\gamma$ as indicated in the
legend; each attenuation curve has $A_{\rm V}$ equal to unity, except for the
black curve, which shows an unattenuated SED with $A_{\rm V}=0.$ From the top
panel we can see that dust attenuation reduces the observed flux to levels
that become more pronounced at shorter wavelengths, and that smaller (more
negative) values of $\gamma$ produce generally stronger levels of attenuation,
particularly at shorter wavelengths.

In the middle panel of Fig. 8, we show the i-z color in the observer frame for
the same galaxy observed at different times. We can see that smaller values of
$\gamma$ correspond to redder colors, as expected from the top panel, and also
that the color history of the galaxy has a complex time dependence produced by
different spectral features that redshift in and out of the wavelength range
of the i-z filter. In the bottom panel of Fig. 8, we show the gradient of the
color of the galaxy with respect to $\gamma,$ which we have calculated for a
fiducial value of $\gamma_{\rm fid}=-0.25.$ The sign of the gradient is
negative, again since larger values of $\gamma$ produce optical colors that
are less reddened by dust obscuration. The SED of the galaxy is generally
bluer at earlier times when specific star formation rates are larger, and so
we see an overall stronger effect of attenuation on the color of the galaxy at
earlier times. Along the same lines, the effect of $\gamma$ is more pronounced
through the bluer g-r filter relative to the redder i-z filter.

####  3.4.1 $\tau_{\rm age}$-dependent attenuation

Some models of dust attenuation account for the natural physical expectation
that the column depth and/or grain-size distribution of dust surrounding star-
forming regions is distinct from the rest of the interstellar medium (e.g.,
Charlot & Fall, 2000). In this section, we adapt the attenuation model
described above to incorporate this effect.

Figure 9: Colors and gradients of $\tau_{\rm age}$-dependent dust. Analogous
to Fig. 8, only here $A_{\rm V}(\tau_{\rm age})$ and $\gamma(\tau_{\rm age})$
depend upon stellar age. The top panel shows the $\tau_{\rm age}$-dependence
of the attenuation curve, $F_{\rm att}(\lambda|\tau_{\rm age}),$ for the same
main sequence galaxy illustrated in the previous sections. The middle panel
shows the i-z color history of the galaxy for different combinations of
$A_{\rm V}^{\rm young}$ and $A_{\rm V}^{\rm old}$, as indicated in the legend.
The bottom panel shows the autodiff-based gradient of r-i color with respect
to each of the four parameters controlling $A_{\rm V}(\tau_{\rm age})$ and
$\gamma(\tau_{\rm age}).$

If the attenuation curve, $F_{\rm att}(\lambda),$ depends upon stellar age,
then our equation for the composite SED of a galaxy becomes:

$\displaystyle L_{\rm CSP}$ $\displaystyle(\lambda|t,\theta_{\rm
SFH},\theta_{\rm Z})=$ $\displaystyle M_{\star}^{\rm
surv}(t)\cdot\sum_{i,j}L_{\rm SSP}(\lambda|\tau_{\rm age}^{i},Z^{j})F_{\rm
att}(\lambda|\tau_{\rm age}^{i},\theta_{\rm dust})$ $\displaystyle\times
P_{\rm SFH}(\tau_{\rm age}^{i}|t,\theta_{\rm SFH})\cdot P_{\rm
met}(Z^{j}|\theta_{\rm Z}).$

For our model of $F_{\rm att}(\lambda|\tau_{\rm age},\theta_{\rm dust}),$ we
use the same parameterization defined by Eqs. 15-16, only we allow the
parameters $A_{\rm V}$ and $\gamma$ to take on different values for young and
old stars. For the functional form capturing the $\tau_{\rm age}$-dependence,
we use a triweight error function, $\mathcal{T}_{\rm erf}(\log_{10}\tau_{\rm
age}),$ to smoothly transition the value of $A_{\rm V}(\tau_{\rm age})$ and
$\gamma(\tau_{\rm age})$ from distinct asymptotic values for young and old
stars (see Appendix B for the definition of $\mathcal{T}_{\rm erf}$).555Note
that in §3.2.2 we used essentially the same technique based on the triweight
error function to capture the $\tau_{\rm age}$-dependence of the metallicity
distribution function.

Figure 9 illustrates how $\tau_{\rm age}$-dependent attenuation impacts the
color history of our fiducial main-sequence galaxy. The top panel shows how
the attenuation curve, $F_{\rm att}(\lambda|\tau_{\rm age}),$ has joint
dependence upon wavelength and stellar age. The middle panel of Fig. 9 shows
the history of i-z color in the observer frame for the same main-sequence
galaxy shown in Fig. 8; color histories of galaxies with different
combinations of $A_{\rm V}^{\rm young}$ and $A_{\rm V}^{\rm old}$ are shown
with different curves, as indicated in the legend. In the bottom panel of Fig.
9, we show the gradients of r-i color with respect to each of the four
parameters that control the $\tau_{\rm age}$-dependence of dust attenuation.
From both the middle and bottom panels we can see that broad-band optical
color is more sensitive to the values of $A_{\rm V}$ and $\gamma$ in the dust
surrounding young populations relative to old; this is sensible, since younger
stars emit bluer light that is more susceptible to obscuration by dust
relative to the redder light emitted by older stars. As discussed further in
§4, in future work we will leverage the flexibility of $\tau_{\rm
age}$-dependent dust models to capture the statistical correlations between
attenuation and physical galaxy properties such as stellar mass and SFR, which
are widely reported in observations (see Salim & Narayanan, 2020, for a
review), and predicted by numerous hydrodynamical simulations of cosmological
volumes (e.g., Hahn et al., 2021).

##  4 Discussion

###  4.1 Computational Benefits of JAX Implementation

In this paper, we have developed DSPS: a JAX-based implementation of many of
the standard calculations of stellar population synthesis. The individual
model evaluations in DSPS are quite efficient, as our code inherits the
performance benefits of JAX, which is one of the highest-performance libraries
used in contemporary deep learning. Using the automated vectorization features
of JAX, we have carried out a simple benchmarking experiment motivated by the
simulation-based forward-modeling applications that are our principal science
target. We performed the following operations in each iteration of our
benchmarking experiment, repeating the sequence ten times and taking the
median of the runtimes as our timing estimate:

  1. 1.

We used diffstar to generate star formation histories for $10^{5}$ galaxies at
$z=0;$

  2. 2.

we used DSPS to calculate photometric fluxes through six LSST-like optical
bands;

  3. 3.

we applied a simple color cut to select galaxies with red colors in the
resulting population;

  4. 4.

we evaluated a histogram on $r$-band magnitude of the color-cut sample.

This set of computational steps mimics the basic operations used in a
simulation-based forward model of the luminosity function of a color-selected
galaxy sample. On the Swing machine at Argonne, the computation with DSPS took
3.1 CPU-minutes on a single Intel Xeon E5 processor; the same computation
requires $\sim 85$ CPU-minutes using FSPS (Conroy et al., 2009b; Conroy &
Gunn, 2010). Repeating the calculation using an NVIDIA A100 GPU on Swing
required only 2.4 seconds of wall-clock time. Thus relative to standard SPS
codes, using DSPS provides a factor of $\sim 20$ speedup on a CPU, and a
factor of over 1000 on a modern GPU.666In this benchmarking exercise, we have
used tabulations of $L_{\rm SSP}(\lambda)$ based on MILES (Falcón-Barroso et
al., 2011), but we note that these absolute runtimes depend on the wavelength
resolution of the underlying collection of SSPs.

The availability of gradient information is another performance benefit of the
JAX-based implementation of DSPS. This feature should be useful in a wide
range of applications. As we have already mentioned, when fitting the SEDs of
individual galaxies, gradient-based optimization and inference algorithms
substantially outperform other methods, particularly in high dimensions.
Gradient information is also extremely useful when conducting a sensitivity
analysis; for example, tracing the trajectory of exact gradients simplifies
the task of identifying degeneracy directions across the parameter space of a
model, and the availability of second-order gradients enables methods using
the Fisher information matrix. We also expect that the property of
differentiability will play a key role in simulation-based forward models of
the galaxy--halo connection. In such applications, the target data vectors are
summary statistics of an entire galaxy population, whereas in this work we
have only demonstrated gradients pertaining to individual galaxies. In
forward-modeling applications targeting large-scale structure observables, we
will use the techniques presented in Hearin et al. (2021b) to propagate the
gradients of individual SEDs through to the one-point estimators used to
measure the luminosity function and the distribution of galaxy colors, as well
as the two-point estimators of the clustering and lensing of galaxy samples
selected by their photometry and/or emission lines.

Of course, direct reimplementation of SPS computations in JAX is not the only
method by which differentiable SED predictions can be achieved. For example,
in Alsing et al. (2020), the authors trained several neural networks (referred
to collectively as SPECULATOR) to approximate the SED and photometry
predictions of two different SPS models, achieving a surrogate function that
is both differentiable and provides a factor of $10^{3}-10^{4}$ speedup
relative to exact evaluation with standard SPS codes. In this approach to
accelerating the computations of stellar population synthesis, an existing SPS
library is treated as a black box that supplies training data quantifying how
the galaxy SED responds to the input model parameters. The advantage of this
approach is that the end result represents essentially the extreme limit of
the computational efficiency that can be achieved for the specific model
prediction that is emulated. The disadvantage is that the human effort
required to identify the appropriate architecture and optimize the parameters
of the network can be quite considerable, and the associated labor must be
repeated each time there is even a small change to the configuration of the
analysis. Since the calculations implemented in DSPS are not emulated, but are
the same computations implemented in other SPS libraries, our performance
gains derive from the one-time effort of our JAX-based reimplementation; for
the same reason, SED calculations with DSPS use the exact SSPs and SFHs, and
are not reliant on approximations based on principal components and Gaussian
mixtures.

Of course, there is no conflict between the use of emulation methods and our
reimplementation of SPS in an autodiff library. In fact, quite the opposite is
true. In generating $\sim 10^{6}$ points of training data for an SPS emulator,
the performance gains quoted above translate into the reduction of wall-clock
times from an entire day to only a few minutes on a GPU. Beyond this simple
translation of our speedup factor into practical terms, the purely functional
software design of the JAX library enables a shift to the workflow of
emulation. Since the entirety of DSPS is implemented in JAX, it is
straightforward to replace specific components of the prediction pipeline with
a dedicated emulator. For example, in profiling DSPS as part of this work, we
identified one of the dominant bottlenecks of our SED predictions to be the
calculation of $P(\tau_{\rm age}|\theta_{\rm sfh}),$ the probability
distribution of stellar ages as a function of the SFH parameters (see §3.1 for
details). This implies that a neural network emulator of $P(\tau_{\rm
age}|\theta_{\rm sfh})$ has potential to provide a considerable performance
enhancement. We highlight that this alternative approach to emulation does not
require sacrificing the flexibility of the analysis. Whereas an end-to-end
photometry emulator must be retrained each time one varies the filters, by
contrast, once a surrogate function for $P(\tau_{\rm age}|\theta_{\rm sfh})$
has been trained, the performance gains apply directly to all downstream SED
predictions. The focused nature of this approach to emulation also simplifies
training, since simpler target functions are less challenging to emulate.

###  4.2 Current Limitations and Future Work

As described in §2, DSPS takes as its starting point a data block such as
$L_{\rm SSP}(\lambda|\tau_{\rm age}^{i},Z^{j})$ that represents the SEDs of a
collection of simple stellar populations. All of the calculations demonstrated
in §3 were based on SSPs for which a fixed IMF had been assumed, so that each
element of the input data block provides an interpolation table for the SED of
a homogeneous population of stars with age, $\tau_{\rm age}^{i},$ and
metallicity, $Z^{j}.$ The differentiable techniques in DSPS can work equally
well with SSP templates that include additional dependencies beyond
metallicity and age; for example, the calculations in §3.3 included additional
dependence of the SSPs on the ionization state of the nebular gas, $U_{\rm
s},$ so that the data block of SSPs has an extra dimension, $L_{\rm
SSP}(\lambda|\tau_{\rm age}^{i},Z^{j},U_{\rm s}^{k}).$

Depending on the physics under consideration, it may or may not be necessary
for the dependence of the SSP upon some parameter, $\theta,$ to appear in the
input data block. Consider, for example, $\theta_{\rm dust},$ the parameters
of the attenuation curves appearing in §3.4; in this case, the effect on the
SED is calculable analytically in JAX, and so there is no need for the data
block of the input SSPs to include a discretized pre-tabulation of the effect
of $\theta_{\rm dust}.$ For modeling ingredients in which it is
straightforward to implement the effect on the SED directly in JAX, this is
generally preferred, because with each new variable upon which the SSPs
depend, there is a multiplicative increase of the memory footprint required by
the SED calculation. In all cases, the required template spectra are not
provided by DSPS, and so users of DSPS must acquire these from another library
that supplies them.

There is no strict need for differentiable SED calculations to rely on fixed-
IMF SSPs; one could instead begin with an input data block that provides a
discretized tabulation of single-star SEDs, $L_{\rm SS}(\lambda|,\tau_{\rm
age}^{i},Z^{j},M_{\star}^{k}).$ Beginning with $L_{\rm SS}(\lambda)$ as the
starting point creates the opportunity to forward model the effect of
parameters encoding uncertainty in the IMF, or even stellar rotation and the
fraction of stellar binaries. The additional complicating factor in this case
is that while $\tau_{\rm age}$ is the natural independent variable appearing
in SPS calculations, for single-star spectra the natural variable is the
equivalent evolutionary phase (EEP, see, e.g., Dotter, 2016), and the
$\tau_{\rm age}$\--EEP relationship is highly non-trivial. Publicly available
libraries such as isochrones (Morton, 2015) can be used to calculate how EEP
depends upon $\tau_{\rm age}$ for stars of different mass and metallicity;
these calculations can in turn be used to train a targeted emulator for the
$\tau_{\rm age}$\--EEP relationship, which is all that would be required to
perform differentiable SED calculations using tabulations of $L_{\rm
SS}(\lambda)$ as a starting point. While this is beyond the scope of the
present paper, we will present ongoing results in this direction in a follow-
up paper that extends the functionality of the DSPS library.

In §3.2, we presented a simple model for stellar metallicity that captures
both diversity in the chemical composition of the stars in a galaxy, as well
as potential correlations between the metallicity and age of stars within a
composite population. Detections of such correlations have been reported in
observational data (Poggianti et al., 2001; Gallazzi et al., 2005), and are
expected on basic physical grounds. Of course, due to the well-known age--
metallicity degeneracy (Worthey, 1994), it is challenging to disentangle these
two influences on the SED, and so it may not be possible to constrain the
chemical makeup of a galaxy beyond its average metallicity, particularly when
the measurements are limited in resolution and/or wavelength range.
Nonetheless, the model in §3.2 improves upon the widely-used assumption of
constant metallicity, requires only one or two additional parameters, and
comes with a practically negligible loss of computational efficiency.

Future extensions of DSPS will include physical models for the chemical
evolution of a galaxy that complement the empirical approach taken in §3.2. We
are currently pursuing two distinct approaches to these extensions. The first
is motivated by Weinberg et al. (2017), who showed that the differential
equations governing the chemical evolution of a galaxy admit analytical
solutions for a wide range of functional forms of star formation history, even
when assumptions of instantaneous recycling are relaxed. The computational
advantages of this analytical approach are attractive, particularly in light
of the diversity in metallicity distribution functions (MDF) that can be
captured by the solution space. In the second approach, we directly
parameterize the MDF evolution of individual galaxies, seeking a functional
form that is sufficiently flexible to capture what is seen in the merger trees
of galaxies in hydrodynamical simulations. This essentially mirrors the effort
in Alarcon et al. (2021) to build the diffstar parameterization of star
formation history, and has the advantage of relaxing the equilibrium
assumptions used to derive the analytical solutions of the differential
equations of chemical evolution.

Throughout this paper, we focused on the diffstar model for demonstrating
autodiff-based gradients of galaxy SEDs; as shown in Alarcon et al. (2021),
this model is flexible enough to give an unbiased description of the SFHs in
the UniverseMachine and IllustrisTNG simulations. However, we note that
diffstar is not as flexible as piecewise-defined models (e.g., Leja et al.,
2019a), since diffstar essentially imposes constraints on the shape of SFHs
derived from basic scaling relations of the galaxy--halo connection. Since the
influence of SFH on broad-band optical colors is relatively simple and largely
insensitive to complex patterns of burstiness (Chaves-Montero & Hearin, 2020,
2021), then smooth models such as diffstar may be sufficient for purposes of
predicting and interpreting the data from large cosmological imaging surveys.
However, recent results indicate that the freedom offered by piecewise-defined
models offers significant benefits over some widely-used parametric SFH forms
(Lower et al., 2020), and so it remains to be seen whether a smooth model such
as diffstar can improve upon these limitations of traditional parametric
forms. Fortunately, there is no technical obstacle to implementing a range of
piecewise-defined models in a differentiable fashion; the behavior of such
models is essentially defined by some form of interpolation on a grid, which
can be naturally implemented with an autodiff library whether this operation
is carried out with linear interpolation (e.g., Chauke et al., 2018; Leja et
al., 2019b; Johnson et al., 2021), or via an ML algorithm such as Gaussian
Process (as in Iyer et al., 2019).

Some SPS models use truly non-parametric formulations of SFH that are defined
by randomly sampling from a pre-computed library of SFHs supplied by a galaxy
formation simulation (Finlator et al., 2007; Pacifici et al., 2012, 2015).
This class of models can also be formulated in a differentiable fashion by
adapting the GalSampler technique (Hearin et al., 2020): the simulated SFHs in
the library are selected with a parametrized function that controls the PDF
used in the weighted sampling. We are currently developing this SFH model
based on galaxy libraries generated with the Galacticus SAM (Benson, 2012),
and will explore this further in future work.

The current version of DSPS does not include several physical ingredients that
are commonly implemented in other widely-used libraries for stellar population
synthesis. For example, the models presented in §3.4 only capture how dust
absorbs optical and UV starlight, but not how the absorbed energy is reemitted
by the dust, and so DSPS does not yet have capability to make physically
realistic predictions for the SEDs of dusty star-forming galaxies in the
infrared. Exploring these effects is beyond the scope of this paper, since our
present focus is on demonstrating the scientific potential of autodiff-based
implementations of stellar population synthesis, but we will include these and
other physical ingredients in future releases of DSPS.

##  5 Summary

We conclude by summarizing our primary results:

  1. 1.

We have developed DSPS, https://github.com/ArgonneCPAC/dsps, a stellar
population synthesis code written in the JAX library for automatic
differentiation. Our software is available for installation with conda or pip.

  2. 2.

In applications of simulation-based forward modeling the galaxy--halo
connection, DSPS improves upon the computational performance of standard SPS
codes by a factor of 20 on a CPU, and by over a factor of 1000 on a GPU.

  3. 3.

We have detailed a set of techniques for formulating models of SPS in a manner
that makes the predictions naturally differentiable with autodiff,
demonstrating worked examples of gradients of the parameters of models of star
formation history in §3.1, metallicity in §3.2, nebular emission in §3.3, and
dust attenuation in §3.4.

  4. 4.

In Appendix A, we have developed new fitting functions for the fraction of
stellar mass that survives as a function of time, $F_{\rm surv}(t),$ and also
the fraction of a population locked up in stellar remnants, $F_{\rm rem}(t);$
our fitting functions accurately capture the IMF-dependence of $F_{\rm
surv}(t)$ and $F_{\rm rem}(t),$ and so their standalone implementations in
DSPS may be useful in other SPS applications.

## Acknowledgements

Special thanks to Ben Johnson, whose active support of FSPS has created a
public record on GitHub that has been an invaluable resource in developing
this work. APH thanks José Feliciano for Feliz Navidad.

We thank the developers of NumPy (Van Der Walt et al., 2011), SciPy (Jones et
al., 2001-2016), Jupyter (Ragan-Kelley et al., 2014), IPython (Pérez &
Granger, 2007), scikit-learn (Pedregosa et al., 2011), JAX (Bradbury et al.,
2018), conda-forge (conda-forge community, 2015), and Matplotlib (Hunter,
2007) for their extremely useful free software. While writing this paper we
made extensive use of the Astrophysics Data Service (ADS) and arXiv preprint
repository.

Work done at Argonne was supported under the DOE contract DE-AC02-06CH11357.
We gratefully acknowledge use of the Bebop cluster in the Laboratory Computing
Resource Center at Argonne National Laboratory.

## References

  * Abadi et al. (2015) Abadi, M., Agarwal, A., Barham, P., et al. 2015, TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems.  https://www.tensorflow.org/
  * Alarcon et al. (2021) Alarcon, A., Hearin, A., Becker, M., & Chaves-Montero, J. 2021, in prep
  * Alsing et al. (2020) Alsing, J., Peiris, H., Leja, J., et al. 2020, ApJS, 249, 5, doi: 10.3847/1538-4365/ab917f
  * Ambikasaran et al. (2014) Ambikasaran, S., Foreman-Mackey, D., Greengard, L., Hogg, D. W., & O’Neil, M. 2014
  * Andrews et al. (2017) Andrews, B. H., Weinberg, D. H., Schönrich, R., & Johnson, J. A. 2017, ApJ, 835, 224, doi: 10.3847/1538-4357/835/2/224
  * Aricò et al. (2021) Aricò, G., Angulo, R. E., Contreras, S., et al. 2021, MNRAS, 506, 4070, doi: 10.1093/mnras/stab1911
  * Arimoto & Yoshii (1987) Arimoto, N., & Yoshii, Y. 1987, A&A, 173, 23
  * Astropy Collaboration et al. (2013) Astropy Collaboration, Robitaille, T. P., Tollerud, E. J., et al. 2013, A&A, 558, A33, doi: 10.1051/0004-6361/201322068
  * Astropy Collaboration et al. (2018) Astropy Collaboration, Price-Whelan, A. M., Sipőcz, B. M., et al. 2018, AJ, 156, 123, doi: 10.3847/1538-3881/aabc4f
  * Audouze & Tinsley (1976) Audouze, J., & Tinsley, B. M. 1976, ARA&A, 14, 43, doi: 10.1146/annurev.aa.14.090176.000355
  * Baugh et al. (1996) Baugh, C. M., Cole, S., & Frenk, C. S. 1996, MNRAS, 283, 1361, doi: 10.1093/mnras/283.4.1361
  * Baydin et al. (2015) Baydin, A. G., Pearlmutter, B. A., Radul, A. A., & Siskind, J. M. 2015, JMLR, 18, 1
  * Benson (2012) Benson, A. J. 2012, New A, 17, 175, doi: 10.1016/j.newast.2011.07.004
  * Blanton & Roweis (2007) Blanton, M. R., & Roweis, S. 2007, AJ, 133, 734, doi: 10.1086/510127
  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., et al. 2018
  * Bradley et al. (2020) Bradley, L., Sipőcz, B., Robitaille, T., et al. 2020, astropy/photutils: 1.0.0, 1.0.0, Zenodo, doi: 10.5281/zenodo.4044744
  * Brinchmann & Ellis (2000) Brinchmann, J., & Ellis, R. S. 2000, ApJ, 536, L77, doi: 10.1086/312738
  * Broyden (1970) Broyden, C. G. 1970, IMA Journal of Applied Mathematics, 6, 76, doi: 10.1093/imamat/6.1.76
  * Bruzual & Charlot (2003) Bruzual, G., & Charlot, S. 2003, MNRAS, 344, 1000, doi: 10.1046/j.1365-8711.2003.06897.x
  * Bruzual A. (1983) Bruzual A., G. 1983, ApJ, 273, 105, doi: 10.1086/161352
  * Buzzoni (1989) Buzzoni, A. 1989, ApJS, 71, 817, doi: 10.1086/191399
  * Calzetti et al. (2000) Calzetti, D., Armus, L., Bohlin, R. C., et al. 2000, ApJ, 533, 682, doi: 10.1086/308692
  * Chabrier (2003) Chabrier, G. 2003, PASP, 115, 763, doi: 10.1086/376392
  * Charlot & Fall (2000) Charlot, S., & Fall, S. M. 2000, ApJ, 539, 718, doi: 10.1086/309250
  * Chauke et al. (2018) Chauke, P., van der Wel, A., Pacifici, C., et al. 2018, ApJ, 861, 13, doi: 10.3847/1538-4357/aac324
  * Chaves-Montero & Hearin (2020) Chaves-Montero, J., & Hearin, A. 2020, MNRAS, 495, 2088, doi: 10.1093/mnras/staa1230
  * Chaves-Montero & Hearin (2021) \---. 2021, MNRAS, 506, 2373, doi: 10.1093/mnras/stab1831
  * Choi et al. (2016) Choi, J., Dotter, A., Conroy, C., et al. 2016, ApJ, 823, 102, doi: 10.3847/0004-637X/823/2/102
  * conda-forge community (2015) conda-forge community. 2015, The conda-forge Project: Community-based Software Distribution Built on the conda Package Format and Ecosystem, Zenodo, doi: 10.5281/zenodo.4774216
  * Conroy (2013) Conroy, C. 2013, ARA&A, 51, 393, doi: 10.1146/annurev-astro-082812-141017
  * Conroy & Gunn (2010) Conroy, C., & Gunn, J. E. 2010, ApJ, 712, 833, doi: 10.1088/0004-637X/712/2/833
  * Conroy et al. (2009a) Conroy, C., Gunn, J. E., & White, M. 2009a, ApJ, 699, 486, doi: 10.1088/0004-637X/699/1/486
  * Conroy et al. (2009b) \---. 2009b, ApJ, 699, 486, doi: 10.1088/0004-637X/699/1/486
  * Diemer (2018) Diemer, B. 2018, ApJS, 239, 35, doi: 10.3847/1538-4365/aaee8c
  * Dotter (2016) Dotter, A. 2016, ApJS, 222, 8, doi: 10.3847/0067-0049/222/1/8
  * Drakos et al. (2021) Drakos, N. E., Villasenor, B., Robertson, B. E., et al. 2021, arXiv:2110.10703, arXiv:2110.10703.  https://arxiv.org/abs/2110.10703
  * Duane et al. (1987) Duane, S., Kennedy, A. D., Pendleton, B. J., & Roweth, D. 1987, Physics Letters B, 195, 216, doi: 10.1016/0370-2693(87)91197-X
  * Eldridge et al. (2017) Eldridge, J. J., Stanway, E. R., Xiao, L., et al. 2017, PASA, 34, e058, doi: 10.1017/pasa.2017.51
  * Euclid Collaboration et al. (2019) Euclid Collaboration, Knabenhans, M., Stadel, J., et al. 2019, MNRAS, 484, 5509, doi: 10.1093/mnras/stz197
  * Falcón-Barroso et al. (2011) Falcón-Barroso, J., Sánchez-Blázquez, P., Vazdekis, A., et al. 2011, A&A, 532, A95, doi: 10.1051/0004-6361/201116842
  * Finlator et al. (2007) Finlator, K., Davé, R., & Oppenheimer, B. D. 2007, MNRAS, 376, 1861, doi: 10.1111/j.1365-2966.2007.11578.x
  * Fioc & Rocca-Volmerange (1999) Fioc, M., & Rocca-Volmerange, B. 1999, arXiv e-prints, astro.  https://arxiv.org/abs/astro-ph/9912179
  * Fletcher (1970) Fletcher, R. 1970, The Computer Journal, 13, 317, doi: 10.1093/comjnl/13.3.317
  * Foreman-Mackey (2020) Foreman-Mackey, D. 2020, python-fsps v0.4.1rc1, Zenodo, doi: 10.5281/zenodo.4737461
  * Gallazzi et al. (2005) Gallazzi, A., Charlot, S., Brinchmann, J., White, S. D. M., & Tremonti, C. A. 2005, MNRAS, 362, 41, doi: 10.1111/j.1365-2966.2005.09321.x
  * Goldfarb (1970) Goldfarb, D. 1970, Math. Comp., 23, doi: 10.1090/S0025-5718-1970-0258249-6
  * GPy (2012) GPy. 2012, GPy: A Gaussian process framework in python, http://github.com/SheffieldML/GPy
  * Hahn et al. (2021) Hahn, C., Starkenburg, T. K., Anglés-Alcázar, D., et al. 2021, arXiv:2106.09741, arXiv:2106.09741.  https://arxiv.org/abs/2106.09741
  * Harnois-Déraps et al. (2019) Harnois-Déraps, J., Giblin, B., & Joachimi, B. 2019, A&A, 631, A160, doi: 10.1051/0004-6361/201935912
  * Hearin et al. (2020) Hearin, A., Korytov, D., Kovacs, E., et al. 2020, MNRAS, 495, 5040, doi: 10.1093/mnras/staa1495
  * Hearin et al. (2021a) Hearin, A. P., Chaves-Montero, J., Becker, M. R., & Alarcon, A. 2021a, The Open Journal of Astrophysics, 4, 7, doi: 10.21105/astro.2105.05859
  * Hearin et al. (2021b) Hearin, A. P., Ramachandra, N., Becker, M. R., & DeRose, J. 2021b, arXiv:2112.08423, arXiv:2112.08423.  https://arxiv.org/abs/2112.08423
  * Heitmann et al. (2006) Heitmann, K., Higdon, D., Nakhleh, C., & Habib, S. 2006, ApJ, 646, L1, doi: 10.1086/506448
  * Hoffman & Gelman (2014) Hoffman, M., & Gelman, A. 2014, Journal of Machine Learning Research, 15, 1593
  * Hunter (2007) Hunter, J. D. 2007, Computing In Science & Engineering, 9, 90, doi: 10.1109/MCSE.2007.55
  * Iyer et al. (2019) Iyer, K. G., Gawiser, E., Faber, S. M., et al. 2019, ApJ, 879, 116, doi: 10.3847/1538-4357/ab2052
  * Johnson et al. (2021) Johnson, B. D., Leja, J., Conroy, C., & Speagle, J. S. 2021, ApJS, 254, 22, doi: 10.3847/1538-4365/abef67
  * Johnson & Weinberg (2020) Johnson, J. W., & Weinberg, D. H. 2020, MNRAS, 498, 1364, doi: 10.1093/mnras/staa2431
  * Jones et al. (2001-2016) Jones, E., Oliphant, T., Peterson, P., et al. 2001-2016, http://www.scipy.org
  * Kauffmann et al. (1999) Kauffmann, G., Colberg, J. M., Diaferio, A., & White, S. D. M. 1999, MNRAS, 303, 188, doi: 10.1046/j.1365-8711.1999.02202.x
  * Kingma & Ba (2014) Kingma, D. P., & Ba, J. 2014, arXiv:1412.6980, arXiv:1412.6980.  https://arxiv.org/abs/1412.6980
  * Kobayashi et al. (2020) Kobayashi, Y., Nishimichi, T., Takada, M., Takahashi, R., & Osato, K. 2020, Phys. Rev. D, 102, 063504, doi: 10.1103/PhysRevD.102.063504
  * Kochkov et al. (2021) Kochkov, D., Smith, J. A., Alieva, A., et al. 2021, arXiv e-prints, arXiv:2102.01010.  https://arxiv.org/abs/2102.01010
  * Kokron et al. (2021) Kokron, N., DeRose, J., Chen, S.-F., White, M., & Wechsler, R. H. 2021, MNRAS, 505, 1422, doi: 10.1093/mnras/stab1358
  * Korytov et al. (2019) Korytov, D., Hearin, A., Kovacs, E., et al. 2019, ApJS, 245, 26, doi: 10.3847/1538-4365/ab510c
  * Kriek & Conroy (2013) Kriek, M., & Conroy, C. 2013, ApJ, 775, L16, doi: 10.1088/2041-8205/775/1/L16
  * Kriek et al. (2009) Kriek, M., van Dokkum, P. G., Labbé, I., et al. 2009, ApJ, 700, 221, doi: 10.1088/0004-637X/700/1/221
  * Kroupa et al. (1993) Kroupa, P., Tout, C. A., & Gilmore, G. 1993, MNRAS, 262, 545, doi: 10.1093/mnras/262.3.545
  * Le Borgne et al. (2004) Le Borgne, D., Rocca-Volmerange, B., Prugniel, P., et al. 2004, A&A, 425, 881, doi: 10.1051/0004-6361:200400044
  * Leitherer et al. (2002) Leitherer, C., Li, I. H., Calzetti, D., & Heckman, T. M. 2002, ApJS, 140, 303, doi: 10.1086/342486
  * Leja et al. (2019a) Leja, J., Carnall, A. C., Johnson, B. D., Conroy, C., & Speagle, J. S. 2019a, ApJ, 876, 3, doi: 10.3847/1538-4357/ab133c
  * Leja et al. (2019b) Leja, J., Johnson, B. D., Conroy, C., et al. 2019b, ApJ, 877, 140, doi: 10.3847/1538-4357/ab1d5a
  * Lower et al. (2020) Lower, S., Narayanan, D., Leja, J., et al. 2020, ApJ, 904, 33, doi: 10.3847/1538-4357/abbfa7
  * LSST Science Collaboration et al. (2009) LSST Science Collaboration, Abell, P. A., Allison, J., et al. 2009, arXiv:0912.0201, arXiv:0912.0201
  * Maraston (2005) Maraston, C. 2005, MNRAS, 362, 799, doi: 10.1111/j.1365-2966.2005.09270.x
  * McClintock et al. (2019) McClintock, T., Rozo, E., Becker, M. R., et al. 2019, ApJ, 872, 53, doi: 10.3847/1538-4357/aaf568
  * McCulloch & Pitts (1943) McCulloch, W. S., & Pitts, W. 1943, The bulletin of mathematical biophysics, 5, 115
  * Modi et al. (2021) Modi, C., Lanusse, F., & Seljak, U. 2021, Astronomy and Computing, 37, 100505, doi: 10.1016/j.ascom.2021.100505
  * Morton (2015) Morton, T. D. 2015, isochrones: Stellar model grid package.  http://ascl.net/1503.010
  * Nishimichi et al. (2019) Nishimichi, T., Takada, M., Takahashi, R., et al. 2019, ApJ, 884, 29, doi: 10.3847/1538-4357/ab3719
  * Oke & Sandage (1968) Oke, J. B., & Sandage, A. 1968, ApJ, 154, 21, doi: 10.1086/149737
  * Pacifici et al. (2012) Pacifici, C., Charlot, S., Blaizot, J., & Brinchmann, J. 2012, MNRAS, 421, 2002, doi: 10.1111/j.1365-2966.2012.20431.x
  * Pacifici et al. (2015) Pacifici, C., da Cunha, E., Charlot, S., et al. 2015, MNRAS, 447, 786, doi: 10.1093/mnras/stu2447
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., et al. 2019, in Advances in Neural Information Processing Systems 32, ed. H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, & R. Garnett (Curran Associates, Inc.), 8024--8035
  * Paxton et al. (2011) Paxton, B., Bildsten, L., Dotter, A., et al. 2011, ApJS, 192, 3, doi: 10.1088/0067-0049/192/1/3
  * Paxton et al. (2013) Paxton, B., Cantiello, M., Arras, P., et al. 2013, ApJS, 208, 4, doi: 10.1088/0067-0049/208/1/4
  * Paxton et al. (2015) Paxton, B., Marchant, P., Schwab, J., et al. 2015, ApJS, 220, 15, doi: 10.1088/0067-0049/220/1/15
  * Pedregosa et al. (2011) Pedregosa, F., Varoquaux, G., Gramfort, A., et al. 2011, Journal of Machine Learning Research, 12, 2825
  * Pérez & Granger (2007) Pérez, F., & Granger, B. E. 2007, Computing in Science and Engineering, 9, 21, doi: 10.1109/MCSE.2007.53
  * Poggianti et al. (2001) Poggianti, B. M., Bridges, T. J., Mobasher, B., et al. 2001, ApJ, 562, 689, doi: 10.1086/323217
  * Ragan-Kelley et al. (2014) Ragan-Kelley, M., Perez, F., Granger, B., et al. 2014, in American Geophysical Union Fall Meeting Abstracts, Vol. D7
  * Ramachandra et al. (2021) Ramachandra, N., Valogiannis, G., Ishak, M., Heitmann, K., & LSST Dark Energy Science Collaboration. 2021, Phys. Rev. D, 103, 123525, doi: 10.1103/PhysRevD.103.123525
  * Rasmussen & Williams (2006) Rasmussen, C. E., & Williams, C. K. I. 2006, Gaussian Processes for Machine Learning
  * Salim et al. (2018) Salim, S., Boquien, M., & Lee, J. C. 2018, ApJ, 859, 11, doi: 10.3847/1538-4357/aabf3c
  * Salim & Narayanan (2020) Salim, S., & Narayanan, D. 2020, ARA&A, 58, 529, doi: 10.1146/annurev-astro-032620-021933
  * Salim et al. (2007) Salim, S., Rich, R. M., Charlot, S., et al. 2007, ApJS, 173, 267, doi: 10.1086/519218
  * Salpeter (1955) Salpeter, E. E. 1955, ApJ, 121, 161, doi: 10.1086/145971
  * Sawicki & Yee (1998) Sawicki, M., & Yee, H. K. C. 1998, AJ, 115, 1329, doi: 10.1086/300291
  * Schoenholz & Cubuk (2019) Schoenholz, S. S., & Cubuk, E. D. 2019, arXiv e-prints, arXiv:1912.04232.  https://arxiv.org/abs/1912.04232
  * Shanno (1970) Shanno, D. 1970, Math. Comp., 647, doi: 10.1090/S0025-5718-1970-0274029-X
  * Somerville & Primack (1999) Somerville, R. S., & Primack, J. R. 1999, MNRAS, 310, 1087, doi: 10.1046/j.1365-8711.1999.03032.x
  * Tinsley (1978) Tinsley, B. M. 1978, ApJ, 222, 14, doi: 10.1086/156116
  * Van Der Walt et al. (2011) Van Der Walt, S., Colbert, S. C., & Varoquaux, G. 2011, ArXiv:1102.1523
  * van Dokkum (2008) van Dokkum, P. G. 2008, ApJ, 674, 29, doi: 10.1086/525014
  * Villaescusa-Navarro et al. (2020) Villaescusa-Navarro, F., Anglés-Alcázar, D., Genel, S., et al. 2020, arXiv:2010.00619, arXiv:2010.00619.  https://arxiv.org/abs/2010.00619
  * Weinberg et al. (2017) Weinberg, D. H., Andrews, B. H., & Freudenburg, J. 2017, ApJ, 837, 183, doi: 10.3847/1538-4357/837/2/183
  * Wiener (1938) Wiener, N. 1938, American Journal of Mathematics, 60, 897–936, doi: 10.2307/2371268
  * Worthey (1994) Worthey, G. 1994, ApJS, 95, 107, doi: 10.1086/192096
  * Xiu (2010) Xiu, D. 2010, Numerical Methods for Stochastic Computations: A Spectral Method Approach

##  Appendix A Fitting Functions for Stellar Lifetimes and Remnants

Figure A1: Surviving stellar mass fraction. For a simple stellar population
that formed in a single burst, the quantity plotted on the y-axis, $F_{\rm
surv}(\tau_{\rm age}),$ describes the fraction of the initial mass of the
population that survives until the time $\tau_{\rm age}$ following the burst.
Due to the $M_{\star}$-dependence of stellar lifetimes, $F_{\rm
surv}(\tau_{\rm age})$ depends on the IMF in a manner shown by the different
colored curves in the figure. Solid curves show results calculated with the
FSPS library; dashed curves show the corresponding fitting function used in
DSPS. The bottom panel shows the residual error in the approximation supplied
by the fitting function. Figure A2: Remnant mass fraction. The quantity
$F_{\rm rem}(\tau_{\rm age})$ describes the fraction of the initial mass of a
simple stellar population that has been transformed into stellar remnants at
the time $\tau_{\rm age}.$ As in Figure A1, different colored curves show
results for different assumed IMFs as indicated in the legend; solid curves
show results derived with FSPS; dashed curves show the corresponding fitting
functions implemented in DSPS; residual errors are shown in the bottom panel.

The lifetime of a star depends sensitively on its initial mass. A massive O
star with $M_{\star}\gtrsim 50{\rm M}_{\odot}$ will explode as a core-collapse
supernova after a few million years, whereas an M star with mass
$M_{\star}\approx 0.1{\rm M}_{\odot}$ will continue to burn hydrogen for
trillions of years, far longer than a Hubble time. Thus after a time
$\tau_{\rm age}$ following some initial burst of star formation, only some
fraction of the stars that form continue to survive, $F_{\rm surv}(\tau_{\rm
age}),$ and the remaining mass is returned to the interstellar medium. Due to
the mass-dependence of stellar lifetimes, $F_{\rm surv}(\tau_{\rm age})$
depends upon the IMF.

In Figure A1, we show the general behavior of $F_{\rm surv}(\tau_{\rm age}),$
including its IMF-dependence, which we have calculated using the python-fsps
wrapper (Foreman-Mackey, 2020) of the FSPS library (Conroy et al., 2009b;
Conroy & Gunn, 2010). We show results for four different widely-used IMFs:
Salpeter (Salpeter, 1955), Kroupa (Kroupa et al., 1993), Chabrier (Chabrier,
2003), and van Dokkum (van Dokkum, 2008). Each colored curve is accompanied by
a dashed curve providing a parametric fit implemented in DSPS; the bottom
panel of Figure A1 shows the difference between the fitting function and the
result calculated by FSPS; the functional form we use to approximate $F_{\rm
surv}(\tau_{\rm age})$ is defined in terms of $\mathcal{T}_{\rm erf}(x),$ with
parameters specified in the surviving_mstar module of the DSPS package.

Most stars leave behind a remnant when they die, which is either a white
dwarf, a neutron star, or a black hole; the physical nature of the remnant
depends on the initial mass of the star. In Figure A2, we show $F_{\rm
rem}(\tau_{\rm age}),$ the fraction of stellar mass locked up in remnants,
plotted as a function of time for the same four IMFs shown in Figure A1. Each
solid each curve in Figure A2 is accompanied by a dashed curve illustrating a
triweight-based parameterized fitting function for $F_{\rm rem}(\tau_{\rm
age})$ defined in the remnant_mass module of the DSPS package.

##  Appendix B Triweight Kernel

Many of the calculations in this paper involve probability-weighted summations
defined by a clipped Gaussian PDF. In this section, we review an alternative
distribution based on the triweight kernel, $\mathcal{T},$ defined as:

$\displaystyle\mathcal{T}(x|\mu,\sigma)\equiv\begin{cases}\frac{35}{96}\left[1-(z/3)^{2}\right]^{3},&|z|\leq
3\\\ 0,&\text{otherwise}\end{cases}$ (19)

where $z\equiv(x-\mu)/\sigma.$ The comparison to a Gaussian is shown in Figure
B1. The two distributions have the same first and second moments, $\mu$ and
$\sigma,$ but differ in their higher-order moments. The coefficients appearing
in Eq. 19 are defined so that the function $\mathcal{T}(x)$ is $C^{\infty}$
across the real line; points with $|x-\mu|>3\sigma$ contribute formally zero
weight, and so using a triweight is essentially equivalent to using a clipped
Gaussian. However, there is a significant computational advantage to the
triweight: evaluating a Gaussian requires a special function evaluation that
that can be far slower on a GPU accelerator device in comparison to the small
number of elementary arithmetical operations required to evaluate Eq. 19.

Figure B1: Triweight kernel. The blue curve shows the probability density
function of a normal distribution centered at $\mu=0$ with second moment
$\sigma=1.$ The orange curve shows the triweight kernel, $\mathcal{T},$
defined in Eq. 19, again with $\mu=0$ and $\sigma=1$. The triweight function
$\mathcal{T}(x)$ vanishes at points beyond $|x-\mu|>3\sigma,$ is continuously
differentiable for all $x,$ and is highly performant on GPUs due to its
formulation in terms of elementary arithmetical operations.

Some calculations in this paper are formulated in terms of $\mathcal{T}_{\rm
erf}(x),$ the cumulative integral of Eq. 19, such as the $\tau_{\rm
age}$-dependent metallicity distribution function shown in §3.2.2. The
$\mathcal{T}_{\rm erf}$ function is also defined by a simple algebraic
expression that is also performant on GPUs:

$\displaystyle\mathcal{T}_{\rm
erf}(x|\mu,\sigma)\equiv\begin{cases}p(z)&|z|\leq 3\\\
0,&\text{otherwise}\end{cases}$ (20) $\displaystyle
p(z)=\frac{1}{2}+\frac{35}{96}z-\frac{35}{864}z^{3}+\frac{7}{2592}z^{5}-\frac{5}{69984}z^{7}$