metalforge 0.1.4

//! Simulated Annealing for arbitrary combinatorial / local-search problems.
//!
//! Unlike the continuous optimizers in this crate, SA does not assume a
//! real-vector search space: it works through the [`Anneal`] trait, which lets a
//! problem define its own state, neighbor moves, and — crucially — an
//! *incremental* energy delta per move. That incremental hook is what lets a
//! caller keep `O(touched)` move evaluation instead of recomputing the whole
//! objective each step (the performance core of the sibling engine **Anvil**,
//! whose spatial-conservation SA this module generalizes).
//!
//! The driver owns the metaheuristic: geometric cooling, an optional
//! sampling-calibrated adaptive schedule, the Metropolis acceptance rule, and
//! independent restarts. It minimizes energy and is deterministic for a seed.
//!
//! The acceptance test short-circuits — `delta < 0 || rng.uniform() < p` — so
//! the RNG is consumed identically to a hand-written annealer that does the
//! same, which keeps byte-for-byte parity when an existing SA is ported onto it.

use crate::rng::Rng;

/// A problem amenable to simulated annealing.
///
/// `State` is the mutable solution (plus any incremental bookkeeping the problem
/// keeps to make [`Anneal::delta`] cheap). `Move` is a proposed local change.
pub trait Anneal {
    /// The annealing state: solution plus any incremental evaluation cache.
    type State: Clone;
    /// A proposed neighbor move.
    type Move;

    /// Builds the starting state. May consume `rng` for a randomized start;
    /// restarts share one RNG, so a deterministic start still diversifies runs.
    fn initial(&self, rng: &mut Rng) -> Self::State;

    /// Total energy of `state` (to **minimize**). Called once per run to seed
    /// the energy; thereafter the driver tracks energy through `delta`.
    fn energy(&self, state: &Self::State) -> f64;

    /// Proposes a random neighbor move without mutating `state`.
    fn propose(&self, state: &Self::State, rng: &mut Rng) -> Self::Move;

    /// Energy change if `mv` were applied to `state` — incremental, no mutation.
    fn delta(&self, state: &Self::State, mv: &Self::Move) -> f64;

    /// Applies `mv` to `state` in place (the move was proposed for this state).
    fn apply(&self, state: &mut Self::State, mv: Self::Move);

    /// Moves a replica makes between swap attempts in parallel tempering (one
    /// "sweep"). Defaults to 1; problems should override with their
    /// neighborhood size (e.g. the number of mutable variables). Unused by the
    /// plain [`Sa`] driver.
    fn sweep_len(&self) -> usize {
        1
    }
}

/// Temperature schedule for the geometric cooling `T_k = t_start · ratio^k`.
#[derive(Debug, Clone, Copy)]
pub enum Schedule {
    /// Fixed start/end temperatures set by the caller.
    Geometric { t_start: f64, t_end: f64 },
    /// Start/end temperatures calibrated by sampling the energy landscape
    /// (Connolly 1990, the "adaptive annealing schedule"). `fallback` is used
    /// when the landscape shows no delta variation (a degenerate problem).
    Adaptive { fallback: (f64, f64) },
}

impl Schedule {
    /// An adaptive schedule with the conventional `(1.0, 1e-3)` fallback.
    pub fn adaptive() -> Self {
        Schedule::Adaptive {
            fallback: (1.0, 1e-3),
        }
    }
}

/// Simulated annealing configuration.
#[derive(Debug, Clone, Copy)]
pub struct Sa {
    /// Iterations (proposed moves) per restart.
    pub iterations: usize,
    /// Number of independent restarts; the best run is kept.
    pub restarts: usize,
    /// Cooling schedule.
    pub schedule: Schedule,
    /// RNG seed; same seed + same problem ⇒ same result.
    pub seed: u64,
}

impl Default for Sa {
    fn default() -> Self {
        Sa {
            iterations: 10_000,
            restarts: 8,
            schedule: Schedule::adaptive(),
            seed: 0,
        }
    }
}

/// Outcome of an annealing run.
#[derive(Debug, Clone)]
pub struct SaResult<S> {
    /// Best state found.
    pub best: S,
    /// Energy of `best`.
    pub energy: f64,
}

impl Sa {
    /// Minimizes the energy of `problem` by simulated annealing.
    pub fn optimize<A: Anneal>(&self, problem: &A) -> SaResult<A::State> {
        let mut rng = Rng::new(self.seed);

        let (t_start, t_end) = match self.schedule {
            Schedule::Geometric { t_start, t_end } => (t_start, t_end),
            Schedule::Adaptive { fallback } => adaptive_schedule(problem, fallback, &mut rng),
        };

        let mut best: Option<(A::State, f64)> = None;
        for _ in 0..self.restarts.max(1) {
            let (state, energy) = one_run(problem, self.iterations, t_start, t_end, &mut rng);
            if best.as_ref().is_none_or(|(_, b)| energy < *b) {
                best = Some((state, energy));
            }
        }
        let (best, energy) = best.expect("at least one restart runs");
        SaResult { best, energy }
    }
}

