$\texttt{SwiftC}_\ell$: fast differentiable angular power spectra beyond Limber

Abstract

The upcoming stage IV wide-field surveys will provide high precision measurements of the large-scale structure (LSS) of the universe. Their interpretation requires fast and accurate theoretical predictions including large scales. For this purpose, we introduce $\texttt{SwiftC}_\ell$, a fast, accurate and differentiable $\texttt{JAX}$-based pipeline for the computation of the angular power spectrum beyond the Limber approximation. It uses a new FFTLog-based method which can reach arbitrary precision and includes interpolation along $k$, allowing for $k$-dependent growth factor and biases. $\texttt{SwiftC}_\ell$ includes a wide range of probes and effects such as galaxy clustering, including magnification bias, redshift-space distortions and primordial non-Gaussianity, weak lensing, including intrinsic alignment, cosmic microwave background (CMB) lensing and CMB integrated Sachs-Wolfe effect. We compare our pipeline to the other available beyond-Limber codes within the N5K challenge from the Rubin Observatory Legacy Survey of Space and Time (LSST) Dark Energy Science Collaboration. $\texttt{SwiftC}_\ell$ computes the 120 different angular power spectra over 103 $\ell$-multipoles in 5 ms on one GPU core while the computation of the gradient is approximately 4$\times$ slower. Using a pre-calculation, $\texttt{SwiftC}_\ell$ is thus about 40$\times$ faster than the winner of the N5K challenge with comparable accuracy. Furthermore, all outputs are auto-differentiable, facilitating gradient-based sampling and robust and accurate Fisher forecasts. We showcase a Markov Chain Monte Carlo, a Hamiltonian Monte Carlo and a Fisher forecast on an LSST-like survey, illustrating $\texttt{SwiftC}_\ell$'s differentiability, speed and reliability in measuring cosmological parameters. The code is publicly available at https://cosmo-gitlab.phys.ethz.ch/cosmo_public/swiftcl.