metalforge 0.1.4

//! Parallel tempering (replica exchange) for combinatorial problems.
//!
//! Instead of independent restarts, `K` replicas run on a geometric temperature
//! ladder from cold (exploits) to hot (explores). Each replica advances by
//! Metropolis moves at its temperature; periodically, adjacent replicas attempt
//! to **exchange** configurations with probability `min(1, exp((β_i − β_j)(E_i −
//! E_j)))`. Good configurations found at high temperature flow down to the cold
//! replica, which escapes local optima better than annealing with restarts. The
//! whole ladder is itself slowly cooled over the run (*annealed* replica
//! exchange), so the cold replica freezes and refines while the hot ones keep
//! feeding it.
//!
//! Built on the same [`Anneal`] trait as [`Sa`](super::Sa). Determinism under
//! parallelism: each replica carries its own RNG seeded independently of
//! execution order, and the exchange phase is sequential with a separate RNG —
//! so the result is identical whether replicas advance on one thread or many,
//! and whether or not the `rayon` feature is enabled.

use super::sa::{adaptive_schedule, Anneal, Schedule};
use super::SaResult;
use crate::rng::{mix_seed, Rng};

/// Parallel tempering configuration.
#[derive(Debug, Clone, Copy)]
pub struct ParallelTempering {
    /// Number of replicas (clamped to ≥ 2).
    pub replicas: usize,
    /// Total Metropolis moves per replica (split into sweeps of
    /// [`Anneal::sweep_len`]).
    pub iterations: usize,
    /// Temperature schedule; `t_start` is the hot end, `t_end` the cold end.
    pub schedule: Schedule,
    /// RNG seed; same seed + same problem ⇒ same result.
    pub seed: u64,
}

impl Default for ParallelTempering {
    fn default() -> Self {
        ParallelTempering {
            replicas: 8,
            iterations: 100_000,
            schedule: Schedule::adaptive(),
            seed: 0,
        }
    }
}

/// Outcome of a parallel-tempering run.
#[derive(Debug, Clone)]
pub struct TemperingResult<S> {
    /// Best result across all replicas (lowest energy; ties to the colder).
    pub best: SaResult<S>,
    /// Each replica's own best, ordered cold → hot. Useful for a per-replica
    /// post-processing quench that exploits the diversity of basins explored.
    pub replicas: Vec<SaResult<S>>,
}

/// One replica: state at a (time-varying) temperature, with its own RNG.
struct Replica<S> {
    state: S,
    energy: f64,
    temp: f64,
    rng: Rng,
    best_state: S,
    best_energy: f64,
}

impl ParallelTempering {
    /// Runs parallel tempering on `problem`, returning every replica's best plus
    /// the global best.
    pub fn optimize<A>(&self, problem: &A) -> TemperingResult<A::State>
    where
        A: Anneal + Sync,
        A::State: Send,
    {
        let k = self.replicas.max(2);
        // Separate RNG for adaptive sampling and exchanges (deterministic order).
        let mut swap_rng = Rng::new(self.seed);
        let (t_hot, t_cold) = match self.schedule {
            Schedule::Geometric { t_start, t_end } => (t_start, t_end),
            Schedule::Adaptive { fallback } => adaptive_schedule(problem, fallback, &mut swap_rng),
        };

        // Geometric ladder: temps[0] = cold … temps[k-1] = hot.
        let temps: Vec<f64> = (0..k)
            .map(|i| {
                let frac = i as f64 / (k - 1) as f64;
                t_cold * (t_hot / t_cold).powf(frac)
            })
            .collect();

        let mut replicas: Vec<Replica<A::State>> = (0..k)
            .map(|i| {
                // Per-replica seed, well separated and independent of scheduling.
                let mut rng = Rng::new(mix_seed(self.seed, i as u64 + 1));
                let state = problem.initial(&mut rng);
                let energy = problem.energy(&state);
                Replica {
                    best_state: state.clone(),
                    best_energy: energy,
                    state,
                    energy,
                    temp: temps[i],
                    rng,
                }
            })
            .collect();

        let sweep_len = problem.sweep_len();
        if sweep_len > 0 && self.iterations > 0 {
            let n_sweeps = (self.iterations / sweep_len).max(1);
            // Annealed replica exchange: cool the whole ladder by `FLOOR` over
            // the run so the cold replica ends frozen (fine refinement).
            const FLOOR: f64 = 1e-2;
            let cool = FLOOR.powf(1.0 / n_sweeps as f64);
            let mut scale = 1.0;
            for _ in 0..n_sweeps {
                for (r, &base) in replicas.iter_mut().zip(&temps) {
                    r.temp = base * scale;
                }
                advance_all(problem, &mut replicas, sweep_len);
                attempt_swaps(&mut replicas, &mut swap_rng);
                scale *= cool;
            }
        }

        let replica_results: Vec<SaResult<A::State>> = replicas
            .into_iter()
            .map(|r| SaResult {
                best: r.best_state,
                energy: r.best_energy,
            })
            .collect();

        // Global best: lowest energy, ties broken toward the colder replica
        // (lower index) so the choice is order-independent.
        let best_idx = (0..replica_results.len())
            .min_by(|&a, &b| {
                replica_results[a]
                    .energy
                    .partial_cmp(&replica_results[b].energy)
                    .unwrap_or(std::cmp::Ordering::Equal)
                    .then(a.cmp(&b))
            })
            .expect("at least two replicas");
        let best = replica_results[best_idx].clone();

        TemperingResult {
            best,
            replicas: replica_results,
        }
    }
}

