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

TL;DR

SwiftC_l introduces a fast, differentiable pipeline for computing beyond-Limber angular power spectra using a novel FFTLog-based method with k-dependent growth, enabling accurate forecasts for stage IV LSS surveys. It supports a wide suite of probes (galaxy clustering with magnification bias, RSD, PNG; weak lensing with intrinsic alignment; CMB lensing and ISW) and yields auto-differential outputs suitable for gradient-based sampling. On LSST-like data and N5K benchmarks, SwiftC_l achieves ~5 ms GPU runtimes for 120 spectra across 103 multipoles, with ~40× speed gains over the challenge winner and subpercent agreement with independent codes, while maintaining a differentiable workflow for MCMC, HMC, and Fisher analyses. The work demonstrates practical utility for rapid parameter inference and Fisher forecasting, with public code availability to facilitate adoption in future cosmological analyses.

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.

Figures (7)

• Figure 1: All probes available in $\texttt{SwiftC}_\ell$ compared to the package CCL when possible. In CCL, we use FKEM for $\ell < 1000$ and then switch to the Limber approximation. We find that the pipelines agree to subpercent level.
• Figure 2: Accuracy of the galaxy clustering auto- and cross-correlated angular power spectra for half of the bins of the N5K challenge for all the entries and $\texttt{SwiftC}_\ell$ against the benchmarks. The uncertainties $\sigma_\ell$ represent a simple Gaussian covariance and the grey region represents the $\ell>200$ range that is not part of the challenge.
• Figure 3: Relative difference in the galaxy-galaxy lensing cross-correlated angular power spectra for some of the bins of the N5K challenge for all the entries and $\texttt{SwiftC}_\ell$ against the benchmarks. The uncertainties $\sigma_\ell$ represent a simple Gaussian covariance and the grey region represents the $\ell>200$ range that is not part of the challenge.
• Figure 4: Relative difference of the weak lensing auto- and cross-correlated angular power spectra for all of the bins of the N5K challenge for all the entries and $\texttt{SwiftC}_\ell$ against the benchmarks. The uncertainties $\sigma_\ell$ represent a simple Gaussian covariance and the grey region represents the $\ell>200$ range that is not part of the challenge.
• Figure 5: Ensemble sampling as implemented in emcee, HMC and a Fisher forecast with data vector computed using CCL with FKEM unless stated otherwise. As can be seen, the parameters are fully recovered in the ensemble sampling when using FKEM, in the HMC and in the Fisher forecast whereas a slight shift in some parameters can be observed in the ensemble sampling when using the Limber approximation. The Fisher information matrix is computed using the auto-differentiation from the JAX package.