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

Technische Samenvatting: Gaussische Processen op Ray-Guided Transformed Uniform Grids

Probleemstelling
Sterke gravitationele lensing (SGL) is een cruciaal instrument voor het beperken van kosmologie en donkere materie-distributies, maar robuuste inferentie is sterk afhankelijk van nauwkeurige bronreconstructie. Een fundamentele uitdaging in SGL is de niet-uniforme informatie-inhoud over het bronvlak: regio's die meerdere beelden produceren (bijv. vier beelden) hebben hogere bemonsteringsfrequenties dan regio's die minder beelden produceren, en de vergroting varieert continu. Traditionele adaptieve meshing-technieken, zoals Delaunay-tessellaties of Voronoi-binning, proberen dit aan te pakken door resolutie te plaatsen waar dat nodig is. Echter, deze methoden vertrouwen vaak op discontinue operaties (bijv. het wisselen van mesh-randen tussen vertices naarmate massamodellen veranderen), wat de differentieerbaarheid kan breken die vereist is voor moderne inferentiekaders en de keuzes voor regularisatie kan beperken. Daartegenover staan standaard Gaussische Proces (GP) priors die flexibele regularisatie bieden, maar doorgaans regelmatige roosters vereisen om Fast Fourier Transforms (FFTs) voor computationele efficiëntie te kunnen gebruiken, wat een spanning creëert tussen adaptiviteit en differentieerbaarheid.

Methodologie
De auteurs stellen een nieuwe bronreconstructie-aanpak voor met een Ray-Guided Transformed Uniform (RTU) grid. Deze methode beoogt de computationele efficiëntie van FFT-gebaseerde GPs op uniforme roosters te combineren met de adaptieve resolutie die vereist is door lensing-data.

1. Constructie van het RTU-grid:

   • De bron wordt gemodelleerd op een uniform rooster in een getransformeerde coördinatenruimte $\mathbf{y} = \mathbf{g}(\mathbf{x})$.
   • De transformatie $\mathbf{g}$ wordt afgeleid van de empirische cumulatieve distributiefuncties (eCDFs) van stralen (rays) die vanuit het beeldvlak terug naar het bronvlak worden getraceerd.
   • Specifiek projecteren de auteurs de posities van de stralen op elke coörataxis, sorteren ze om 1D eCDFs te vormen, en passen zij kubische splines toe op deze distributies. Dit zorgt ervoor dat de transformatie ten minste twee keer continu differentieerbaar is.
   • De eCDFs worden geconstrueerd met behulp van enkel stralen die binnen een boogmasker vallen (gedefinieerd door signaal-ruis drempelwaarden), waardoor resolutie alleen wordt toegewezen aan regio's waar de bron daadwerkelijk wordt waargenomen.
   • De transformatie vervormt de coördinaten van het bronvlak effectief zodat de straaldichtheid ongeveer uniform wordt in de nieuwe $\mathbf{y}$-coördinaten.

2. Bronmodellering:

   • De helderheid van de bron wordt gedefinieerd als een Gaussisch Proces op dit getransformeerde uniforme rooster.
   • De GP wordt gerealiseerd door witte ruis in het Fourier-domein te genereren, te vermenigvuldigen met de wortel van een vermogensspectrum (specifiek een flexibel Matérn-vermogensspectrum), en een inverse FFT toe te passen.
   • Positiviteit wordt afgedwongen via een softplus-functie, en de uiteindelijke helderheid van de bron wordt geschaald door een machtswet om scherpere kenmerken mogelijk te maken.
   • De mapping terug naar het fysieke bronvlak wordt uitgevoerd via $s_{source}(\mathbf{x}) = s^*_{source}(\mathbf{g}(\mathbf{x}))$.

3. Inferentiekader:

   • De aanpak is geïmplementeerd binnen Herculens, een auto-differentieerbaar forward modeling framework geschreven in JAX.
   • Inferentie wordt uitgevoerd met behulp van Stochastic Variational Inference (SVI) met Numpyro, waarbij de Evidence Lower Bound (ELBO) wordt geoptimaliseerd.
   • De methode ondersteunt willekeurige vermogensspectrum-keuzes voor regularisatie, waardoor de flexibiliteit van GPs behouden blijft terwijl de snelheid van FFT wordt benut.

Belangrijkste Bijdragen

• Differentiabele Adaptiviteit: Het RTU-grid biedt een adaptief resolutieschema dat continu differentieerbaar is, waarmee het de discontinuïteitsproblemen van Delaunay/Voronoi-meshes overwint terwijl de snelheid van FFT-gebaseerde GPs behouden blijft.
• Efficiëntie Vuistregel: De auteurs stellen een heuristiek voor het selecteren van het aantal bronpixels ($N_{src}$) op basis van de straaldichtheid, waarbij zij suggereren dat ongeveer 4 stralen per pixel nodig zijn om de vrijheidsgraden van bilineaire interpolatie te evenaren.
• Implementatie: De methode is volledig geïntegreerd in het Herulens-framework, waardoor het compatibel is met state-of-the-art sampling-algoritmen en auto-differentieerbare workflows.

Resultaten
De auteurs hebben het RTU-grid getest op drie mock-datasets (die de kwaliteit van JWST NIRCAM simuleren) en echte data van het lensysteem SDSS J0946+1006.

• Fitkwaliteit en Pixel-efficiëntie: Het RTU-grid bereikt een beeldvlak-fitkwaliteit ($\chi^2_{img}$) die vergelijkbaar is met uniforme roosters, maar vereist ongeveer de helft van het aantal pixels per dimensie (een factor $\simeq 4$ minder totale pixels) om de ruis-niveau residuen te bereiken.
• ELBO-prestaties: Voor hetzelfde aantal pixels levert het RTU-grid over het algemeen hogere Evidence Lower Bounds (ELBO's) op vergeleken met uniforme roosters, wat wijst op een betere fit bij de data onder de variationele benadering.
• Gevoeligheid voor Substructuur: Het gebruik van het RTU-grid heeft slechts een milde invloed op het verschil in ELBO tussen modellen met en zonder substructuren. De methode onderscheidt succesvol modellen met en zonder substructuren (bijv. in Mock 1 en SDSS J0946+1006) zonder de voorkeur voor substructuur kunstmatig op te blazen door grid-artefacten.
• Nauwkeurigheid van Bronreconstructie: Wanneer het onderliggende massamodel accuraat is, produceert het RTU-grid nauwkeurigere bronreconstructies (lagere $\Delta^2_{src}$) met minder pixels dan de uniforme rooster-aanpak.

Betekenis en Claims
Het artikel beweert dat de RTU-grid methode de trade-off tussen realisme, flexibiliteit, auto-differentieerbaarheid en computationele efficiëntie oplost. Door ervoor te zorgen dat de bronpixels een meer uniforme hoeveelheid stralen bevatten, maakt de methode hoogwaardige bronreconstructie mogelijk met aanzienlijk verminderde computationele kosten. De auteurs stellen dat deze aanpak bijzonder gunstig is voor het analyseren van grote steekproeven van lenssystemen, zoals verwacht bij de Euclid-survey, waar robuuste en flexibele bronmodellering cruciaal is voor het minimaliseren van biases in kosmologische en donkere materie-beperkingen. De methode wordt gepresenteerd als een praktische verbetering ten opzichte van bestaande adaptieve meshing-technieken, die een snelle, flexibele en auto-differentieerbare alternatief biedt zonder de capaciteit om substructuren te detecteren te compromitteren.