Gaussian processes on ray-guided transformed uniform grids for fast, flexible, and auto-differentiable adaptive source reconstruction in lens modelling

Abstract

Strong gravitational lensing constrains cosmology and dark matter, but robust inference requires accurate source reconstruction. The achievable source resolution is highly position-dependent. Adaptive meshes can place resolution where needed, but typically rely on discontinuous operations, such as Delaunay tessellations or Voronoi binning, which can restrict regularization choices and break differentiability. In this paper, we present a novel approach for modelling the source on a ray-guided transformed uniform grid (RTU grid), that is adaptive to the lens mass model, auto-differentiable and flexible with respect to the regularization by allowing for an arbitrary choice of power spectrum. We achieve this by defining the source as a Gaussian process on a uniform grid, which is then transformed based on the cumulative distributions of rays traced back to the source plane. This approach ensures that source pixels contain a more uniform number of rays. The approach is fast by leveraging the fast Fourier transform to describe the Gaussian process in Fourier space. We apply this new approach to mock data and show that it achieves comparable fit quality with fewer source pixels, typically corresponding to about a factor of two fewer pixels per dimension, and increases Evidence Lower Bounds (ELBOs) for the same number of pixels. Using the RTU grid only mildly affects the difference in ELBO for models with and without substructures within lens galaxies. A fast, flexible, and auto-differentiable source reconstruction can greatly benefit the analysis of large samples of lens systems, e.g. those found within the Euclid survey