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

Resumen Técnico: Procesos Gaussianos en Rejillas Uniformes Transformadas Guiadas por Rayos

Planteamiento del Problema
El lente gravitacional fuerte (SGL, por sus siglas en inglés) es una herramienta crítica para restringir la cosmología y las distribuciones de materia oscura, pero la inferencia robusta depende en gran medida de una reconstrucción precisa de la fuente. Un desafío fundamental en el SGL es el contenido de información no uniforme en el plano de la fuente: las regiones que producen múltiples imágenes (por ejemplo, cuatro imágenes) tienen tasas de muestreo más altas que aquellas que producen menos imágenes, y la magnificación varía continuamente. Las técnicas tradicionales de mallado adaptativo, como las teselaciones de Delaunay o el agrupamiento de Voronoi, intentan abordar esto colocando resolución donde se necesita. Sin embargo, estos métodos suelen depender de operaciones discontinuas (por ejemplo, el intercambio de vértices en los bordes de la malla a medida que cambian los modelos de masa), lo que puede romper la diferenciabilidad requerida por los marcos de inferencia modernos y restringir las opciones de regularización. Por el contrario, los modelos previos de Procesos Gaussianos (GP) estándar ofrecen una regularización flexible, pero típicamente requieren rejillas regulares para utilizar Transformadas Rápidas de Fourier (FFT) para la eficiencia computacional, lo que crea una tensión entre la adaptabilidad y la diferenciabilidad.

Metodología
Los autores proponen un nuevo enfoque de reconstrucción de fuentes utilizando una rejilla Uniforme Transformada Guiada por Rayos (RTU). Este método tiene como objetivo combinar la eficiencia computacional de los GP basados en FFT en rejillas uniformes con la resolución adaptativa requerida por los datos de lentes.

1. Construcción de la Rejilla RTU:

   • La fuente se modela en una rejilla uniforme en un espacio de coordenadas transformado $\mathbf{y} = \mathbf{g}(\mathbf{x})$.
   • La transformación $\mathbf{g}$ se deriva de las funciones de distribución acumulativa empírica (eCDF) de los rayos trazados desde el plano de la imagen hacia el plano de la fuente.
   • Específicamente, los autores proyectan las posiciones de los rayos sobre cada eje de coordenadas, los ordenan para formar eCDFs unidimensionales y ajustan splines cúbicos a estas distribuciones. Esto asegura que la transformación sea al menos dos veces continuamente diferenciable.
   • Las eCDFs se construyen utilizando únicamente los rayos que caen dentro de una máscara de arco (definida por umbrales de relación señal-ruido), asegurando que la resolución se asigne solo a las regiones donde la fuente es observada.
   • La transformación estira efectivamente las coordenadas del plano de la fuente de modo que la densidad de rayos sea aproximadamente uniforme en las nuevas coordenadas $\mathbf{y}$.

2. Modelado de la Fuente:

   • El brillo de la fuente se define como un Proceso Gaussiano en esta rejilla uniforme transformada.
   • El GP se realiza generando ruido blanco en el espacio de Fourier, multiplicándolo por la raíz cuadrada de un espectro de potencia (específicamente un espectro de potencia de Matérn flexible) y aplicando una FFT inversa.
   • La positividad se impone mediante una función softplus, y el brillo final de la fuente se escala mediante una ley de potencia para permitir características más nítidas.
   • El mapeo de vuelta al plano de la fuente física se realiza mediante $s_{source}(\mathbf{x}) = s^*_{source}(\mathbf{g}(\mathbf{x}))$.

3. Marco de Inferencia:

   • El enfoque está implementado dentro de Herculens, un marco de modelado hacia adelante auto-diferenciable escrito en JAX.
   • La inferencia se realiza mediante Inferencia Variacional Estocástica (SVI) con Numpyro, optimizando el Límite Inferior de Evidencia (ELBO).
   • El método soporta elecciones arbitrarias de espectro de potencia para la regularización, manteniendo la flexibilidad de los GP mientras se aprovecha la velocidad de las FFT.

Contribuciones Clave

• Adaptabilidad Diferenciable: La rejilla RTU proporciona un esquema de resolución adaptativa que es continuamente diferenciable, superando los problemas de discontinuidad de las mallas de Delaunay/Voronoi mientras retiene la velocidad de los GP basados en FFT.
• Regla de Eficiencia Heurística: Los autores proponen una heurística para seleccionar el número de píxeles de la fuente ($N_{src}$) basada en la densidad de rayos, sugiriendo aproximadamente 4 rayos por píxel para igualar los grados de libertad en la interpolación bilineal.
• Implementación: El método está totalmente integrado en el marco Herculens, lo que lo hace compatible con algoritmos de muestreo de vanguardia y flujos de trabajo auto-diferenciables.

Resultados
Los autores probaron la rejilla RTU en tres conjuntos de datos simulados (simulando la calidad de JWST NIRCAM) y datos reales del sistema de lente SDSS J0946+1006.

• Calidad del Ajuste y Eficiencia de Píxeles: La rejilla RTU logra una calidad de ajuste en el plano de la imagen ($\chi^2_{img}$) comparable a las rejillas uniformes, pero requiere aproximadamente la mitad del número de píxeles por dimensión (un factor de $\simeq 4$ menos de píxeles totales) para alcanzar los residuos del nivel de ruido.
• Rendimiento del ELBO: Para el mismo número de píxeles, la rejilla RTU generalmente produce Límites Inferiores de Evidencia (ELBO) más altos en comparación con las rejillas uniformes, lo que indica un mejor ajuste a los datos bajo la aproximación variacional.
• Sensibilidad a la Subestructura: El uso de la rejilla RTU solo afecta levemente la diferencia de ELBO entre modelos con y sin subestructuras. El método distingue con éxito entre modelos con y sin subestructuras (por ejemplo, en Mock 1 y SDSS J0946+1006) sin inflar artificialmente la preferencia por la subestructura debido a artefactos de la rejilla.
• Precisión de la Reconstrucción de la Fuente: Cuando el modelo de masa subyacente es preciso, la rejilla RTU produce reconstrucciones de la fuente más exactas (menor $\Delta^2_{src}$) con menos píxeles que el enfoque de la rejilla uniforme.

Significancia y Reivindicaciones
El artículo afirma que el método de la rejilla RTU resuelve el compromiso entre realismo, flexibilidad, auto-diferenciabilidad y eficiencia computacional en el modelado de lentes. Al asegurar que los píxeles de la fuente contengan un número más uniforme de rayos, el método permite una reconstrucción de la fuente de alta fidelidad con un costo computacional significativamente reducido. Los autores sostienen que este enfoque es particularmente beneficioso para analizar grandes muestras de sistemas de lentes, como los que se esperan de la encuesta Euclid, donde un modelado de fuentes robusto y flexible es crítico para minimizar los sesgos en las restricciones cosmológicas y de materia oscura. El método se presenta como una mejora práctica sobre las técnicas existentes de mallado adaptativo, ofreciendo una alternativa rápida, flexible y auto-diferenciable que no compromete la capacidad de detectar subestructuras.