Ensemble Sampler¶

We include an ensemble sampler that can be used to sample from a target distribution. This sampler is based on the EnsembleSampler implementation in emcee.

class hemcee.samplers.ensemble.EnsembleSampler(total_chains: int, dim: int, log_prob: ~typing.Callable, move=<function stretch_move>, backend: ~hemcee.backend.backend.Backend = None)¶

Bases: BaseSampler

Affine-invariant ensemble sampler wrapper.

get_acceptance_prob()¶

get_autocorr(discard, thin)¶

get_chain(discard: int = 0, thin: int = 1, flat: bool = False) → Array¶

get_logprob(discard: int = 0, thin: int = 1, flat: bool = False) → Array¶

run_mcmc(key: PRNGKey, initial_state: Array, num_samples: int, warmup: int = 1000, thin_by: int = 1, batch_size: int = None, show_progress: bool = False, **kwargs) → Tuple[Array, dict]¶

Run the vanilla ensemble sampler.

Args:

key (jax.random.PRNGKey): Random number generator key. initial_state (jnp.ndarray): Initial ensemble state with shape

(total_chains, dim).

num_samples (int): Number of post-warmup samples to retain. warmup (int): Number of warmup iterations. Defaults to 1000. thin_by (int): Keep every thin_by sample. Defaults to 1. show_progress (bool): Whether to display a progress bar. Defaults

to False.

**kwargs: Additional keyword arguments passed to the proposal move.

Returns:

tuple[jnp.ndarray, dict]: Post-warmup samples and diagnostic: information including acceptance rates.