Skip to content

Models and priors

Four objects make one differentiable log density:

  1. A Theory turns parameters into a provider of quantities (hubble(z), comoving_distance(z), sound_horizon_drag(), pk_cb(k, z), ...).
  2. Likelihood objects declare what they require and score the provider.
  3. Priors say which parameters are sampled.
  4. A Model checks that all of it fits together.
import jax

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

from cosmonuts import Model, Uniform, Fixed
from cosmonuts.theories import W0WaCDM
from cosmonuts.likelihoods import DESIDR2BAO

model = Model(W0WaCDM(rdrag="param"), [DESIDR2BAO()], {
    "h": Fixed(0.7),
    "omega_b": Fixed(0.0224),
    "omega_cdm": Uniform(0.05, 0.2),
    "w0": Uniform(-2.0, 0.0),
    "wa": Uniform(-3.0, 2.0),
    "rdrag": Uniform(120.0, 170.0),
})

The model raises ValueError at construction if a likelihood requires a quantity the theory does not provide, or a parameter has no prior.

Priors

Prior Meaning
Uniform(lower, upper) Flat on an interval
Normal(mean, std) Gaussian
Fixed(value) Held constant, not sampled
Defined(fn) A function of other parameters

The samplers move in an unconstrained space; the model maps to physical values and adds the Jacobian, so model.log_density(u) is a pure function for jit, grad and vmap.

Reparametrise with Defined

Changes of variables go in the priors dict, never in theory code. Sample what you want and define what the theory reads:

import jax.numpy as jnp
from cosmonuts import Defined, Uniform

priors = {
    "logA": Uniform(1.61, 3.91),
    "A_s": Defined(lambda logA: 1e-10 * jnp.exp(logA)),
    # ...
}

Defined parameters are saved as derived (*) columns in the chains. A parameter that only feeds a Defined one counts as used. Do not wrap a theory to change its inputs.

Labels and derived quantities

LaTeX labels for new names go in labels; without $:

model = Model(theory, likelihoods, priors, labels={"A_s": r"A_\mathrm{s}"})

derived={"name": fn(provider, params)} records extra quantities with the samples. Each likelihood's own derived is recorded too.

Bundled theories

Theory Provides Needs
W0WaCDM Background: \(H(z)\), distances, \(r_d\), \(\theta_*\) (\(\Lambda\)CDM with \(w_0=-1\), \(w_a=0\) fixed) core
ClassETheory Background and CMB spectra from class-e emulators extra classe
EFT4PyTheory Background, growth and \(P(k)\) of Horndeski EFT models extra eft4py
Theories Several theories combined; each quantity from the one that provides it none

Pass a dict {name: theory} to Model to combine theories, and provides={quantity: name} to choose which one serves a quantity several provide.