/// Advances every replica `sweep_len` moves. Replicas are disjoint and each
/// uses only its own RNG, so the outcome is independent of execution order —
/// parallel (with `rayon`) or sequential give identical results.
#[cfg(feature = "rayon")]
fn advance_all<A>(problem: &A, replicas: &mut [Replica<A::State>], sweep_len: usize)
where
    A: Anneal + Sync,
    A::State: Send,
{
    use rayon::prelude::*;
    replicas
        .par_iter_mut()
        .for_each(|r| advance(problem, r, sweep_len));
}

#[cfg(not(feature = "rayon"))]
fn advance_all<A: Anneal>(problem: &A, replicas: &mut [Replica<A::State>], sweep_len: usize) {
    for r in replicas.iter_mut() {
        advance(problem, r, sweep_len);
    }
}

/// Metropolis moves at the replica's fixed temperature.
fn advance<A: Anneal>(problem: &A, r: &mut Replica<A::State>, n_moves: usize) {
    for _ in 0..n_moves {
        let mv = problem.propose(&r.state, &mut r.rng);
        let delta = problem.delta(&r.state, &mv);
        if delta < 0.0 || r.rng.uniform() < (-delta / r.temp).exp() {
            problem.apply(&mut r.state, mv);
            r.energy += delta;
            if r.energy < r.best_energy {
                r.best_energy = r.energy;
                r.best_state = r.state.clone();
            }
        }
    }
}

/// Sequential exchange sweep over adjacent replica pairs (cold→hot), accepting
/// each swap with probability `min(1, exp((β_cold − β_hot)(E_cold − E_hot)))`.
fn attempt_swaps<S>(replicas: &mut [Replica<S>], swap_rng: &mut Rng) {
    for c in 0..replicas.len() - 1 {
        let beta_cold = 1.0 / replicas[c].temp;
        let beta_hot = 1.0 / replicas[c + 1].temp;
        let arg = (beta_cold - beta_hot) * (replicas[c].energy - replicas[c + 1].energy);
        if arg >= 0.0 || swap_rng.uniform() < arg.exp() {
            let (left, right) = replicas.split_at_mut(c + 1);
            std::mem::swap(&mut left[c].state, &mut right[0].state);
            std::mem::swap(&mut left[c].energy, &mut right[0].energy);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Number partitioning (same as the SA tests): minimize the squared
    /// difference of the two group sums.
    struct Partition {
        weights: Vec<f64>,
    }
    impl Partition {
        fn diff(&self, x: &[bool]) -> f64 {
            self.weights
                .iter()
                .zip(x)
                .map(|(w, &b)| if b { *w } else { -*w })
                .sum()
        }
    }
    impl Anneal for Partition {
        type State = Vec<bool>;
        type Move = usize;
        fn initial(&self, _r: &mut Rng) -> Vec<bool> {
            vec![false; self.weights.len()]
        }
        fn energy(&self, x: &Vec<bool>) -> f64 {
            self.diff(x).powi(2)
        }
        fn propose(&self, x: &Vec<bool>, r: &mut Rng) -> usize {
            r.index(x.len())
        }
        fn delta(&self, x: &Vec<bool>, &i: &usize) -> f64 {
            let d = self.diff(x);
            let step = if x[i] { -2.0 } else { 2.0 } * self.weights[i];
            (d + step).powi(2) - d.powi(2)
        }
        fn apply(&self, x: &mut Vec<bool>, i: usize) {
            x[i] = !x[i];
        }
        fn sweep_len(&self) -> usize {
            self.weights.len()
        }
    }

    #[test]
    fn finds_perfect_partition() {
        let p = Partition {
            weights: vec![8.0, 7.0, 6.0, 5.0, 4.0, 9.0, 3.0, 2.0],
        };
        let pt = ParallelTempering {
            replicas: 6,
            iterations: 8000,
            schedule: Schedule::adaptive(),
            seed: 1,
        };
        let res = pt.optimize(&p);
        assert!(res.best.energy < 1e-9, "energy {}", res.best.energy);
        assert_eq!(res.replicas.len(), 6);
    }

    #[test]
    fn is_deterministic() {
        let p = Partition {
            weights: vec![5.0, 3.0, 9.0, 7.0, 1.0, 8.0, 4.0],
        };
        let pt = ParallelTempering {
            replicas: 5,
            iterations: 5000,
            schedule: Schedule::Geometric {
                t_start: 50.0,
                t_end: 0.01,
            },
            seed: 7,
        };
        let a = pt.optimize(&p);
        let b = pt.optimize(&p);
        assert_eq!(a.best.best, b.best.best);
        assert_eq!(a.best.energy, b.best.energy);
    }
}