⚛️ general relativity

HIcosmo: a differentiable JAX-based framework for cosmology inference

This paper introduces HIcosmo, an open-source, JAX-based framework for efficient cosmological inference that leverages automatic differentiation and GPU acceleration to achieve significantly faster sampling speeds than existing tools like Cobaya while supporting a wide range of Stage IV survey likelihoods and cosmological models.

Original authors: Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Published 2026-06-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the universe as a giant, complex puzzle that scientists are trying to solve. For decades, they've been using a standard set of rules (called the $\Lambda$ CDM model) to figure out how the universe began, how it's expanding, and what it's made of. But recently, the pieces of the puzzle haven't been fitting together perfectly. Measurements from the early universe (like the Cosmic Microwave Background) disagree with measurements from the nearby universe (like exploding stars called Supernovae). This is known as the "Hubble Tension."

To solve this, new, massive telescopes are being built that will collect billions of data points. The problem? The computer programs scientists use to solve these puzzles are too slow and clunky to handle this new flood of information. They are like trying to count grains of sand on a beach using a single, slow calculator.

This paper introduces HIcosmo, a brand-new, super-fast computer tool designed to solve these cosmic puzzles efficiently. Here is how it works, explained simply:

1. The "Smart Calculator" (JAX and Differentiability)

Traditional tools are like a chef who has to manually measure every ingredient, write down the result, and then manually calculate the next step. If they need to know how changing one ingredient affects the final taste, they have to guess and check, which takes forever.

HIcosmo is different. It is built on a technology called JAX, which acts like a "smart calculator" that understands the entire recipe at once.

Automatic Differentiation: Instead of guessing, HIcosmo knows exactly how the result changes if you tweak a single number. It's like having a GPS that instantly tells you the best route to take without needing to drive every possible road to see which is fastest.
Just-in-Time Compilation: The first time you ask HIcosmo to do a calculation, it takes a split second to "warm up" and translate the instructions into super-optimized machine code. After that, it runs at lightning speed, bypassing the slow "translator" (the Python interpreter) that slows down older tools.

2. The Three-Layer Cake

The authors built HIcosmo like a three-layer cake, where each layer has a specific job:

The Bottom Layer (The Model): This is the "theory" part. It takes the rules of the universe (like how much dark matter or dark energy there is) and calculates what the universe should look like.
The Middle Layer (The Likelihood): This is the "comparison" part. It takes the theory from the bottom layer and compares it against real data from telescopes (like the Pantheon+ supernova data or DESI galaxy maps). It calculates how "wrong" the theory is.
The Top Layer (The Inference): This is the "solver" part. It uses the information from the middle layer to find the best possible answer. It uses a smart method called NUTS (No-U-Turn Sampler) that uses the "smart calculator" gradients to zoom straight to the correct answer, rather than wandering around aimlessly like older tools do.

