HIcosmo: a differentiable JAX-based framework for cosmology inference

Abstract

The Stage IV cosmological surveys, such as Euclid, LSST, DESI, and SKA, will deliver observational data of unprecedented volume, calling for efficient and reliable inference tools. This paper presents HIcosmo (High-performance Inference for Cosmology), an open-source JAX-based framework for cosmology inference. In HIcosmo, the forward model, distance integrals, likelihood evaluations, posterior sampling, and Fisher forecasts are all built from JAX primitives, so that gradients and Hessians of the log-likelihood are obtained directly by automatic differentiation, without any finite-difference approximation. The framework implements the $\Lambda$CDM, $w$CDM, $w_0 w_a$CDM, and interacting dark-energy models, and provides likelihoods for Type Ia supernovae (Pantheon+, DES-SN5YR, Union3), baryon acoustic oscillations (DESI DR1/DR2, SDSS), Planck 2018 distance priors, local $H_0$ measurements, and strong-lensing time delays. Its scope is restricted to background cosmology, with Boltzmann solvers and full perturbation-level likelihoods left to external tools. We validate HIcosmo against the reference implementation of each likelihood and against Cobaya. $\chi^2$ values agree to absolute differences of $10^{-6}$-$10^{-2}$, and the marginalized constraints from the two codes differ by less than $0.2\sigma$ in every analysis tested. Leveraging just-in-time compilation and automatic differentiation, HIcosmo achieves about $8.7\times$ the end-to-end sampling throughput of Cobaya on CPU. As the dataset grows to survey scale, GPU acceleration over CPU reaches up to $20\times$. As applications, we present multi-probe $\Lambda$CDM joint constraints, dark-energy equation-of-state constraints, and Fisher forecasts for six 21 cm intensity-mapping surveys, including SKA1, MeerKAT, BINGO, Tianlai, and CHIME.