use burn::tensor::backend::Backend;
use rand::{Rng, RngExt};
use crate::fitness::FitnessFn;
use crate::local_search::{BudgetedEval, LocalSearch, clamp_vec, sanitize_fitness};
use rlevo_core::bounds::Bounds;
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CoolingSchedule {
Geometric {
factor: f32,
},
Linear {
delta: f32,
},
}
#[derive(Debug, Clone)]
pub struct SimulatedAnnealingParams {
bounds: Bounds,
max_iters: usize,
initial_temp: f32,
cooling: CoolingSchedule,
min_temp: f32,
step_size: f32,
}
impl SimulatedAnnealingParams {
#[must_use]
pub fn default_for(bounds: Bounds) -> Self {
let (lo, hi): (f32, f32) = bounds.into();
debug_assert!(
(hi - lo) > 0.0,
"SimulatedAnnealingParams::default_for: zero-width bounds yields step_size 0 (search cannot move)"
);
Self {
bounds,
max_iters: 200,
initial_temp: 1.0,
cooling: CoolingSchedule::Geometric { factor: 0.95 },
min_temp: 1e-6,
step_size: 0.1 * (hi - lo),
}
}
#[must_use]
pub fn with_bounds(mut self, bounds: Bounds) -> Self {
self.bounds = bounds;
self
}
#[must_use]
pub fn with_max_iters(mut self, max_iters: usize) -> Self {
assert!(
max_iters >= 1,
"SimulatedAnnealingParams::with_max_iters: max_iters must be >= 1"
);
self.max_iters = max_iters;
self
}
#[must_use]
pub fn with_initial_temp(mut self, initial_temp: f32) -> Self {
assert!(
initial_temp.is_finite() && initial_temp > 0.0,
"SimulatedAnnealingParams::with_initial_temp: initial_temp must be finite and > 0"
);
self.initial_temp = initial_temp;
self
}
#[must_use]
pub fn with_cooling(mut self, cooling: CoolingSchedule) -> Self {
match cooling {
CoolingSchedule::Geometric { factor } => assert!(
factor.is_finite() && factor > 0.0 && factor < 1.0,
"SimulatedAnnealingParams::with_cooling: geometric factor must be in (0, 1)"
),
CoolingSchedule::Linear { delta } => assert!(
delta.is_finite() && delta > 0.0,
"SimulatedAnnealingParams::with_cooling: linear delta must be finite and > 0"
),
}
self.cooling = cooling;
self
}
#[must_use]
pub fn with_min_temp(mut self, min_temp: f32) -> Self {
assert!(
min_temp.is_finite() && min_temp >= 0.0,
"SimulatedAnnealingParams::with_min_temp: min_temp must be finite and >= 0"
);
self.min_temp = min_temp;
self
}
#[must_use]
pub fn with_step_size(mut self, step_size: f32) -> Self {
assert!(
step_size.is_finite() && step_size > 0.0,
"SimulatedAnnealingParams::with_step_size: step_size must be finite and > 0"
);
self.step_size = step_size;
self
}
#[must_use]
pub fn bounds(&self) -> Bounds {
self.bounds
}
#[must_use]
pub fn max_iters(&self) -> usize {
self.max_iters
}
#[must_use]
pub fn initial_temp(&self) -> f32 {
self.initial_temp
}
#[must_use]
pub fn cooling(&self) -> CoolingSchedule {
self.cooling
}
#[must_use]
pub fn min_temp(&self) -> f32 {
self.min_temp
}
#[must_use]
pub fn step_size(&self) -> f32 {
self.step_size
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct SimulatedAnnealing;
impl SimulatedAnnealing {
fn refine_impl(
params: &SimulatedAnnealingParams,
genome: Vec<f32>,
known: Option<f32>,
fitness_fn: &mut dyn FitnessFn<Vec<f32>>,
rng: &mut dyn Rng,
) -> (Vec<f32>, f32) {
assert!(
params.max_iters >= 1,
"SimulatedAnnealingParams::max_iters must be >= 1 (the input genome \
is always evaluated once to seed the best-so-far tracker)"
);
let mut budget: BudgetedEval = BudgetedEval::new(fitness_fn, params.max_iters);
let initial_fit: f32 = if let Some(f) = known {
sanitize_fitness(f)
} else {
let Some(f) = budget.eval(&genome) else {
unreachable!("budget of >= 1 cannot be exhausted before the first eval");
};
f
};
let mut current: Vec<f32> = genome;
let mut current_fit: f32 = initial_fit;
let mut best: Vec<f32> = current.clone();
let mut best_fit: f32 = current_fit;
let dim: usize = current.len();
if dim == 0 {
return (best, best_fit);
}
let mut temp: f32 = params.initial_temp;
loop {
let mut candidate: Vec<f32> = current.clone();
for x in &mut candidate {
*x += params.step_size * crate::sampling::standard_normal(rng);
}
clamp_vec(&mut candidate, params.bounds);
let Some(cand_fit) = budget.eval(&candidate) else {
break;
};
if cand_fit > best_fit {
best_fit = cand_fit;
best.clone_from(&candidate);
}
let delta: f32 = cand_fit - current_fit;
let accept: bool = delta >= 0.0 || rng.random::<f32>() < (delta / temp).exp();
if accept {
current = candidate;
current_fit = cand_fit;
}
match params.cooling {
CoolingSchedule::Geometric { factor } => temp *= factor,
CoolingSchedule::Linear { delta } => temp = (temp - delta).max(0.0),
}
if temp < params.min_temp {
break;
}
}
(best, best_fit)
}
}
impl<B: Backend> LocalSearch<B> for SimulatedAnnealing {
type Params = SimulatedAnnealingParams;
fn refine(
&self,
params: &SimulatedAnnealingParams,
genome: Vec<f32>,
fitness_fn: &mut dyn FitnessFn<Vec<f32>>,
rng: &mut dyn Rng,
) -> (Vec<f32>, f32) {
Self::refine_impl(params, genome, None, fitness_fn, rng)
}
fn refine_with_known_fitness(
&self,
params: &SimulatedAnnealingParams,
genome: Vec<f32>,
known_fitness: f32,
fitness_fn: &mut dyn FitnessFn<Vec<f32>>,
rng: &mut dyn Rng,
) -> (Vec<f32>, f32) {
Self::refine_impl(params, genome, Some(known_fitness), fitness_fn, rng)
}
}
#[cfg(test)]
mod tests {
use super::*;
use burn::backend::Flex;
use rand::SeedableRng;
use rand::rngs::StdRng;
type TestBackend = Flex;
#[test]
fn with_setters_override_defaults() {
let sa = SimulatedAnnealingParams::default_for(Bounds::new(-2.0, 2.0))
.with_max_iters(50)
.with_initial_temp(3.0)
.with_cooling(CoolingSchedule::Linear { delta: 0.1 })
.with_min_temp(0.01)
.with_step_size(0.5);
assert_eq!(sa.max_iters(), 50);
assert!((sa.initial_temp() - 3.0).abs() < 1e-6);
assert_eq!(sa.cooling(), CoolingSchedule::Linear { delta: 0.1 });
assert!((sa.min_temp() - 0.01).abs() < 1e-6);
assert!((sa.step_size() - 0.5).abs() < 1e-6);
}
#[test]
#[should_panic(expected = "geometric factor must be in (0, 1)")]
fn with_cooling_rejects_out_of_range_geometric_factor() {
let _ = SimulatedAnnealingParams::default_for(Bounds::new(-2.0, 2.0))
.with_cooling(CoolingSchedule::Geometric { factor: 1.5 });
}
struct NegSphere;
impl FitnessFn<Vec<f32>> for NegSphere {
fn evaluate_one(&mut self, x: &Vec<f32>) -> f32 {
-x.iter().map(|v| v * v).sum::<f32>()
}
}
struct NegRosenbrock;
impl FitnessFn<Vec<f32>> for NegRosenbrock {
fn evaluate_one(&mut self, x: &Vec<f32>) -> f32 {
let a = 1.0 - x[0];
let b = x[1] - x[0] * x[0];
-(a * a + 100.0 * b * b)
}
}
struct Flat;
impl FitnessFn<Vec<f32>> for Flat {
fn evaluate_one(&mut self, _x: &Vec<f32>) -> f32 {
1.0
}
}
struct Counting<'a> {
inner: &'a mut dyn FitnessFn<Vec<f32>>,
calls: usize,
}
impl<'a> Counting<'a> {
fn new(inner: &'a mut dyn FitnessFn<Vec<f32>>) -> Self {
Self { inner, calls: 0 }
}
}
impl FitnessFn<Vec<f32>> for Counting<'_> {
fn evaluate_one(&mut self, x: &Vec<f32>) -> f32 {
self.calls += 1;
self.inner.evaluate_one(x)
}
}
struct Recording<'a> {
inner: &'a mut dyn FitnessFn<Vec<f32>>,
fitnesses: Vec<f32>,
}
impl<'a> Recording<'a> {
fn new(inner: &'a mut dyn FitnessFn<Vec<f32>>) -> Self {
Self {
inner,
fitnesses: Vec::new(),
}
}
}
impl FitnessFn<Vec<f32>> for Recording<'_> {
fn evaluate_one(&mut self, x: &Vec<f32>) -> f32 {
let f = self.inner.evaluate_one(x);
self.fitnesses.push(f);
f
}
}
const BOUNDS: Bounds = Bounds::new(-5.12, 5.12);
fn random_start(rng: &mut StdRng, dim: usize, bounds: Bounds) -> Vec<f32> {
let (lo, hi): (f32, f32) = bounds.into();
(0..dim)
.map(|_| lo + (hi - lo) * rng.random::<f32>())
.collect()
}
#[test]
fn sphere_d2_improves_substantially() {
let searcher = SimulatedAnnealing;
let params = SimulatedAnnealingParams::default_for(BOUNDS);
let mut fitness = NegSphere;
let mut rng = StdRng::seed_from_u64(1);
let start = random_start(&mut rng, 2, BOUNDS);
let start_fit: f32 = -start.iter().map(|v| v * v).sum::<f32>();
let (_g, fit) =
LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
assert!(
fit > 0.1 * start_fit,
"sphere D=2 should improve substantially: best={fit}, start={start_fit}"
);
}
#[test]
fn sphere_d10_strictly_improves() {
let searcher = SimulatedAnnealing;
let params = SimulatedAnnealingParams::default_for(BOUNDS);
let mut fitness = NegSphere;
let mut rng = StdRng::seed_from_u64(2);
let start = random_start(&mut rng, 10, BOUNDS);
let start_fit: f32 = -start.iter().map(|v| v * v).sum::<f32>();
let (_g, fit) =
LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
assert!(fit > start_fit, "expected improvement: {fit} > {start_fit}");
}
#[test]
fn output_len_equals_input_len() {
let searcher = SimulatedAnnealing;
let params = SimulatedAnnealingParams::default_for(BOUNDS);
let mut fitness = NegSphere;
let mut rng = StdRng::seed_from_u64(3);
for dim in [1_usize, 2, 5, 10] {
let start = random_start(&mut rng, dim, BOUNDS);
let (g, _f) = LocalSearch::<TestBackend>::refine(
&searcher,
¶ms,
start,
&mut fitness,
&mut rng,
);
assert_eq!(g.len(), dim);
}
}
#[test]
fn returned_fitness_matches_fresh_eval() {
let searcher = SimulatedAnnealing;
let params = SimulatedAnnealingParams::default_for(BOUNDS);
let mut fitness = NegSphere;
let mut rng = StdRng::seed_from_u64(4);
let start = random_start(&mut rng, 4, BOUNDS);
let (g, fit) =
LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
let fresh = fitness.evaluate_one(&g);
approx::assert_relative_eq!(fit, fresh, epsilon = 1e-6);
}
#[test]
fn rosenbrock_monotone_non_worsening() {
let searcher = SimulatedAnnealing;
let params = SimulatedAnnealingParams::default_for(BOUNDS);
let mut rng = StdRng::seed_from_u64(5);
for _ in 0..6 {
let start = random_start(&mut rng, 2, BOUNDS);
let mut fitness = NegRosenbrock;
let start_fit = fitness.evaluate_one(&start);
let (_g, fit) = LocalSearch::<TestBackend>::refine(
&searcher,
¶ms,
start,
&mut fitness,
&mut rng,
);
assert!(fit >= start_fit, "monotone: {fit} >= {start_fit}");
}
}
#[test]
#[allow(clippy::float_cmp)]
fn flat_landscape_terminates_within_budget() {
let searcher = SimulatedAnnealing;
let mut params = SimulatedAnnealingParams::default_for(BOUNDS);
params.max_iters = 37;
let mut base = Flat;
let mut counting = Counting::new(&mut base);
let mut rng = StdRng::seed_from_u64(6);
let start = vec![1.0_f32, 2.0, 3.0];
let (g, fit) = LocalSearch::<TestBackend>::refine(
&searcher,
¶ms,
start.clone(),
&mut counting,
&mut rng,
);
assert!(
counting.calls <= params.max_iters,
"evals {} must not exceed budget {}",
counting.calls,
params.max_iters
);
assert_eq!(g, start);
assert_eq!(fit, 1.0);
}
#[test]
#[allow(clippy::float_cmp)]
fn same_seed_bit_identical_different_seed_differs() {
let searcher = SimulatedAnnealing;
let params = SimulatedAnnealingParams::default_for(BOUNDS);
let start = vec![2.0_f32, -3.0, 1.5];
let mut fitness_a = NegSphere;
let mut rng_a = StdRng::seed_from_u64(123);
let (g_a, f_a) = LocalSearch::<TestBackend>::refine(
&searcher,
¶ms,
start.clone(),
&mut fitness_a,
&mut rng_a,
);
let mut fitness_b = NegSphere;
let mut rng_b = StdRng::seed_from_u64(123);
let (g_b, f_b) = LocalSearch::<TestBackend>::refine(
&searcher,
¶ms,
start.clone(),
&mut fitness_b,
&mut rng_b,
);
assert_eq!(g_a, g_b);
assert_eq!(f_a, f_b);
let mut fitness_c = NegSphere;
let mut rng_c = StdRng::seed_from_u64(999);
let (g_c, _f_c) = LocalSearch::<TestBackend>::refine(
&searcher,
¶ms,
start,
&mut fitness_c,
&mut rng_c,
);
assert_ne!(g_a, g_c);
}
#[test]
fn min_temp_early_stop_below_budget() {
let searcher = SimulatedAnnealing;
let mut params = SimulatedAnnealingParams::default_for(BOUNDS);
params.max_iters = 1000;
params.initial_temp = 1e-3;
params.min_temp = 1e-1;
params.cooling = CoolingSchedule::Geometric { factor: 0.5 };
let mut base = NegSphere;
let mut counting = Counting::new(&mut base);
let mut rng = StdRng::seed_from_u64(7);
let start = vec![1.0_f32, -1.0];
let _ =
LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut counting, &mut rng);
assert!(
counting.calls < params.max_iters,
"min_temp early stop: evals {} should be < budget {}",
counting.calls,
params.max_iters
);
}
#[test]
fn boundary_start_stays_within_bounds() {
let searcher = SimulatedAnnealing;
let mut params = SimulatedAnnealingParams::default_for(BOUNDS);
params.step_size = 4.0;
let mut fitness = NegSphere;
let mut rng = StdRng::seed_from_u64(8);
let start = vec![BOUNDS.hi(); 4];
let (g, _f) =
LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
for &x in &g {
assert!(
x >= BOUNDS.lo() && x <= BOUNDS.hi(),
"coord {x} out of bounds {BOUNDS:?}"
);
}
}
#[test]
fn uphill_moves_accepted_at_high_temperature() {
let searcher = SimulatedAnnealing;
let mut params = SimulatedAnnealingParams::default_for(BOUNDS);
params.max_iters = 200;
params.initial_temp = 1e9;
params.min_temp = 1e-9;
params.step_size = 0.5;
let mut base = NegSphere;
let mut recording = Recording::new(&mut base);
let mut rng = StdRng::seed_from_u64(11);
let start = vec![0.05_f32, -0.05];
let _ =
LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut recording, &mut rng);
let mut running_best = f32::NEG_INFINITY;
let mut worse_than_best = 0_usize;
for &f in &recording.fitnesses {
if f < running_best {
worse_than_best += 1;
}
if f > running_best {
running_best = f;
}
}
assert!(
worse_than_best >= 3,
"expected sustained worsening exploration at high temperature, saw {worse_than_best} \
worse-than-best evaluations"
);
}
#[test]
fn known_fitness_skips_exactly_the_seeding_eval() {
let searcher = SimulatedAnnealing;
let mut params = SimulatedAnnealingParams::default_for(BOUNDS);
params.max_iters = 10_000;
let start = vec![1.0_f32, 2.0, 3.0];
let refine_evals = {
let mut base = Flat;
let mut counting = Counting::new(&mut base);
let mut rng = StdRng::seed_from_u64(31);
let _ = LocalSearch::<TestBackend>::refine(
&searcher,
¶ms,
start.clone(),
&mut counting,
&mut rng,
);
counting.calls
};
let hint_evals = {
let mut base = Flat;
let mut counting = Counting::new(&mut base);
let mut rng = StdRng::seed_from_u64(31);
let _ = LocalSearch::<TestBackend>::refine_with_known_fitness(
&searcher,
¶ms,
start.clone(),
1.0, &mut counting,
&mut rng,
);
counting.calls
};
assert_eq!(
hint_evals + 1,
refine_evals,
"hint path must skip exactly the seeding eval ({hint_evals} vs {refine_evals})"
);
}
#[test]
fn nan_hint_does_not_propagate() {
let searcher = SimulatedAnnealing;
let params = SimulatedAnnealingParams::default_for(BOUNDS);
let mut fitness = NegSphere;
let mut rng = StdRng::seed_from_u64(32);
let start = vec![2.0_f32, -1.0];
let (g, fit) = LocalSearch::<TestBackend>::refine_with_known_fitness(
&searcher,
¶ms,
start,
f32::NAN,
&mut fitness,
&mut rng,
);
assert!(fit.is_finite(), "NaN hint must be sanitized, got {fit}");
let fresh = fitness.evaluate_one(&g);
approx::assert_relative_eq!(fit, fresh, epsilon = 1e-6);
}
}