HIcosmo: uma estrutura diferenciável baseada em JAX para inferência cosmológica

A arquitetura do HIcosmo é robusta e flexível, implementando modelos cosmológicos fundamentais como $Λ$CDM, $w$CDM, $w_0 w_a$CDM, além de modelos interativos de energia escura. Para a inferência, a estrutura integra probabilidades de diversas fontes observacionais, incluindo dados de supernovas Tipo Ia (Pantheon+, DES-SN5YR, Union3), oscilações acústicas de bárions (DESI DR1/DR2, SDSS), e medições da distância anterior do Planck 2018, entre outras. É crucial ressaltar que o escopo atual do HIcosmo é restrito à cosmologia de fundo, com os solucionadores de Boltzmann e as probabilidades completas de nível de perturbação sendo delegados a ferramentas externas, permitindo um foco otimizado em sua especialidade.

A validação rigorosa do HIcosmo foi conduzida comparando seus resultados com implementações de referência para cada probabilidade e com o Cobaya, um software de inferência cosmológica amplamente reconhecido. Os testes revelaram uma excelente concordância: os valores de $χ^2$ obtidos pelo HIcosmo e pelas referências apresentam diferenças absolutas na ordem de $10^{-6}$ a $10^{-2}$. Similarmente, as restrições marginalizadas derivadas dos dois códigos divergem em menos de $0, 2σ$ em todas as análises testadas. Essa precisão e consistência são fundamentais para assegurar a confiabilidade e a robustez dos resultados gerados pelo HIcosmo, confirmando sua adequação para análises cosmológicas de alta precisão.

Um dos maiores diferenciais do HIcosmo reside em sua notável eficiência computacional. Aproveitando as capacidades de compilação just-in-time (JIT) e diferenciação automática inerentes ao JAX, a estrutura alcança um rendimento de amostragem ponta a ponta na CPU aproximadamente 8, 7 vezes superior ao do Cobaya. Essa otimização é vital para processar o vasto volume de dados esperado das futuras pesquisas cosmológicas. Adicionalmente, em cenários onde o conjunto de dados cresce para atingir a escala das grandes pesquisas, a aceleração proporcionada pela GPU em relação à CPU pode chegar a impressionantes 20 vezes, evidenciando o potencial do HIcosmo para lidar com os desafios computacionais mais exigentes da cosmologia de Estágio IV.

Para demonstrar suas amplas capacidades, o HIcosmo foi aplicado em diversas frentes. Foram apresentadas restrições conjuntas $Λ$CDM multi-sonda, restrições da equação de estado da energia escura e previsões de Fisher para seis importantes pesquisas de mapeamento de intensidade de 21 cm. Essas pesquisas incluem projetos de grande relevância internacional como SKA1, MeerKAT, BINGO, Tianlai e CHIME. Tais aplicações não apenas ilustram a versatilidade da estrutura em abordar diferentes aspectos da inferência cosmológica, mas também sublinham seu papel como uma ferramenta poderosa para a caracterização de modelos cosmológicos fundamentais e a previsão de resultados para futuras observações.