3. What It Can Do (and What It Can't)

What it does:
HIcosmo is currently the master of background cosmology. Think of this as measuring the "big picture" of the universe: how fast it's expanding, how far away things are, and the overall shape of space. It handles data from:

Exploding stars (Supernovae).
Galaxy clustering (BAO).
The afterglow of the Big Bang (CMB distance priors).
Time delays in gravitational lenses.

It has been tested against the current gold-standard software (called Cobaya) and found to be statistically identical in its answers but much faster.

On a standard computer processor (CPU), it is about 8.7 times faster than the old tool.
On a powerful graphics card (GPU), which is great for handling massive amounts of data at once, it can be up to 20 times faster as the data gets bigger.

What it doesn't do:
The paper is very clear that HIcosmo does not currently handle the "fine details" or "perturbations" of the universe, such as the complex wiggles in the cosmic web or the detailed power spectra of the early universe. Those require a different, very heavy engine (called a Boltzmann solver) that HIcosmo leaves to other tools. Think of HIcosmo as the expert on the shape of the balloon, while other tools are needed to analyze the texture of the rubber.

4. Real-World Results

The authors didn't just build it; they used it to solve real problems:

Joint Constraints: They combined data from supernovae, galaxy surveys, and the CMB to pin down the universe's expansion rate ( $H_0$ ) and matter density ( $\Omega_m$ ) with incredible precision.
Future Forecasting: They used HIcosmo to predict how well future radio telescopes (like SKA, MeerKAT, and CHIME) will be able to measure the "equation of state" of dark energy (a fancy way of asking: "Is dark energy constant, or does it change over time?"). They found that these new telescopes could measure this with amazing accuracy.

The Bottom Line

HIcosmo is a modern, high-speed engine for cosmology. It replaces the slow, manual methods of the past with a streamlined, automatic system that can handle the massive data floods coming from the next generation of telescopes. It doesn't change the laws of physics, but it changes how fast and efficiently we can test them against the universe.

Technical Summary: HIcosmo – A Differentiable JAX-Based Framework for Cosmology Inference

1. Problem Statement

The upcoming Stage IV cosmological surveys (e.g., Euclid, LSST, DESI, SKA) will generate observational data of unprecedented volume. While background-level cosmological inference (involving distances, distance ratios, and distance moduli) typically operates in low-dimensional parameter spaces ( $d \sim 5\text{--}7$ ), the computational cost of inference remains a bottleneck due to the sheer number of likelihood evaluations required by Markov Chain Monte Carlo (MCMC) methods.

Traditional pipelines (e.g., CosmoMC, Cobaya, MontePython) rely on random-walk Metropolis–Hastings (MH) samplers. These methods suffer from slow mixing times that scale linearly with dimensionality ( $d$ ) and require finite-difference approximations for gradients, which are computationally expensive ( $O(d)$ extra evaluations) and sensitive to step-size choices. Furthermore, existing mature codes are not implemented as differentiable programs, preventing the use of gradient-based samplers like the No-U-Turn Sampler (NUTS) and automatic differentiation (AD) for Fisher forecasting. While JAX-based tools have emerged for perturbation-level analyses (e.g., jax-cosmo, CosmoPower), there was no systematically validated, end-to-end differentiable framework for background-level cosmology that integrates likelihoods, NUTS sampling, and AD-based Fisher forecasts within a single architecture.

2. Methodology and Architecture

The paper presents HIcosmo (High-performance Inference for Cosmology), an open-source framework built entirely on JAX. The framework is designed around the principle of end-to-end differentiability, ensuring that every operation from cosmological parameters to the scalar log-likelihood is traceable by JAX's automatic differentiation system.

2.1 Layered Architecture

HIcosmo consists of three primary layers:

Model Layer: Computes theoretical predictions. It implements four background models: $\Lambda$ CDM, $w$ CDM, $w_0w_a$ CDM (CPL parametrization), and an interacting dark-energy model ( $I\Lambda$ CDM). Crucially, it replaces adaptive integrators (which are non-differentiable) with fixed-grid trapezoidal quadrature on dense redshift grids. This allows for exact gradient computation via the chain rule without finite differences.
Likelihood Layer: Compares model predictions with observational data. It supports:
- Type Ia Supernovae: Pantheon+, DES-SN5YR, and Union3.
- Baryon Acoustic Oscillations (BAO): DESI DR1/DR2, SDSS, BOSS, and 6dFGS.
- CMB: Planck 2018 distance priors (compressed into $R$ , $l_A$ , $\omega_b$ ).
- Local $H_0$ : SH0ES distance ladder.
- Strong Lensing: H0LiCOW time delays.
- All likelihoods are implemented in pure JAX, supporting Just-In-Time (JIT) compilation and vectorization.
Inference Layer: Consumes the log-likelihood and its derivatives. It provides:
- NUTS Sampling: Implemented via NumPyro, utilizing gradients directly from JAX to improve mixing times (scaling as $d^{1/4}$ ).
- Fisher Forecasting: Computes the Hessian matrix via nested automatic differentiation, eliminating the need for step-size selection and finite-difference errors.

2.2 Implementation Details

Differentiability: By using fixed-step integration and precomputing parameter-independent quantities (e.g., inverse covariance matrices), the framework ensures that gradients and Hessians are obtained directly via reverse-mode AD.
Hardware Acceleration: The framework leverages JAX's unified CPU/GPU/TPU execution. JIT compilation optimizes the likelihood into machine code, bypassing the Python interpreter overhead during MCMC runs.
Calibration: To ensure physical accuracy, the framework calibrates derived quantities (like the sound horizon $r_d$ ) against the Boltzmann solver CAMB, reducing residuals to $\sim 0.01\%$ .

3. Key Results

3.1 Validation and Accuracy

HIcosmo was systematically validated against reference implementations (Cobaya) and official likelihood codes across $\Lambda$ CDM, $w$ CDM, and $w_0w_a$ CDM models.

Likelihood Precision: The $\chi^2$ values agree with reference codes to absolute differences of $10^{-6}$ to $10^{-2}$ .
Parameter Constraints: Marginalized constraints from HIcosmo and Cobaya differ by less than 0.2 $\sigma$ in all tested analyses. For joint SN+BAO+CMB analyses, the agreement is within 0.04 $\sigma$ .
Consistency: The framework reproduces established constraints, such as $H_0 = 68.24 \pm 0.38$ km/s/Mpc and $\Omega_m = 0.307 \pm 0.005$ for the joint analysis, confirming that the differentiable approximations do not introduce significant bias.

3.2 Performance Benchmarks

Speedup over Scipy: For distance calculations, HIcosmo achieves a 140 $\times$ speedup over scipy adaptive quadrature after JIT compilation. For the full Pantheon+ likelihood pipeline, it is 16 $\times$ faster on CPU.
Sampling Throughput: Compared to Cobaya (using MH samplers) on an 8-core CPU, HIcosmo (using NUTS) achieves an 8.7 $\times$ higher Effective Sample Size (ESS) throughput (400 ESS/s vs. 45.8 ESS/s). This is due to both the faster likelihood evaluation and the superior mixing of NUTS ( $\tau \approx 1.2$ vs. $\tau \approx 11.4$ ).
GPU Acceleration: While single evaluations on small datasets (Pantheon+) show modest GPU gains due to overhead, the GPU acceleration scales with dataset size. For synthetic datasets with $N_{SN} = 20,000$ , the GPU achieves a 20.8 $\times$ speedup over the CPU.
Fisher Matrix: AD-based Fisher matrix computation is stable and requires no step-size tuning, avoiding the floating-point round-off errors that plague finite-difference methods at small step sizes.

3.3 Applications

The paper demonstrates the framework's utility through:

Multi-probe Joint Constraints: Combining Pantheon+, DESI BAO, and Planck priors to tighten constraints on $H_0$ and $\Omega_m$ .
Dark Energy Constraints: Deriving constraints on $w_0$ and $w_a$ consistent with DESI 2024 results.
21 cm Intensity Mapping Forecasts: Performing Fisher forecasts for six surveys (SKA1, MeerKAT, BINGO, Tianlai, CHIME). The framework predicts that SKA1 Band 1 could constrain the dark energy equation of state to $\sigma(w) \approx 0.046$ .

4. Significance and Claims

The authors claim that HIcosmo fills a critical gap in the cosmological inference landscape by providing the first systematically validated, end-to-end differentiable JAX-native pipeline for background cosmology.

Efficiency: It demonstrates that gradient-based sampling (NUTS) combined with JIT-compiled likelihoods can significantly outperform traditional random-walk MCMC in terms of both computational speed and statistical efficiency (ESS).
Scalability: The framework is uniquely positioned to handle the massive datasets of future surveys (Stage IV) by leveraging GPU acceleration, where speedups can reach 20 $\times$ over CPU execution.
Modularity and Extensibility: By separating the model, likelihood, and inference layers, the framework allows for easy addition of new models or datasets without rewriting derivative code.
Foundation for Perturbation-Level Inference: While currently restricted to background cosmology (excluding full Boltzmann solvers), the paper argues that the successful end-to-end validation at the background level establishes a robust foundation. The architecture is designed to be extended to perturbation-level analyses by replacing the model layer with differentiable Boltzmann emulators (e.g., CosmoPower-JAX), thereby enabling gradient-based inference for high-dimensional Stage IV problems.

In summary, HIcosmo offers a reproducible, high-performance alternative to existing tools, proving that automatic differentiation and modern hardware acceleration can be effectively applied to standard cosmological inference tasks.