FlowSN: Normalising Flows for Simulation-Based Inference under Realistic Selection Effects applied to Supernova Cosmology

Abstract

We present FlowSN, a statistical framework using simulation-based inference with normalising flows to account for selection effects in observational astronomy. Failure to account for selection effects can lead to biased inference on global parameters. An example is Malmquist bias, where detection limits result in a sample skewed towards brighter objects. In Type Ia supernova (SN Ia) cosmology, these selection effects can systematically shift the inferred posterior distributions of cosmological parameters, necessitating the development of robust statistical frameworks to account for the biases. Simulation-based inference enables us to implicitly learn probability distributions that are analytically intractable to calculate. In this work, we introduce a novel approach that employs a normalising flow to learn the non-analytic selected SN likelihood for a given survey from forward simulations, independent of the assumed cosmological model. The resulting likelihood approximation is incorporated into a hierarchical Bayesian framework and posterior sampling is performed using Hamiltonian Monte Carlo to obtain constraints on cosmological parameters conditioned on the observed data. The modular learnt likelihood approximation can be reused without retraining to evaluate different cosmological models, providing a key advantage over other simulation-based inference approaches. We demonstrate the performance of this methodology by training and testing the simulation-based inference technique using realistic LSST-like SNANA simulations for the first time. Our FlowSN approach yields accurate posterior estimates on cosmological parameters, including the dark energy equation of state $w_0$, that are an order of magnitude less biased than those obtained with conventional techniques and also exhibit improved frequentist calibration.

Figures (11)

• Figure 1: Plot illustrating the impact of selection effects in realistic SNANA survey simulations, modelling a spectroscopic LSST supernova sample. At high redshift, Malmquist bias causes the selected supernovae to systematically lie below the true underlying cosmology relation curve. The histograms are normalised to unit area, and the dashed curve represents the spectroscopic selection efficiency, corresponding to the mock 4MOST selection function used in plasticc.
• Figure 2: Graphical model illustrating the dependencies between global parameters $\bm{\Theta}=(\bm{C},M_0,\alpha,\beta,\bm{\theta}_c,\bm{\theta}_x)$ to be inferred, supernova-level latent variables $(\mu_s,z_s,m_0^s,m_s,c_s,x_s)$ and observed data shaded in grey $(\hat{d}_s,\hat{z}^{\text{hel}}_s,I_s)$. The blue dashed area demonstrates the dependencies learnt with the FlowSN normalising flow. Each parameter is clearly defined in Section \ref{['sec:method']}. Note we only have observed data for selected SNe with $I_s=1$ and train the FlowSN normalising flow using only these examples.
• Figure 3: Schematic showing how the parameters interface with the normalising flow that models $q_\phi (\bm{\hat{d}}_s|\,m^s_0,\hat{z}^{\text{hel}}_s\bm{\Sigma}_s,\bm{\Theta}_{\text{SN}})$ in order to approximate the auxiliary density $g$ used in the observed data likelihood. The three $u_0$ parameters are each sampled from unit normal distributions and propagated through $L$ transforms $f_l$ with learnable parameters $\bm{\phi}$. These transforms are conditioned on the expected apparent magnitude $m_0^s$ for a given supernova (at zero stretch and colour), observed heliocentric redshift $\hat{z}^{\text{hel}}_s$, the measurement covariances $\bm{\Sigma}_s$, and the supernova population parameters $\bm{\Theta}_{\text{SN}}$. The outputs of $f_L$ are the three observed summary statistics corresponding to the observed apparent magnitude $\hat{m}_s$, peak $B-V$ colour $\hat{c}_s$ and stretch $\hat{x}_s$.
• Figure 4: Learnt FlowSN three-dimensional likelihood approximation in comparison to the analytical solution and naive likelihood solution (that does not account for selection effects) for the simplified forward model. The conditional information is equal to $\bm{\psi}=(m_0,\hat{z}^{\text{hel}},\bm{\Sigma},\bm{\Theta}_{\text{SN}})$. We see as latent apparent magnitude $m_0$ gets dimmer, the analytical and naive likelihood solutions differ more in all three dimensions. Both the analytical and naive solutions are derived in Appendix \ref{['app:toy_proof']}.
• Figure 5: Posteriors on flat $w$CDM cosmological parameters and SN population parameters from 2000 simulated spectroscopic SNe using the simplified forward model. The top right plot shows constraints on different non-flat $\Lambda$CDM cosmological parameters inferred without retraining the FlowSN normalising flow. We plot the FlowSN posterior, the analytical likelihood solution posterior and the naive solution posterior that does not account for selection effects. Both the analytical and naive solutions are derived in Appendix \ref{['app:toy_proof']}. The true global parameter values used in the simulations are illustrated with straight lines. We see that the FlowSN and analytical posteriors agree very closely, while the naive posterior is biased, particularly for the cosmological parameters. Both the analytical and naive solutions are derived in Appendix \ref{['app:toy_proof']}. Full posteriors across eleven global parameters are shown in Figure \ref{['fig:toy_post_full']}.