Efficient Bayesian analysis of kilonovae and GRB afterglows with fiesta

Abstract: Gamma-ray burst (GRB) afterglows and kilonovae (KNe) are electromagnetic transients that can accompany binary neutron star (BNS) mergers. Therefore, studying their emission processes is of general interest for constraining cosmological parameters or the behavior of ultra-dense matter. One common method to analyze electromagnetic data from BNS mergers is to sample a Bayesian posterior over the parameters of a physical model for the transient. However, sampling the posterior is computationally costly, and because of the many likelihood evaluations needed in this process, detailed models are too expensive to be used directly in inference. In the present article, we address this problem by introducing fiesta, a python package to train machine learning (ML) surrogates for GRB afterglow and kilonova models that can accelerate likelihood evaluations. Specifically, we introduce extensive ML surrogates for the state-of-the-art GRB afterglow models afterglowpy and pyblastafterglow, as well as a new surrogate for the KN emission as modeled by the possis code. Our surrogates enable evaluation of the lightcurve posterior within minutes. We also provide built-in posterior sampling capabilities in fiesta that rely on the flowMC package which scale efficiently to higher dimensions when adding up to tens of nuisance sampling parameters. Because of its use of the JAX framework, fiesta also allows for GPU acceleration during both surrogate training and posterior sampling. We apply our framework to reanalyze AT2017gfo/GRB170817A and GRB211211A using our surrogates, thus employing the new pyblastafterglow model for the first time in Bayesian inference.