/// One annealing run with geometric cooling between `t_start` and `t_end`.
fn one_run<A: Anneal>(
    problem: &A,
    iterations: usize,
    t_start: f64,
    t_end: f64,
    rng: &mut Rng,
) -> (A::State, f64) {
    let mut current = problem.initial(rng);
    let mut current_e = problem.energy(&current);
    let mut best = current.clone();
    let mut best_e = current_e;

    if iterations == 0 {
        return (best, best_e);
    }

    let ratio = (t_end / t_start).powf(1.0 / iterations as f64);
    let mut t = t_start;

    for _ in 0..iterations {
        let mv = problem.propose(&current, rng);
        let delta = problem.delta(&current, &mv);
        // Short-circuit: a downhill move never draws from the RNG (parity hook).
        let accept = delta < 0.0 || rng.uniform() < (-delta / t).exp();
        if accept {
            problem.apply(&mut current, mv);
            current_e += delta;
            if current_e < best_e {
                best_e = current_e;
                best = current.clone();
            }
        }
        t *= ratio;
    }

    (best, best_e)
}

/// Calibrates the start/end temperatures by sampling (Connolly 1990).
///
/// Walks randomly (accepting every move) for a fixed number of samples and
/// measures the magnitude of energy deltas. From the acceptance rule
/// `P = exp(−Δ/T)` it solves `T = Δ / (−ln P)` and sets `t_start` from the mean
/// delta (so a typical move is accepted with probability `P_START`) and `t_end`
/// from the minimum delta (so even the smallest move freezes at `P_END`). Using
/// the minimum for `t_end` matters for multiscale objectives, where a mean-based
/// end temperature would leave the fine scale thermal. Falls back to `fallback`
/// if the landscape is flat.
pub(crate) fn adaptive_schedule<A: Anneal>(
    problem: &A,
    fallback: (f64, f64),
    rng: &mut Rng,
) -> (f64, f64) {
    const SAMPLES: usize = 2_000;
    const P_START: f64 = 0.8;
    const P_END: f64 = 0.01;

    let mut current = problem.initial(rng);
    let mut sum_abs = 0.0;
    let mut min_abs = f64::INFINITY;
    let mut count = 0usize;
    for _ in 0..SAMPLES {
        let mv = problem.propose(&current, rng);
        let d = problem.delta(&current, &mv).abs();
        if d > 0.0 {
            sum_abs += d;
            min_abs = min_abs.min(d);
            count += 1;
        }
        problem.apply(&mut current, mv); // random walk: accept everything
    }

    if count == 0 {
        return fallback;
    }
    let mean = sum_abs / count as f64;
    let t_start = mean / (-P_START.ln());
    let t_end = min_abs / (-P_END.ln());
    if t_end > 0.0 && t_end < t_start {
        (t_start, t_end)
    } else {
        fallback
    }
}

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

    /// Number partitioning as an SA problem: split `weights` into two groups
    /// minimizing the squared difference of their sums. Energy 0 ⇔ a perfect
    /// partition exists. `state[i]` puts weight i in group B; flipping moves it.
    struct Partition {
        weights: Vec<f64>,
    }

    impl Partition {
        fn diff(&self, x: &[bool]) -> f64 {
            let mut d = 0.0;
            for (w, &b) in self.weights.iter().zip(x) {
                d += if b { *w } else { -*w };
            }
            d
        }
    }

    impl Anneal for Partition {
        type State = Vec<bool>;
        type Move = usize;

        fn initial(&self, _rng: &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>, rng: &mut Rng) -> usize {
            rng.index(x.len())
        }
        fn delta(&self, x: &Vec<bool>, &i: &usize) -> f64 {
            // Closed-form delta of the squared difference for a single flip.
            let d = self.diff(x);
            let step = if x[i] {
                -2.0 * self.weights[i]
            } 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];
        }
    }

    #[test]
    fn finds_perfect_partition() {
        // {8,7,6,5,4} splits as {8,7} vs {6,5,4} ⇒ difference 0.
        let p = Partition {
            weights: vec![8.0, 7.0, 6.0, 5.0, 4.0],
        };
        let sa = Sa {
            iterations: 4000,
            restarts: 8,
            schedule: Schedule::adaptive(),
            seed: 1,
        };
        let res = sa.optimize(&p);
        assert!(res.energy < 1e-9, "energy {}", res.energy);
        assert!((p.diff(&res.best)).abs() < 1e-9);
    }

    #[test]
    fn incremental_delta_matches_full_energy() {
        // The closed-form delta must equal the difference of full energies.
        let p = Partition {
            weights: vec![3.0, 1.0, 4.0, 1.0, 5.0, 9.0, 2.0],
        };
        let mut rng = Rng::new(42);
        let mut x = p.initial(&mut rng);
        for _ in 0..500 {
            let mv = p.propose(&x, &mut rng);
            let predicted = p.delta(&x, &mv);
            let before = p.energy(&x);
            p.apply(&mut x, mv);
            let after = p.energy(&x);
            assert!((predicted - (after - before)).abs() < 1e-9);
        }
    }

    #[test]
    fn is_deterministic() {
        let p = Partition {
            weights: vec![5.0, 3.0, 9.0, 7.0, 1.0, 8.0],
        };
        let sa = Sa {
            iterations: 2000,
            schedule: Schedule::Geometric {
                t_start: 10.0,
                t_end: 0.01,
            },
            seed: 7,
            restarts: 4,
        };
        let a = sa.optimize(&p);
        let b = sa.optimize(&p);
        assert_eq!(a.best, b.best);
        assert_eq!(a.energy, b.energy);
    }
}