Cobaya: emulators in place of CLASS¶

class_e.cobaya.ClassE is a Cobaya theory that serves a class-e bundle through the same interface as Cobaya's classy. Replace the classy block and the likelihoods run unchanged.

pip install 'class-e[cobaya]'

Replace classy¶

class-eclassy

theory:
  class_e.cobaya.ClassE:
    bundle: bundles/w0wacdm-mnu-theta.cemu   # one bundle: every emulator of the model
    out_of_bounds: reject                    # reject (default) | warn | ignore
    prior_check: 1000                        # prior samples for the coverage warning (0: off)
    input_map: {}                            # {bundle input: sampled name or expression}
likelihood:
  planck_2018_highl_plik.TTTEEE_lite_native: null
  planck_2018_lowl.TT: null
  planck_2018_lowl.EE: null
  bao.desi_dr2: null
  sn.union3: null
params:
  logA: {prior: {min: 2.95, max: 3.15}, drop: true, latex: \ln(10^{10} A_\mathrm{s})}
  A_s: {value: "lambda logA: 1e-10*np.exp(logA)"}
  theta_s_100: {prior: {min: 1.0390, max: 1.0430}, drop: true}
  100*theta_s: {value: "lambda theta_s_100: theta_s_100", derived: false}
  # ... omega_b, omega_cdm, n_s, tau_reio, w0_fld, wa_fld, A_planck as for classy
  H0: {latex: H_0}
  rdrag: {latex: r_\mathrm{drag}}

theory:
  classy:
    extra_args: {N_ncdm: 1, m_ncdm: 0.06, N_ur: 2.0308, non_linear: hmcode, ...}
likelihood:
  # ... the same
params:
  # ... the same

ClassE is a BoltzmannBase, the base of Cobaya's classy and camb. It has the same getters (get_Cl, get_Hubble, get_angular_diameter_distance, ...), the same units, and Cobaya's classy renames (rdrag → rs_drag, tau → tau_reio, omegabh2 → omega_b, ...).

CLASS settings such as extra_args have no effect here: the bundle was trained with fixed settings.

What the bundle must contain¶

One bundle per model, holding every emulator for that parameter space. ClassE predicts only what the likelihoods request: a BAO-only run never evaluates the CMB emulators.

Likelihood	Needs	Bundle observables
`planck_2018_highl_plik.TTTEEE_lite_native`	lensed `Cl` tt, te, ee to $\ell = 2508$	`TT`, `TE`, `EE` on every $\ell$ from 2 to $\geq 2508$
`planck_2018_lowl.TT`, `.EE`	`Cl` tt or ee to $\ell = 29$	`TT`, `EE`
`bao.desi_dr2`	$D_A(z)$, $H(z)$, derived `rdrag`	`Hz` (distance emulator), `rs_drag`
`sn.union3` and other SNe	$D_A(z)$	`Hz`

Spectra are returned from $\ell = 0$, with $C_0 = C_1 = 0$, in CLASS's dimensionless convention. get_Cl applies the units (FIRASmuK2, muK2, ...) and $\ell$ factors exactly as classy does. A bundle of unlensed spectra provides unlensed_Cl instead.
Distances come from the distance emulator of Hz, with the curvature of its background, for any $z$ in its grid. See Distances and expansion history, including how to train one bundle for CMB and background in a model sampled in $\theta_s$.
Derived parameters come from the bundle's scalar observables (rs_drag, H0, z_drag, theta_star_100, ...), plus h = H0 / 100, T_cmb, and any sampled input passed through.

Anything else is refused when the model is built, with a list of what the bundle provides. That includes matter power, sigma8, fsigma8, BB, lensed_scal_Cl, an $\ell$ beyond the bundle's grid, and a redshift beyond its $z$ grid.

Parameters¶

ClassE resolves each bundle input once, when the model is built. It tries these in order:

input_map, which you set: {bundle input: sampled name or expression};
the input's own name;
its alias in the generation YAML (stored in the bundle), e.g. 100*theta_s for theta_s_100;
a known conversion.

Bundle input	Read from
`logA`	`log(1e10 * A_s)`
`A_s`	`1e-10 * exp(logA)`
`h`	`H0 / 100`
`H0`	`100 * h`
`theta_s_100`	`100*theta_s`
`omega_cdm` / `omega_c`	each other
`w0`, `wa`	`w0_fld`, `wa_fld`
`m_nu`	`m_ncdm` (one massive species)

Expressions use class_e.utils.expressions (arithmetic, exp, log, sqrt, ...). They are evaluated with JAX when the emulators are, so the mapping stays differentiable.

theory:
  class_e.cobaya.ClassE:
    bundle: bundles/w0wa.cemu
    input_map: {s: "w0_fld + wa_fld"}   # the bundle was trained on s = w0 + wa

Every parameter that reaches the theory must be read by some input. The exception is a CLASS setting fixed at exactly the value the bundle was trained with, e.g. N_ur: 2.0308 in both the bundle and the Cobaya params.

Anything else is an error: the emulators were trained without it and cannot follow it. Sampling mnu with a bundle trained at m_ncdm = 0.06 is one such error.

Priors and the training box¶

Emulators do not extrapolate. With out_of_bounds: reject, the default, a point outside the bundle's training box, or one breaking a sampling constraint stored in the bundle (e.g. w0 + wa < 0), gets zero likelihood.

