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

TL;DR

This paper introduces a novel Hamiltonian ray tracing algorithm for gravitational lensing that is accurate, efficient, and compatible with various structure formation simulations, accounting for post-Born effects and three-dimensional geometry.

Contribution

The paper presents the first on-the-fly, post-Born, three-dimensional Hamiltonian ray tracing algorithm integrated with existing gravity solvers, enabling precise and efficient cosmological light propagation simulations.

Findings

Achieves sub-percent accuracy in deflection angles for point-mass lenses

Successfully tests in cosmological simulations with convergence maps

Compatible with GPU acceleration and automatic differentiation

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.