mixle.inference.mcmc package

Small MCMC utilities over mixle log-density objects.

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.

On top of that machinery this module also provides a high-level parameter posterior API (sample_parameter_posterior(), sample_conjugate_posterior()). Given a prototype distribution that fixes the model family, a dataset, and a prior over parameters, it samples p(theta | data) proportional to exp(sum_i log p(x_i | theta)) * prior(theta) by running Metropolis-Hastings or Hamiltonian Monte Carlo in an unconstrained reparameterization (log for positive scales, stick-breaking for probability simplices) and mapping the retained samples back to parameter space (or to rebuilt distribution objects).

class AdaptiveCovarianceProposal(initial_covariance=1.0, scale=None, regularization=1.0e-6, adapt_after=2, adapt_during_burn_in_only=True)[source]

Bases: Proposal

Full-covariance Gaussian random walk with burn-in covariance learning.

Parameters:

  • initial_covariance (Any)

  • scale (float | None)

  • regularization (float)

  • adapt_after (int)

  • adapt_during_burn_in_only (bool)

sample(current, rng)[source]

Draw a full-covariance Gaussian random-walk proposal.

Parameters:
Return type:

Any

log_density(proposed, current)[source]

Return the current adaptive Gaussian proposal log density.

Parameters:

  • proposed (Any)

  • current (Any)

Return type:

float

adapt(current, proposed, accepted, step, in_burn_in)[source]

Update the empirical covariance estimate during adaptation.

Parameters:
Return type:

None

class AdaptiveRandomWalkProposal(scale, target_acceptance=0.44, adaptation_rate=0.05, adapt_during_burn_in_only=True, min_scale=1.0e-12, max_scale=1.0e12)[source]

Bases: RandomWalkProposal

Gaussian random walk with Robbins-Monro scale adaptation.

By default adaptation only runs during burn-in, which preserves the stationary post-burn chain used for retained samples.

Parameters:

  • scale (Any)

  • target_acceptance (float)

  • adaptation_rate (float)

  • adapt_during_burn_in_only (bool)

  • min_scale (float)

  • max_scale (float)

adapt(current, proposed, accepted, step, in_burn_in)[source]

Adjust the random-walk scale toward the target acceptance rate.

Parameters:
Return type:

None

class BlockProposal(keys, proposal)[source]

Bases: Proposal

Lift a proposal on one or more mapping fields to a full-record proposal.

Parameters:

  • keys (Any)

  • proposal (Proposal)

sample(current, rng)[source]

Propose a replacement for the configured mapping field block.

Parameters:
Return type:

dict[Any, Any]

log_density(proposed, current)[source]

Return the block proposal density inside the full mapping state.

Parameters:
Return type:

float

adapt(current, proposed, accepted, step, in_burn_in)[source]

Forward adaptation to the wrapped block proposal.

Parameters:
Return type:

None

class IndependentProposal(sampler, log_density=None)[source]

Bases: Proposal

Independence proposal from a sampler plus optional log density.

Parameters:

  • sampler (Callable[[np.random.RandomState], Any])

  • log_density (Callable[[Any], float] | None)

sample(current, rng)[source]

Draw a state independent of current from the supplied sampler.

Parameters:
Return type:

Any

log_density(proposed, current)[source]

Return the proposal log density when one was supplied.

Parameters:

  • proposed (Any)

  • current (Any)

Return type:

float

class LangevinProposal(step_size, grad_log_target)[source]

Bases: Proposal

Metropolis-adjusted Langevin proposal for scalar/vector states.

Parameters:

  • step_size (float)

  • grad_log_target (Callable[[Any], Any])

sample(current, rng)[source]

Draw a MALA proposal centered at the Langevin drift mean.

Parameters:
Return type:

Any

log_density(proposed, current)[source]

Return the asymmetric Langevin proposal log density.

Parameters:

  • proposed (Any)

  • current (Any)

Return type:

float

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]

class MixtureProposal(proposals, weights=None)[source]

Bases: Proposal

Mixture of proposal kernels with exact mixture proposal density.

Parameters:

  • proposals (Sequence[Proposal])

  • weights (Sequence[float] | None)

