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

Resumo Técnico: Processos Gaussianos em Grades Uniformes Transformadas por Guia de Raios

Enunciado do Problema
O lenticulamento gravitacional forte (SGL) é uma ferramenta crítica para restringir a cosmologia e as distribções de matéria escura, mas a inferência robusta depende fortemente de uma reconstrução precisa da fonte. Um desafio fundamental no SGL é o conteúdo de informação não uniforme no plano da fonte: regiões que produzem múltiplas imagens (por exemplo, quatro imagens) possuem taxas de amostragem mais altas do que aquelas que produzem menos imagens, e a magnificação varia continuamente. Técnicas tradicionais de malha adaptativa, como as tesselações de Delaunay ou a binarização de Voronoi, tentam abordar isso posicionando a resolução onde ela é necessária. No entanto, esses métodos frequentemente dependem de operações descontínuas (por exemplo, troca de vérticos nas arestas da malha conforme os modelos de massa mudam), o que pode quebrar a diferenciabilidade necessária para os frameworks de inferência modernos e restringir as escolhas de regularização. Por outro lado, os priors de Processos Gaussianos (GP) padrão oferecem regularização flexível, mas tipicamente exigem grades regulares para utilizar Transformadas Rápidas de Fourier (FFTs) para eficiência computacional, criando uma tensão entre adaptatividade e diferenciabilidade.

Metodologia
Os autores propõem uma nova abordagem de reconstrução de fonte usando uma grade Uniforme Transformada por Guia de Raios (RTU). Este método visa combinar a eficiência computacional de GPs baseados em FFT em grades uniformes com a resolução adaptativa exigida pelos dados de lenticulamento.

1. Construção da Grade RTU:

   • A fonte é modelada em uma grade uniforme em um espaço de coordenadas transformado $\mathbf{y} = \mathbf{g}(\mathbf{x})$.
   • A transformação $\mathbf{g}$ é derivada das funções de distribuição cumulativa empíricas (eCDFs) dos raios traçados de volta do plano da imagem para o plano da fonte.
   • Especificamente, os autores projetam as posições dos raios em cada eixo de coordenadas, ordenam-nos para formar eCDFs 1D e ajustam splines cúbicas a essas distribuições. Isso garante que a transformação seja pelo menos duas vezes continuamente diferenciável.
   • As eCDFs são construídas usando apenas raios que caem dentro de uma máscara de arco (definida por limiares de relação sinal-ruído), garantindo que a resolução seja alocada apenas em regiões onde a fonte é observada.
   • A transformação efetivamente estica as coordenadas do plano da fonte para que a densidade de raios se torne aproximadamente uniforme nas novas coordenadas $\mathbf{y}$.

2. Modelagem da Fonte:

   • O brilho da fonte é definido como um Processo Gaussiano nesta grade uniforme transformada.
   • O GP é realizado gerando ruído branco no espaço de Fourier, multiplicando-o pela raiz quadrada de um espectro de potência (especificamente um espectro de potência Matérn flexível) e aplicando uma FFT inversa.
   • A positividade é imposta via uma função softplus, e o brilho final da fonte é escalonado por uma lei de potência para permitir características mais nítidas.
   • O mapeamento de volta para o plano físico da fonte é realizado via $s_{source}(\mathbf{x}) = s^*_{source}(\mathbf{g}(\mathbf{x}))$.

3. Framework de Inferência:

   • A abordagem é implementada no Herculens, um framework de modelagem direta auto-diferenciável escrito em JAX.
   • A inferência é realizada usando Inferência Variacional Estocástica (SVI) com Numpyro, otimizando o Limite Inferior de Evidência (ELBO).
   • O método suporta escolhas arbitrárias de espectro de potência para regularização, mantendo a flexibilidade dos GPs enquanto aproveita a velocidade da FFT.

Principais Contribuições

• Adaptatividade Diferenciável: A grade RTU fornece um esquema de resolução adaptativa que é continuamente diferenciável, superando os problemas de descontinuidade das malhas de Delaunay/Voronoi, mantendo a velocidade dos GPs baseados em FFT.
• Regra de Bolso de Eficiência: Os autores propõem uma heurística para selecionar o número de pixels da fonte ($N_{src}$) com base na densidade de raios, sugerindo aproximadamente 4 raios por pixel para corresponder aos graus de liberdade na interpolação bilinear.
• Implementação: O método está totalmente integrado ao framework Herculens, tornando-o compatível com algoritmos de amostragem de última geração e fluxos de trabalho auto-diferenciáveis.

Resultados
Os autores testaram a grade RTU em três conjuntos de dados simulados (simulando qualidade JWST NIRCAM) e dados reais do sistema de lente SDSS J0946+1006.

• Qualidade do Ajuste e Eficiência de Pixels: A grade RTU alcança uma qualidade de ajuste no plano da imagem ($\chi^2_{img}$) comparável às grades uniformes, mas requer aproximadamente metade do número de pixels por dimensão (um fator de $\simeq 4$ menos pixels no total) para atingir os resíduos de nível de ruído.
• Desempenho do ELBO: Para o mesmo número de pixels, a grade RTU geralmente produz Limites Inferiores de Evidência (ELBOs) mais altos em comparação com grades uniformes, indicando um melhor ajuste aos dados sob a aproximação variacional.
• Sensibilidade à Subestrutura: O uso da grade RTU afeta apenas levemente a diferença no ELBO entre modelos com e sem subestruturas. O método distingue com sucesso modelos com e sem subestruturas (por exemplo, no Mock 1 e SDSS J0946+1006) sem inflar artificialmente a preferência por subestrutura devido a artefatos de grade.
• Precisão da Reconstrução da Fonte: Quando o modelo de massa subjacente é preciso, a grade RTU produz reconstruções de fonte mais precisas (menor $\Delta^2_{src}$) com menos pixels do que a abordagem de grade uniforme.

Significância e Alegações
O artigo afirma que o método da grade RTU resolve o compromisso entre realismo, flexibilidade, auto-diferenciabilidade e eficiência computacional no modelamento de lentes. Ao garantir que os pixels da fonte contenham um número mais uniforme de raios, o método permite uma reconstrução de fonte de alta fidelidade com redução significativa do custo computacional. Os autores afirmam que esta abordagem é particularmente benéfica para analisar grandes amostras de sistemas de lentes, como as esperadas para o levantamento Euclid, onde o modelamento de fontes robusto e flexível é crítico para minimizar vieses em restrições cosmológicas e de matéria escura. O método é apresentado como uma melhoria prática sobre as técnicas existentes de malha adaptativa, oferecendo uma alternativa rápida, flexível e auto-diferenciável que não compromete a capacidade de detectar subestruturas.