mixle.inference.mcmc.samplers module

Generic MCMC drivers over user-supplied log targets and proposals.

The low-level functions here deliberately operate on user-supplied log-target callables and proposal objects. That keeps the transition machinery orthogonal to the distribution, estimator, and compute-engine protocols while still making ordinary dist.log_density(x) models easy to sample from.

class MCMCResult(samples, log_probs, accepted, transition_labels=None)[source]

Bases: object

Samples and diagnostics returned by an MCMC run.

Parameters:
samples: list[Any]
log_probs: ndarray
accepted: ndarray
transition_labels: tuple[str, ...] | None = None
property acceptance_rate: float

Return the overall fraction of accepted transitions.

property acceptance_rate_by_label: dict[str, float]

Return acceptance rates for labelled transition kernels.

sample_array()[source]

Return numeric samples as an ndarray for diagnostics.

Return type:

ndarray

effective_sample_size(max_lag=None)[source]

Estimate effective sample size using positive autocorrelation lags.

Scalar samples return a float. Vector samples return one ESS value per trailing dimension.

Parameters:

max_lag (int | None)

Return type:

Any

summary(max_lag=None)[source]

Return basic numeric chain diagnostics.

The summary intentionally stays dependency-free and returns plain numbers/arrays: sample count, mean, variance, Monte Carlo standard error estimate, ESS, and acceptance diagnostics.

Parameters:

max_lag (int | None)

Return type:

dict[str, Any]

distribution_log_target(dist, evidence=None)[source]

Return log_target(x) = dist.log_density(x) + evidence(x).

Parameters:
Return type:

Callable[[Any], float]

metropolis_hastings(log_target, initial, proposal, num_samples, burn_in=0, thin=1, rng=None)[source]

Run a generic Metropolis-Hastings chain.

Parameters:

  • log_target (Callable[[Any], float]) – Callable returning an unnormalized log target.

  • initial (Any) – Initial Markov-chain state.

  • proposal (Proposal) – Proposal object with sample and optional log_density methods.

  • num_samples (int) – Number of post-burn/thinned states to return.

  • burn_in (int) – Number of initial transitions to discard.

  • thin (int) – Keep one sample every thin transitions.

  • rng (RandomState | None) – Optional RandomState.

Returns:

MCMCResult with retained samples, retained log probabilities, and the accept/reject indicator for every transition.

Return type:

MCMCResult

affine_invariant_ensemble(log_target, p0, num_samples, burn_in=0, thin=1, a=2.0, rng=None)[source]

Goodman & Weare affine-invariant ensemble sampler (the “stretch move”).

A population of W walkers explores the target jointly; each walker is proposed along the line to a randomly chosen complementary walker, so the sampler is invariant to affine rescalings of the target and needs no per-dimension step tuning. It mixes far better than random-walk Metropolis on correlated/poorly-scaled posteriors and, because every proposal is one log-target evaluation, delivers very high ESS/sec on low/medium-dimensional models.

Parameters:

  • log_target (Callable[[ndarray], float]) – unnormalized log target for a single walker state (d,).

  • p0 (ndarray) – initial ensemble, shape (W, d) with W even and W >= 2*d + 2.

  • num_samples (int) – retained sweeps; each sweep contributes all W walker states.

  • burn_in (int) – sweeps to discard. thin: keep one sweep in thin.

  • a (float) – stretch scale (>1; 2.0 is the standard default).

  • rng (RandomState | None) – optional RandomState.

  • thin (int)

Returns:

MCMCResult whose samples are the pooled walker states (sweep-major), so its diagnostics see W * num_samples / thin draws.

Return type:

MCMCResult

metropolis_within_gibbs(log_target, initial, proposals, num_samples, burn_in=0, thin=1, rng=None)[source]

Cycle labelled proposal kernels and accept/reject each against one target.

This is useful for record/dict states where each proposal updates a field or a small block while the full joint log target still owns all model math. Retained samples are recorded after complete sweeps through all proposals.

Parameters:
Return type:

MCMCResult

hamiltonian_monte_carlo(log_target, grad_log_target, initial, num_samples, step_size, num_steps, mass=1.0, burn_in=0, thin=1, rng=None)[source]

Run Hamiltonian Monte Carlo for scalar/vector numeric states.

log_target may be unnormalized. grad_log_target must return the gradient of that log target with respect to the numeric state. Both callables stay user/model-owned; this utility only owns the transition mechanics.

Parameters:
Return type:

MCMCResult

reflective_hmc(log_target, grad_log_target, initial, lower, upper, num_samples, step_size, num_steps, mass=1.0, burn_in=0, thin=1, rng=None)[source]

Hamiltonian Monte Carlo on a box [lower, upper] by reflecting trajectories off the walls.

Samples a target constrained to a hyper-rectangle without distorting it: the leapfrog trajectory bounces specularly off each boundary (mirror the position, flip that momentum component), which is volume-preserving and time-reversible, so the usual Metropolis correction leaves the box-restricted target invariant. log_target (may be unnormalized) and grad_log_target are defined on the box; lower/upper broadcast to the state shape. Returns an MCMCResult whose samples all lie in the box – the constrained-HMC answer for box/simplex-style bounds (WS-1 constraints).

Parameters:
Return type:

MCMCResult

particle_filter(observations, propagate, log_likelihood, initial_particles, *, resample=True, rng=None)[source]

Bootstrap particle filter (sequential Monte Carlo) for a general state-space model.

