Resumo Técnico: HIcosmo – Um Framework Diferenciável Baseado em JAX para Inferência Cosmológica

1. Definição do Problema

Os próximos levantamentos cosmológicos de Estágio IV (ex: Euclid, LSST, DESI, SKA) gerarão dados observacionais de volume sem precedentes. Embora a inferência cosmológica de nível de fundo (envolvendo distâncias, razões de distância e módulos de distância) opere tipicamente em espaços de parâmetros de baixa dimensionalidade ($d \sim 5\text{--}7$), o custo computacional da inferência permanece um gargalo devido ao enorme número de avaliações de verossimilhança (likelihood) exigidas pelos métodos de Cadeia de Markov Monte Carlo (MCMC).

Pipelines tradicionais (ex: CosmoMC, Cobaya, MontePython) dependem de amostradores de caminhada aleatória de Metropolis–Hastings (MH). Esses métodos sofrem com tempos de mistura (mixing) lentos que escalam linearmente com a dimensionalidade ($d$) e requerem aproximações de diferença finita para gradientes, que são computacionalmente caras ($O(d)$ avaliações extras) e sensíveis às escolhas de tamanho de passo. Além disso, os códigos existentes maduros não são implementados como programas diferenciáveis, impedindo o uso de amostradores baseados em gradiente, como o No-U-Turn Sampler (NUTS), e diferenciação automática (AD) para previsões de Fisher. Embora ferramentas baseadas em JAX tenham surgido para análises de nível de perturbação (ex: jax-cosmo, CosmoPower), não havia um framework sistematicamente validado e de ponta a ponta para a cosmologia de nível de fundo que integrasse verossimilhanças, amostragem NUTS e previsões de Fisher baseadas em AD dentro de uma única arquitetura.

2. Metodologia e Arquitetura

O artigo apresenta o HIcosmo (High-performance Inference for Cosmology), um framework de código aberto construído inteiramente sobre JAX. O framework é desenhado em torno do princípio da diferenciabilidade de ponta a ponta, garantindo que cada operação, desde os parâmetros cosmológicos até o logaritmo da verossimilhança escalar, seja rastreável pelo sistema de diferenciação automática do JAX.

2.1 Arquitetura em Camadas

O HIcosmo consiste em três camadas primárias:

1. Camada de Modelo: Computa previsões teóricas. Implementa quatro modelos de fundo: $\Lambda$CDM, $w$CDM, $w_0w_a$CDM (parametrização CPL) e um modelo de energia escura interagente ($I\Lambda$CDM). Crucialmente, substitui integradores adaptativos (que não são diferenciáveis) por quadratura trapezoidal em grade fixa em grades de redshift densas. Isso permite o cálculo exato de gradientes via regra da cadeia sem o uso de diferenças finitas.
2. Camada de Verossimilhança (Likelihood): Compara as previsões do modelo com dados observacionais. Suporta:
   • Supernovas Tipo Ia: Pantheon+, DES-SN5YR e Union3.
   • Oscilações Acústicas de Bárions (BAO): DESI DR1/DR2, SDSS, BOSS e 6dFGS.
   • CMB: Priores de distância do Planck 2018 (comprimidos em $R$, $l_A$, $\omega_b$).
   • $H_0$ Local: Escada de distância SH0ES.
   • Lenteamento Forte: Atrasos de tempo H0LiCOW.
   • Todas as verossimilhanças são implementadas em JAX puro, suportando compilação Just-In-Time (JIT) e vetorização.
3. Camada de Inferência: Consome o logaritmo da verossimilhança e suas derivadas. Fornece:
   • Amostragem NUTS: Implementada via NumPyro, utilizando gradientes diretamente do JAX para melhorar os tempos de mistura (escalando como $d^{1/4}$).
   • Previsão de Fisher: Computa a matriz Hessiana via diferenciação automática aninhada, eliminando a necessidade de seleção de tamanho de passo e erros de diferença finita.

2.2 Detalhes de Implementação

• Diferenciabilidade: Ao utilizar integração de passo fixo e pré-computar quantidades independentes de parâmetros (ex: matrizes de covariância inversa), o framework garante que gradientes e Hessianas sejam obtidos diretamente via AD de modo reverso.
• Aceleração de Hardware: O framework aproveita a execução unificada CPU/GPU/TPU do JAX. A compilação JIT otimiza a verossimilhança em código de máquina, contornando o overhead do interpretador Python durante as execuções MCMC.
• Calibração: Para garantir a precisão física, o framework calibra quantidades derivadas (como o horizonte sonoro $r_d$) contra o solver de Boltzmann CAMB, reduzindo os resíduos para $\sim 0,01\%$.

3. Principais Resultados

3.1 Validação e Precisão

O HIcosmo foi sistematicamente validado contra implementações de referência (Cobaya) e códigos de verossimilhança oficiais através dos modelos $\Lambda$CDM, $w$CDM e $w_0w_a$CDM.

