Quick Start

This guide will get you up and running with HIcosmo for cosmological parameter estimation.

Inference in 3 Lines of Code

HIcosmo’s design philosophy is 3 lines of code for core functionality:

from hicosmo import hicosmo

inf = hicosmo(cosmology="LCDM", likelihood="sn", free_params=["H0", "Omega_m"])
samples = inf.run()

After completion, you will obtain posterior distribution samples for \(H_0\) and \(\Omega_m\).

Tip

Parallel acceleration: Add the following code at the very beginning of your script to enable multi-device parallelism:

import hicosmo as hc
hc.init(8)  # 8 parallel devices = 8 parallel chains = ~800% CPU

This must be called before importing any other modules!

Basic Examples

Single-Probe Analysis (Supernovae)

from hicosmo import hicosmo

# Constrain LCDM model using Pantheon+ supernova data
inf = hicosmo(
    cosmology="LCDM",
    likelihood="sn",           # Pantheon+ data
    free_params=["H0", "Omega_m"],
    num_samples=4000,          # Total number of samples
    num_chains=4               # Number of parallel chains
)

# Run MCMC sampling
samples = inf.run()

# View results summary
inf.summary()

# Generate corner plot
inf.corner_plot('corner_sn.pdf')

Multi-Probe Joint Analysis

from hicosmo import hicosmo

# Joint SN + BAO + CMB constraints
inf = hicosmo(
    cosmology="LCDM",
    likelihood=["sn", "bao", "cmb"],  # Multi-probe joint analysis
    free_params=["H0", "Omega_m"],
    num_samples=8000,
    num_chains=4
)

samples = inf.run()
inf.summary()
inf.corner_plot('corner_joint.pdf')

Supported Likelihood Strings

HIcosmo supports the following convenient likelihood strings:

String

Corresponding Class

Description

sn

SN_likelihood("pantheon+")

Pantheon+ 1701 SNe Ia

sn_shoes

SN_likelihood("pantheon+shoes")

With Cepheid calibration

bao

BAO_likelihood("desi2024")

DESI 2024 BAO

bao_sdss_dr12

BAO_likelihood("sdss_dr12")

SDSS DR12 BAO

cmb

Planck2018DistancePriorsLikelihood

Planck 2018 compressed likelihood

h0licow

H0LiCOWLikelihood

Strong lensing time delay

tdcosmo

TDCOSMOLikelihood

TDCOSMO hierarchical Bayesian

sh0es

SH0ESLikelihood

SH0ES local \(H_0\)

Extended Dark Energy Models

wCDM model (constant dark energy equation of state):

inf = hicosmo(
    cosmology="wCDM",
    likelihood=["sn", "bao", "cmb"],
    free_params=["H0", "Omega_m", "w0"],
    num_samples=10000
)
samples = inf.run()

CPL parameterization (time-varying dark energy):

\[w(a) = w_0 + w_a (1 - a)\]
inf = hicosmo(
    cosmology="CPL",
    likelihood=["sn", "bao", "cmb"],
    free_params=["H0", "Omega_m", "w0", "wa"],
    num_samples=16000
)
samples = inf.run()

Custom Likelihood Functions

HIcosmo can be used not only for cosmological analysis but also for arbitrary parameter fitting problems. The following example demonstrates how to fit a quadratic polynomial \(y = ax^2 + bx + c\):

#!/usr/bin/env python
"""
Simple MCMC Example: Polynomial Fitting (y = a*x² + b*x + c)

"""
import hicosmo as hc
hc.init()
import numpy as np
import jax.numpy as jnp
from pathlib import Path

from hicosmo.samplers import MCMC
from hicosmo.visualization import Plotter

# Load data
data_path = Path(__file__).parent / 'data' / 'sim_data.txt'
x, y_obs, y_err = np.loadtxt(data_path, unpack=True)
x, y_obs, y_err = jnp.asarray(x), jnp.asarray(y_obs), jnp.asarray(y_err)


def log_likelihood(a, b, c):
    """Log-likelihood: -0.5 * chi²"""
    y_th = a * x**2 + b * x + c
    return -0.5 * jnp.sum((y_obs - y_th)**2 / y_err**2)


# Parameters: {name: (initial, min, max, latex)}
params = {
    'a': (3.5, 0.0, 10.0, r'$a$'),
    'b': (1.0, 0.0, 4.0,  r'$b$'),
    'c': (1.0, 0.0, 3.0,  r'$c$'),
}