sample(current, rng)[source]

Draw from one mixture component proposal and remember its index.

Parameters:
Return type:

Any

log_density(proposed, current)[source]

Return the log mixture density of all component proposals.

Parameters:

  • proposed (Any)

  • current (Any)

Return type:

float

adapt(current, proposed, accepted, step, in_burn_in)[source]

Forward adaptation to the component used for the last proposal.

Parameters:
Return type:

None

class ParameterBridge(dim, to_unconstrained, from_unconstrained, log_abs_det_jacobian, build, param_names, initial_theta=None)[source]

Bases: object

Map between a distribution’s parameters and an unconstrained vector.

Parameters:
dim

Length of the unconstrained vector phi.

Type:

int

to_unconstrained

theta -> phi for the prototype’s parameters.

Type:

collections.abc.Callable[[Any], numpy.ndarray]

from_unconstrained

phi -> theta (parameter-space value).

Type:

collections.abc.Callable[[numpy.ndarray], Any]

log_abs_det_jacobian

phi -> log|det dtheta/dphi| so a flat prior in theta-space maps to the right density in phi-space.

Type:

collections.abc.Callable[[numpy.ndarray], float]

build

theta -> distribution rebuilds a model from parameters.

Type:

collections.abc.Callable[[Any], Any]

param_names

Human-readable names for each theta block (diagnostics).

Type:

tuple[str, …]

initial_theta

The prototype’s parameters in theta-space (chain start).

Type:

Any

dim: int
to_unconstrained: Callable[[Any], ndarray]
from_unconstrained: Callable[[ndarray], Any]
log_abs_det_jacobian: Callable[[ndarray], float]
build: Callable[[Any], Any]
param_names: tuple[str, ...]
initial_theta: Any = None
class Proposal[source]

Bases: object

Base proposal protocol for Metropolis-Hastings kernels.

sample(current, rng)[source]

Draw a proposed state given the current state.

Parameters:
Return type:

Any

log_density(proposed, current)[source]

Return log q(proposed | current) for Hastings correction.

Parameters:

  • proposed (Any)

  • current (Any)

Return type:

float

adapt(current, proposed, accepted, step, in_burn_in)[source]

Optional adaptation hook called after each transition.

Parameters:
Return type:

None

class RandomWalkProposal(scale)[source]

Bases: Proposal

Symmetric Gaussian random-walk proposal for scalar/vector states.

Parameters:

scale (Any)

sample(current, rng)[source]

Draw a Gaussian random-walk proposal centered at current.

Parameters:
Return type:

Any

log_density(proposed, current)[source]

Return the Gaussian random-walk transition log density.

Parameters:

  • proposed (Any)

  • current (Any)

Return type:

float

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

build_parameter_bridge(prototype)[source]

Build a ParameterBridge for a prototype distribution.

Supported mixle.stats families:

  • Gaussian (mu, sigma2) -> (mu, log sigma2)

  • Gamma (k, theta) -> (log k, log theta)

  • Exponential beta/lam (positive scalar) -> log

  • Poisson lam -> log lam

  • Bernoulli p -> logit p

  • Beta (a, b) -> (log a, log b)

  • Categorical probability map -> stick-breaking over the simplex

Raises:

NotImplementedError – if the family or a parameter shape is unsupported.

Parameters:

prototype (Any)

Return type:

ParameterBridge

distribution_log_target(dist, evidence=None)[source]

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

Parameters:
Return type:

Callable[[Any], float]

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

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

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

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]

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

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

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

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

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]

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]

sample_conjugate_posterior(dist, data, draws=1000, seed=None, return_distributions=False)[source]

Draw exact posterior parameter samples for a conjugate mixle.stats leaf.

For mixle.stats distributions carrying a closed-form conjugate prior, the posterior over parameters is available analytically. This runs the distribution’s own conjugate estimator over data to obtain the posterior hyperparameters (read back via the fitted model’s get_prior()), then draws iid parameter samples from that posterior. This is an exact alternative to sample_parameter_posterior().

