jwspecmcmc.samplers

MCMC sampler wrappers for emcee, nautilus, and NumPyro NUTS.

All wrappers accept a LikelihoodSpec and PriorSet, run the sampler, and return a common result dict.

Functions

run_emcee(spec, prior_set, p0_free, *[, ...])

Run the emcee ensemble sampler.

run_nautilus(spec, prior_set, *[, n_live, ...])

Run the nautilus nested sampler.

run_nuts(spec, prior_set, p0_free, *[, ...])

Run NumPyro NUTS (No-U-Turn Sampler) with JAX-accelerated likelihood.
jwspecmcmc.samplers.run_emcee(spec, prior_set, p0_free, *, n_walkers='auto', n_steps=2000, n_burn=None, progress=True, seed=42, moves=None)[source]

Run the emcee ensemble sampler.

Parameters:

  • spec (LikelihoodSpec) – Cached data for likelihood evaluation.

  • prior_set (PriorSet) – Prior distributions.

  • p0_free (np.ndarray) – MLE estimate in free-parameter space (used to initialise walkers).

  • n_walkers (int or "auto") – Number of walkers. "auto" (default) picks a value based on n_dim and the number of CPU cores.

  • n_steps (int) – Number of MCMC steps (default 2000).

  • n_burn (int or None) – Burn-in steps to discard. If None, estimated from the integrated autocorrelation time.

  • progress (bool) – Show a progress bar.

  • seed (int) – Random seed.

  • moves (optional) – Custom emcee moves. If None, uses the default StretchMove.

Returns:

Keys: flat_chains (n_samples, n_dim), flat_log_prob (n_samples,), chains (n_walkers, n_steps_kept, n_dim), log_prob_chains (n_walkers, n_steps_kept), n_burn (int), sampler_name (str), sampler_meta (dict).

Return type:

dict

jwspecmcmc.samplers.run_nautilus(spec, prior_set, *, n_live=2000, n_eff=10000, progress=True, seed=42)[source]

Run the nautilus nested sampler.

Parameters:

  • spec (LikelihoodSpec) – Cached data for likelihood evaluation.

  • prior_set (PriorSet) – Prior distributions.

  • n_live (int) – Number of live points (default 2000).

  • n_eff (int) – Target effective sample size (default 10000).

  • progress (bool) – Show a progress bar.

  • seed (int) – Random seed.

Returns:

Same keys as run_emcee(), except chains and log_prob_chains are None (nautilus does not produce walker chains).

Return type:

dict

jwspecmcmc.samplers.run_nuts(spec, prior_set, p0_free, *, n_warmup=500, n_samples=2000, n_chains=6, progress=True, seed=42, target_accept_prob=0.8, max_tree_depth=10)[source]

Run NumPyro NUTS (No-U-Turn Sampler) with JAX-accelerated likelihood.

Uses automatic differentiation through the emission-line model for gradient-based HMC sampling, requiring far fewer likelihood evaluations than ensemble samplers like emcee.

Parameters:

  • spec (LikelihoodSpec) – Cached data for likelihood evaluation.

  • prior_set (PriorSet) – Prior distributions.

  • p0_free (np.ndarray) – MLE estimate in free-parameter space (used to initialise).

  • n_warmup (int) – Number of warmup (adaptation) steps (default 500).

  • n_samples (int) – Number of posterior samples per chain (default 2000).

  • n_chains (int) – Number of independent chains (default 1).

  • progress (bool) – Show a progress bar.

  • seed (int) – Random seed.

  • target_accept_prob (float) – Target acceptance probability for NUTS adaptation (default 0.8).

  • max_tree_depth (int) – Maximum tree depth for NUTS (default 10).

Returns:

Same keys as run_emcee().

Return type:

dict