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 =
90                comparable_fitness(candidate.fitness) - comparable_fitness(self.current.fitness);
91            let accepted = if delta > 0.0 {
92                true
93            } else {
94                let p = (delta / self.temperature.max(1e-9)).exp();
95                // Deterministic acceptance using eval_counter as jitter source
96                let threshold = ((self.eval_counter % 1000) as f64) / 1000.0;
97                p > threshold
98            };
99            if accepted {
100                self.current = candidate;
101                if comparable_fitness(self.current.fitness) > comparable_fitness(self.best.fitness)
102                {
103                    self.best = self.current.clone();
104                }
105            }
106        }
107        self.generation += 1;
108        self.temperature = (self.temperature * self.cooling_rate).max(self.min_temperature);
109    }
110
111    fn should_terminate(&self, stats: &SearchStats, budget: &Budget) -> bool {
112        stats.evaluations >= budget.max_requests
113            || stats.generation >= budget.max_generations
114            || stats.stagnation_counter >= budget.stagnation_limit
115            || self.temperature <= self.min_temperature
116    }
117
118    fn best(&self) -> Option<&Chromosome> {
119        Some(&self.best)
120    }
121
122    fn checkpoint(&self) -> Result<Vec<u8>, EvolutionError> {
123        serde_json::to_vec(self).map_err(EvolutionError::SerializationFailed)
124    }
125
126    fn restore(&mut self, bytes: &[u8]) -> Result<(), EvolutionError> {
127        if bytes.len() > crate::types::MAX_CHECKPOINT_BYTES {
128            return Err(EvolutionError::OversizedData {
129                context: "sim_anneal checkpoint restore".into(),
130                size: bytes.len(),
131                max: crate::types::MAX_CHECKPOINT_BYTES,
132            });
133        }
134        *self = serde_json::from_slice(bytes).map_err(EvolutionError::DeserializationFailed)?;
135        Ok(())
136    }
137
138    fn clone_box(&self) -> Box<dyn SearchAlgorithm> {
139        Box::new(self.clone())
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use rand::SeedableRng;
147
148    #[test]
149    fn non_finite_fitness_is_refused_in_acceptance() {
150        let mut alg = SimulatedAnnealing::new();
151        let pool = GenePool::default_wafrift();
152        let mut rng = StdRng::seed_from_u64(9);
153        alg.initialize(vec![Chromosome::new(vec![])], &pool, &mut rng);
154
155        alg.submit_evaluations(vec![(
156            1,
157            OracleVerdict {
158                passed: false,
159                status_delta: 0,
160                body_delta: 0,
161                latency_ms: 0,
162                confidence: f64::NAN,
163                triggered_rules: 1,
164            },
165        )]);
166        let best_after_nan = comparable_fitness(alg.best().expect("best must exist").fitness);
167
168        alg.submit_evaluations(vec![(2, OracleVerdict::from_bool(true))]);
169        let best_after_valid = comparable_fitness(alg.best().expect("best must exist").fitness);
170        assert!(best_after_valid > best_after_nan);
171    }
172}