A Hamiltonian, post-Born, three-dimensional, on-the-fly ray tracing algorithm for gravitational lensing

Abstract

The analyses of the next generation cosmological surveys demand an accurate, efficient, and differentiable method for simulating the universe and its observables across cosmological volumes. We present Hamiltonian ray tracing (HRT) -- the first post-Born (accounting for lens-lens coupling and without relying on the Born approximation), three-dimensional (without assuming the thin-lens approximation), and on-the-fly (applicable to any structure formation simulations) ray tracing algorithm based on the Hamiltonian formalism. HRT performs symplectic integration of the photon geodesics in a weak gravitational field, and can integrate tightly with any gravity solver, enabling co-evolution of matter particles and light rays with minimal additional computations. We implement HRT in the particle-mesh library $\texttt{pmwd}$, leveraging hardware accelerators such as GPUs and automatic differentiation capabilities based on $\texttt{JAX}$. When tested on a point-mass lens, HRT achieves sub-percent accuracy in deflection angles above the resolution limit across both weak and moderately strong lensing regimes. We also test HRT in cosmological simulations on the convergence maps and their power spectra.