$ exttt{SwiftC}_ll$: fast differentiable angular power spectra beyond Limber

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.

💡 Research Summary

The paper introduces SwiftCℓ, a novel, JAX‑based pipeline designed to compute angular power spectra (Cℓ) for large‑scale‑structure (LSS) analyses beyond the Limber approximation. Upcoming Stage IV surveys such as LSST, Euclid, DESI, and SphereX will deliver percent‑level measurements over a wide range of angular scales, including the very large scales (ℓ ≲ 30) where the Limber approximation fails. Accurate theoretical predictions on these scales are essential for probing effects that only appear at low ℓ, such as primordial non‑Gaussianity (PNG) and the Integrated Sachs‑Wolfe (ISW) effect.

Core methodology
SwiftCℓ builds on the FFTLog algorithm, which performs a fast Fourier transform in logarithmic space. The matter power spectrum is split as

P(k, χ, χ′) = P(k, z = z_fid) · D(k, z(χ)) · D(k, z(χ′)),

where D(k, z) = P(k, z)/P(k, z_fid) captures the (potentially scale‑dependent) growth factor. This split isolates most of the k‑dependence in the first factor, allowing a mild residual D(k, z) to be treated efficiently.

For each probe the line‑of‑sight kernel W_i(χ) · D(k, z(χ)) is expanded with FFTLog:

W(χ) D(χ, k) ≈ ∑_n c_n(k) χ^{b+iη_n}.

The integral over χ of χ^p j_ℓ(kχ) can then be evaluated analytically using gamma‑function ratios, yielding a closed‑form expression that depends only on ℓ and the FFTLog exponent p. Consequently, the ℓ‑loop reduces to a simple sum over pre‑computed gamma factors.

The remaining k‑integral is performed numerically. Rather than evaluating FFTLog coefficients at every k‑sample, SwiftCℓ computes them at a modest number N_interp of interpolation points (typically ~64) and interpolates the rest using splines. This dramatically reduces the computational cost while preserving high accuracy because D(k, z) varies smoothly with k. The two main hyper‑parameters are N_FFT (the number of logarithmic points in the FFTLog transform) and N_interp (the number of k‑points with explicit coefficients). Convergence tests in the appendix show that N_FFT ≈ 256 and N_interp ≈ 64 achieve relative errors below 10⁻⁴ for all relevant ℓ.

Automatic differentiation
All operations are implemented with JAX primitives, making the entire pipeline auto‑differentiable. Gradients with respect to cosmological parameters, bias parameters, or any user‑defined function can be obtained at a modest overhead (≈ 4× the forward evaluation). This enables gradient‑based samplers such as Hamiltonian Monte Carlo (HMC) and facilitates rapid Fisher‑matrix calculations.

Supported probes
SwiftCℓ provides a unified framework for a comprehensive set of probes:

• Galaxy clustering (δ_g) with linear bias b₁, magnification bias C_g, linear redshift‑space distortions (RSD), and local PNG (f_NL).
• Weak gravitational lensing (γ) using the standard convergence kernel and the Non‑Linear Alignment (NLA) model for intrinsic alignments.
• CMB lensing (κ) with integration up to the last‑scattering surface.
• CMB ISW effect, modeled via the time derivative of the gravitational potential.

Each probe is expressed as a source function Δ_iℓ(k) that follows the generic form Δ_iℓ(k) = c(ℓ) ∫ dχ W_i(χ) j_ℓ^{(n)}(kχ) · D(k, z(χ)), allowing all cross‑spectra C_ijℓ to be computed with the same FFTLog machinery.

Performance and validation
The authors benchmark SwiftCℓ against the N5K challenge, which gathered a suite of beyond‑Limber codes for LSST. Computing 120 distinct angular power spectra over 103 multipoles takes ≈ 5 ms on a single GPU core; computing the full gradient costs about 20 ms (≈ 4× slower). Compared with the challenge winner (FKEM), SwiftCℓ is roughly 40× faster while delivering comparable accuracy (relative errors ≲ 10⁻⁴).

Validation includes:

• Internal convergence tests varying N_FFT and N_interp.
• Comparison with other public codes (FKEM, matter, Levin, AngPow, Blast.jl).
• End‑to‑end cosmological inference on an LSST‑like mock: an MCMC run, an HMC run, and a Fisher‑matrix forecast are presented, demonstrating that the auto‑differentiable pipeline yields unbiased parameter constraints and dramatically reduces wall‑clock time.

Implications
SwiftCℓ solves three long‑standing bottlenecks in LSS theory:

1. Speed – GPU‑accelerated FFTLog reduces the evaluation of oscillatory Bessel integrals from seconds to milliseconds.
2. Accuracy on large scales – By abandoning the Limber approximation and retaining full k‑dependence, it reliably models low‑ℓ observables crucial for PNG, ISW, and tests of modified gravity.
3. Differentiability – Full JAX integration enables gradient‑based sampling, opening the door to efficient Hamiltonian Monte Carlo, variational inference, and rapid Fisher forecasts for multi‑probe analyses.

The code is open‑source (GitLab link provided) and designed to be extensible: users can add custom bias models, non‑linear corrections, or alternative window functions without breaking the auto‑diff pipeline. SwiftCℓ therefore represents a significant step toward the next generation of cosmological data analysis pipelines, where speed, precision, and differentiability are all required simultaneously.