Using Neural Emulators and Hamiltonian Monte Carlo to constrain the Epoch of Reionization's History with the Ly$α$ Forest Power Spectrum

Resumo Técnico: Utilizando Emuladores Neurais e Monte Carlo Hamiltoniano para Restringir a História da Época da Reionização

Definição do Problema
A Época da Reionização (EoR) representa uma transição crítica na história cósmica, mas sua cronologia detalhada e estrutura inominada permanecem mal restringidas. Embora as medições da Radiação Cósmica de Fundo em Micro-ondas (CMB) forneçam uma restrição integrada da profundidade óptica, elas são insensíveis à duração e topologia específicas da reionização. As restrições atuais a partir de espectros de quasares de alto redshift (por exemplo, asas de amortecimento, frações de pixels escuros) oferecem precisão limitada e frequentemente dependem de interpretações dependentes de modelos. O floresta de Lyman-$\alpha$ (Ly$\alpha$) em $z \sim 5$ oferece uma promissora sonda dos estados térmicos e de ionização do meio intergaláctico (IGM), mas extrair histórias detalhadas de reionização requer o aproveitamento de toda a distribuição de opacidades de Ly$\alpha$.

Uma barreira significativa à inferência quantitativa é o custo computacional. Gerar simulações hidrodinâmicas com resolução suficiente para capturar estruturas de pequena escala relevantes para a EoR requer milhares de horas-GPU por modelo. Consequentemente, estudos anteriores foram limitados a comparações qualitativas ou espaços de parâmetros restritos. Além disso, a estimativa robusta de parâmetros requer matrizes de covariância precisas, que são difíceis de derivar empiricamente a partir das amostras esparsas de quasares de alto-$z$ disponíveis. Abordagens tradicionais frequentemente assumem estruturas de covariância fixas ou dependem de conjuntos de simulações proibitivamente caros para computar covariâncias dependentes do modelo.

Metodologia
Este trabalho apresenta um novo framework de inferência baseado em JAX que integra observáveis da floresta de Ly$\alpha$ modelados via forward-modeling com emuladores neurais diferenciáveis e Monte Carlo Hamiltoniano (HMC). A metodologia procede através dos seguintes passos:

1. Construção do Conjunto de Dados de Simulação: Os autores geraram um conjunto de dados de 501 simulações hidrodinâmicas de baixa resolução ($256^3$ voxels, $L_{box} = 20 \, h^{-1}$ Mpc) usando o código Nyx. Essas simulações foram impulsionadas por histórias de reionização geradas pelo AMBER, um código semi-numérico que mapeia parâmetros definidos pelo usuário para campos de ionização. O espaço de parâmetros inclui o ponto médio da reionização ($z_{mid}$), duração ($\Delta z$), assimetria ($A_z$), temperatura de aquecimento máxima ($\Delta T_{re}$) e fluxo médio ($\langle F \rangle$).
2. Modelagem Direta (Forward Modeling): A partir dessas simulações, os autores geraram espias (skewers) de fluxo 1D de Ly$\alpha$, aplicando efeitos observacionais como resolução de espectrógrafo ($R=30.000$) e ruído ($SNR=30$). Eles computaram o espectro de potência 1D médio ($\langle P_{Ly\alpha} \rangle$) e matrizes de covariância dependentes do modelo ($\Sigma_{P_{Ly\alpha}}$) para cada modelo.
3. Emulação Neural: Dois emuladores neurais independentes foram treinados:
   • Emulador de Espectro de Potência: Uma rede feed-forward totalmente conectada que prevê $\log_{10}\langle P_{Ly\alpha} \rangle$ a partir dos parâmetros de entrada.
   • Emulador de Matriz de Covariância: Para garantir que a saída permaneça simétrica e definida positiva, os autores empregaram a decomposição de Cholesky ($\Sigma = LL^T$). Uma rede neural prevê o fator de Cholesky inferior achatado ($L_{flat}$), com uma ativação exponencial aplicada aos elementos diagonais para impor positividade.
