JAX-bandflux: Differentiable Supernova Modelling

• JAX-bandflux is a differentiable, GPU-accelerated framework implementing SALT models for supernova light curves, enabling rapid cosmological inference.
• The framework leverages gradient-based optimizers and vectorized parallelism with JAX’s JIT and vmap to accelerate parameter estimation.
• It reproduces and extends SNCosmo functionality while ensuring scalability to high-volume survey data and improved quantification of systematic uncertainties.

JAX-bandflux is a differentiable framework for supernova SALT modelling, implemented using JAX to enable efficient, GPU-accelerated parameter inference for cosmological surveys. By reproducing and extending the core functionality of established tools such as SNCosmo, JAX-bandflux offers a computationally efficient and parallelizable approach to modeling supernova light curves, with particular focus on facilitating gradient-based optimization workflows central to contemporary cosmological analyses.

1. Supernova Light Curve Modelling and SALT Formalism

At the core of JAX-bandflux lies the implementation of critical routines from the SNCosmo library, recast in a differentiable programming paradigm. The principal physical model is the SALT3 representation of Type Ia supernova light curves, parameterized as:

$$F(p, \lambda) = x_0 \big[ M_0(p, \lambda) + x_1 M_1(p, \lambda) + \ldots \big] \cdot \exp \left[c \cdot CL(\lambda)\right]$$

where $x_0$ is the overall scale (normalization), $x_1$ encodes stretch (light curve shape), $t_0$ is the reference time of maximum brightness, $c$ represents color variation, and $M_0, M_1$ are empirically derived flux surfaces. The parameter vector $p$ encodes dependencies on redshift and phase ($t-2$), while $CL(\lambda)$ parameterizes the wavelength-dependent color law.

Observed bandflux is computed via wavelength integration over the specified bandpass:

$$\mathrm{bandflux} = \int_{\lambda_{\mathrm{min}}}^{\lambda_{\mathrm{max}}} F(\lambda) \cdot T(\lambda) \cdot \frac{\lambda}{hc}\, d\lambda$$

where $T(\lambda)$ is the instrument's transmission function, with factors of $h$ (Planck’s constant) and $c$ (speed of light) converting energy into photon counts.

2. Differentiable Computational Framework

JAX-bandflux leverages JAX’s autodifferentiation engine, enabling end-to-end differentiation of the supernova modeling pipeline. This is essential for propagating gradients through both the light curve model and the wavelength integration operations. In contrast to legacy frameworks which primarily depend on non-differentiable code paths or grid-based fitting procedures, this design enables the following:

• Utilization of gradient-based optimizers (e.g., L-BFGS-B, nested sampling) for parameter estimation.
• Large-scale, simultaneous fitting of multiple supernovae, with gradient information capturing the sensitivity of the model to observational and calibration uncertainties.
• Direct integration of the model within probabilistic inference pipelines, improving both computational throughput and accuracy in parameter estimation.

The differentiable implementation thus enables more precise and rapid cosmological inference workflows, particularly valuable for fitting SALT parameters and propagating uncertainties into cosmological measurements.

3. GPU-Accelerated Parallelization

The computational demands of evaluating light curve models and performing integrations across high-dimensional parameter spaces are substantial, especially when working with the large datasets typical of contemporary astronomical surveys. JAX-bandflux addresses this by utilizing JAX's just-in-time (JIT) compilation and vectorized parallelization patterns (e.g., vmap). These mechanisms enable:

• Parallel evaluation of flux calculations across supernovae, bandpasses, and observational epochs.
• Substantial reduction in time-to-solution for inference problems, as most computationally intensive operations are offloaded to GPUs.
• Scalability to the requirements of high-cadence, high-volume survey data, where traditional CPU-bound routines become prohibitive.

A plausible implication is that this design paradigm will become increasingly essential as future surveys continue to expand in scope and data volume.

4. SALT Parameter Inference and Cosmological Utility

The SALT parameters—$x_0$, $x_1$, $t_0$, and $c$—are central to the construction of standardized candles for cosmological distance determination. Their physical interpretation is as follows:

Parameter	Role	Description
$x_0$	Flux normalization	Sets the brightness scale
$x_1$	Stretch	Encodes light curve shape variability
$t_0$	Peak time	Reference epoch of maximum brightness
$c$	Color	Captures color-dependent variation

Accurate inference of these parameters transforms observed light curves into distance indicators, crucial for measuring cosmic expansion and placing constraints on cosmological parameters. JAX-bandflux supports fully differentiable computation of SALT parameters, integrating them with optimization routines to facilitate both speed and precision.

5. Comparison with SNCosmo and Computational Advancements

While both JAX-bandflux and SNCosmo target the modeling of supernova light curves using SALT models, there are several key distinctions:

• Differentiability: JAX-bandflux is designed for end-to-end gradient propagation, in contrast to SNCosmo’s primarily non-differentiable architecture.
• Hardware Acceleration: By leveraging GPU and TPU hardware via JAX, JAX-bandflux offers substantial acceleration and scalability beyond SNCosmo’s CPU-bound implementations.
• Parallelized Evaluation: JAX-bandflux efficiently vectorizes and parallelizes calculations for large datasets, making contemporary, large-scale analyses tractable.

Despite these differences, JAX-bandflux maintains functional parity with SNCosmo in terms of core modeling capabilities and realism of its flux calculations.

6. Broader Impact and Methodological Implications

JAX-bandflux enables more robust and efficient cosmological analyses by improving computational efficiency, scalability, and precision in parameter estimation. The adoption of differentiable programming and GPU acceleration in core astrophysical modeling routines reflects a trend toward end-to-end, high-throughput analysis pipelines. This methodological advance stands to improve the calibration of supernova samples, the quantification of systematic uncertainties, and the ultimate precision of dark energy and cosmic expansion measurements.

A plausible implication is that similar differentiable modeling frameworks will propagate to other domains within astrophysics, enabling analysis of increasingly complex and voluminous datasets as observational capabilities expand. The JAX-bandflux architecture thus serves as a template for high-performance, differentiable scientific computing in cosmology and beyond.