jax-SNANA: A GPU-accelerated Supernova Framework for the SBI Era
Abstract
We present jax-SNANA, a JAX reimplementation of the supernova analysis framework SNANA. jax-SNANA enables GPU-accelerated transient simulations and light curve fitting, supernova standardisation and cosmological inference. Running in single precision on an NVIDIA A100 GPU, jax-SNANA simulates more than 100,000 LSST-observed supernovae per second and fits 5,000 light curves per second after JIT compilation, a two orders-of-magnitude increase over current tools, essential for analysing samples at the scale expected from LSST. jax-SNANA currently supports the SALT and BayeSN spectral energy distribution (SED) models as well as physical models from redback-jax and a variety of SN SED templates; light curve fitting is supported with both photometric and spectroscopic redshifts. We fully validate the accuracy of jax-SNANA by comparing to SNANA, finding excellent consistency between the two across a variety of tasks. The code provides a modular JAX end-to-end supernova cosmology pipeline, user-extensible to support new features including different SED templates and photometric classification algorithms. jax-SNANA has been designed with future applications of simulation-based inference in mind, providing fast simulations and enabling fully compiled on-the-fly simulations to be generated within the training loop of downstream SBI tasks. jax-SNANA extends beyond SNANA by supporting simulations of generic SN-host galaxy relationships and multiple sub-populations of events, for example originating from young and old progenitors. jax-SNANA is publicly available at https://github.com/mattgrayling/jax-snana and installable through pip.
Related
- Detectability of 100 TeV Gamma Rays from PeV Cosmic Rays Accelerated in the Early Phase of a Supernova Shock Wave
- Type Ia supernova population machine – I. A unified cosmological simulation–binary synthesis framework establishing non-universal delay-time distributions and cosmic progenitor-channel dominance crossover
- IViS: Interferometric Visibility-domain inversion Software. A GPU-accelerated Python framework for joint deconvolution with ASKAP
- A Two-Parameter Framework for the Diffuse Supernova Neutrino Background
- A Framework for Analysing Nebular-Phase JWST Supernova Spectra with Application to Type Ic SN 2024aecx
BibTeX
@misc{grayling2026jaxsnana,
title = {{jax-SNANA: A GPU-accelerated Supernova Framework for the SBI Era}},
author = {Matthew Grayling},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.19000},
eprint = {2610.19000},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.19000},
}
RIS
TY - UNPB
TI - jax-SNANA: A GPU-accelerated Supernova Framework for the SBI Era
AU - Matthew Grayling
PY - 2026
DA - 2026/10/13
PB - arXiv
DO - 10.48550/arxiv.2610.19000
UR - https://arxiv.org/abs/2610.19000
AB - We present jax-SNANA, a JAX reimplementation of the supernova analysis framework SNANA. jax-SNANA enables GPU-accelerated transient simulations and light curve fitting, supernova standardisation and cosmological inference. Running in single precision on an NVIDIA A100 GPU, jax-SNANA simulates more than 100,000 LSST-observed supernovae per second and fits 5,000 light curves per second after JIT compilation, a two orders-of-magnitude increase over current tools, essential for analysing samples at the scale expected from LSST. jax-SNANA currently supports the SALT and BayeSN spectral energy distribution (SED) models as well as physical models from redback-jax and a variety of SN SED templates; light curve fitting is supported with both photometric and spectroscopic redshifts. We fully validate the accuracy of jax-SNANA by comparing to SNANA, finding excellent consistency between the two across a variety of tasks. The code provides a modular JAX end-to-end supernova cosmology pipeline, user-extensible to support new features including different SED templates and photometric classification algorithms. jax-SNANA has been designed with future applications of simulation-based inference in mind, providing fast simulations and enabling fully compiled on-the-fly simulations to be generated within the training loop of downstream SBI tasks. jax-SNANA extends beyond SNANA by supporting simulations of generic SN-host galaxy relationships and multiple sub-populations of events, for example originating from young and old progenitors. jax-SNANA is publicly available at https://github.com/mattgrayling/jax-snana and installable through pip.
ER -