1use std::f32::consts::PI;
38use std::marker::PhantomData;
39
40use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
41use rand::Rng;
42use rand::RngExt;
43
44use rlevo_core::bounds::Bounds;
45use rlevo_core::config::{self, ConfigError, Validate};
46
47use super::len_matches_pop;
48use crate::ops::selection::argmax_host;
49use crate::rng::{SeedPurpose, seed_stream};
50use crate::strategy::{Strategy, StrategyMetrics};
51
52const MAX_SPIRAL_EXP: f32 = 80.0; fn spiral_factor(l: f32, b: f32) -> f32 {
69 (b * l).clamp(-MAX_SPIRAL_EXP, MAX_SPIRAL_EXP).exp() * (2.0 * PI * l).cos()
70}
71
72#[derive(Debug, Clone)]
74pub struct WoaConfig {
75 pub pop_size: usize,
77 pub genome_dim: usize,
79 pub bounds: Bounds,
81 pub max_generations: usize,
83 pub b: f32,
85}
86
87impl WoaConfig {
88 #[must_use]
90 pub fn default_for(pop_size: usize, genome_dim: usize) -> Self {
91 Self {
92 pop_size,
93 genome_dim,
94 bounds: Bounds::new(-5.12, 5.12),
95 max_generations: 500,
96 b: 1.0,
97 }
98 }
99}
100
101impl Validate for WoaConfig {
102 fn validate(&self) -> Result<(), ConfigError> {
103 const C: &str = "WoaConfig";
104 config::at_least(C, "pop_size", self.pop_size, 1)?;
105 config::nonzero(C, "genome_dim", self.genome_dim)?;
106 config::at_least(C, "max_generations", self.max_generations, 1)?;
107 config::positive(C, "b", f64::from(self.b))?;
108 Ok(())
109 }
110}
111
112#[derive(Debug, Clone)]
114pub struct WoaState<B: Backend> {
115 positions: Tensor<B, 2>,
117 fitness: Vec<f32>,
119 best_genome: Option<Tensor<B, 2>>,
121 best_fitness: f32,
123 generation: usize,
125}
126
127impl<B: Backend> WoaState<B> {
128 pub fn try_new(
135 positions: Tensor<B, 2>,
136 fitness: Vec<f32>,
137 best_genome: Option<Tensor<B, 2>>,
138 best_fitness: f32,
139 generation: usize,
140 ) -> Result<Self, ConfigError> {
141 let pop = positions.dims()[0];
142 config::nonzero("WoaState", "pop_size", pop)?;
143 len_matches_pop("WoaState", "fitness", pop, fitness.len())?;
144 Ok(Self {
145 positions,
146 fitness,
147 best_genome,
148 best_fitness,
149 generation,
150 })
151 }
152
153 #[must_use]
155 pub fn positions(&self) -> &Tensor<B, 2> {
156 &self.positions
157 }
158
159 #[must_use]
161 pub fn fitness(&self) -> &[f32] {
162 &self.fitness
163 }
164
165 #[must_use]
167 pub fn best_genome(&self) -> Option<&Tensor<B, 2>> {
168 self.best_genome.as_ref()
169 }
170
171 #[must_use]
173 pub fn best_fitness(&self) -> f32 {
174 self.best_fitness
175 }
176
177 #[must_use]
179 pub fn generation(&self) -> usize {
180 self.generation
181 }
182}
183
184#[derive(Debug, Clone, Copy, Default)]
207pub struct WhaleOptimization<B: Backend> {
208 _backend: PhantomData<fn() -> B>,
209}
210
211impl<B: Backend> WhaleOptimization<B> {
212 #[must_use]
214 pub fn new() -> Self {
215 Self {
216 _backend: PhantomData,
217 }
218 }
219}
220
221impl<B: Backend> Strategy<B> for WhaleOptimization<B>
222where
223 B::Device: Clone,
224{
225 type Params = WoaConfig;
226 type State = WoaState<B>;
227 type Genome = Tensor<B, 2>;
228
229 fn init(
233 &self,
234 params: &WoaConfig,
235 rng: &mut dyn Rng,
236 device: &<B as burn::tensor::backend::BackendTypes>::Device,
237 ) -> WoaState<B> {
238 debug_assert!(
239 params.validate().is_ok(),
240 "invalid WoaConfig reached init: {params:?}"
241 );
242 let (lo, hi): (f32, f32) = params.bounds.into();
243 let pop = params.pop_size;
248 let genome_dim = params.genome_dim;
249 let mut stream = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
250 let mut position_rows = Vec::with_capacity(pop * genome_dim);
251 for _ in 0..pop * genome_dim {
252 position_rows.push(lo + (hi - lo) * stream.random::<f32>());
253 }
254 let positions =
255 Tensor::<B, 2>::from_data(TensorData::new(position_rows, [pop, genome_dim]), device);
256 WoaState {
257 positions,
258 fitness: Vec::new(),
259 best_genome: None,
260 best_fitness: f32::NEG_INFINITY,
261 generation: 0,
262 }
263 }
264
265 #[allow(clippy::many_single_char_names)]
279 fn ask(
280 &self,
281 params: &WoaConfig,
282 state: &WoaState<B>,
283 rng: &mut dyn Rng,
284 device: &<B as burn::tensor::backend::BackendTypes>::Device,
285 ) -> (Tensor<B, 2>, WoaState<B>) {
286 if state.fitness.is_empty() {
288 return (state.positions.clone(), state.clone());
289 }
290
291 let pop_size = params.pop_size;
292 let genome_dim = params.genome_dim;
293
294 #[allow(clippy::cast_precision_loss)]
296 let t = state.generation as f32;
297 #[allow(clippy::cast_precision_loss)]
298 let max_t = params.max_generations.max(1) as f32;
299 let a = 2.0 * (1.0 - (t / max_t).min(1.0));
300
301 let mut stream = seed_stream(rng.next_u64(), state.generation as u64, SeedPurpose::Other);
304 let mut rand_idx: Vec<i64> = Vec::with_capacity(pop_size);
305 let mut a_scalar: Vec<f32> = Vec::with_capacity(pop_size);
306 let mut c_scalar: Vec<f32> = Vec::with_capacity(pop_size);
307 let mut p_scalar: Vec<f32> = Vec::with_capacity(pop_size);
308 let mut l_scalar: Vec<f32> = Vec::with_capacity(pop_size);
309 let mut abs_a_lt_one: Vec<i64> = Vec::with_capacity(pop_size);
310 let mut p_lt_half: Vec<i64> = Vec::with_capacity(pop_size);
311 for i in 0..pop_size {
312 let r_a: f32 = stream.random::<f32>();
313 let r_c: f32 = stream.random::<f32>();
314 let p: f32 = stream.random::<f32>();
315 let l: f32 = 2.0 * stream.random::<f32>() - 1.0;
316 let a_val = 2.0 * a * r_a - a;
317 let c_val = 2.0 * r_c;
318 a_scalar.push(a_val);
319 c_scalar.push(c_val);
320 p_scalar.push(p);
321 l_scalar.push(l);
322 abs_a_lt_one.push(i64::from(a_val.abs() < 1.0));
323 p_lt_half.push(i64::from(p < 0.5));
324 let mut r = stream.random_range(0..pop_size);
326 if r == i {
327 r = (r + 1) % pop_size;
328 }
329 #[allow(clippy::cast_possible_wrap)]
330 rand_idx.push(r as i64);
331 }
332
333 let a_row = Tensor::<B, 1>::from_data(TensorData::new(a_scalar, [pop_size]), device)
334 .unsqueeze_dim::<2>(1)
335 .expand([pop_size, genome_dim]);
336 let c_row = Tensor::<B, 1>::from_data(TensorData::new(c_scalar, [pop_size]), device)
337 .unsqueeze_dim::<2>(1)
338 .expand([pop_size, genome_dim]);
339 let rand_idx_t =
340 Tensor::<B, 1, Int>::from_data(TensorData::new(rand_idx, [pop_size]), device);
341 let x_rand = state.positions.clone().select(0, rand_idx_t);
342
343 let x_best = state
344 .best_genome
345 .as_ref()
346 .expect("best_genome populated after the first tell")
347 .clone()
348 .expand([pop_size, genome_dim]);
349
350 let enc_best = x_best.clone()
352 - a_row
353 .clone()
354 .mul((c_row.clone().mul(x_best.clone()) - state.positions.clone()).abs());
355 let enc_rand =
357 x_rand.clone() - a_row.mul((c_row.mul(x_rand) - state.positions.clone()).abs());
358 let dist = (x_best.clone() - state.positions.clone()).abs();
363 let factor_host: Vec<f32> = l_scalar
364 .iter()
365 .map(|&l| spiral_factor(l, params.b))
366 .collect();
367 let factor = Tensor::<B, 1>::from_data(TensorData::new(factor_host, [pop_size]), device);
368 let factor_mat = factor.unsqueeze_dim::<2>(1).expand([pop_size, genome_dim]);
369 let spiral = dist.mul(factor_mat) + x_best;
370
371 let m_abs_a_lt_one =
373 Tensor::<B, 1, Int>::from_data(TensorData::new(abs_a_lt_one, [pop_size]), device)
374 .equal_elem(1)
375 .unsqueeze_dim::<2>(1)
376 .expand([pop_size, genome_dim]);
377 let m_p_lt_half =
378 Tensor::<B, 1, Int>::from_data(TensorData::new(p_lt_half, [pop_size]), device)
379 .equal_elem(1)
380 .unsqueeze_dim::<2>(1)
381 .expand([pop_size, genome_dim]);
382
383 let encircle = enc_rand.mask_where(m_abs_a_lt_one, enc_best);
384 let new_positions = spiral.mask_where(m_p_lt_half, encircle);
385
386 let (lo, hi): (f32, f32) = params.bounds.into();
387 let new_positions = new_positions.clamp(lo, hi);
388
389 let mut next = state.clone();
390 next.positions.clone_from(&new_positions);
391 (new_positions, next)
392 }
393
394 fn tell(
400 &self,
401 _params: &WoaConfig,
402 population: Tensor<B, 2>,
403 fitness: Tensor<B, 1>,
404 mut state: WoaState<B>,
405 _rng: &mut dyn Rng,
406 ) -> (WoaState<B>, StrategyMetrics) {
407 let fitness_host = fitness
408 .into_data()
409 .into_vec::<f32>()
410 .expect("fitness tensor must be readable as f32");
411 state.fitness.clone_from(&fitness_host);
412 state.positions.clone_from(&population);
413 let best_idx = argmax_host(&fitness_host);
414 if fitness_host[best_idx] > state.best_fitness {
415 state.best_fitness = fitness_host[best_idx];
416 let device = population.device();
417 #[allow(clippy::cast_possible_wrap)]
418 let idx = Tensor::<B, 1, Int>::from_data(
419 TensorData::new(vec![best_idx as i64], [1]),
420 &device,
421 );
422 state.best_genome = Some(population.select(0, idx));
423 }
424 state.generation += 1;
425 let m =
426 StrategyMetrics::from_host_fitness(state.generation, &fitness_host, state.best_fitness);
427 state.best_fitness = m.best_fitness_ever();
428 (state, m)
429 }
430
431 fn best(&self, state: &WoaState<B>) -> Option<(Tensor<B, 2>, f32)> {
434 state
435 .best_genome
436 .as_ref()
437 .map(|g| (g.clone(), state.best_fitness))
438 }
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444 use crate::fitness::FromFitnessEvaluable;
445 use crate::strategy::EvolutionaryHarness;
446 use burn::backend::Flex;
447 use rand::SeedableRng;
448 use rand::rngs::StdRng;
449 use rlevo_core::fitness::FitnessEvaluable;
450
451 type TestBackend = Flex;
452
453 #[allow(clippy::trivially_copy_pass_by_ref)] fn finite_fitness(
457 n: usize,
458 device: &<TestBackend as burn::tensor::backend::BackendTypes>::Device,
459 ) -> Tensor<TestBackend, 1> {
460 #[allow(clippy::cast_precision_loss)]
461 let vals: Vec<f32> = (0..n).map(|i| -(i as f32) - 1.0).collect();
462 Tensor::<TestBackend, 1>::from_data(TensorData::new(vals, [n]), device)
463 }
464
465 #[test]
466 fn try_new_checks_fitness_length() {
467 let device = Default::default();
468 let pos = Tensor::<TestBackend, 2>::zeros([3, 2], &device);
469 assert!(WoaState::try_new(pos.clone(), vec![1.0; 3], None, 1.0, 1).is_ok());
470 assert!(WoaState::try_new(pos.clone(), vec![], None, f32::MIN, 0).is_ok());
471 assert!(WoaState::try_new(pos, vec![1.0; 2], None, 1.0, 1).is_err());
472 let empty = Tensor::<TestBackend, 2>::zeros([0, 2], &device);
473 assert!(WoaState::try_new(empty, vec![], None, 1.0, 0).is_err());
474 }
475
476 #[test]
477 fn default_config_validates() {
478 assert!(WoaConfig::default_for(30, 10).validate().is_ok());
479 }
480
481 #[test]
482 fn rejects_nonpositive_b() {
483 let mut cfg = WoaConfig::default_for(30, 10);
484 cfg.b = 0.0;
485 assert_eq!(cfg.validate().unwrap_err().field, "b");
486 }
487
488 struct Sphere;
489 struct SphereFit;
490 impl FitnessEvaluable for SphereFit {
491 type Individual = Vec<f64>;
492 type Landscape = Sphere;
493 fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
494 x.iter().map(|v| v * v).sum()
495 }
496 }
497
498 #[test]
499 fn woa_converges_on_sphere_d10() {
500 let device = Default::default();
505 let strategy = WhaleOptimization::<TestBackend>::new();
506 let params = WoaConfig::default_for(32, 10);
507 let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
508 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
509 strategy, params, fitness_fn, 5, device, 600,
510 )
511 .expect("valid params");
512 harness.reset();
513 while !harness.step(()).done {}
514 let best = harness.latest_metrics().unwrap().best_fitness_ever();
515 assert!(best < 1e-4, "WOA D10 best={best}");
516 }
517
518 #[test]
519 fn spiral_factor_stays_finite_under_overflow() {
520 for &(l, b) in &[
533 (0.75_f32, 200.0_f32),
534 (0.5, 200.0),
535 (0.9, 200.0),
536 (0.75, 500.0),
537 ] {
538 let unguarded: f32 = (b * l).exp() * (2.0 * PI * l).cos();
540 assert!(
541 !unguarded.is_finite(),
542 "test setup: expected overflow for l={l}, b={b}, got {unguarded}"
543 );
544 let guarded: f32 = spiral_factor(l, b);
545 assert!(
546 guarded.is_finite(),
547 "spiral_factor non-finite for l={l}, b={b}: {guarded}"
548 );
549 }
550 }
551
552 #[test]
553 fn spiral_factor_matches_unguarded_when_in_range() {
554 for &(l, b) in &[(0.3_f32, 1.0_f32), (-0.7, 2.0), (0.25, 10.0)] {
557 let expected: f32 = (b * l).exp() * (2.0 * PI * l).cos();
558 let got: f32 = spiral_factor(l, b);
559 approx::assert_relative_eq!(got, expected, epsilon = 1e-6);
560 }
561 }
562
563 #[test]
564 fn best_is_none_until_first_tell() {
565 let device = Default::default();
566 let strategy = WhaleOptimization::<TestBackend>::new();
567 let params = WoaConfig::default_for(4, 3);
568 let mut rng = StdRng::seed_from_u64(0);
569 let state = strategy.init(¶ms, &mut rng, &device);
570 assert!(strategy.best(&state).is_none());
571 let (pop, state) = strategy.ask(¶ms, &state, &mut rng, &device);
572 let fitness = finite_fitness(4, &device);
573 let (state, _m) = strategy.tell(¶ms, pop, fitness, state, &mut rng);
574 assert!(strategy.best(&state).is_some());
575 }
576
577 #[test]
578 fn ask_keeps_positions_in_bounds() {
579 let device = Default::default();
582 let strategy = WhaleOptimization::<TestBackend>::new();
583 let params = WoaConfig::default_for(6, 4);
584 let (lo, hi): (f32, f32) = params.bounds.into();
585 for seed in 0..32 {
586 let mut rng = StdRng::seed_from_u64(seed);
587 let state = strategy.init(¶ms, &mut rng, &device);
588 let (pop1, state) = strategy.ask(¶ms, &state, &mut rng, &device);
589 let fitness = finite_fitness(6, &device);
590 let (state, _m) = strategy.tell(¶ms, pop1, fitness, state, &mut rng);
591 let (pop2, _state) = strategy.ask(¶ms, &state, &mut rng, &device);
592 let values = pop2.into_data().into_vec::<f32>().unwrap();
593 for v in values {
594 assert!(
595 v >= lo - 1e-4 && v <= hi + 1e-4,
596 "seed {seed}: position {v} out of bounds [{lo}, {hi}]"
597 );
598 }
599 }
600 }
601
602 #[test]
603 fn pop_size_two_runs() {
604 let device = Default::default();
607 let strategy = WhaleOptimization::<TestBackend>::new();
608 let params = WoaConfig::default_for(2, 3);
609 let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
610 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
611 strategy, params, fitness_fn, 0, device, 5,
612 )
613 .expect("valid params");
614 harness.reset();
615 while !harness.step(()).done {}
616 assert!(
617 harness
618 .latest_metrics()
619 .unwrap()
620 .best_fitness_ever()
621 .is_finite()
622 );
623 }
624}