• Precisão da Verossimilhança: Os valores de $\chi^2$ concordam com os códigos de referência com diferenças absoltas de $10^{-6}$ a $10^{-2}$.
• Restrições de Parâmetros: As restrições marginalizadas do HIcosmo e do Cobaya diferem por menos de 0,2$\sigma$ em todas as análises testadas. Para análises conjuntas SN+BAO+CMB, o acordo está dentro de 0,04$\sigma$.
• Consistência: O framework reproduz restrições estabelecidas, como $H_0 = 68,24 \pm 0,38$ km/s/Mpc e $\Omega_m = 0,307 \pm 0,005$ para a análise conjunta, confirmando que as aproximações diferenciáveis não introduzem viés significativo.

3.2 Benchmarks de Desempenho

• Aceleração sobre o Scipy: Para cálculos de distância, o HIcosmo alcança uma aceleração de 140$\times$ sobre a quadratura adaptativa do scipy após a compilação JIT. Para o pipeline completo de verossimilhança Pantheon+, ele é 16$\times$ mais rápido em CPU.
• Taxa de Amostragem (Throughput): Comparado ao Cobaya (usando amostradores MH) em uma CPU de 8 núcleos, o HIcosmo (usando NUTS) atinge um throughput de Tamanho de Amostra Efetiva (ESS) 8,7$\times$ maior (400 ESS/s vs 45,8 ESS/s). Isso se deve tanto à avaliação mais rápida da verossimilhança quanto à melhor mistura do NUTS ($\tau \approx 1,2$ vs $\tau \approx 11,4$).
• Aceleração de GPU: Embora as avaliações individuais em conjuntos de dados pequenos (Pantheon+) mostrem ganhos modestos de GPU devido ao overhead, a aceleração de GPU escala com o tamanho do conjunto de dados. Para conjuntos de dados sintéticos com $N_{SN} = 20.000$, a GPU atinge uma aceleração de 20,8$\times$ sobre a CPU.
• Matriz de Fisher: O cálculo da matriz de Fisher baseado em AD é estável e não requer ajuste de tamanho de passo, evitando os erros de arredondamento de ponto flutuante que assolam os métodos de diferença finita em pequenos tamanhos de passo.

3.3 Aplicações

O artigo demonstra a utilidade do framework através de:

• Restrições Conjuntas de Múltiplos Probes: Combinando Pantheon+, DESI BAO e priores de Planck para estreitar as restrições em $H_0$ e $\Omega_m$.
• Restrições de Energia Escura: Derivando restrições sobre $w_0$ e $w_a$ consistentes com os resultados do DESI 2024.
• Previsões de Mapeamento de Intensidade de 21 cm: Realizando previsões de Fisher para seis levantamentos (SKA1, MeerKAT, BINGO, Tianlai, CHIME). O framework prevê que o SKA1 Banda 1 poderia restringir a equação de estado da energia escura para $\sigma(w) \approx 0,046$.

4. Significância e Alegações

Os autores afirmam que o HIcosmo preenche uma lacuna crítica no cenário da inferência cosmológica ao fornecer o primeiro pipeline de cosmologia de fundo nativo em JAX, sistematicamente validado e de ponta a ponta.

• Eficiência: Demonstra que a amostragem baseada em gradiente (NUTS) combinada com verossimilhanças compiladas via JIT pode superar significativamente o MCMC tradicional de caminhada aleatória tanto em termos de velocidade computacional quanto de eficiência estatística (ESS).
• Escalabilidade: O framework está unicamente posicionado para lidar com os massivos conjuntos de dados de futuros levantamentos (Estágio IV), aproveitando a aceleração de GPU, onde os ganhos de velocidade podem chegar a 20$\times$ sobre a execução em CPU.
• Modularidade e Extensibilidade: Ao separar as camadas de modelo, verossimilhança e inferência, o framework permite a adição fácil de novos modelos ou conjuntos de dados sem a necessidade de reescrever o código de derivada.
• Fundação para Inferência de Nível de Perturbação: Embora atualmente restrito à cosmologia de fundo (excluindo solvers de Boltzmann completos), o artigo argumenta que a validação bem-sucedida de ponta a ponta no nível de fundo estabelece uma base robusta. A arquitetura foi projetada para ser estendida para análises de nível de perturbação, substituindo a camada de modelo por emuladores de Boltzmann diferenciáveis (ex: CosmoPower-JAX), permitindo assim a inferência baseada em gradiente para problemas de alta dimensionalidade do Estágio IV.

Em resumo, o HIcosmo oferece uma alternativa reprodutível e de alto desempenho às ferramentas existentes, provando que a diferenciação automática e a aceleração de hardware moderna podem ser aplicadas eficazmente a tarefas padrão de inferência cosmológica.