mod mcmc_sa

module mcmc_sa

MCMC-style simulated annealing with Gelman-Rubin convergence termination at each epoch.

Classical SA runs a fixed number of inner-loop steps K per temperature epoch. MCMC tradition runs until a between-chain convergence diagnostic crosses a threshold (typically Rhat < 1.1 across M >= 2 chains). The two are mathematically identical at fixed temperature – both implement a Metropolis-Hastings kernel with Boltzmann-Gibbs stationary distribution – but differ in loop control. This module ships the MCMC-style point of the typed algebra.

This is the same fixed-temperature kernel used by parallel tempering; only the epoch stopping rule changes.

Structs and Unions

struct GelmanRubin

Gelman-Rubin convergence diagnostic on a multi-chain trace.

For M chains of length N, this computes the standard chain means, within-chain variance, between-chain variance, pooled variance estimate, and Rhat = sqrt(Var / W).

M >= 2 and N >= 2 are required; otherwise returns f64::INFINITY. We track per-coordinate Rhat and return the maximum across coords; SA practice converges on the worst-mixing coordinate.

Implementations

impl GelmanRubin

Functions

fn compute(traces: &[Vec<Vec<f64>>]) -> f64

Computes max-per-coordinate Rhat from a (n_chains, n_draws, dim) view passed as a flat slice with row-major (chain, draw, dim) layout.

struct MultiChainResult

Aggregated multi-chain history: one History per chain plus the per-epoch Rhat trace.

chain_histories: Vec<History>

One History per chain.

epoch_rhat: Vec<f64>

Final Rhat at each epoch.

epoch_steps: Vec<usize>

Inner-loop step count used at each epoch (varies; the MCMC payoff).

struct MultiChainSampler<S: Sampler<f64>>

MCMC-style multi-chain SA driver.

At each epoch, run k_min steps per chain, compute Rhat, advance when it is below the threshold, otherwise run k_check more steps and recheck until k_max caps the epoch.

When sparse_straggler_only = true, the additional batches step only the chains farthest from the pooled mean (the “stragglers”) rather than all chains. This is the skip-connection design: chains that have already mixed get skipped on subsequent batches, reducing total fevals at the cost of a slightly slower Rhat computation (the converged chains’ history grows more slowly so the pooled mean estimate is noisier on early checks). For Rastrigin 5D this reduces fevals by ~30 percent at unchanged variance reduction.

sampler: S

The per-chain sampler (shared across chains; each chain owns its State and RNG).

n_chains: usize

Number of chains.

k_min: usize

Minimum inner-loop steps per epoch before checking Rhat.

k_check: usize

Step batch size between Rhat checks after k_min.

k_max: usize

Hard cap on inner-loop steps per epoch.

rhat_threshold: f64

Convergence threshold (typical: 1.1).

sparse_straggler_only: bool

Skip-connection mode: in phase 2, step only the straggler_top_k chains farthest from the pooled mean instead of all chains.

straggler_top_k: usize

How many stragglers to step per phase-2 batch when sparse mode is on. Setting to 0 falls back to full multi-chain stepping.

Implementations

impl<S: Sampler<f64>> MultiChainSampler<S>

Functions

fn run(&self, cooling: &dyn crate::cool::Cooling<f64>, n_epochs: usize, seed: u64) -> MultiChainResult

Drives the multi-chain MCMC-SA loop. Each chain’s RNG is seeded from seed.wrapping_add(chain_idx as u64) for reproducibility.

struct MultiChainState

Per-chain state for a multi-chain MCMC-SA run.

chains: Vec<State>

One per-chain State.

seeds: Vec<u64>

Per-chain RNG state, advanced independently.