4. Pipeline de Inferência: Os emuladores foram combinados com o Monte Carlo Hamiltoniano (HMC) via pacote Numpyro. A diferenciabilidade dos emuladores baseados em JAX permite o cálculo dos gradientes necessários para o HMC, acelerando significamente a exploração do espaço de parâmetros em comparação com algoritmos padrão de Metropolis-Hastings.

Principais Contribuições

• Emulação de Covariância Diferenciável: Este trabalho apresenta a primeira tentativa de emular matrizes de covariência da floresta de Ly$\alpha$ usando redes neurais. Ao prever o fator de Cholesky, o método garante validade física (simetria e definição positiva) enquanto captura correlações dependentes de parâmetros.
• Framework Diferenciável de Ponta a Ponta: A integração de modelagem direta, emulação neural e HMC dentro do ecossistema JAX cria um pipeline totalmente diferenciável para inferência de parâmetros cosmológicos.
• Validação Quantitativa: O framework é validado contra um conjunto de 100 observações fictícias (mock), demonstrando a capacidade de recuperar os parâmetros reais e avaliando a validade estatística das distribuições posteriores resultantes.

Resultos

• Precisão do Emulador: O emulador de espectro de potência alcançou um erro percentual absoluto médio (MAPE) mediano de 0,15% no conjunto de teste, com 95% dos erros permanecendo abaixo de 2%. O emulador de matriz de covariância alcançou um erro mediano de 0,42% (relativo ao produto dos elementos diagonais), com erros geralmente permanecendo abaixo de 5% em todas as escalas. Ambos os emuladores performaram com precisão bem abaixo dos níveis de ruído observacional esperados assumidos na modelagem direta.
• Recuperação de Parâmetros: Quando aplicado a observações fictícias, o framework HMC recuperou com sucesso os parâmetros de entrada reais. As distribuições posteriores foram visualmente consistentes com a verdade fundamental (ground truth).
• Validade Estatística: Um teste de cobertura realizado em 100 observações fictícias revelou que as distribuições posteriores são ligeiramente excessivamente confiantes (a probabilidade de cobertura estimada reside marginalmente abaixo da linha ideal). Os autores atribuem isso primariamente à suposição de uma função de verossimilhança multivariada Gaussiana, observando que a floresta de Ly$\alpha$ é inerentemente um campo aleatório não-Gaussiano.
• Eficiência de Dados: Os autores testaram o método com um conjunto de treinamento reduzido de apenas 100 simulações. Embora os erros de emulação tenham aumentado (o MAPE do espectro de potência subiu para ~1,2%, os erros de covariância aumentaram), as distribuições posteriores resultantes e a cobertura estatística permaneceram comparáveis às derivadas do conjunto de dados completo. Isso sugere que o método é robusto mesmo com dados de treinamento limitados, um fator crucial para futuras simulações de alta resolução.

Significância e Limitações
O artigo afirma que este framework fornece um caminho viável para restringir quantitativamente a história da EoR usando espectros de potência da floresta de Ly$\alpha$. A principal significância reside em demonstrar que emuladores neurais podem substituir eficientemente modelos diretos caros e que matrizes de covariência dependentes de modelo podem ser emuladas, permitindo uma inferência Bayesiana mais rápida e rigorosa.

No entanto, os autores declaram explicitamente que os resultados atuais ainda não podem ser aplicados a dados observacionais reais devido à baixa resolução das simulações utilizadas, que suprime o poder em pequenas escalas. Além disso, o modelo de reionização depende de prescrições simplificadas de injeção de calor e carece de flutuações espaciais no fundo ultravioleta (UVB). Os autores concluem que trabalhos futuros devem incorporar simulações de maior resolução, tratamentos térmicos mais motivados fisicamente e modelos de UVB flutuante. Eles também sugerem que a leve excessiva confiança observada nas posteriores pode ser abordada movendo-se em direção a métodos de inferência baseados em simulação (SBI) que não dependem de suposições de verossimilhança Gaussiana.