🔭 astrophysics

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

This paper introduces a novel, auto-differentiable source reconstruction method for gravitational lensing that utilizes Gaussian processes on ray-guided transformed uniform grids to achieve fast, flexible, and high-resolution adaptive modeling with improved statistical evidence compared to traditional adaptive mesh techniques.

Original authors: Wolfgang J. R. Enzi, Coleman M. Krawczyk, Tian Li, Thomas E. Collett

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Wolfgang J. R. Enzi, Coleman M. Krawczyk, Tian Li, Thomas E. Collett

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to take a high-resolution photograph of a distant, distorted galaxy. This galaxy's light has been bent by a massive object (like another galaxy) sitting in front of it, acting like a cosmic magnifying glass. This phenomenon is called gravitational lensing.

The problem is that this "magnifying glass" doesn't magnify everything equally. Some parts of the background galaxy are stretched out and magnified a lot, while other parts are barely magnified at all.

The Old Way: The "One-Size-Fits-All" Grid

Traditionally, astronomers tried to reconstruct the image of that background galaxy by dividing it into a grid of tiny squares (pixels), like a chessboard. They would assume every square on the board needed the same amount of detail.

The Problem: This is inefficient. In the highly magnified areas, the grid squares are too big to see the fine details (like a low-resolution photo). In the barely magnified areas, the grid squares are tiny and full of empty space, wasting computer power.
The Alternative: Some scientists used "adaptive" grids that changed shape, like a net made of triangles that could stretch and shrink. However, these shapes often changed in jerky, discontinuous ways when the math was tweaked, making it hard for computers to find the perfect solution smoothly.

The New Solution: The "Smart, Stretchy Map" (RTU Grid)

The authors of this paper propose a clever new method called the Ray-Guided Transformed Uniform (RTU) grid.

Here is how it works, using an analogy:

Imagine you are a cartographer trying to draw a map of a country where the population density varies wildly. In the cities, people are packed shoulder-to-shoulder; in the deserts, they are miles apart.

The Old Map: You draw a grid where every square represents exactly 1 square mile. In the city, you have to cram thousands of people into one square (losing detail), and in the desert, your square is empty.
The RTU Map: Instead of forcing a rigid grid, you stretch and warp your map paper. You stretch the paper over the cities so that one square on your map now covers a smaller physical area, capturing the crowd. You shrink the paper over the deserts so one square covers a huge empty space.
- The Result: Now, every single square on your map contains roughly the same number of people. You get perfect detail everywhere without needing a million tiny squares.

How They Did It (The "Magic" Trick)

The scientists used a mathematical tool called a Gaussian Process (think of it as a very flexible, smooth way to guess what an image looks like based on patterns). Usually, this tool works best on a perfect, regular grid (like a chessboard).

Trace the Rays: They traced the path of light rays from the telescope back to the source galaxy.
Count the Rays: They noticed that in some areas, many rays land on the same spot (high magnification), and in others, few rays land (low magnification).
The Transformation: They created a mathematical "stretching" rule. They rearranged the coordinates of their regular grid so that the areas with many rays got stretched out, and areas with few rays got squished together.
The Result: On this new, warped grid, the light rays are distributed evenly. This allows them to use the fast, smooth math of the Gaussian Process while still getting high-resolution details exactly where they are needed.

What They Found

The team tested this new method on fake data (simulations) and real data from a famous galaxy system (SDSS J0946+1006).

Efficiency: They found that the new "stretched" grid could produce images just as clear as the old "rigid" grid, but it needed half as many pixels in each direction. Since a grid is 2D, this means they used about four times fewer total pixels to get the same quality.
Speed: Because they used fewer pixels, the computer calculations were much faster.
Accuracy: The new method didn't miss any important details, like small clumps of dark matter (substructures) inside the lensing galaxy. It was just as good at spotting these hidden features as the old method, but much more efficient.

Why This Matters

The paper concludes that this method is a "fast, flexible, and auto-differentiable" way to reconstruct images. In plain English: it's a smart, smooth, and speedy way to un-distort images of the universe.

This is particularly important because upcoming space telescopes (like the Euclid mission mentioned in the paper) will find thousands of these lensing events. Astronomers need a method that is fast enough to process all of them without getting bogged down by inefficient math. This new "stretched map" approach provides that speed without sacrificing the quality of the cosmic pictures.

Technical Summary: Gaussian Processes on Ray-Guided Transformed Uniform Grids

