LINX: A Fast, Differentiable, and Extensible Big Bang Nucleosynthesis Package

TL;DR

LINX delivers an accurate, fast, and fully differentiable Big Bang Nucleosynthesis code built on JAX, enabling gradient-based inference and seamless coupling with differentiable CMB emulators for joint cosmological analyses. By modularizing background evolution and nuclear networks, LINX achieves sub-percent agreement with established codes while offering extensive network flexibility and a path to beyond-Standard-Model physics. Gradient-assisted inference, including neural-transport reparameterization and HMC, demonstrates efficient joint CMB+BBN analyses with practical resource needs, illustrating LINX’s potential to broaden exploration of early-Universe physics. The work highlights significant speedups, extensibility, and the practical impact of differentiable BBN in modern cosmology.

Abstract

We introduce LINX (Light Isotope Nucleosynthesis with JAX), a new differentiable public Big Bang Nucleosynthesis (BBN) code designed for fast parameter estimation. By leveraging JAX, LINX achieves both speed and differentiability, enabling the use of Bayesian inference, including gradient-based methods. We discuss the formalism used in LINX for rapid primordial elemental abundance predictions and give examples of how LINX can be used. When combined with differentiable Cosmic Microwave Background (CMB) power spectrum emulators, LINX can be used for joint CMB and BBN analyses without requiring extensive computational resources, including on personal hardware.

Figures (11)

• Figure 1: A code block diagram illustrating the functions of different modules of LINX. Blue indicates user inputs, yellow indicates modules the user would typically call, and green indicates modules that the code calls internally. Blue arrows indicate the flow of user inputs, green arrows indicate internal calls, and black arrows indicate the flow of outputs. Two additional modules, const and special_funcs, are not shown; they provide auxiliary constants and special functions to all of the other modules. Notation is defined in the main text. Refer to code documentation for more information about the specific inputs and outputs of each module.
• Figure 2: A diagram representing the "key" network, showing the nuclei that are tracked and the reactions between them. Shown in red are deuterium reactions that dominate the uncertainty in the prediction of D/H Pisanti:2020efzPitrou:2021vqr.
• Figure 3: Residuals between PRyMordial and LINX for four primordial element abundances as a function of the baryon-to-photon ratio $\eta$. For this comparison, we set $\tau_n=879.4s$ in both codes, use the key_PRIMAT_2018 network in LINX, and set smallnet_flag=True in PRyMordial. The agreement between the two codes is $\ll 1\%$ throughout the range of $\eta$ for D/H and $Y_\mathrm{P}$. The jagged behavior is likely due to numerical noise associated with the differential equation solvers in either code.
• Figure 4: Scaling of evaluation time with batch size when LINX is run on a CPU versus a GPU. For small batch sizes, the CPU is more efficient, but for batch sizes $\gtrsim 550$, the GPU becomes more efficient due to the more favorable scaling.
• Figure 5: Dependence of the deuterium (left) and helium-4 (right) predicted abundances on key reaction-rate uncertainties in the network. The deuterium prediction depends strongly on the rate of $d(p,\gamma)^3$He (as well as $d(d,p)t$ and $d(d,n)^3$He), while the helium-4 prediction depends strongly on the neutron lifetime. The $x$-axes show the scaling of the $d(p,\gamma)^3$He rate (left, see text) or the value of the neutron lifetime (right), and the $y$-axes show the corresponding predictions for D/H (left) and $Y_\mathrm{P}$ (right). The green regions have a height corresponding to $1\sigma$ the uncertainty of the measurements of D/H Cooke_2018 and $Y_\mathrm{P}$Aver_2015 and a width corresponding to the $1\sigma$ uncertainty of the quantities on the $x$-axes (defined to be 1 for the left plot or determined by measurement for the right plot Czarnecki:2018okw). The helium-4 prediction does not depend significantly on the choice of reaction network.