Propagates a cloud of weighted particles through a user-supplied state-space model and conditions on each observation in turn – the nonlinear/non-Gaussian generalization of the Kalman filter (and a member of the SMC family). initial_particles is an (N, d) array from the prior; propagate( particles, rng) returns the particles advanced one step through the transition (including process noise); log_likelihood(particles, y) returns the per-particle observation log-density. Each step reweights by the likelihood, records the weighted-mean filtered state, and (by default) multinomially resamples to fight weight degeneracy. Returns (filtered_means, log_likelihood) where the second value is the SMC estimate of the model’s marginal log-likelihood log p(y_1:T) (an unbiased evidence estimate, usable for parameter inference). For a linear-Gaussian model it converges to the exact Kalman filter as N -> infinity.

Parameters:
Return type:

tuple[ndarray, float]

dense_mass_hmc(log_target, grad_log_target, initial, num_samples, step_size, num_steps, warmup=500, rng=None)[source]

HMC with a warmup-adapted dense mass matrix (the Stan-style Euclidean metric).

Runs a warmup phase with an identity metric, estimates the posterior covariance Sigma from the warmup draws, and runs the sampling phase with mass matrix M = Sigma^{-1} (momentum ~ N(0, M), position step M^{-1} p). Matching the metric to the posterior shape decorrelates the parameters, so strongly-correlated/ill-conditioned targets mix far better than identity-mass HMC – the standard fix that makes HMC competitive on realistic posteriors. log_target may be unnormalized; grad_log_target returns its gradient. Returns an MCMCResult from the sampling phase.

Parameters:
Return type:

MCMCResult

nuts(log_target=None, grad_log_target=None, initial=None, num_samples=0, warmup=1000, mass=1.0, target_accept=0.8, max_tree_depth=10, thin=1, rng=None, *, value_and_grad=None, adapt_mass=False)[source]

No-U-Turn Sampler (Hoffman & Gelman 2014, efficient NUTS with dual-averaging step size).

Auto-tunes the leapfrog trajectory length (recursive tree doubling, U-turn termination) and, during warmup, the step size to hit target_accept — so unlike fixed-step HMC it needs no manual tuning and mixes well on correlated / higher-dimensional posteriors. mass is a diagonal mass matrix.

Two equivalent target interfaces (back-compatible):

  • nuts(log_target, grad_log_target, initial, ...) — separate value/gradient callables (the historical signature). grad_log_target returns the gradient of the (unnormalized) log target.

  • nuts(value_and_grad=fn, initial=..., ...) — a fused callable returning (logp, grad) in one shot. This halves forward passes (the value NUTS already needs for the slice/Metropolis criterion is shared with the gradient) and lets the sampler cache (logp, grad) at every trajectory endpoint so shared leapfrog/tree nodes are never re-evaluated — typically ~2-3x fewer target evaluations than the split path.

Parameters:
Return type:

MCMCResult

sample_distribution(dist, initial, proposal, num_samples, burn_in=0, thin=1, rng=None, evidence=None)[source]

Sample from a distribution’s log-density, optionally with evidence.

Parameters:
Return type:

MCMCResult

posterior_predictive(samples, sampler, rng=None, size=None)[source]

Draw posterior predictive samples from retained MCMC states.

sampler is called as sampler(state, rng) or sampler(state, rng, size). It can build a mixle distribution, evaluate arbitrary simulation code, or call a model-specific predictive function; the MCMC utility only handles iteration and RNG plumbing.

Parameters:
Return type:

list[Any]

gelman_rubin(chains)[source]

Gelman-Rubin potential scale reduction factor (R-hat) across independent chains.

R-hat compares the variance between chains to the variance within chains for each parameter. Values near 1.0 indicate the chains have mixed and are sampling a common target; values noticeably above 1.0 (a common threshold is 1.01-1.1) flag non-convergence – chains stuck in different regions, too short a run, or poor mixing. This is the standard multi-chain convergence check (Gelman & Rubin 1992) and the multi-chain complement to MCMCResult.effective_sample_size().

Parameters:

chains (Sequence[Any]) – Two or more chains, each an MCMCResult or an array-like of states. Chains may differ in length; all are truncated to the shortest common length.

Returns:

A float for scalar parameters, otherwise an array of R-hat values shaped like a single sampled state (one R-hat per parameter dimension).

Return type:

Any

run_chains(sampler, num_chains, initials, rng=None, **sampler_kwargs)[source]

Run several independent chains and report their Gelman-Rubin R-hat.

Each chain is given its own initial state (from initials) and its own RNG seeded deterministically from rng for reproducibility. The chains are otherwise independent, so this is the multi-chain convergence harness: overdisperse the initials, run, and check the returned R-hat is near 1.0 before trusting the pooled samples.

Parameters:

  • sampler (Callable[[...], MCMCResult]) – Callable invoked as sampler(initial=..., rng=..., **sampler_kwargs) and returning an MCMCResult (e.g. metropolis_hastings(), nuts()).

  • num_chains (int) – Number of independent chains to run (>= 2 for a meaningful R-hat).

  • initials (Sequence[Any] | Callable[[RandomState], Any]) – Either a sequence of per-chain initial states (length num_chains) or a callable initials(rng) -> state that draws an overdispersed start per chain.

  • rng (RandomState | None) – Optional RandomState used to seed the per-chain RNGs.

  • **sampler_kwargs (Any) – Forwarded to sampler (e.g. proposal, num_samples).

Returns:

(results, rhat) – the list of per-chain results and their R-hat.

Return type:

tuple[list[MCMCResult], Any]