Problem Statement
Strong gravitational lensing (SGL) is a critical tool for constraining cosmology and dark matter distributions, but robust inference relies heavily on accurate source reconstruction. A fundamental challenge in SGL is the non-uniform information content across the source plane: regions producing multiple images (e.g., four images) have higher sampling rates than those producing fewer images, and magnification varies continuously. Traditional adaptive meshing techniques, such as Delaunay tessellations or Voronoi binning, attempt to address this by placing resolution where needed. However, these methods often rely on discontinuous operations (e.g., mesh edges swapping vertices as mass models change), which can break the differentiability required for modern inference frameworks and restrict regularization choices. Conversely, standard Gaussian Process (GP) priors offer flexible regularization but typically require regular grids to utilize Fast Fourier Transforms (FFTs) for computational efficiency, creating a tension between adaptivity and differentiability.

Methodology
The authors propose a novel source reconstruction approach using a Ray-Guided Transformed Uniform (RTU) grid. This method aims to combine the computational efficiency of FFT-based GPs on uniform grids with the adaptive resolution required by lensing data.

RTU Grid Construction:
- The source is modeled on a uniform grid in a transformed coordinate space $\mathbf{y} = \mathbf{g}(\mathbf{x})$ .
- The transformation $\mathbf{g}$ is derived from the empirical cumulative distribution functions (eCDFs) of rays traced back from the image plane to the source plane.
- Specifically, the authors project ray positions onto each coordinate axis, sort them to form 1D eCDFs, and fit cubic splines to these distributions. This ensures the transformation is at least twice continuously differentiable.
- The eCDFs are constructed using only rays falling within an arc mask (defined by signal-to-noise thresholds), ensuring resolution is allocated only to regions where the source is observed.
- The transformation effectively stretches the source plane coordinates so that the ray density becomes approximately uniform in the new $\mathbf{y}$ coordinates.
Source Modeling:
- The source brightness is defined as a Gaussian Process on this transformed uniform grid.
- The GP is realized by generating white noise in Fourier space, multiplying by the square root of a power spectrum (specifically a flexible Matérn power spectrum), and applying an inverse FFT.
- Positivity is enforced via a softplus function, and the final source brightness is scaled by a power law to allow for sharper features.
- The mapping back to the physical source plane is performed via $s_{source}(\mathbf{x}) = s^*_{source}(\mathbf{g}(\mathbf{x}))$ .
Inference Framework:
- The approach is implemented within Herculens, an auto-differentiable forward modeling framework written in JAX.
- Inference is performed using Stochastic Variational Inference (SVI) with Numpyro, optimizing the Evidence Lower Bound (ELBO).
- The method supports arbitrary power spectrum choices for regularization, maintaining the flexibility of GPs while leveraging FFT speed.

Key Contributions

Differentiable Adaptivity: The RTU grid provides an adaptive resolution scheme that is continuously differentiable, overcoming the discontinuity issues of Delaunay/Voronoi meshes while retaining the speed of FFT-based GPs.
Efficiency Rule of Thumb: The authors propose a heuristic for selecting the number of source pixels ( $N_{src}$ ) based on ray density, suggesting approximately 4 rays per pixel to match the degrees of freedom in bilinear interpolation.
Implementation: The method is fully integrated into the Herculens framework, making it compatible with state-of-the-art sampling algorithms and auto-differentiable workflows.

Results
The authors tested the RTU grid on three mock datasets (simulating JWST NIRCAM quality) and real data from the lens system SDSS J0946+1006.

Fit Quality and Pixel Efficiency: The RTU grid achieves image-plane fit quality ( $\chi^2_{img}$ ) comparable to uniform grids but requires approximately half the number of pixels per dimension (a factor of $\simeq 4$ fewer total pixels) to reach the noise-level residuals.
ELBO Performance: For the same number of pixels, the RTU grid generally yields higher Evidence Lower Bounds (ELBOs) compared to uniform grids, indicating a better fit to the data under the variational approximation.
Substructure Sensitivity: The use of the RTU grid only mildly affects the difference in ELBO between models with and without substructures. The method successfully distinguishes between models with and without substructures (e.g., in Mock 1 and SDSS J0946+1006) without artificially inflating the preference for substructure due to grid artifacts.
Source Reconstruction Accuracy: When the underlying mass model is accurate, the RTU grid produces more accurate source reconstructions (lower $\Delta^2_{src}$ ) with fewer pixels than the uniform grid approach.

Significance and Claims
The paper claims that the RTU grid method resolves the trade-off between realism, flexibility, auto-differentiability, and computational efficiency in lens modeling. By ensuring source pixels contain a more uniform number of rays, the method allows for high-fidelity source reconstruction with significantly reduced computational cost. The authors assert that this approach is particularly beneficial for analyzing large samples of lens systems, such as those expected from the Euclid survey, where robust and flexible source modeling is critical for minimizing biases in cosmological and dark matter constraints. The method is presented as a practical improvement over existing adaptive meshing techniques, offering a fast, flexible, and auto-differentiable alternative that does not compromise on the ability to detect substructures.