ClassE logs the count of such points at 1, 2, 4, 8, ... so long chains stay quiet. With warn or ignore, the emulators extrapolate instead.

A prior wider than the box therefore truncates the posterior silently. Check an input before running it:

class-e cobaya check chains/lcdm.yaml && cobaya-run chains/lcdm.yaml

                  Bundle inputs over the prior (4000 samples)
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┓
┃ input     ┃ from            ┃ prior range          ┃ bundle box         ┃ outside ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━┩
│ h         │ H0 / 100        │ [0.640021, 0.709991] │ [0.6466, 0.7006]   │   23.1% │
│ omega_b   │ omega_b         │ [0.0216004, 0.023099]│ [0.02162, 0.02312] │    1.4% │
│ omega_cdm │ omega_cdm       │ [0.114001, 0.125992] │ [0.114, 0.126]     │       - │
│ logA      │ log(1e10 * A_s) │ [2.95004, 3.14991]   │ [2.974, 3.114]     │   29.9% │
│ n_s       │ n_s             │ [0.920009, 1]        │ [0.9439, 0.9859]   │   48.2% │
│ tau_reio  │ tau_reio        │ [0.0211929, 0.08249] │ [0.0189, 0.0899]   │       - │
└───────────┴─────────────────┴──────────────────────┴────────────────────┴─────────┘
warning: In total 72.5% of the prior volume is outside the emulators' domain (zero likelihood) ...
error: Derived parameters ['sigma8'] cannot be computed from the bundle (it gives ['T_cmb']).

check works without likelihood data and runs nothing. It uses Cobaya's own parameterization, so lambda-defined inputs, renames, drop and the likelihoods' nuisance parameters are handled as in a real run. It does three things:

samples the prior;
maps each sample onto the bundle's inputs as ClassE would;
reports the volume outside the box (per input), the volume removed by the stored constraints, and every error that would stop the model from building.

It exits with status 1 on errors. From Python, class_e.cobaya.check.check(info) returns the same report.

ClassE also does the coverage part itself when the model is built: from prior_check samples of the run's prior (1000 by default; 0 turns it off), it logs a warning for each input whose prior reaches beyond the box, for each failing constraint, and for the total volume outside the domain. A wide prior is then visible in the run's log, even without class-e cobaya check.

bundle: also takes the name of a published emulator (bundle: lcdm@jax, see Published emulators). It is downloaded on first use. On nodes without network, set offline: true or $CLASS_E_OFFLINE=1 to use only the download cache.

Accuracy against CLASS¶

Take a bundle that reproduces a CLASS run exactly at one point. There, Cobaya with ClassE and Cobaya with classy give the same $\chi^2$, which confirms the units, the multipole indexing, $T_{\rm cmb}$ and the distances:

Likelihood	$\chi^2$ (classy)	ClassE − classy
Planck plik lite TTTEEE	583.5035	$4\times10^{-8}$
Planck low-$\ell$ TT, EE	23.5036, 396.1429	0
DESI DR2 BAO	29.9528	$3\times10^{-4}$
Union3	26.3718	$9\times10^{-4}$

Away from such a point, the difference is the emulators' error. dev/benchmarks/cobaya_dchi2.py maps $\Delta\chi^2$ over a bundle's test set.

$\Delta\chi^2$ grows with the distance to the data. A spectrum error of $10^{-2}\,\sigma_{\rm CV}$ gives $|\Delta\chi^2| \sim 1$ at points with $\chi^2 \sim 2500$, but $\sim 0.2$ near the posterior. So judge a bundle on the points near the data.

From Python¶

from cobaya.model import get_model

model = get_model({
    "theory": {"class_e.cobaya.ClassE": {"bundle": "bundles/lcdm.cemu"}},
    "likelihood": {"planck_2018_lowl.TT": None},
    "params": {"h": 0.6736, "omega_b": 0.02237, "omega_cdm": 0.12, "logA": 3.044,
               "n_s": 0.9649, "tau_reio": 0.0544, "A_planck": 1.0},
})
model.logposterior({})

The theory has no Cobaya-specific logic: class_e.theory.ParameterAdapter and class_e.theory.EmulatedCosmology map parameters and give the products. Other samplers can build on them.

Likelihood	Needs	Bundle observables
`planck_2018_highl_plik.TTTEEE_lite_native`	lensed `Cl` tt, te, ee to \(\ell = 2508\)	`TT`, `TE`, `EE` on every \(\ell\) from 2 to \(\geq 2508\)
`planck_2018_lowl.TT`, `.EE`	`Cl` tt or ee to \(\ell = 29\)	`TT`, `EE`
`bao.desi_dr2`	\(D_A(z)\), \(H(z)\), derived `rdrag`	`Hz` (distance emulator), `rs_drag`
`sn.union3` and other SNe	\(D_A(z)\)	`Hz`

Likelihood	\(\chi^2\) (classy)	ClassE − classy
Planck plik lite TTTEEE	583.5035	\(4\times10^{-8}\)
Planck low-\(\ell\) TT, EE	23.5036, 396.1429	0
DESI DR2 BAO	29.9528	\(3\times10^{-4}\)
Union3	26.3718	\(9\times10^{-4}\)