Supported leaves: Gaussian (NormalGamma posterior, samples (mu, sigma2)), Poisson (Gamma posterior, samples lam), Exponential (Gamma posterior over the rate, samples the scale beta), and Bernoulli, Binomial, and Geometric (Beta posterior, samples p). Binomial draws keep the prototype trial count and support shift fixed.

Parameters:

  • dist (Any) – A mixle.stats distribution; if it carries no conjugate prior a non-informative default for the family is attached automatically.

  • data (Any) – Observations for the family.

  • draws (int) – Number of iid posterior samples.

  • seed (int | None) – Seed for the RandomState.

  • return_distributions (bool) – Return rebuilt distributions instead of parameters.

Returns:

MCMCResult with iid samples (all accepted, no autocorrelation).

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

sample_parameter_posterior(prototype_dist, data, prior=None, sampler='mh', steps=2000, burn_in=500, thin=1, seed=None, proposal=None, initial=None, step_size=0.05, num_steps=20, grad_log_target=None, return_distributions=False)[source]

Sample the parameter posterior p(theta | data) of a distribution.

The model family is fixed by prototype_dist; its parameters define the sampled space. The unnormalized log target is the data log-likelihood (rebuilding the distribution per proposed theta and summing seq_log_density) plus a prior log-density and the reparameterization Jacobian. Sampling runs in an unconstrained space (see build_parameter_bridge()) so proposals stay in-domain, and retained samples are mapped back to parameter space.

Parameters:

  • prototype_dist (Any) – A distribution instance fixing the model family.

  • data (Any) – Observations accepted by the family’s encoder.

  • prior (Any) – None (flat), a callable theta -> log p(theta), or a distribution with log_density over the parameter representation.

  • sampler (str) – 'mh' (Metropolis-Hastings), 'hmc' (Hamiltonian Monte Carlo), or 'nuts' (No-U-Turn Sampler, self-tuning). HMC/NUTS use grad_log_target if given, else a finite-difference gradient.

  • steps (int) – Number of retained posterior samples.

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

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

  • seed (int | None) – Seed for the RandomState.

  • proposal (Proposal | None) – Optional MH proposal in the unconstrained space; defaults to a random walk (MH only).

  • initial (Any) – Optional starting theta; defaults to the prototype’s parameters.

  • step_size (float) – HMC leapfrog controls (NUTS self-tunes step size).

  • num_steps (int) – HMC leapfrog controls (NUTS self-tunes step size).

  • grad_log_target (Callable[[Any], Any] | None) – Optional exact gradient of log_target in the unconstrained space (e.g. from mixle.inference.mcmc.gradients.torch_gradient()). Replaces the finite-difference gradient for HMC/NUTS – one backward pass per step instead of O(dim) target evaluations.

  • return_distributions (bool) – If True, samples are rebuilt distribution objects instead of parameter-space values.

Returns:

MCMCResult whose samples are parameter-space values (or rebuilt distributions) and whose diagnostics come from the underlying driver.

Return type:

MCMCResult

torch_available()[source]

Return True if Torch can be imported (autograd gradients are available).

Return type:

bool

torch_gradient(log_target_torch, dtype='float64')[source]

Build an exact grad_log_target from a Torch-valued log-target via autograd.

Parameters:

  • log_target_torch (Callable[[Any], Any]) – Callable mapping a 1-D Torch tensor of parameters to a scalar Torch tensor (the unnormalized log target). Differentiation flows through whatever Torch ops it uses, so the gradient is exact – one backward pass regardless of dimension.

  • dtype (str) – Torch float dtype for the parameter tensor (“float64” by default for MCMC numerics; use “float32” to match a single-precision model).

Returns:

grad(x) -> gradient accepting and returning numpy (a float for scalar x, otherwise an array shaped like x). Suitable as the grad_log_target argument to hamiltonian_monte_carlo(), nuts(), or a Langevin proposal.

Return type:

Callable[[Any], Any]

value_and_torch_gradient(log_target_torch, dtype='float64')[source]

Like torch_gradient() but returns (value, gradient) in one backward pass.

Useful for samplers that need both the log-target and its gradient at the same point (HMC/NUTS leapfrog), avoiding a redundant forward evaluation.

Parameters:
Return type:

Callable[[Any], tuple[float, Any]]

Submodules