API reference

The public API, as exported by cosmonuts. Docstrings use the Google style.

Model

cosmonuts.Model

Model(theory: Theory | Mapping[str, Theory], likelihoods: Likelihood | Sequence[Likelihood], priors: Mapping[str, Any], derived: Mapping[str, Callable] | None = None, labels: Mapping[str, str] | None = None, provides: Mapping[str, str] | None = None)

Bases: Module

Posterior of a theory under likelihoods and priors.

Parameters:

Name	Type	Description	Default
theory	Theory | Mapping[str, Theory]	The theory, or several by name, {name: theory}, combined into Theories: each quantity comes from the theory that provides it, and a theory parameter without a prior that another theory provides (e.g. an emulated h) is read from that theory.	required
likelihoods	Likelihood | Sequence[Likelihood]	One or more likelihoods.	required
priors	Mapping[str, Any]	For every parameter the theory and likelihoods read, a prior (Uniform, Normal, ...), a Fixed value, or a Defined function of other parameters (a change of variables, e.g. "A_s": Defined(lambda logA: 1e-10 * jnp.exp(logA)) with "logA" sampled). Parameters that only feed a Defined one count as used.	required
derived	Mapping[str, Callable] | None	Extra quantities to record with samples, name -> fn(provider, params): a scalar, or a 1-D array saved as the columns name_0, name_1, ... (each likelihood's derived is recorded too).	None
labels	Mapping[str, str] | None	LaTeX labels (no $) overriding the theory's and likelihoods' param_labels (and Defined.label), e.g. for a reparametrised variable.	None
provides	Mapping[str, str] | None	With several theories, which one serves a quantity that several provide, {quantity: name} (Cobaya's provides:).	None

Raises:

Type	Description
ValueError	If a likelihood requires a quantity the theory does not provide, a parameter has no prior, a prior is unused, the Defined parameters read an unknown name or depend on each other in a cycle, or several theories cannot be combined (see Theories).

Attributes:

Name	Type	Description
sampled	tuple[str, ...]	Names of the sampled parameters, in the order of \(u\).
fixed	tuple[str, ...]	Names of the fixed parameters.
defined	tuple[str, ...]	Names of the Defined parameters, in the order they are computed.

ndim property

ndim: int

Number of sampled parameters.

reference property

reference: dict[str, Any]

A central point: the middle of uniform priors, the mean of normal ones.

labels property

labels: dict[str, str]

LaTeX labels (no $) of all parameters and derived quantities.

From the theory's and likelihoods' param_labels; unlabelled names label themselves.

log_normalization property

log_normalization: float

Sum of the likelihoods' log_normalization: what \(\log L\) leaves out.

Add it to an evidence computed from log_likelihood for the absolute \(\ln Z\); Bayes factors between models of the same data do not need it.

Raises:

Type	Description
NotImplementedError	If a likelihood does not define it.
ValueError	If a likelihood cannot compute it as configured.

complete

complete(params: Mapping[str, Any]) -> dict[str, Any]

params with the fixed and defined parameters added (given values are kept).

to_physical

to_physical(u: Any) -> tuple[dict[str, Any], Any]

All parameters (sampled, fixed and defined) at \(u\), and \(\log|dx/du|\).

to_unconstrained

to_unconstrained(params: Mapping[str, Any]) -> Any

\(u\) for the sampled parameters in params (fixed ones are ignored).

log_prior

log_prior(params: Mapping[str, Any]) -> Any

Sum of the sampled parameters' log prior densities.

log_likelihoods

log_likelihoods(params: Mapping[str, Any]) -> tuple[Any, Any]

Each likelihood's \(\log L\) at physical params, and provider.ok.

params needs only the sampled parameters: fixed and defined ones are added.

evaluate

evaluate(params: Mapping[str, Any]) -> tuple[Any, Any, dict[str, Any]]

log_likelihoods and the quantities recorded with samples, from one theory call.

Returns:

Name	Type	Description
	Any	Each likelihood's \(\log L\), provider.ok, and the outputs by name, in column
order	Any	the saved Defined parameters, derived, then each likelihood's
	dict[str, Any]	derived.

output_columns

output_columns() -> list[tuple[str, str, int | None]]

The chain columns of the outputs of evaluate.

Each is (column, output, index): a scalar output is one column of its name (index None), a 1-D output name of length \(n\) the columns name_0 to name_{n-1}. The shapes come from tracing the theory once (jax.eval_shape), without computing it.

Raises:

Type	Description
ValueError	If an output has more than one dimension, or two columns (or a column and a parameter) share a name.

log_likelihood

log_likelihood(params: Mapping[str, Any]) -> Any

Total \(\log L\) at physical params; \(-\infty\) where the theory fails.

log_posterior

log_posterior(params: Mapping[str, Any]) -> Any

\(\log L + \log\pi\) at physical params (no Jacobian): what a MAP fit maximises.

log_prior_density

log_prior_density(u: Any) -> Any

The prior as a density in the unconstrained space, \(\log\pi(x(u)) + \log|dx/du|\).

What nested sampling draws from (inside its likelihood contours): logistic in \(u\) for a Uniform prior, standard normal for a Normal one.

log_density

log_density(u: Any) -> Any

Log density in the unconstrained space, \(\log L + \log\pi + \log|dx/du|\).

derived_values

derived_values(params: Mapping[str, Any]) -> dict[str, Any]

The derived quantities at physical params.

sample_prior

sample_prior(key: Any, n: int) -> Any

n draws of \(u\) from the prior, shape (n, ndim).

cosmonuts.Theory

Bases: Module

Base class of theories: parameters in, a Provider out.

Subclasses define two static fields and __call__.

Attributes:

Name	Type	Description
params	AbstractVar[tuple[str, ...]]	Names of the parameters the theory reads from the dictionary it is called with.
provides	AbstractVar[frozenset[str]]	Names of the provider methods it implements, from QUANTITIES.
param_labels	dict[str, str]	LaTeX labels (no $) of its parameters, for saved chains.
l_max	int | None	Largest multipole the provider's cl_* methods serve, for theories that provide CMB spectra (None otherwise). Likelihoods compare it with the multipoles they need in Likelihood.check.

__call__ abstractmethod

__call__(params: dict[str, Any]) -> Provider

Compute the provider at a parameter point.

Parameters:

Name	Type	Description	Default
params	dict[str, Any]	Parameter values by name; contains at least self.params.	required

Returns:

Type	Description
Provider	The provider. It must be computable inside jax.jit and
Provider	differentiable in the parameters.

cosmonuts.Provider

Bases: Protocol

Standard interface of the object a theory returns.

A provider implements a subset of these methods, the ones its theory lists in provides; each takes and returns JAX arrays. Units are fixed (see QUANTITIES): lengths in Mpc, \(c = 1\).

Attributes:

Name	Type	Description
ok	Any	Scalar boolean, False if the theory computation failed (e.g. a solver did not converge or the model is unstable); the model's log density is then \(-\infty\).

hubble

hubble(z: Any) -> Any

\(H(z)\) in Mpc\(^{-1}\).

comoving_distance

comoving_distance(z: Any) -> Any

Comoving distance in Mpc.

angular_diameter_distance

angular_diameter_distance(z: Any) -> Any

Angular diameter distance in Mpc.

sound_horizon_drag

sound_horizon_drag() -> Any

\(r_d\) in Mpc.

theta_star

theta_star() -> Any

\(\theta_* = r_s(z_*)/D_M(z_*)\) in radians, \(z_*\) where the optical depth of recombination reaches 1.

growth_rate

growth_rate(z: Any) -> Any

Linear growth rate \(f(z)\).

pk_cb

pk_cb(k: Any, z: Any) -> Any

Linear \(P_{cb}(k, z)\) in Mpc\(^3\) at \(k\) in Mpc\(^{-1}\), with shape k.shape.

h

h() -> Any

Dimensionless Hubble rate today.

Omega_b

Omega_b() -> Any

Baryon density parameter today.

Omega_cb

Omega_cb() -> Any

Baryon + CDM density parameter today.

n_s

n_s() -> Any

Scalar spectral index.

cl_tt

cl_tt(ell: Any) -> Any

Lensed \(D_\ell^{TT} = \ell(\ell+1)C_\ell/2\pi\) in \(\mu{\rm K}^2\) at integer ell, shape ell.shape.

cl_te

cl_te(ell: Any) -> Any

Lensed \(D_\ell^{TE}\) in \(\mu{\rm K}^2\) at integer ell.

cl_ee

cl_ee(ell: Any) -> Any

Lensed \(D_\ell^{EE}\) in \(\mu{\rm K}^2\) at integer ell.

cl_bb

cl_bb(ell: Any) -> Any

Lensed \(D_\ell^{BB}\) in \(\mu{\rm K}^2\) at integer ell.

cl_pp

cl_pp(ell: Any) -> Any

\([L(L+1)]^2 C_L^{\phi\phi}/2\pi\) at integer ell (\(= L\)).

cosmonuts.Likelihood

Bases: Module

Base class of likelihoods.

Attributes:

Name	Type	Description
requires	AbstractVar[frozenset[str]]	Provider quantities it uses (see QUANTITIES).
params	AbstractVar[tuple[str, ...]]	Names of the parameters it reads itself (nuisance parameters); their priors are given to the Model.
install_name	str | None	Name for cosmonuts install, for likelihoods with data files.
data_sources	tuple[Source, ...]	Files to download, into the directory install_name of the packages path.
files	tuple[str, ...]	Files (or directories) that must exist once installed, relative to that directory, when they differ from data_sources (e.g. the contents of a downloaded archive that postprocess extracts and removes).
param_labels	dict[str, str]	LaTeX labels (no $) of its nuisance parameters, for saved chains (not labels, which the DESI DR1 likelihoods use for tracers).

derived_labels property

derived_labels: dict[str, str]

LaTeX labels (no $) of the derived quantities.

has_derived property

has_derived: bool

Whether the class records derived quantities.

log_normalization property

log_normalization: float

Constant to add to log_likelihood for the normalised density of the data.

Only absolute evidences need it. Subclasses that drop a constant implement it; a likelihood that is already normalised returns 0.0.

Raises:

Type	Description
NotImplementedError	If the likelihood does not define it.

log_likelihood abstractmethod

log_likelihood(provider: Any, params: dict[str, Any]) -> Any

\(\log L\) at a provider and parameter point (a JAX scalar).

check

check(theory: Any) -> None

Check, once when the model is built, that theory can serve this likelihood.

requires is checked by the model itself; override this for what it cannot express, such as the multipole range of CMB spectra (theory.l_max).

Raises:

Type	Description
ValueError	If the theory cannot serve the likelihood.

derived

derived(provider: Any, params: dict[str, Any]) -> dict[str, Any]

Quantities to record with each sample, by name (none by default).

Each is a scalar, or a 1-D array saved as the columns name_0, name_1, ... They are computed after sampling, for the saved draws only, from the same theory call as the chains' \(\chi^2\); names must not clash with parameters or other outputs. Their LaTeX labels go in derived_labels.

postprocess classmethod

postprocess(directory: Path) -> None

Hook run by cosmonuts install after downloading, e.g. to convert formats.

Parameters:

Name	Type	Description	Default
directory	Path	The likelihood's data directory.	required

installed_files classmethod

installed_files() -> tuple[str, ...]

What must exist in the data directory once installed: files, else the sources.

data_directory classmethod

data_directory(path: str | Path | None = None) -> Path

The installed data directory, checking that every file is there.

Parameters:

Name	Type	Description	Default
path	str | Path | None	Packages path; default from packages_path.	None

Returns:

Type	Description
Path	The directory.

Raises:

Type	Description
DataNotFoundError	If a file is missing, with the install command.

cosmonuts.GaussianLikelihood

Bases: Likelihood

\(\log L = -\tfrac12\,\Delta^T C^{-1} \Delta\) with \(\Delta = t(\theta) - d\).

Subclasses provide data, precision and theory_vector.

Attributes:

Name	Type	Description
data	AbstractVar[Any]	Data vector \(d\).
precision	AbstractVar[Any]	Inverse covariance \(C^{-1}\).

log_normalization property

log_normalization: float

\(-\tfrac12\ln|2\pi C| = \tfrac12\ln|C^{-1}| - \tfrac n2\ln 2\pi\), from precision.

Raises:

Type	Description
ValueError	If precision is not positive definite (e.g. a nuisance parameter was marginalised into it: the subclass must then supply the constant itself).

theory_vector abstractmethod

theory_vector(provider: Any, params: dict[str, Any]) -> Any

The model prediction \(t(\theta)\) for data.

chi2

chi2(provider: Any, params: dict[str, Any]) -> Any

\(\chi^2 = \Delta^T C^{-1} \Delta\).

log_likelihood

log_likelihood(provider: Any, params: dict[str, Any]) -> Any

Return \(-\tfrac12\chi^2\) (without the normalisation).

cosmonuts.DataNotFoundError

Bases: FileNotFoundError

A likelihood's data files are not installed.

Priors

cosmonuts.Prior

Bases: Module

Base class of one-dimensional priors.

log_prob abstractmethod

log_prob(x: Any) -> Any

Log prior density at the physical value x.

to_physical abstractmethod

to_physical(u: Any) -> tuple[Any, Any]

Physical value at the unconstrained u, and \(\log|dx/du|\).

to_unconstrained abstractmethod

to_unconstrained(x: Any) -> Any

Unconstrained value at the physical x.

sample abstractmethod

sample(key: Any) -> Any

One draw of the physical value.

cosmonuts.Uniform

Bases: Prior

Flat prior on \([\text{lower}, \text{upper}]\), sampled through \(x = l + (u_p - l)\,\sigma(u)\).

Attributes:

Name	Type	Description
lower	float	Lower bound.
upper	float	Upper bound.

log_prob

log_prob(x: Any) -> Any

Return the log of the prior density at the physical value x (\(-\infty\) outside the support).

to_physical

to_physical(u: Any) -> tuple[Any, Any]

Map the unconstrained \(u\) to the physical value, returning it with \(\log|dx/du|\).

to_unconstrained

to_unconstrained(x: Any) -> Any

Map a physical value x to the unconstrained \(u\) the samplers move in.

sample

sample(key: Any) -> Any

Draw a physical value from the prior with the jax.random key key.

cosmonuts.Normal

Bases: Prior

Gaussian prior \(\mathcal{N}(\mu, \sigma^2)\), sampled standardised: \(x = \mu + \sigma u\).

Without the scaling, a tight prior (e.g. BBN's \(\sigma(\omega_b) = 5.5\times10^{-4}\)) would make \(u\) hundreds of times narrower than the other parameters, and NUTS's warm-up would start with a tiny step size and very long trajectories.

Attributes:

Name	Type	Description
mean	float	Mean \(\mu\).
std	float	Standard deviation \(\sigma\).

log_prob

log_prob(x: Any) -> Any

Return the log of the prior density at the physical value x.

to_physical

to_physical(u: Any) -> tuple[Any, Any]

Map the unconstrained \(u\) to the physical value, returning it with \(\log|dx/du|\).

to_unconstrained

to_unconstrained(x: Any) -> Any

Map a physical value x to the unconstrained \(u\) the samplers move in.

sample

sample(key: Any) -> Any

Draw a physical value from the prior with the jax.random key key.

cosmonuts.Fixed

Bases: Module

A parameter held at a fixed value, not sampled.

Attributes:

Name	Type	Description
value	float	The value.

cosmonuts.Defined

Defined(fn: Callable, *, save: bool = True, label: str | None = None)

Bases: Module

A parameter defined as a function of others (sampled, fixed or defined), not sampled.

The function's argument names are the parameters it reads, as in Cobaya's value: "lambda logA: 1e-10*np.exp(logA)". It carries no prior of its own: the prior is on the parameters it is computed from. The theory and likelihoods read it like any other parameter, so a change of variables never touches theory code.

Parameters:

Name	Type	Description	Default
fn	Callable	The function; written with jax.numpy, since it runs inside the sampler.	required
save	bool	Record its value in saved chains, as a derived (*) column.	True
label	str | None	LaTeX label (no $) for saved chains and plots.	None

Examples:

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

Samplers

cosmonuts.sample

sample(model: Model, key: Any, method: Literal['nuts', 'hmc', 'metropolis', 'nested'] = 'nuts', device: Any = None, **kwargs: Any) -> NUTSResult | MetropolisResult | NestedResult

Sample the model with method on device: one call for every sampler.

The keyword arguments are the method's (nuts, hmc, metropolis or nested), so moving a script from CPU NUTS to GPU nested sampling changes only method and device.

Parameters:

Name	Type	Description	Default
model	Model	The model.	required
key	Any	A jax.random key.	required
method	Literal['nuts', 'hmc', 'metropolis', 'nested']	"nuts", "hmc", "metropolis" or "nested".	'nuts'
device	Any	None (JAX's default device), "cpu", "gpu" or a jax.Device.	None
**kwargs	Any	The method's own arguments.	{}

Returns:

Type	Description
NUTSResult | MetropolisResult | NestedResult	The method's result: NUTSResult for NUTS and HMC,
NUTSResult | MetropolisResult | NestedResult	MetropolisResult for Metropolis,
NUTSResult | MetropolisResult | NestedResult	NestedResult for nested sampling.

Raises:

Type	Description
ValueError	For an unknown method.
TypeError	For an argument the method does not take.

cosmonuts.nuts

nuts(model: Model, key: Any, *, num_warmup: int = 500, num_samples: int = 1000, num_chains: int = 4, initial: dict[str, Any] | None = None, jitter: float = 0.05, output: str | PathLike | None = None, overwrite: bool = False, progress: bool | int = False, mass_matrix: Literal['diagonal', 'dense'] = 'diagonal', target_ess: float | None = None, check_every: int | None = None, rhat_method: RhatMethod = 'rank', adaptation: str | PathLike | None = None, resume: bool = False, max_num_doublings: int = 10, mass_matrix_seed: Mapping[str, Any] | None = None, mass_matrix_seed_weight: float = 0.0, inverse_mass_matrix: Any = None, parallel: Literal['auto', 'vmap', 'shard', 'process', 'mpi'] = 'auto', device: Any = None) -> NUTSResult

Sample the model's posterior with NUTS.

Parameters:

Name	Type	Description	Default
model	Model	The model.	required
key	Any	A jax.random key.	required
num_warmup	int	Adaptation steps: of one chain, shared by all, with parallel="vmap"; of each chain otherwise.	500
num_samples	int	Draws per chain; with target_ess, the most per chain.	1000
num_chains	int	Chains, run in parallel (see parallel).	4
initial	dict[str, Any] | None	Physical starting point of the sampled parameters; default model.reference.	None
jitter	float	Standard deviation of the chains' starting offsets from the adapted so the offsets respect tight and correlated directions).	0.05
output	str | PathLike | None	If given, a path root: write the chains there as they are drawn, every check_every draws, with a ChainWriter (<output>.1.txt, ..., .paramnames, .ranges; .adaptation.npz after warm-up; .covmat and .updated.yaml at the end), so a crash loses at most one chunk. If writing fails, sampling carries on and all draws end up in <output>.draws.npz.	None
overwrite	bool	Replace existing chains at output instead of raising.	False
progress	bool | int	Report progress in warm-up and in sampling: draws done, time (compilation excluded) and ETA, leapfrog steps per draw, divergences so far, and the step size in warm-up. True: a live bar in a terminal, or a line every 5% of the draws when the output is not one (a log file); an integer N: a line every N draws; False: silent. Fed by jax.debug.callback inside the compiled loops, one host call per draw, which costs little next to a gradient.	False
mass_matrix	Literal['diagonal', 'dense']	Inverse mass matrix learnt in warm-up: "diagonal", one scale per parameter, or "dense", the full covariance. Dense follows correlations (e.g. \(w_0\)-\(w_a\)), so NUTS needs fewer leapfrog steps per draw there, but it estimates \(d(d+1)/2\) entries: give it more warm-up posteriors with \(d \lesssim 100\).	'diagonal'
target_ess	float | None	Stop sampling early, once the smallest effective sample size over the parameters (all chains together) reaches this and \(\hat R < 1.01\) (see rhat_method); checked every check_every draws. None: always num_samples draws.	None
check_every	int | None	Draws per chain between ESS checks and chain writes (default 100): the run can overshoot target_ess by up to this many draws. Each check is cheap (ESS and \(\hat R\) on the draws so far); each write evaluates the likelihoods once per new draw.	None
rhat_method	RhatMethod	The \(\hat R\) of the stopping rule: "rank", the rank-normalised split-\(\hat R\) (Vehtari et al. 2021), which also catches chains that differ in scale only; or "gelman-rubin", the classic one. The result reports both.	'rank'
adaptation	str | PathLike | None	A .adaptation.npz written by an earlier run with output (or by ChainWriter.write_adaptation): skip warm-up and sample with its step size and inverse mass matrix, from its final position (plus jitter). initial and num_warmup are then unused, and the saved matrix sets mass_matrix (1-D: diagonal, 2-D: dense). Its parameter names must be the model's.	None
resume	bool	Continue the chains at output instead of starting new ones: each chain restarts from its last saved draw with the saved step size and mass matrix (<output>.adaptation.npz, or adaptation), without warm-up or jitter, and new draws are appended to the same files. num_samples then counts all draws per chain, old ones included, and target_ess is checked on all of them: nothing is sampled if the saved chains already meet it. The result holds every draw (the old ones' log_density rebuilt from their minuslogpost); divergences and num_steps cover this session only. The new draws use the random numbers that follow the saved ones' (their keys are indexed by draw), never theirs.	False
max_num_doublings	int	Most doublings of a NUTS trajectory: at most \(2^n - 1\) leapfrog steps per draw, in warm-up and sampling (BlackJAX's default, 10, allows 1023). A lower cap bounds the cost of a draw where the step size is small or the metric poor (early warm-up, a funnel); draws that hit it are shorter than NUTS would make them, which lowers the ESS per draw but keeps the chains correct.	10
mass_matrix_seed	Mapping[str, Any] | None	Physical draws of the sampled parameters, name -> array (e.g. an earlier run's chains, mapped to this model's parameters): their covariance in the sampler's unconstrained space (its diagonal with mass_matrix="diagonal") is the inverse mass matrix warm-up starts from, instead of the identity. That speeds up the first windows, which otherwise take many small steps when the posterior widths differ a lot. Draws outside the priors' support are dropped. Unused with adaptation or resume.	None
mass_matrix_seed_weight	float	How long the seed persists: a pseudo-count of draws blending it into each warm-up window's estimate (BlackJAX's imm_shrinkage_to_previous). 0: the seed only sets the first window, later ones use their own draws alone (Stan's rule); about the window size (25-100): half and half. It weights inverse_mass_matrix the same way.	0.0
inverse_mass_matrix	Any	The inverse mass matrix warm-up starts from, given directly: (d,) or (d, d) in the sampler's unconstrained coordinates, in model.sampled order, e.g. the Laplace covariance at a best fit, covariance(model, best.params, space="unconstrained") (covariance). Its diagonal is used with mass_matrix="diagonal". Excludes mass_matrix_seed; unused with adaptation or resume.	None
parallel	Literal['auto', 'vmap', 'shard', 'process', 'mpi']	How the chains run in parallel. "vmap": batched with jax.vmap on one device, after one warm-up whose step size and mass matrix all chains share. A batch moves in lockstep: every NUTS draw costs the longest trajectory of the batch, so on a CPU several chains take several times as long as one. "shard": spread over the devices with jax.shard_map (several CPU devices with COSMONUTS_CPU_DEVICES=n in the environment, or several GPUs), each device running its own chains concurrently, and each chain warming up on its own from initial plus jitter (Stan's scheme), so warm-up is parallel too. .adaptation.npz then holds each chain's step size, mass matrix and position; the result's step_size and inverse_mass_matrix are their means. "process": one process per chain, as Cobaya with MPI, for theories that cannot be batched or sharded too. Each chain warms up on its own, and the chains meet every check_every draws, where the chain files are written and target_ess is checked on all of them. The model goes to the workers with cloudpickle, and the calling script needs the if __name__ == "__main__": guard. Each worker runs XLA on one thread; the progress shown is chain 0's. "mpi": the same with one MPI rank per chain: run the script with mpirun -n <num_chains> python -m mpi4py script.py (needs mpi4py). Each rank builds its own model, rank 0 writes the files, and every rank returns the result. With "process" and "mpi", chain i draws from fold_in(key, i). "auto": "shard" when there are more chains than one and more devices than one, else "vmap".	'auto'
device	Any	Where to run: None (JAX's default device), "cpu", "gpu", or a jax.Device. NUTS gains from a GPU mainly when the theory itself is heavy; nested is the method that parallelises.	None

Returns:

Type	Description
NUTSResult	The samples and diagnostics.

Raises:

Type	Description
FileExistsError	If chains exist at output and overwrite is false (checked before sampling).
ValueError	If mass_matrix is not "diagonal" or "dense", if num_samples, num_chains, target_ess, check_every or max_num_doublings is not positive, if parallel is unknown, if a per-chain adaptation has another number of chains, if mass_matrix_seed lacks a sampled parameter or has too few draws inside the priors (at most the number of parameters), if inverse_mass_matrix has the wrong shape or is not positive definite, or comes with mass_matrix_seed, if adaptation does not match the model (parameter names, shapes) or holds another sampler's (HMC's) adaptation, or if resume is asked without output, with overwrite, or with chains that do not match (number, parameters).
FileNotFoundError	With resume, if there are no chains or no .adaptation.npz at output.
RuntimeError	If device names a platform JAX does not have.

cosmonuts.hmc

hmc(model: Model, key: Any, *, num_warmup: int = 500, num_samples: int = 1000, num_chains: int = 4, initial: dict[str, Any] | None = None, jitter: float = 0.05, output: str | PathLike | None = None, overwrite: bool = False, progress: bool | int = False, mass_matrix: Literal['diagonal', 'dense'] = 'diagonal', target_ess: float | None = None, check_every: int | None = None, rhat_method: RhatMethod = 'rank', adaptation: str | PathLike | None = None, resume: bool = False, mass_matrix_seed: Mapping[str, Any] | None = None, mass_matrix_seed_weight: float = 0.0, inverse_mass_matrix: Any = None, num_integration_steps: int | None = None, parallel: Literal['auto', 'vmap', 'shard', 'process', 'mpi'] = 'auto', device: Any = None) -> HMCResult

Sample with standard fixed-trajectory HMC.

Accepts the options of nuts (warm-up, chains, output, resume, mass matrix, ESS stopping, device, ...), with num_integration_steps in place of max_num_doublings.

Parameters:

Name	Type	Description	Default
model	Model	The model.	required
key	Any	A jax.random key.	required
num_warmup	int	As in nuts.	500
num_samples	int	As in nuts.	1000
num_chains	int	As in nuts.	4
initial	dict[str, Any] | None	As in nuts.	None
jitter	float	As in nuts.	0.05
output	str | PathLike | None	As in nuts.	None
overwrite	bool	As in nuts.	False
progress	bool | int	As in nuts.	False
mass_matrix	Literal['diagonal', 'dense']	As in nuts.	'diagonal'
target_ess	float | None	As in nuts.	None
check_every	int | None	As in nuts.	None
rhat_method	RhatMethod	As in nuts.	'rank'
adaptation	str | PathLike | None	As in nuts.	None
resume	bool	As in nuts.	False
mass_matrix_seed	Mapping[str, Any] | None	As in nuts.	None
mass_matrix_seed_weight	float	As in nuts.	0.0
inverse_mass_matrix	Any	As in nuts.	None
parallel	Literal['auto', 'vmap', 'shard', 'process', 'mpi']	As in nuts.	'auto'
device	Any	As in nuts.	None
num_integration_steps	int | None	Leapfrog steps per trajectory; default 10. With adaptation or resume, the saved count is used: the step size was adapted for it.	None

Returns:

Type	Description
HMCResult	The samples and diagnostics, as for NUTS (settings records
HMCResult	num_integration_steps and sampler="hmc").

Raises:

Type	Description
ValueError	As nuts; also if num_integration_steps is not positive, differs from that of a loaded adaptation, or if the adaptation is NUTS's.

cosmonuts.metropolis

metropolis(model: Model, key: Any, *, num_samples: int = 10000, num_chains: int = 4, initial: dict[str, Any] | None = None, jitter: float = 0.05, output: str | PathLike | None = None, overwrite: bool = False, resume: bool = False, progress: bool | int = False, target_ess: float | None = None, check_every: int | None = None, max_learning_steps: int = 20000, learn_rhat_tolerance: float = 0.05, rhat_method: RhatMethod = 'rank', target_acceptance: float | None = None, proposal_seed: Mapping[str, Any] | None = None, device: Any = None) -> MetropolisResult

Sample the model's posterior with an adaptive random-walk Metropolis algorithm.

The walk runs in the unconstrained space on Model.log_density (the Jacobian included), one posterior evaluation per step and chain, with the chains vectorised. It has two phases.

• Learning: in chunks of check_every steps, the proposal covariance is re-estimated from the pooled draws of all chains (the first 30% dropped) and its scale tuned towards target_acceptance. It ends once \(\hat R - 1\) of the recent draws (each chain split in two, with rhat_method) is below learn_rhat_tolerance, the covariance estimate has stopped drifting (rms change below 0.25 between the two halves of the draws) and the acceptance is within a factor 2 of its target; or, with a warning, at max_learning_steps.
• Sampling: the proposal is then frozen and the chains continue from where learning left them; the learning draws are discarded, so the kept chains are those of a plain Metropolis algorithm. Every step is kept (no thinning).

No gradient is needed, so it suits cheap or non-differentiable likelihoods and low-dimensional problems; NUTS wins as the dimension and the cost of a posterior evaluation grow. Evaluations per chain are learning_steps + num_draws.

Parameters:

Name	Type	Description	Default
model	Model	The model.	required
key	Any	A jax.random key.	required
num_samples	int	Draws per chain after learning; with target_ess, the most per chain.	10000
num_chains	int	Chains, run in parallel (at least 2 are recommended for the diagnostics).	4
initial	dict[str, Any] | None	Physical starting point; default model.reference.	None
jitter	float	Scale of the chains' random starting offsets, in units of the Cholesky factor of the starting proposal covariance (the proposal_seed covariance, or the identity in the unconstrained coordinates).	0.05
output	str | PathLike | None	If given, a path root: write the chains there as they are drawn, like nuts.	None
overwrite	bool	Replace existing chains at output instead of raising.	False
resume	bool	Not supported by this sampler; must be False.	False
progress	bool | int	Print a line after every chunk of learning and of sampling (True), or every progress steps (rounded up to the next chunk end) and at the end of sampling; False: silent.	False
target_ess	float | None	Stop sampling once the smallest ESS (all chains together) reaches this and \(\hat R < 1.01\), checked every check_every draws (from 200 draws per chain on). A budget heuristic, not a proof of convergence.	None
check_every	int | None	Steps per chunk of learning and of sampling (default min(250, num_samples)); also how often chains are written and ESS checked.	None
max_learning_steps	int	Most learning steps per chain.	20000
learn_rhat_tolerance	float	Learning may end once \(\hat R - 1\) of the recent draws is below this (see the learning phase above for the other conditions).	0.05
rhat_method	RhatMethod	The \(\hat R\) of the learning stop and of the target_ess rule: "rank", the rank-normalised split-\(\hat R\) (Vehtari et al. 2021), or "gelman-rubin", the classic one. The result reports both.	'rank'
target_acceptance	float | None	Acceptance rate the scale is tuned for; default 0.44 for one parameter, falling linearly to 0.234 at five and above.	None
proposal_seed	Mapping[str, Any] | None	Physical draws of the sampled parameters, name -> array (e.g. an earlier run's chains): their covariance in the unconstrained space is the proposal covariance learning starts from, instead of the identity.	None
device	Any	Where to run: None (JAX's default device), "cpu", "gpu", or a jax.Device.	None

Returns:

Type	Description
MetropolisResult	The samples and diagnostics.

Raises:

Type	Description
ValueError	For non-positive num_samples, num_chains, check_every or target_ess; a target_acceptance outside (0, 1); an unknown rhat_method; resume; a proposal_seed that lacks a sampled parameter or has too few draws inside the priors; or chains that cannot start at a finite log density.
FileExistsError	If chains exist at output and overwrite is false.

Warns:

Type	Description
UserWarning	If learning does not meet its stopping rule within max_learning_steps; sampling then uses the last proposal.

cosmonuts.nested

nested(model: Model, key: Any, *, nlive: int | None = None, num_delete: int | None = None, num_inner_steps: int | None = None, frac_remain: float = 0.001, kernel: Literal['slice', 'swig'] = 'slice', max_steps: int = 10, max_shrinkage: int = 100, check_every: int = 10, max_dead: int | None = None, output: str | PathLike | None = None, overwrite: bool = False, resume: bool = False, progress: bool | int | Literal['step'] = False, device: Any = None, parallel: Literal['auto', 'vmap', 'shard'] = 'auto', batching: Batching = 'auto') -> NestedResult

Nested slice sampling of the model: the evidence and weighted posterior samples.

Runs BlackJAX's nested slice sampler (NestedSampler) in the model's unconstrained space. Each step replaces the num_delete lowest-likelihood live points by as many constrained slice chains run in parallel, which is what fills a GPU; no gradients are needed. It stops when the live points hold less than frac_remain of the evidence; their contribution is then added.

The evidence is that of log_likelihood (\(\chi^2\) convention), enough for Bayes factors between models of the same data; logZ_normalized adds the likelihoods' log_normalization (see Likelihood), with a warning and None when one cannot give it (e.g. DESDovekie without magnitude_prior_width).

Parameters:

Name	Type	Description	Default
model	Model	The model.	required
key	Any	A jax.random key.	required
nlive	int | None	Live points (default 500, or the saved run's with resume). The error on \(\ln Z\) scales as \(\sqrt{H/n_{\rm live}}\), \(H\) the information.	None
num_delete	int | None	Points replaced per step, by parallel slice chains. Default: nlive // 4 on an accelerator, nlive // 10 on a CPU (a large batch only coarsens a CPU run). Above nlive // 2 it warns: the error on \(\ln Z\) grows sharply beyond (Yallup, Kroupa & Handley 2026, Fig. 2). With parallel="shard", take a multiple of the number of devices: the chains are spread over the largest number of devices that divides it, and several chains per device even out the wait for the slowest device at each step.	None
num_inner_steps	int | None	Slice steps per new point (default max(5, 3 d); BlackJAX advises at least 2 d: fewer can bias \(\ln Z\) upwards). The right number depends on the problem: compare \(\ln Z\) at two values before lowering it (in a 7-parameter eft4py model, \(\ln Z\) rose with it from d to 3 d by more than its error).	None
frac_remain	float	Stop when \(Z_{\rm live}/Z\) falls below this (PolyChord's precision_criterion). May differ from the saved run's with resume, e.g. to continue a finished run to a tighter tolerance.	0.001
kernel	Literal['slice', 'swig']	"slice": hit-and-run along directions from the live points' covariance; "swig": slice-within-Gibbs along the axes, for near axis-aligned posteriors.	'slice'
max_steps	int	Most stepping-out expansions per slice.	10
max_shrinkage	int	Most shrinkage steps per slice.	100
check_every	int	Steps per compiled chunk: convergence is checked, progress reported and the files written between chunks.	10
max_dead	int | None	Stop after this many dead points even if not converged (converged is then false and \(\ln Z\) incomplete): a budget, or a run to continue with resume.	None
output	str | PathLike | None	Path root for files written as it runs: <output>_dead-birth.txt (each dead point as it dies) and <output>.nested_state.npz (to resume), then at the end the weighted chain <output>.1.txt and its files, as save_chains writes them.	None
overwrite	bool	Replace existing chains at output instead of raising.	False
resume	bool	Continue the run saved at output (its live points, evidence and random key), so it ends as if it had never been interrupted; the given key is not used. Settings left None are the saved ones; given ones must match them.	False
progress	bool | int | Literal['step']	Print a line every chunk (True) or every progress chunks: dead points, \(\ln Z\), \(\ln(Z_{\rm live}/Z)\) against its target, likelihood calls, seconds per step and likelihood calls per second (compilation excluded; the first line gives its time), and the slice expansions and shrinkages per slice (mean/max) and the share of slices that ran out of shrinkages, to tune max_steps and max_shrinkage. "step": also a line per step, from inside the compiled chunk, so check_every can stay large.	False
device	Any	Where to run: None (JAX's default device), "cpu", "gpu" or a jax.Device.	None
parallel	Literal['auto', 'vmap', 'shard']	How the num_delete slice chains of a step run: "vmap", batched on one device, or "shard", spread over the platform's devices (several GPUs, or CPU devices from COSMONUTS_CPU_DEVICES=n, best with XLA_FLAGS=--xla_cpu_multi_thread_eigen=false for one device per core), each device running its chains with its own loops, so that it waits only for the slowest of its own chains and likelihood calls. "auto": "shard" when there are several devices. The run does not depend on the number of devices (the same keys go to the same chains), and a run resumes onto any number of them.	'auto'
batching	Batching	How the chains on one device run (and the initial live points are evaluated): "vmap", in lockstep, every chain waiting at each loop for the slowest; "sequential", one after another (lax.map); "auto": "sequential" on a CPU when the likelihood has a while_loop (an adaptive ODE solve, a shooting), whose cost varies from point to point, else "vmap" (always on accelerators). For eft4py on a CPU, sequential is two orders of magnitude faster than a vmapped batch of 50. "vmap" and "sequential", and runs on different numbers of devices, can differ in the last bits of log L (about \(10^{-13}\) relative), not in the points. On a many-core CPU, with COSMONUTS_CPU_DEVICES=n XLA_FLAGS=--xla_cpu_multi_thread_eigen=false and a num_delete a multiple of n, each core then runs num_delete / n chains one after another. Points of log L \(= -\infty\) (non-viable) are counted in the progress lines among the dead points and the initial ones; during a step, rejected proposals count as shrinkages whatever their log L.	'auto'

Returns:

Type	Description
NestedResult	The weighted points and the evidence.

Raises:

Type	Description
ValueError	For bad settings (including an unknown parallel or batching), resume without output or with overwrite, or settings or parameters that do not match the saved run.
FileExistsError	If chains exist at output and neither overwrite nor resume.
FileNotFoundError	With resume, if there is no saved run at output.
RuntimeError	If device names a platform JAX does not have.

cosmonuts.maximize

maximize(model: Model, initial: dict[str, Any] | list[dict[str, Any]] | None = None, *, key: Any = None, num_starts: int = 1, ignore_prior: bool = False, max_steps: int = 500, tol: float = 1e-06, progress: bool | int = False) -> MaximizeResult

Maximum a posteriori (or maximum likelihood) point, by L-BFGS.

The objective is \(\log L + \log\pi\) (only \(\log L\) with ignore_prior), optimised over the unconstrained \(u\); without the Jacobian, so the optimum is that of the physical-space posterior. For uniform priors both options give the best fit. Starts run in parallel.

Parameters:

Name	Type	Description	Default
model	Model	The model.	required
initial	dict[str, Any] | list[dict[str, Any]] | None	Physical starting point(s) of the sampled parameters.	None
key	Any	A jax.random key; with num_starts > len(initial) the extra starts are prior draws.	None
num_starts	int	Total number of starts.	1
ignore_prior	bool	Maximise the likelihood only.	False
max_steps	int	L-BFGS iterations per start.	500
tol	float	Stop when the gradient norm falls below this.	1e-06
progress	bool | int	Print a line every progress iterations (True: every 5% of max_steps): iterations, time (compilation excluded), the best objective so far, how many starts have converged and the largest gradient norm among the others; then a summary. False: silent.	False

Returns:

Type	Description
MaximizeResult	The best of the starts.

Raises:

Type	Description
ValueError	If no start ends at a finite value.

Results and output

cosmonuts.NUTSResult

Bases: NamedTuple

Output of nuts.

Attributes:

Name	Type	Description
samples	dict[str, ndarray]	Physical samples by name, each of shape (chains, draws).
log_density	ndarray	Log density of each sample, (chains, draws).
rhat	dict[str, float]	Rank-normalised split-\(\hat R\) by parameter (Vehtari et al. 2021; see rhat).
ess	dict[str, float]	Effective sample size by parameter.
divergences	int	Number of divergent transitions.
num_steps	float	Mean number of leapfrog steps per draw.
step_size	float	Adapted step size.
inverse_mass_matrix	ndarray	Adapted inverse mass matrix: (d,) (its diagonal) with mass_matrix="diagonal", (d, d) with "dense".
settings	dict[str, Any]	The run's arguments: num_warmup, num_samples, num_chains, mass_matrix, target_ess, check_every, jitter, initial (the physical starting point), seed (the key's data), max_num_doublings and mass_matrix_seed (the number of seed draws used, or None); and num_draws, the draws per chain actually taken (fewer than num_samples after an early stop).
wall_time	float	Seconds for warm-up and sampling, compilation included.
derived	dict[str, ndarray]	The quantities recorded with each sample, by chain column, (chains, draws): saved Defined parameters, the model's derived (a 1-D one as name_0, ...) and the likelihoods' derived.
rhat_gelman_rubin	dict[str, float]	Classic potential scale reduction by parameter (Gelman & Rubin 1992), nan for one chain.

sampler property

sampler: str

The method that drew the samples: "nuts" or "hmc".

posterior

posterior() -> dict[str, dict[str, float]]

Mean, standard deviation, \(\hat R\) and ESS of each sampled parameter.

diagnostics

diagnostics() -> dict[str, Any]

The run's diagnostics: the results section of .updated.yaml.

summary

summary() -> dict[str, Any]

The headline numbers of the run, formatted for display.

cosmonuts.HMCResult module-attribute

HMCResult = NUTSResult

cosmonuts.MetropolisResult

Bases: NamedTuple

Output of metropolis.

Attributes:

Name	Type	Description
samples	dict[str, ndarray]	Physical samples by name, each of shape (chains, draws).
log_density	ndarray	Log density of each sample in the unconstrained space, (chains, draws).
rhat	dict[str, float]	Rank-normalised split-\(\hat R\) by parameter (Vehtari et al. 2021; see rhat).
ess	dict[str, float]	Effective sample size by parameter.
acceptance_rate	ndarray	Acceptance rate of each chain in the sampling phase, (chains,).
proposal_covariance	ndarray	The frozen proposal covariance, in the unconstrained space (proposal_scale**2 times the learned posterior covariance).
proposal_scale	float	Scale of the proposal relative to the learned covariance (the optimal \(2.38/\sqrt{d}\) is the starting value).
settings	dict[str, Any]	The run's arguments, plus learning_steps, learning_converged and num_draws (below num_samples if target_ess stopped the run early).
wall_time	float	Seconds for learning and sampling, compilation included.
derived	dict[str, ndarray]	The quantities recorded with each sample, by chain column, (chains, draws).
rhat_gelman_rubin	dict[str, float]	Classic potential scale reduction by parameter (Gelman & Rubin 1992), nan for one chain.

sampler property

sampler: str

The method that drew the samples: "metropolis".

posterior

posterior() -> dict[str, dict[str, float]]

Mean, standard deviation, \(\hat R\) and ESS of each sampled parameter.

diagnostics

diagnostics() -> dict[str, Any]

The run's diagnostics: the results section of .updated.yaml.

summary

summary() -> dict[str, Any]

The headline numbers of the run, formatted for display.

cosmonuts.NestedResult

Bases: NamedTuple

Output of nested.

Every dead point, in the order they died, then the final live points, each with its posterior weight. GetDist and cplots take the weights; resample gives equal-weight draws.

Attributes:

Name	Type	Description
samples	dict[str, ndarray]	Physical sampled parameters by name, each of shape (n,).
weights	ndarray	Posterior weights, (n,), summing to 1.
loglike	ndarray	\(\log L\) of each point (Model.log_likelihood: without the likelihoods' normalisation constants, as the \(\chi^2\) of the chains).
loglike_birth	ndarray	\(\log L\) of the contour each point was drawn inside (nan for the first live points, drawn from the prior).
logZ	float	\(\ln Z\) of loglike (mean over simulated prior volumes): Bayes factors between models of the same data.
logZ_err	float	Its standard deviation over the simulations.
logZ_normalized	float | None	The absolute evidence, logZ + log_normalization, or None if a likelihood cannot give its constant.
log_normalization	float | None	model.log_normalization, or None.
information	float	Kullback-Leibler divergence from prior to posterior, in nats.
ess	float	Kish effective sample size of the weights.
converged	bool	Whether the run reached frac_remain (rather than max_dead).
num_dead	int	Dead points: the first num_dead entries; the rest were live at the end.
num_loglike	int	Likelihood calls (approximate: initial points, slice expansions and shrinkages).
settings	dict[str, Any]	The run's settings: sampler ("nested"), nlive, num_delete, num_inner_steps, frac_remain, kernel, max_steps, max_shrinkage, check_every, seed, device and resumed_dead (dead points from an earlier session).
wall_time	float	Seconds, compilation included (all sessions, when resumed).
derived	dict[str, ndarray]	The quantities recorded with each point, by chain column, (n,): saved Defined parameters, the model's derived (a 1-D one as name_0, ...) and the likelihoods' derived. Evaluated after the run, for every point.

sampler property

sampler: str

The method that drew the samples: "nested".

posterior

posterior() -> dict[str, dict[str, float]]

Weighted mean and standard deviation of each sampled parameter.

diagnostics

diagnostics() -> dict[str, Any]

The evidence and the run's counters: the results section of .updated.yaml.

summary

summary() -> dict[str, Any]

The headline numbers of the run, formatted for display.

resample

resample(key: Any, n: int) -> dict[str, np.ndarray]

n equal-weight posterior draws (with replacement), physical, by name (derived too).

cosmonuts.MaximizeResult

Bases: NamedTuple

Output of maximize.

Attributes:

Name	Type	Description
params	dict[str, float]	Best-fit physical parameters (sampled and fixed).
value	float	Objective at the best fit (log posterior, or log likelihood).
chi2	float	\(-2\log L\) at the best fit.
converged	bool	Whether the best start met the tolerance.
all_values	ndarray	Objective at every start's end point (-inf for dropped starts).

cosmonuts.save_chains

save_chains(root: str | PathLike, result: Any, model: Model, *, labels: Mapping[str, str] | None = None, overwrite: bool = False) -> list[Path]

Write chains, parameter names and run settings under a path root.

Each draw's likelihoods, prior and model.derived quantities are recomputed (one likelihood call per draw, batched with jax.vmap) for the chi2 and derived columns. nuts(..., output=root) writes the same files as it samples, with a ChainWriter.

A NestedResult gives one weighted chain, <root>.1.txt (without the points of negligible weight, below \(10^{-30}\) of the largest), and every point in <root>_dead-birth.txt (parameters, \(\log L\) and the \(\log L\) at birth, \(-\infty\) for the first live points: PolyChord's format, which anesthetic reads); .updated.yaml then records the evidence.

Parameters:

Name	Type	Description	Default
root	str | PathLike	Path root, e.g. "chains/w0wa"; its directory is created.	required
result	Any	A NUTSResult or NestedResult of model.	required
model	Model	The model that was sampled.	required
labels	Mapping[str, str] | None	LaTeX labels (no $) overriding model.labels.	None
overwrite	bool	Replace existing chains at root (all of them, so no stale chain is left for GetDist to load).	False

Returns:

Type	Description
list[Path]	The files written.

Raises:

Type	Description
FileExistsError	If chains exist at root and overwrite is false.
ValueError	If a derived quantity is neither a scalar nor 1-D.

cosmonuts.rhat

rhat(draws: Any, method: RhatMethod = 'rank') -> np.ndarray

\(\hat R\) by parameter of draws (chains, n, d), shape (d,).

Parameters:

Name	Type	Description	Default
draws	Any	The draws, chains first.	required
method	RhatMethod	"rank", the rank-normalised split-\(\hat R\), or "gelman-rubin", the classic potential scale reduction.	'rank'

Returns:

Type	Description
ndarray	\(\hat R\) of each parameter. "rank" splits each chain in two, so it works on one
ndarray	chain too, and is nan with fewer than 4 draws per chain; "gelman-rubin" is nan
ndarray	with one chain.

Raises:

Type	Description
ValueError	For an unknown method.

cosmonuts.QUANTITIES module-attribute

QUANTITIES = {'hubble': '$H(z)$ in Mpc$^{-1}$ ($c = 1$)', 'comoving_distance': 'Comoving distance $\\chi(z)$ in Mpc (equal to $D_M$ when flat)', 'angular_diameter_distance': '$D_A(z)$ in Mpc', 'sound_horizon_drag': 'Sound horizon at the baryon drag epoch, $r_d$ in Mpc', 'theta_star': "CMB acoustic angle $\\theta_* = r_s(z_*)/D_M(z_*)$ in radians (CAMB's input `thetastar`, not $\\times 100$)", 'growth_rate': 'Linear growth rate $f(z) = d\\ln D/d\\ln a$', 'pk_cb': 'Linear power spectrum of baryons and CDM, $P_{cb}(k, z)$ in Mpc$^3$, $k$ in Mpc$^{-1}$', 'h': '$h = H_0 / (100\\,{\\rm km\\,s^{-1}\\,Mpc^{-1}})$', 'Omega_b': 'Baryon density parameter today', 'Omega_cb': 'Baryon + CDM density parameter today', 'n_s': 'Scalar spectral index', 'cl_tt': 'Lensed $D_\\ell^{TT} = \\ell(\\ell+1)C_\\ell/2\\pi$ in $\\mu{\\rm K}^2$, at integer $\\ell$', 'cl_te': 'Lensed $D_\\ell^{TE}$ in $\\mu{\\rm K}^2$, at integer $\\ell$', 'cl_ee': 'Lensed $D_\\ell^{EE}$ in $\\mu{\\rm K}^2$, at integer $\\ell$', 'cl_bb': 'Lensed $D_\\ell^{BB}$ in $\\mu{\\rm K}^2$, at integer $\\ell$', 'cl_pp': 'Lensing potential $[L(L+1)]^2 C_L^{\\phi\\phi}/2\\pi$ (dimensionless), at integer $L$'}

The standard quantity names, with their units: the methods of a provider.

Theories

cosmonuts.theories.w0wa.W0WaCDM

W0WaCDM(rdrag: Literal['fit', 'param'] = 'fit', T_cmb: float = 2.7255, N_eff: float = 3.044, m_nu: float | tuple[float, ...] = ())

Bases: Theory

Flat \(w_0w_a\)CDM background with radiation.

Parameters:

Name	Type	Description	Default
rdrag	Literal['fit', 'param']	"fit" to integrate \(r_d\) over the expansion up to the fitted drag redshift (drag_redshift_fit), or "param" to read it from the parameter rdrag (in Mpc), e.g. for BAO alone.	'fit'
T_cmb	float	CMB temperature in K.	2.7255
N_eff	float	Effective number of neutrino species, massive ones included.	3.044
m_nu	float | tuple[float, ...]	Masses in eV of the massive neutrino states (e.g. 0.06 or (0.06,) for one state, as CAMB's default); each carries \(N_{\rm eff}/3\). Empty or 0.0: all massless.	()

Attributes:

Name	Type	Description
params	tuple[str, ...]	h, omega_b, omega_cdm, w0, wa, and rdrag if sampled.
provides	frozenset[str]	Background quantities: hubble, comoving_distance, angular_diameter_distance, sound_horizon_drag, theta_star, h, Omega_b, Omega_cb.

cosmonuts.theories.classe.ClassETheory

ClassETheory(bundle: str | Path | EmulatorBundle, *, T_cmb: float = T_CMB, reject_outside_box: bool = True)

Bases: Theory

Lensed CMB spectra from a class-e emulator bundle.

Parameters:

Name	Type	Description	Default
bundle	str | Path | EmulatorBundle	Path to a .cemu file, or an EmulatorBundle. Enable x64 first (jax.config.update("jax_enable_x64", True)): emulators loaded in float32 lose precision.	required
T_cmb	float	CMB temperature in K used to convert \(C_\ell\) to \(\mu{\rm K}^2\).	T_CMB
reject_outside_box	bool	Set provider.ok false outside the bundle's prior box, where the emulators extrapolate, and where the training design's constraints (bundle.meta["constraints"], e.g. "w0 + wa < 0") fail.	True

Attributes:

Name	Type	Description
params	tuple[str, ...]	The bundle's parameter names, in its order.
provides	frozenset[str]	cl_xy for each emulated spectrum (TT, TE, EE, and BB, PP if present), the distances with an Hz distance emulator, and sound_horizon_drag, theta_star, h from the scalars.
l_max	int | None	Largest multipole served (None without spectra).
z_max	float | None	Largest redshift of the distances (None without a distance emulator).
bounds	dict[str, tuple[float, float]]	The emulators' prior box, {name: (min, max)}.

bounds property

bounds: dict[str, tuple[float, float]]

The emulators' prior box by parameter name.

cosmonuts.theories.eft4py.EFT4PyTheory

EFT4PyTheory(background: Any, ln_a: Any, couplings: Sequence[str | None], *, cosmology: Callable[[dict], Cosmology] | None = None, cosmology_params: Sequence[str] | None = None, shoot_guess: Callable[[Cosmology, dict[str, Any]], Any] = _default_guess, shoot_slope: Callable[[Cosmology, dict[str, Any]], Any] | None = _default_slope, shoot_options: Mapping[str, Any] | None = None, rdrag: Literal['computed', 'param'] = 'computed', growth: bool = True, power_spectrum: bool = True, require_stable: bool = True, transfer: Any = None, primordial: Any = None)

Bases: Theory

A Horndeski model solved by eft4py's JaxBackground.

Parameters:

Name	Type	Description	Default
background	Any	The eft4py JaxBackground (model, free couplings, initial conditions).	required
ln_a	Any	Grid in \(\ln a\) for the solve; it must end at 0 and start before \(a \approx 10^{-4}\) for \(r_d\) and \(\theta_*\), and before \(10^{-3}\) for the growth. \(\theta_*\) integrates over the early grid: 200 points between \(a = 10^{-8}\) and \(z = 10\) give \(10^{-6}\) (60 give \(8\times10^{-5}\), a third of Planck's error).	required
couplings	Sequence[str | None]	For each of the background's free couplings, in order, the name of the parameter that sets it, or None for the one tuned by shooting.	required
cosmology	Callable[[dict], Cosmology] | None	params -> Cosmology; default Cosmology.from_physical(h, omega_b, omega_cdm).	None
cosmology_params	Sequence[str] | None	Parameter names cosmology reads; default ("h", "omega_b", "omega_cdm").	None
shoot_guess	Callable[[Cosmology, dict[str, Any]], Any]	(cosmo, params) -> starting guess for the tuned coupling; default \(3\Omega_{\rm DE}H_0^2\), a cosmological constant's value. A guess that knows the model (from the other couplings in params) saves shooting iterations, each a background solve.	_default_guess
shoot_slope	Callable[[Cosmology, dict[str, Any]], Any] | None	(cosmo, params) -> estimate of \(d(H(a=1)/H_0)/ds\) near the guess, \(s\) the tuned coupling's factor: the secant's second point is then a Newton step from the first solve, which saves about one background solve. Default \(\Omega_{\rm DE}/2\), right for a coupling that acts as a cosmological constant (as the default guess assumes); a poor estimate costs iterations, not accuracy; None: eft4py's fixed second point.	_default_slope
shoot_options	Mapping[str, Any] | None	Keyword arguments for eft4py's JaxBackground.shoot_and_solve: method ("secant", eft4py's default, or "newton"), rtol, atol (on \(H(a=1)/H_0 - 1\)), max_steps; default eft4py's.	None
rdrag	Literal['computed', 'param']	"computed" (\(r_d\) from the solved background) or "param" (read from the parameter rdrag, in Mpc).	'computed'
growth	bool	Solve the growth, providing growth_rate.	True
power_spectrum	bool	Provide pk_cb (needs growth), and n_s when the primordial spectrum has it. Without transfer: syren-new at \(z_{\rm ref}\), reading A_s and n_s.	True
transfer	Any	A class-e emulator bundle (path or EmulatorBundle, needs cosmonuts[classe]) of the CDM+baryon transfer function Tk on grids k (Mpc\(^{-1}\)) and z, e.g. from class-e's configs/lcdm-transfer.yaml. Then \(P_{cb} = 2\pi^2/k^3\,\mathcal{P}_\mathcal{R}(k)\,T^2\) evolved by the model's growth (see TransferPower), for any primordial; the bundle's parameters are read by name, and points outside its prior box give ok = False. With massive neutrinos use a bundle with a redshift grid and a cosmology that includes them.	None
primordial	Any	\(\mathcal{P}_\mathcal{R}(k)\) template with transfer (default PowerLaw); its parameters are read.	None
require_stable	bool	Reject models that violate the stability conditions.	True

Attributes:

Name	Type	Description
params	tuple[str, ...]	Parameter names read, from the above.
provides	frozenset[str]	Provider quantities, from the above.

from_yaml classmethod

from_yaml(path: str | Path, ln_a: Any = None, *, solver: Mapping[str, Any] | None = None, **kwargs: Any) -> EFT4PyTheory

The theory of an eft4py run file (eft4py.config.HorndeskiFactory).

The run file gives the model (\(G_i\) and couplings), its initial conditions and the coupling tuned for \(H_0\) (shoot). Every other coupling becomes a sampled parameter of the same name. As in eft4py, the values entry of the tuned coupling is its starting guess: an expression of the other couplings, the cosmology (h, H0, Omega_m, Omega_r, Omega_de) and inputs, used as shoot_guess unless one is given. The rest of values, and the cosmology and output blocks, are not used, since the sampler sets the couplings and the cosmology, and may be left out. In a run file:

theory:
  cosmonuts.theories.eft4py:EFT4PyTheory.from_yaml:
    {path: model.yaml, rdrag: param, growth: false, power_spectrum: false}

Parameters:

Name	Type	Description	Default
path	str | Path	The eft4py run file.	required
ln_a	Any	Grid in \(\ln a\) (see the class); default 200 points from the file's a_ini to \(z = 10\) and 240 from there to today (only the second segment if a_ini is later than \(z = 10\)).	None
solver	Mapping[str, Any] | None	Options of eft4py's JaxBackground, over the file's solver block (whose backend is always "jax"). The ODE's rtol defaults to DEFAULT_RTOL (\(10^{-8}\)), not eft4py's \(10^{-10}\). Its max_steps is eft4py's: viable histories need far fewer steps, and a lower cap (e.g. 5000) makes the non-viable points, which otherwise run to the cap, much cheaper, but check that it leaves log L unchanged for your model (5000 was checked for the SBSS models only). stop_at_ghost_wall: true is not yet safe: it drops viable points.	None
**kwargs	Any	Further arguments of EFT4PyTheory.	{}

Returns:

Type	Description
EFT4PyTheory	The theory.

Raises:

Type	Description
ValueError	If the file has no shoot coupling.

cosmonuts.theories.combined.Theories

Theories(theories: Mapping[str, Theory], provides: Mapping[str, str] | None = None, given: Iterable[str] | None = None)

Bases: Theory

Named theories run together, each quantity served by one of them.

Parameters:

Name	Type	Description	Default
theories	Mapping[str, Theory]	The theories by name (the names appear in errors and saved settings).	required
provides	Mapping[str, str] | None	Which theory serves a quantity that several provide, {quantity: name}.	None
given	Iterable[str] | None	The names that have a prior, fixed value or definition. A theory parameter outside it that another theory provides as a quantity is read from that theory. None (default) reads every parameter that another theory provides; Model passes its priors' names.	None

Raises:

Type	Description
ValueError	For an empty mapping, a provides entry naming an unknown theory or one that does not provide the quantity, a parameter that several other theories could supply and provides does not settle, or theories that read each other's outputs in a cycle.

Attributes:

Name	Type	Description
params	tuple[str, ...]	The theories' parameters, minus those read from other theories.
provides	frozenset[str]	The quantities with an assigned theory.
inputs	tuple[tuple[str, str, str], ...]	(theory, parameter, supplier) for each chained parameter.
order	tuple[str, ...]	The theories' names in the order they run.
l_max	int | None	The smallest l_max of the theories that serve CMB spectra.

param_labels property

param_labels: dict[str, str]

The theories' parameter labels (the first theory wins on a clash).

name property

name: str

The theories' names, for summaries.

Functions as theories and likelihoods

cosmonuts.external.FunctionTheory

FunctionTheory(fn: Callable, provides: Iterable[str], *, name: str = 'theory', source: str | None = None)

Bases: Theory

A theory from a function of the parameters that returns a dictionary of quantities.

The dictionary may hold ok, a boolean that is False where the computation failed.

Parameters:

Name	Type	Description	Default
fn	Callable	The function; its argument names are the parameters it reads.	required
provides	Iterable[str]	The quantity names it returns.	required
name	str	A name for saved settings.	'theory'
source	str | None	The function's source, for saved settings (a run file's lambda string).	None

Raises:

Type	Description
ValueError	If provides is empty.

cosmonuts.external.FunctionLikelihood

FunctionLikelihood(fn: Callable, requires: Iterable[str] = (), *, name: str = 'likelihood', source: str | None = None)

Bases: Likelihood

A likelihood from a function returning \(\log L\).

Parameters:

Name	Type	Description	Default
fn	Callable	The function. Its argument names are the parameters it reads, except provider, which receives the theory's output.	required
requires	Iterable[str]	The theory quantities it reads from the provider.	()
name	str	Its name, for the chains' chi2__<name> column and the saved settings.	'likelihood'
source	str | None	The function's source, for saved settings (a run file's lambda string).	None

Raises:

Type	Description
ValueError	If requires is given and the function takes no provider.

log_likelihood

log_likelihood(provider: Any, params: dict[str, Any]) -> Any

Return fn called with the provider and the parameters it names.

cosmonuts.external.NoTheory

Bases: Theory

No theory: for models whose likelihoods only read parameters.

Data, run files and output files

cosmonuts.installer.install

install(likelihood: Any, path: str | Path | None = None, *, force: bool = False, log: Any = print) -> Path

Download a likelihood's data files and run its postprocess hook.

Parameters:

Name	Type	Description	Default
likelihood	Any	A likelihood class, or a registered dataset name (see REGISTRY).	required
path	str | Path | None	Packages path; default from packages_path.	None
force	bool	Download again even if files exist.	False
log	Any	Called with progress messages; None to silence.	print

Returns:

Type	Description
Path	The likelihood's data directory.

Raises:

Type	Description
ValueError	If the class declares no install_name or no data_sources.

cosmonuts.installer.is_installed

is_installed(likelihood: Any, path: str | Path | None = None) -> bool

Whether every data file of a likelihood (class or registered name) is present.

cosmonuts.installer.paths.packages_path

packages_path(path: str | Path | None = None) -> Path

The packages path (not created).

Parameters:

Name	Type	Description	Default
path	str | Path | None	Explicit path; takes precedence.	None

Returns:

Type	Description
Path	The resolved directory.

cosmonuts.config.run

run(*sources: str | Path | Mapping[str, Any], overrides: Sequence[str] = (), test: bool = False, progress: bool | int = True) -> RunResult

Sample the model of a run file, as cobaya.run: cosmonuts sample from Python.

Parameters:

Name	Type	Description	Default
*sources	str | Path | Mapping[str, Any]	Run files (paths or mappings), merged left to right.	()
overrides	Sequence[str]	key.path=value overrides, applied last.	()
test	bool	Only build the model and evaluate it at the start point.	False
progress	bool | int	The sampler's progress display.	True

Examples:

from cosmonuts.config import run
out = run("w0wa.yaml", overrides=["sampler.nuts.num_samples=200"])

cosmonuts.optimize.covariance

covariance(model: Model, params: dict[str, Any], *, ignore_prior: bool = False, space: str = 'physical') -> np.ndarray

Laplace covariance at a best fit: the inverse of minus the Hessian of the objective.

For a Gaussian posterior (or likelihood) this is its covariance; near a best fit it is the Fisher matrix estimate of the errors. Parameters are in model.sampled order.

Parameters:

Name	Type	Description	Default
model	Model	The model.	required
params	dict[str, Any]	The best fit, e.g. maximize(...).params (fixed parameters are ignored).	required
ignore_prior	bool	Curvature of \(\log L\) only, as maximize(ignore_prior=True) maximises.	False
space	str	"physical", the covariance of the sampled parameters; "unconstrained", that of the sampler's coordinates \(u\), from Model.log_density (Jacobian and prior included, so ignore_prior must be false), e.g. a NUTS inverse mass matrix.	'physical'

Raises:

Type	Description
ValueError	For an unknown space, or a Hessian that is not negative definite (not a maximum: a saddle, a flat direction, or a best fit on a prior edge).

cosmonuts.output.ChainWriter

ChainWriter(root: str | PathLike, model: Model, num_chains: int, *, labels: Mapping[str, str] | None = None, overwrite: bool = False, resume: bool = False)

Chains on disk, written while sampling: rows appended chunk by chunk, like Cobaya.

The chain files (with their # header), .paramnames and .ranges are written at once; append adds rows (the likelihoods, prior and derived quantities are evaluated for the new draws only); write_adaptation stores the adapted step size and mass matrix after warm-up (<root>.adaptation.npz); finalize adds .covmat and .updated.yaml. A crash then loses at most the draws since the last append.

Parameters:

Name	Type	Description	Default
root	str | PathLike	Path root, e.g. "chains/w0wa"; its directory is created.	required
model	Model	The model being sampled.	required
num_chains	int	Number of chains (one file each).	required
labels	Mapping[str, str] | None	LaTeX labels (no $) overriding model.labels.	None
overwrite	bool	Replace existing chains at root instead of raising.	False

Raises:

Type	Description
FileExistsError	If chains exist at root and overwrite is false.

append

append(samples: Mapping[str, Any], weights: Any = None, evaluated: Evaluation | None = None) -> None

Add draws: samples[name] of shape (chains, n), in physical units.

Parameters:

Name	Type	Description	Default
samples	Mapping[str, Any]	samples[name] of shape (chains, n), in physical units.	required
weights	Any	Weights of the same shape, filling the weight column (default 1: unweighted draws).	None
evaluated	Evaluation | None	The draws' _evaluate, flattened over chains, if already done.	None

preload

preload(samples: Mapping[str, Any]) -> None

Record draws already in the files (when resuming), for the final .covmat.

Their outputs (derived) are read back from the files.

derived

derived() -> dict[str, np.ndarray]

The outputs of every draw written (and preloaded), (chains, draws) by column.

write_adaptation

write_adaptation(step_size: Any, inverse_mass_matrix: Any, position: Any, *, sampler: str = 'nuts', num_integration_steps: int | None = None) -> Path

Store the warm-up's adapted step size, inverse mass matrix and final position.

Parameters:

Name	Type	Description	Default
step_size	Any	The adapted step size (per chain, or one).	required
inverse_mass_matrix	Any	The adapted inverse mass matrix.	required
position	Any	The final position.	required
sampler	str	The sampler's name, "nuts" or "hmc".	'nuts'
num_integration_steps	int | None	HMC's trajectory length, saved so a resume reuses it.	None

finalize

finalize(result: Any) -> list[Path]

Write .covmat and .updated.yaml; return every file written.