Samplers

sample is the one entry point: method picks the algorithm, device picks CPU or GPU, and the rest of the call is the method's.

import jax
from cosmonuts import sample

key = jax.random.PRNGKey(0)
result = sample(model, key, output="chains/w0wa", progress=True)                          # NUTS
result = sample(model, key, method="nested", device="gpu", output="chains/w0wa_ns")       # evidence

Method	Function	Use it for
"nuts"	nuts	Default. Posteriors, with adaptive step size and mass matrix
"hmc"	hmc	Fixed-length trajectories (num_integration_steps)
"metropolis"	metropolis	Gradient-free random walk, for models that are not differentiable
"nested"	nested	The evidence \(\ln Z\) and its error; fills a GPU
(best fit)	maximize	L-BFGS from several starts

Each returns a result with samples, derived, posterior() and summary(); NUTS and HMC add \(\hat R\), ESS and divergences.

Chains in parallel

NUTS, HMC and Metropolis take parallel=:

parallel	Chains run	Warm-up
"vmap"	batched on one device	one, shared
"shard"	spread over devices (COSMONUTS_CPU_DEVICES=n makes CPU cores)	one per chain
"process"	one process each	one per chain
"mpi"	one MPI rank each (mpirun -n <chains> python -m mpi4py script.py)	one per chain
"auto"	"shard" with several chains and devices, else "vmap"

CPU or GPU

NUTS gains from a GPU mainly when the theory itself is heavy. Nested sampling is the method that parallelises: its batch of slice chains is what fills a GPU.

Stopping on ESS

target_ess= stops once the smallest effective sample size reaches it and \(\hat R < 1.01\), checked every check_every draws.

Resume

With output, progress is saved as it goes; run again with resume=True to continue. See Your first fit.

Best fit

from cosmonuts import maximize

best = maximize(model, num_starts=4, key=jax.random.PRNGKey(0))
best.params, best.chi2

Use best.params as initial= for a sampler. For constraints that do not depend on prior volume, see Profile likelihoods.