Skip to main content

wafrift_evolution/search/
sim_anneal.rs

1use crate::evolution::crossover::mutation::mutate_with_log;
2use crate::evolution::{Chromosome, GenePool, population::random_chromosome};
3use crate::lineage::Lineage;
4use crate::search::{EvalCandidate, SearchAlgorithm, comparable_fitness, fitness_cmp};
5use crate::types::{Budget, EvolutionError, OracleVerdict, SearchStats};
6use rand::rngs::StdRng;
7use serde::{Deserialize, Serialize};
8
9/// Simulated annealing search.
10///
11/// Uses a temperature schedule to occasionally accept worse candidates,
12/// helping escape local optima.
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct SimulatedAnnealing {
15    current: Chromosome,
16    best: Chromosome,
17    gene_pool: GenePool,
18    generation: u32,
19    eval_counter: u64,
20    temperature: f64,
21    cooling_rate: f64,
22    min_temperature: f64,
23}
24
25impl SimulatedAnnealing {
26    #[must_use]
27    pub fn new() -> Self {
28        Self {
29            current: Chromosome::new(vec![]),
30            best: Chromosome::new(vec![]),
31            gene_pool: GenePool::default_wafrift(),
32            generation: 0,
33            eval_counter: 0,
34            temperature: 1.0,
35            cooling_rate: 0.95,
36            min_temperature: 0.01,
37        }
38    }
39
40    fn neighbor(&self, rng: &mut StdRng) -> Chromosome {
41        let mut child = self.current.clone();
42        let log = mutate_with_log(&mut child, &self.gene_pool, 0.25, rng);
43        child.lineage = Lineage::mutation(&self.current, log, self.generation);
44        child
45    }
46}
47
48impl Default for SimulatedAnnealing {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl SearchAlgorithm for SimulatedAnnealing {
55    fn name(&self) -> &'static str {
56        "simulated_annealing"
57    }
58
59    fn initialize(&mut self, population: Vec<Chromosome>, gene_pool: &GenePool, rng: &mut StdRng) {
60        self.gene_pool = gene_pool.clone();
61        if let Some(best) = population
62            .iter()
63            .max_by(|a, b| fitness_cmp(a.fitness, b.fitness))
64        {
65            self.current = best.clone();
66            self.best = best.clone();
67        } else {
68            self.current = random_chromosome(gene_pool, rng);
69            self.best = self.current.clone();
70        }
71    }
72
73    fn request_evaluations(&mut self, n: usize, rng: &mut StdRng) -> Vec<EvalCandidate> {
74        let mut out = Vec::with_capacity(n);
75        for _ in 0..n {
76            self.eval_counter += 1;
77            out.push(EvalCandidate {
78                id: self.eval_counter,
79                chromosome: self.neighbor(rng),
80            });
81        }
82        out
83    }
84
85    fn submit_evaluations(&mut self, results: Vec<(u64, OracleVerdict)>) {
86        for (_id, verdict) in results {
87            let mut candidate = self.current.clone();
88            candidate.record_verdict(&verdict);
89            let delta = comparable_fitness(candidate.fitness) - comparable_fitness(self.current.fitness);
90            let accepted = if delta > 0.0 {
91                true
92            } else {
93                let p = (delta / self.temperature.max(1e-9)).exp();
94                // Deterministic acceptance using eval_counter as jitter source
95                let threshold = ((self.eval_counter % 1000) as f64) / 1000.0;
96                p > threshold
97            };
98            if accepted {
99                self.current = candidate;
100                if comparable_fitness(self.current.fitness) > comparable_fitness(self.best.fitness) {
101                    self.best = self.current.clone();
102                }
103            }
104        }
105        self.generation += 1;
106        self.temperature = (self.temperature * self.cooling_rate).max(self.min_temperature);
107    }
108
109    fn should_terminate(&self, stats: &SearchStats, budget: &Budget) -> bool {
110        stats.evaluations >= budget.max_requests
111            || stats.generation >= budget.max_generations
112            || stats.stagnation_counter >= budget.stagnation_limit
113            || self.temperature <= self.min_temperature
114    }
115
116    fn best(&self) -> Option<&Chromosome> {
117        Some(&self.best)
118    }
119
120    fn checkpoint(&self) -> Result<Vec<u8>, EvolutionError> {
121        serde_json::to_vec(self).map_err(|e| EvolutionError::SerializationFailed(e.to_string()))
122    }
123
124    fn restore(&mut self, bytes: &[u8]) -> Result<(), EvolutionError> {
125        *self = serde_json::from_slice(bytes)
126            .map_err(|e| EvolutionError::DeserializationFailed(e.to_string()))?;
127        Ok(())
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134    use rand::SeedableRng;
135
136    #[test]
137    fn non_finite_fitness_is_refused_in_acceptance() {
138        let mut alg = SimulatedAnnealing::new();
139        let pool = GenePool::default_wafrift();
140        let mut rng = StdRng::seed_from_u64(9);
141        alg.initialize(vec![Chromosome::new(vec![])], &pool, &mut rng);
142
143        alg.submit_evaluations(vec![(
144            1,
145            OracleVerdict {
146                passed: false,
147                status_delta: 0,
148                body_delta: 0,
149                latency_ms: 0,
150                confidence: f64::NAN,
151                triggered_rules: 1,
152            },
153        )]);
154        let best_after_nan = comparable_fitness(alg.best().expect("best must exist").fitness);
155
156        alg.submit_evaluations(vec![(2, OracleVerdict::from_bool(true))]);
157        let best_after_valid = comparable_fitness(alg.best().expect("best must exist").fitness);
158        assert!(best_after_valid > best_after_nan);
159    }
160}