if __name__ == '__main__':
    # Run MCMC
    mcmc = MCMC(params, log_likelihood, chain_name='polynomial_fit')
    mcmc.run(num_samples=20000)

    # Plot results
    plotter = Plotter('polynomial_fit')
    plotter.corner()
    plotter.report()

Parameter format explanation:

# Dict format: {name: (initial, min, max, latex)}
params = {
    'a': (3.5, 0.0, 10.0, r'$a$'),  # Parameter a: initial 3.5, range [0, 10]
    'b': (1.0, 0.0, 4.0,  r'$b$'),  # Parameter b: initial 1.0, range [0, 4]
    'c': (1.0, 0.0, 3.0,  r'$c$'),  # Parameter c: initial 1.0, range [0, 3]
}

Output:

INFO:hicosmo.samplers.init:✅ HIcosmo: 8 devices ready for parallel MCMC
INFO:hicosmo.samplers.inference:MCMC completed:
INFO:hicosmo.samplers.inference:  Start time: 2026-01-08 16:54:43
INFO:hicosmo.samplers.inference:  End time:   2026-01-08 16:54:45
INFO:hicosmo.samplers.inference:  Duration:   1.94s
INFO:hicosmo.samplers.inference:  Saved to:   mcmc_chains/polynomial_fit.h5

Tip

The generated report (plotter.report()) includes:

  • Parameter constraints (mean, median, 68%/95% confidence intervals)

  • LaTeX-formatted output (ready for papers)

  • MCMC diagnostics (ESS, R-hat)

  • Model selection criteria (AIC, Bayesian evidence)

Using the Class Interface (Advanced)

For scenarios requiring more control, you can use the class interface directly:

# 1. Initialize first (must be at the very beginning!)
import hicosmo as hc
hc.init(8)  # 8 parallel devices

# 2. Import other modules
from hicosmo.models import LCDM
from hicosmo.likelihoods import SN_likelihood, BAO_likelihood
from hicosmo.samplers import MCMC

# 3. Create likelihood objects
sne = SN_likelihood(LCDM, "pantheon+", M_B="marginalize")
bao = BAO_likelihood(LCDM, "desi2024")

# Joint likelihood
joint_likelihood = sne + bao

# 4. Configure parameters
params = {
    "H0": (70.0, 60.0, 80.0),       # (reference, min, max)
    "Omega_m": (0.3, 0.1, 0.5),
}

# 5. Create MCMC and run (num_chains defaults to device count = 8)
mcmc = MCMC(params, joint_likelihood, chain_name="lcdm_sn_bao")
samples = mcmc.run(num_samples=4000)

# 6. Output results
mcmc.print_summary()

Visualizing Results

Corner Plot:

from hicosmo.visualization import Plotter

# Load from saved chains
plotter = Plotter("lcdm_sn_bao")
plotter.corner(["H0", "Omega_m"], filename="corner.pdf")

Multi-chain comparison:

plotter = Plotter(
    ["chain_planck", "chain_shoes", "chain_bao"],
    chain_labels=["Planck", "SH0ES", "DESI BAO"]
)
plotter.corner(["H0", "Omega_m"], filename="h0_tension.pdf")

Saving and Loading Results

# Save results
inf.save_results("results/my_analysis.pkl")

# Load and continue analysis
from hicosmo.visualization import Plotter
plotter = Plotter("results/my_analysis.pkl")
plotter.corner(["H0", "Omega_m"])

YAML Configuration-Driven

For reproducible analyses, YAML configuration files are recommended:

config.yaml:

name: joint_analysis
theory: LCDM

likelihood:
  - name: sn
    dataset: pantheon+
  - name: bao
    dataset: desi2024
  - name: cmb
    dataset: planck2018

params:
  H0:
    prior: {dist: uniform, min: 50, max: 100}
    ref: 70.0
  Omega_m:
    prior: {dist: uniform, min: 0.1, max: 0.5}
    ref: 0.3

sampler:
  name: numpyro
  num_samples: 8000
  num_chains: 4
  num_warmup: 1000

output:
  root: results
  chain_name: joint_lcdm

Running from configuration:

from hicosmo import run_from_config

run_from_config("config.yaml")

Next Steps