How it works

Four objects, one differentiable path

Theoryparameters to a providerhubble(z), pk_cb(k, z)

Likelihooddeclares what it requiresreturns log L

Modelchecks theory, data and priorslog_density(u)

Samplerruns on CPU or GPUnuts · nested · maximize

\(\nabla\)

Gradients everywhere

Theories are Equinox modules, so log_density is ready for jit, grad and vmap.

\(\ln Z\)

Posterior and evidence

BlackJAX NUTS with dense mass matrices, nested slice sampling for \(\ln Z\) and its error, L-BFGS for the best fit.

.txt

Cobaya-style output

Chains, .covmat and updated.yaml that getdist and anesthetic read as they are.

Bundled likelihoods

DESI DR2 BAO DESI DR1 BAO DESI DR1 ShapeFit + BAO DES-Dovekie Pantheon+ Pantheon+ SH0ES Union3 Union3.1 Planck PR4 Compressed\(\theta_*, \omega_\mathrm{b}, \omega_\mathrm{c}\) Planck 2018 low-\(\ell\) TT Commander Planck 2018 low-\(\ell\) EE SimAll Planck 2018 plik-lite TT, TE, EE Any .clik file clipy

Quick example

NUTSEvidence on GPUInstall data

import jax

jax.config.update("jax_enable_x64", True)
from cosmonuts import Model, Uniform, nuts, maximize
from cosmonuts.theories import W0WaCDM
from cosmonuts.likelihoods import DESIDR2BAO, DESDovekie, PlanckPR4EarlyLCDM

model = Model(
    W0WaCDM(),
    [PlanckPR4EarlyLCDM(), DESIDR2BAO(), DESDovekie()],
    {
        "h": Uniform(0.55, 0.85),
        "omega_b": Uniform(0.018, 0.026),
        "omega_cdm": Uniform(0.08, 0.16),
        "w0": Uniform(-2.0, 0.0),
        "wa": Uniform(-3.0, 2.0),
    },
)
best = maximize(model, num_starts=4, key=jax.random.PRNGKey(0))
chains = nuts(model, jax.random.PRNGKey(1), initial=best.params)

# pip install 'cosmonuts[gpu]'  (Linux, CUDA 13)
from cosmonuts import sample

ns = sample(
    model, jax.random.PRNGKey(2), method="nested", device="gpu", output="chains/w0wa_ns", progress=True
)
print(ns.logZ, ns.logZ_err)

cosmonuts install desi_dr2_bao des_dovekie desi_dr1
cosmonuts install pantheon_plus union3 union3.1
cosmonuts install --all --path ~/cosmo_data   # or set COSMONUTS_PACKAGES_PATH

Where next

• Install and data: extras, cosmonuts install, the bundled likelihoods.
• Models and priors: theories, likelihoods, priors and Defined reparametrisations.
• Samplers: NUTS, HMC, Metropolis, nested sampling and the best fit, on CPU or GPU.
• Run files and CLI: the same runs from YAML.
• Outputs and plots: the chain files, diagnostics and figures.
• Extending and the API reference.

Beta: under active development

cosmonuts and this site are still being built. The API, run-file format and output layout may change between releases without deprecation warnings, and some pages are incomplete. Pin the version you use, and check results against an established code before relying on them for publication.