1use burn::tensor::backend::Backend;
34use rand::Rng;
35
36use crate::fitness::FitnessFn;
37use crate::local_search::{BudgetedEval, LocalSearch, clamp_vec, sanitize_fitness};
38use rlevo_core::bounds::Bounds;
39
40#[derive(Debug, Clone)]
42pub struct NelderMeadParams {
43 pub bounds: Bounds,
46 pub max_iters: usize,
50 pub alpha: f32,
52 pub gamma: f32,
54 pub rho: f32,
56 pub sigma: f32,
58 pub initial_step: f32,
63 pub tolerance: f32,
67}
68
69impl NelderMeadParams {
70 #[must_use]
76 pub fn default_for(bounds: Bounds) -> Self {
77 let (lo, hi): (f32, f32) = bounds.into();
78 Self {
79 bounds,
80 max_iters: 200,
81 alpha: 1.0,
82 gamma: 2.0,
83 rho: 0.5,
84 sigma: 0.5,
85 initial_step: 0.05 * (hi - lo),
86 tolerance: 1e-8,
87 }
88 }
89}
90
91#[derive(Debug, Clone, Copy, Default)]
129pub struct NelderMead;
130
131struct Vertex {
133 point: Vec<f32>,
135 fitness: f32,
137}
138
139impl NelderMead {
140 #[allow(clippy::too_many_lines)]
163 fn refine_impl(
164 params: &NelderMeadParams,
165 genome: Vec<f32>,
166 known: Option<f32>,
167 fitness_fn: &mut dyn FitnessFn<Vec<f32>>,
168 ) -> (Vec<f32>, f32) {
169 assert!(
170 params.max_iters >= 1,
171 "NelderMeadParams::max_iters must be >= 1 (the input genome is \
172 always evaluated once to seed the best-so-far tracker)"
173 );
174 let mut budget: BudgetedEval<'_> = BudgetedEval::new(fitness_fn, params.max_iters);
175
176 let dim: usize = genome.len();
177
178 let (lo, hi): (f32, f32) = params.bounds.into();
182
183 let in_bounds: bool = genome.iter().all(|&x| x >= lo && x <= hi);
186
187 let mut vertex0: Vec<f32> = genome;
192 clamp_vec(&mut vertex0, params.bounds);
193 let mut best: Vec<f32> = vertex0.clone();
194 let initial_fit: f32 = match known {
195 Some(f) if in_bounds => sanitize_fitness(f),
196 _ => {
197 let Some(f) = budget.eval(&vertex0) else {
198 unreachable!("budget of >= 1 cannot be exhausted before the first eval");
199 };
200 f
201 }
202 };
203 let mut best_fit: f32 = initial_fit;
204
205 if dim == 0 {
208 return (best, best_fit);
209 }
210
211 let mut simplex: Vec<Vertex> = Vec::with_capacity(dim + 1);
217 simplex.push(Vertex {
218 point: vertex0.clone(),
219 fitness: initial_fit,
220 });
221
222 for j in 0..dim {
223 let mut point: Vec<f32> = vertex0.clone();
224 let forward: f32 = point[j] + params.initial_step;
225 if forward > hi || forward < lo {
226 point[j] -= params.initial_step;
227 } else {
228 point[j] = forward;
229 }
230 clamp_vec(&mut point, params.bounds);
231
232 let Some(fitness) = budget.eval(&point) else {
233 update_best(&mut best, &mut best_fit, &simplex);
236 return (best, best_fit);
237 };
238 if fitness > best_fit {
239 best_fit = fitness;
240 best.clone_from(&point);
241 }
242 simplex.push(Vertex { point, fitness });
243 }
244
245 let n: usize = simplex.len(); loop {
250 simplex.sort_by(|a, b| b.fitness.total_cmp(&a.fitness));
255
256 let f_best: f32 = simplex[0].fitness;
257 let f_worst: f32 = simplex[n - 1].fitness;
258 let f_second_worst: f32 = simplex[n - 2].fitness;
259
260 if f_best - f_worst < params.tolerance {
263 break;
264 }
265 if budget.remaining() == 0 {
266 break;
267 }
268
269 let centroid: Vec<f32> = centroid_excluding_worst(&simplex, dim);
271 let worst_point: &[f32] = &simplex[n - 1].point;
272
273 let reflected: Vec<f32> = affine(
275 ¢roid,
276 ¢roid,
277 worst_point,
278 params.alpha,
279 params.bounds,
280 );
281 let Some(f_reflected) = eval_clamped(&mut budget, &reflected, &mut best, &mut best_fit)
282 else {
283 break;
284 };
285
286 if f_reflected > f_best {
287 let expanded: Vec<f32> = affine(
289 ¢roid,
290 &reflected,
291 ¢roid,
292 params.gamma,
293 params.bounds,
294 );
295 let Some(f_expanded) =
296 eval_clamped(&mut budget, &expanded, &mut best, &mut best_fit)
297 else {
298 break;
299 };
300 if f_expanded > f_reflected {
301 replace_worst(&mut simplex, expanded, f_expanded);
302 } else {
303 replace_worst(&mut simplex, reflected, f_reflected);
304 }
305 } else if f_reflected > f_second_worst {
306 replace_worst(&mut simplex, reflected, f_reflected);
308 } else {
309 let (target, target_fit): (&[f32], f32) = if f_reflected > f_worst {
313 (&reflected, f_reflected)
314 } else {
315 (worst_point, f_worst)
316 };
317 let contracted: Vec<f32> =
318 affine(¢roid, target, ¢roid, params.rho, params.bounds);
319 let Some(f_contracted) =
320 eval_clamped(&mut budget, &contracted, &mut best, &mut best_fit)
321 else {
322 break;
323 };
324
325 if f_contracted > target_fit {
326 replace_worst(&mut simplex, contracted, f_contracted);
327 } else {
328 let best_point: Vec<f32> = simplex[0].point.clone();
331 for v in simplex.iter_mut().skip(1) {
332 let mut shrunk: Vec<f32> = Vec::with_capacity(dim);
333 for (b, c) in best_point.iter().zip(v.point.iter()) {
334 shrunk.push(b + params.sigma * (c - b));
335 }
336 clamp_vec(&mut shrunk, params.bounds);
337 let Some(f_shrunk) =
338 eval_clamped(&mut budget, &shrunk, &mut best, &mut best_fit)
339 else {
340 return (best, best_fit);
345 };
346 v.point = shrunk;
347 v.fitness = f_shrunk;
348 }
349 }
350 }
351 }
352
353 (best, best_fit)
354 }
355}
356
357impl<B: Backend> LocalSearch<B> for NelderMead {
358 type Params = NelderMeadParams;
359
360 fn refine(
364 &self,
365 params: &NelderMeadParams,
366 genome: Vec<f32>,
367 fitness_fn: &mut dyn FitnessFn<Vec<f32>>,
368 _rng: &mut dyn Rng,
369 ) -> (Vec<f32>, f32) {
370 Self::refine_impl(params, genome, None, fitness_fn)
371 }
372
373 fn refine_with_known_fitness(
385 &self,
386 params: &NelderMeadParams,
387 genome: Vec<f32>,
388 known_fitness: f32,
389 fitness_fn: &mut dyn FitnessFn<Vec<f32>>,
390 _rng: &mut dyn Rng,
391 ) -> (Vec<f32>, f32) {
392 Self::refine_impl(params, genome, Some(known_fitness), fitness_fn)
393 }
394}
395
396fn eval_clamped(
401 budget: &mut BudgetedEval<'_>,
402 point: &Vec<f32>,
403 best: &mut Vec<f32>,
404 best_fit: &mut f32,
405) -> Option<f32> {
406 let fitness: f32 = budget.eval(point)?;
407 if fitness > *best_fit {
408 *best_fit = fitness;
409 best.clone_from(point);
410 }
411 Some(fitness)
412}
413
414fn update_best(best: &mut Vec<f32>, best_fit: &mut f32, simplex: &[Vertex]) {
416 for v in simplex {
417 if v.fitness > *best_fit {
418 *best_fit = v.fitness;
419 best.clone_from(&v.point);
420 }
421 }
422}
423
424fn centroid_excluding_worst(simplex: &[Vertex], dim: usize) -> Vec<f32> {
426 let count: usize = simplex.len() - 1;
427 #[allow(clippy::cast_precision_loss)]
428 let inv: f32 = 1.0 / count as f32;
429 let mut centroid: Vec<f32> = vec![0.0; dim];
430 for v in &simplex[..count] {
431 for (c, &p) in centroid.iter_mut().zip(v.point.iter()) {
432 *c += p;
433 }
434 }
435 for c in &mut centroid {
436 *c *= inv;
437 }
438 centroid
439}
440
441fn affine(base: &[f32], a: &[f32], b: &[f32], coeff: f32, bounds: Bounds) -> Vec<f32> {
446 let mut out: Vec<f32> = Vec::with_capacity(base.len());
447 for k in 0..base.len() {
448 out.push(base[k] + coeff * (a[k] - b[k]));
449 }
450 clamp_vec(&mut out, bounds);
451 out
452}
453
454fn replace_worst(simplex: &mut [Vertex], point: Vec<f32>, fitness: f32) {
456 let last: usize = simplex.len() - 1;
457 simplex[last] = Vertex { point, fitness };
458}
459
460#[cfg(test)]
461mod tests {
462 use super::*;
463 use burn::backend::Flex;
464 use rand::rngs::StdRng;
465 use rand::{RngExt, SeedableRng};
466
467 type TestBackend = Flex;
468
469 struct NegSphere;
472 impl FitnessFn<Vec<f32>> for NegSphere {
473 fn evaluate_one(&mut self, x: &Vec<f32>) -> f32 {
474 -x.iter().map(|v| v * v).sum::<f32>()
475 }
476 }
477
478 struct NegRosenbrock;
480 impl FitnessFn<Vec<f32>> for NegRosenbrock {
481 fn evaluate_one(&mut self, x: &Vec<f32>) -> f32 {
482 let a = 1.0 - x[0];
483 let b = x[1] - x[0] * x[0];
484 -(a * a + 100.0 * b * b)
485 }
486 }
487
488 struct Flat;
490 impl FitnessFn<Vec<f32>> for Flat {
491 fn evaluate_one(&mut self, _x: &Vec<f32>) -> f32 {
492 1.0
493 }
494 }
495
496 struct Counting<'a> {
498 inner: &'a mut dyn FitnessFn<Vec<f32>>,
499 calls: usize,
500 }
501 impl<'a> Counting<'a> {
502 fn new(inner: &'a mut dyn FitnessFn<Vec<f32>>) -> Self {
503 Self { inner, calls: 0 }
504 }
505 }
506 impl FitnessFn<Vec<f32>> for Counting<'_> {
507 fn evaluate_one(&mut self, x: &Vec<f32>) -> f32 {
508 self.calls += 1;
509 self.inner.evaluate_one(x)
510 }
511 }
512
513 const BOUNDS: Bounds = Bounds::new(-5.12, 5.12);
514
515 fn random_start(rng: &mut StdRng, dim: usize, bounds: Bounds) -> Vec<f32> {
516 let (lo, hi): (f32, f32) = bounds.into();
517 (0..dim)
518 .map(|_| lo + (hi - lo) * rng.random::<f32>())
519 .collect()
520 }
521
522 #[test]
523 fn sphere_d2_converges_below_threshold() {
524 let searcher = NelderMead;
525 let params = NelderMeadParams::default_for(BOUNDS);
526 let mut fitness = NegSphere;
527 let mut rng = StdRng::seed_from_u64(1);
528 let start = random_start(&mut rng, 2, BOUNDS);
529 let (_g, fit) =
530 LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
531 assert!(fit > -1e-6, "sphere D=2 should converge: best={fit}");
532 }
533
534 #[test]
535 fn sphere_d10_strictly_improves() {
536 let searcher = NelderMead;
537 let params = NelderMeadParams::default_for(BOUNDS);
538 let mut fitness = NegSphere;
539 let mut rng = StdRng::seed_from_u64(2);
540 let start = random_start(&mut rng, 10, BOUNDS);
541 let start_fit: f32 = -start.iter().map(|v| v * v).sum::<f32>();
542 let (_g, fit) =
543 LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
544 assert!(fit > start_fit, "expected improvement: {fit} > {start_fit}");
545 }
546
547 #[test]
548 fn output_len_equals_input_len() {
549 let searcher = NelderMead;
550 let params = NelderMeadParams::default_for(BOUNDS);
551 let mut fitness = NegSphere;
552 let mut rng = StdRng::seed_from_u64(3);
553 for dim in [1_usize, 2, 5, 10] {
554 let start = random_start(&mut rng, dim, BOUNDS);
555 let (g, _f) = LocalSearch::<TestBackend>::refine(
556 &searcher,
557 ¶ms,
558 start,
559 &mut fitness,
560 &mut rng,
561 );
562 assert_eq!(g.len(), dim);
563 }
564 }
565
566 #[test]
567 fn returned_fitness_matches_fresh_eval() {
568 let searcher = NelderMead;
569 let params = NelderMeadParams::default_for(BOUNDS);
570 let mut fitness = NegSphere;
571 let mut rng = StdRng::seed_from_u64(4);
572 let start = random_start(&mut rng, 4, BOUNDS);
573 let (g, fit) =
574 LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
575 let fresh = fitness.evaluate_one(&g);
576 approx::assert_relative_eq!(fit, fresh, epsilon = 1e-6);
577 }
578
579 #[test]
580 fn rosenbrock_monotone_non_worsening() {
581 let searcher = NelderMead;
582 let params = NelderMeadParams::default_for(BOUNDS);
583 let mut rng = StdRng::seed_from_u64(5);
584 for _ in 0..6 {
585 let start = random_start(&mut rng, 2, BOUNDS);
586 let mut fitness = NegRosenbrock;
587 let start_fit = fitness.evaluate_one(&start);
588 let (_g, fit) = LocalSearch::<TestBackend>::refine(
589 &searcher,
590 ¶ms,
591 start,
592 &mut fitness,
593 &mut rng,
594 );
595 assert!(fit >= start_fit, "monotone: {fit} >= {start_fit}");
596 }
597 }
598
599 #[test]
600 fn eval_count_never_exceeds_budget() {
601 let searcher = NelderMead;
602 let mut params = NelderMeadParams::default_for(BOUNDS);
603 params.max_iters = 37;
604 let mut base = Flat;
605 let mut counting = Counting::new(&mut base);
606 let mut rng = StdRng::seed_from_u64(6);
607 let start = vec![1.0_f32, 2.0, 3.0, 4.0];
608 let (g, _f) = LocalSearch::<TestBackend>::refine(
609 &searcher,
610 ¶ms,
611 start.clone(),
612 &mut counting,
613 &mut rng,
614 );
615 assert!(
616 counting.calls <= params.max_iters,
617 "evals {} must not exceed budget {}",
618 counting.calls,
619 params.max_iters
620 );
621 assert_eq!(g.len(), start.len());
622 }
623
624 #[test]
625 fn degenerate_budget_no_worse_than_input() {
626 let searcher = NelderMead;
629 let mut params = NelderMeadParams::default_for(BOUNDS);
630 params.max_iters = 2;
631 let mut fitness = NegSphere;
632 let mut rng = StdRng::seed_from_u64(7);
633 let start = random_start(&mut rng, 5, BOUNDS);
634 let start_fit: f32 = -start.iter().map(|v| v * v).sum::<f32>();
635 let (g, fit) =
636 LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
637 assert_eq!(g.len(), 5, "dimensionality preserved");
638 assert!(
639 fit >= start_fit,
640 "no worse than input: {fit} >= {start_fit}"
641 );
642 }
643
644 #[test]
645 fn tolerance_early_stops_before_budget() {
646 let searcher = NelderMead;
649 let mut params = NelderMeadParams::default_for(BOUNDS);
650 params.max_iters = 1000;
651 let mut base = NegSphere;
652 let mut counting = Counting::new(&mut base);
653 let mut rng = StdRng::seed_from_u64(8);
654 let start = vec![1.0_f32, -0.5];
655 let (_g, _f) =
656 LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut counting, &mut rng);
657 assert!(
658 counting.calls < params.max_iters,
659 "tolerance should early-stop: evals {} < budget {}",
660 counting.calls,
661 params.max_iters
662 );
663 }
664
665 #[test]
666 fn boundary_start_stays_within_bounds() {
667 let searcher = NelderMead;
668 let params = NelderMeadParams::default_for(BOUNDS);
669 let mut fitness = NegSphere;
670 let mut rng = StdRng::seed_from_u64(9);
671 let start = vec![BOUNDS.hi(); 4];
674 let (g, _f) =
675 LocalSearch::<TestBackend>::refine(&searcher, ¶ms, start, &mut fitness, &mut rng);
676 for &x in &g {
677 assert!(
678 x >= BOUNDS.lo() && x <= BOUNDS.hi(),
679 "coord {x} out of bounds {BOUNDS:?}"
680 );
681 }
682 }
683
684 #[test]
685 #[allow(clippy::float_cmp)]
686 fn same_seed_is_bit_identical() {
687 let searcher = NelderMead;
688 let params = NelderMeadParams::default_for(BOUNDS);
689 let start = vec![2.0_f32, -3.0, 1.5];
690
691 let mut fitness_a = NegSphere;
692 let mut rng_a = StdRng::seed_from_u64(123);
693 let (g_a, f_a) = LocalSearch::<TestBackend>::refine(
694 &searcher,
695 ¶ms,
696 start.clone(),
697 &mut fitness_a,
698 &mut rng_a,
699 );
700
701 let mut fitness_b = NegSphere;
702 let mut rng_b = StdRng::seed_from_u64(123);
703 let (g_b, f_b) = LocalSearch::<TestBackend>::refine(
704 &searcher,
705 ¶ms,
706 start,
707 &mut fitness_b,
708 &mut rng_b,
709 );
710
711 assert_eq!(g_a, g_b);
712 assert_eq!(f_a, f_b);
713 }
714
715 #[test]
716 fn known_fitness_skips_seeding_eval_when_in_bounds() {
717 let searcher = NelderMead;
722 let params = NelderMeadParams::default_for(BOUNDS);
723 let start = vec![1.0_f32, 2.0, 3.0]; let refine_evals = {
726 let mut base = Flat;
727 let mut counting = Counting::new(&mut base);
728 let mut rng = StdRng::seed_from_u64(41);
729 let _ = LocalSearch::<TestBackend>::refine(
730 &searcher,
731 ¶ms,
732 start.clone(),
733 &mut counting,
734 &mut rng,
735 );
736 counting.calls
737 };
738 let hint_evals = {
739 let mut base = Flat;
740 let mut counting = Counting::new(&mut base);
741 let mut rng = StdRng::seed_from_u64(41);
742 let _ = LocalSearch::<TestBackend>::refine_with_known_fitness(
743 &searcher,
744 ¶ms,
745 start.clone(),
746 1.0, &mut counting,
748 &mut rng,
749 );
750 counting.calls
751 };
752 assert_eq!(
753 hint_evals + 1,
754 refine_evals,
755 "in-bounds hint must skip exactly the vertex-0 eval ({hint_evals} vs {refine_evals})"
756 );
757 }
758
759 #[test]
760 fn out_of_bounds_hint_is_ignored() {
761 let searcher = NelderMead;
766 let params = NelderMeadParams::default_for(BOUNDS);
767 let start = vec![100.0_f32, 100.0]; let refine_evals = {
770 let mut base = NegSphere;
771 let mut counting = Counting::new(&mut base);
772 let mut rng = StdRng::seed_from_u64(42);
773 let _ = LocalSearch::<TestBackend>::refine(
774 &searcher,
775 ¶ms,
776 start.clone(),
777 &mut counting,
778 &mut rng,
779 );
780 counting.calls
781 };
782 let (g, fit, hint_evals) = {
783 let mut base = NegSphere;
784 let mut counting = Counting::new(&mut base);
785 let mut rng = StdRng::seed_from_u64(42);
786 let (g, fit) = LocalSearch::<TestBackend>::refine_with_known_fitness(
787 &searcher,
788 ¶ms,
789 start.clone(),
790 999.0, &mut counting,
792 &mut rng,
793 );
794 (g, fit, counting.calls)
795 };
796 assert_eq!(
797 hint_evals, refine_evals,
798 "out-of-bounds hint must fall back to evaluating vertex 0"
799 );
800 let mut fresh_fn = NegSphere;
801 let fresh = fresh_fn.evaluate_one(&g);
802 approx::assert_relative_eq!(fit, fresh, epsilon = 1e-6);
803 assert!(
804 fit <= 0.0,
805 "neg-sphere fitness is non-positive; bogus hint leaked: {fit}"
806 );
807 }
808
809 #[test]
810 fn nan_hint_does_not_propagate() {
811 let searcher = NelderMead;
812 let params = NelderMeadParams::default_for(BOUNDS);
813 let mut fitness = NegSphere;
814 let mut rng = StdRng::seed_from_u64(43);
815 let start = vec![2.0_f32, -1.0]; let (g, fit) = LocalSearch::<TestBackend>::refine_with_known_fitness(
817 &searcher,
818 ¶ms,
819 start,
820 f32::NAN,
821 &mut fitness,
822 &mut rng,
823 );
824 assert!(fit.is_finite(), "NaN hint must be sanitized, got {fit}");
825 let fresh = fitness.evaluate_one(&g);
826 approx::assert_relative_eq!(fit, fresh, epsilon = 1e-6);
827 }
828}