Skip to content

Your first fit: flat \(\Lambda\)CDM

Fit flat \(\Lambda\)CDM to DESI DR2 BAO and Union3.1 supernovae, with a BBN prior on \(\omega_\mathrm{b}\). \(\Lambda\)CDM is the bundled W0WaCDM with \(w_0 = -1\) and \(w_a = 0\) fixed. Download the data once:

cosmonuts install desi_dr2_bao union3.1

Then run the fit from Python or from the command line. The two are equivalent.

lcdm_union31_desi.py
import jax

jax.config.update("jax_enable_x64", True)

from cosmonuts import Defined, Fixed, Model, Normal, Uniform, nuts
from cosmonuts.likelihoods import DESIDR2BAO, Union31
from cosmonuts.theories import W0WaCDM

model = Model(W0WaCDM(m_nu=0.06), [DESIDR2BAO(), Union31()], {
    "h": Uniform(0.5, 0.9),
    "omega_b": Normal(0.02218, 0.00055),     # BBN
    "omega_cdm": Uniform(0.05, 0.25),
    "w0": Fixed(-1.0),                       # LCDM
    "wa": Fixed(0.0),
    "Omega_m": Defined(
        lambda h, omega_b, omega_cdm: (omega_b + omega_cdm + 0.06 / 93.14) / h**2),
})

chains = nuts(model, jax.random.PRNGKey(0), num_samples=1000,
              output="chains/lcdm_union31_desi")
h, Om = chains.samples["h"], chains.derived["Omega_m"]
print(f"h = {h.mean():.4f} ± {h.std():.4f}")
print(f"Omega_m = {Om.mean():.4f} ± {Om.std():.4f}")

python lcdm_union31_desi.py
lcdm_union31_desi.yaml
theory:
  W0WaCDM: {m_nu: 0.06}

likelihood:
  DESIDR2BAO: null
  Union31: null

params:
  h: {prior: {min: 0.5, max: 0.9}, ref: 0.68}
  omega_b: {prior: {dist: norm, loc: 0.02218, scale: 0.00055}}  # BBN
  omega_cdm: {prior: {min: 0.05, max: 0.25}, ref: 0.12}
  w0: -1.0                                                       # LCDM
  wa: 0.0
  Omega_m:
    value: "lambda h, omega_b, omega_cdm: (omega_b + omega_cdm + 0.06 / 93.14) / h**2"

sampler:
  nuts: {num_samples: 1000, num_chains: 4}

output: chains/lcdm_union31_desi

cosmonuts sample lcdm_union31_desi.yaml --test   # build, evaluate once, stop
cosmonuts sample lcdm_union31_desi.yaml          # NUTS, 4 chains

Either way you get

\[ h = 0.686 \pm 0.006, \qquad \Omega_\mathrm{m} = 0.302 \pm 0.008, \]

in about 15 s on a laptop CPU. The chains in chains/ are in Cobaya's format, so getdist.loadMCSamples("chains/lcdm_union31_desi") reads them.

If a run stops

A run with an output saves its progress as it goes. If it is interrupted (a crash, a killed job, a dead laptop), run the same thing again with resume on to continue it:

chains = nuts(model, jax.random.PRNGKey(0), num_samples=1000,
              output="chains/lcdm_union31_desi", resume=True)

cosmonuts sample lcdm_union31_desi.yaml --resume   # or -r, or `resume: true` in the file

What is resumed depends on the sampler:

  • NUTS / HMC. Each chain restarts from its last saved draw with the saved step size and mass matrix (<output>.adaptation.npz), with no warm-up. New draws are appended to the same files. num_samples counts all the draws per chain, old ones included.
  • Nested sampling. The live points, the evidence and the random key are saved after every chunk (<output>.nested_state.npz, next to <output>_dead-birth.txt), so the resumed run ends as if it had never stopped. Settings you leave unset are read from the saved run, and settings you do give must match it (nlive, num_delete, ...). frac_remain may differ, for example to take a finished run to a tighter tolerance.

resume and force (overwrite=True in Python) exclude each other. Without either, a run whose output already exists stops with FileExistsError instead of overwriting it.