1use std::marker::PhantomData;
37
38use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
39use rand::Rng;
40use rand::RngExt;
41
42use rlevo_core::bounds::Bounds;
43use rlevo_core::config::{self, ConfigError, ConstraintKind, Validate};
44
45use super::len_matches_pop;
46use crate::ops::selection::{argmax_host, tournament_indices_host};
47use crate::rng::{SeedPurpose, seed_stream};
48use crate::strategy::{Strategy, StrategyMetrics};
49
50#[derive(Debug, Clone)]
52pub struct AbcConfig {
53 pub pop_size: usize,
56 pub genome_dim: usize,
58 pub bounds: Bounds,
60 pub limit: usize,
63 pub tournament_size: usize,
68}
69
70impl AbcConfig {
71 #[must_use]
77 pub fn default_for(pop_size: usize, genome_dim: usize) -> Self {
78 Self {
79 pop_size,
80 genome_dim,
81 bounds: Bounds::new(-5.12, 5.12),
82 limit: (pop_size * genome_dim) / 2,
83 tournament_size: 3,
84 }
85 }
86}
87
88impl Validate for AbcConfig {
89 fn validate(&self) -> Result<(), ConfigError> {
90 const C: &str = "AbcConfig";
91 config::at_least(C, "pop_size", self.pop_size, 2)?;
92 config::nonzero(C, "genome_dim", self.genome_dim)?;
93 config::at_least(C, "limit", self.limit, 1)?;
94 config::at_least(C, "tournament_size", self.tournament_size, 1)?;
95 if self.tournament_size > 2 * self.pop_size {
96 return Err(ConfigError {
97 config: C,
98 field: "tournament_size",
99 kind: ConstraintKind::Custom("tournament_size must not exceed 2 * pop_size"),
100 });
101 }
102 Ok(())
103 }
104}
105
106#[derive(Debug, Clone)]
112pub struct AbcState<B: Backend> {
113 colony: Tensor<B, 2>,
115 fitness: Vec<f32>,
117 trial: Vec<usize>,
119 target_of_candidate: Vec<usize>,
124 best_genome: Option<Tensor<B, 2>>,
126 best_fitness: f32,
128 generation: usize,
130}
131
132impl<B: Backend> AbcState<B> {
133 #[allow(clippy::too_many_arguments)]
142 pub fn try_new(
143 colony: Tensor<B, 2>,
144 fitness: Vec<f32>,
145 trial: Vec<usize>,
146 target_of_candidate: Vec<usize>,
147 best_genome: Option<Tensor<B, 2>>,
148 best_fitness: f32,
149 generation: usize,
150 ) -> Result<Self, ConfigError> {
151 let pop = colony.dims()[0];
152 config::nonzero("AbcState", "pop_size", pop)?;
153 len_matches_pop("AbcState", "fitness", pop, fitness.len())?;
154 len_matches_pop("AbcState", "trial", pop, trial.len())?;
155 if !target_of_candidate.is_empty() && target_of_candidate.len() != 2 * pop {
156 return Err(ConfigError {
157 config: "AbcState",
158 field: "target_of_candidate",
159 kind: ConstraintKind::Custom("length must equal 2 * pop_size"),
160 });
161 }
162 Ok(Self {
163 colony,
164 fitness,
165 trial,
166 target_of_candidate,
167 best_genome,
168 best_fitness,
169 generation,
170 })
171 }
172
173 #[must_use]
175 pub fn colony(&self) -> &Tensor<B, 2> {
176 &self.colony
177 }
178
179 #[must_use]
181 pub fn fitness(&self) -> &[f32] {
182 &self.fitness
183 }
184
185 #[must_use]
187 pub fn trial(&self) -> &[usize] {
188 &self.trial
189 }
190
191 #[must_use]
194 pub fn target_of_candidate(&self) -> &[usize] {
195 &self.target_of_candidate
196 }
197
198 #[must_use]
200 pub fn best_genome(&self) -> Option<&Tensor<B, 2>> {
201 self.best_genome.as_ref()
202 }
203
204 #[must_use]
206 pub fn best_fitness(&self) -> f32 {
207 self.best_fitness
208 }
209
210 #[must_use]
212 pub fn generation(&self) -> usize {
213 self.generation
214 }
215}
216
217#[derive(Debug, Clone, Copy, Default)]
236pub struct ArtificialBeeColony<B: Backend> {
237 _backend: PhantomData<fn() -> B>,
238}
239
240impl<B: Backend> ArtificialBeeColony<B> {
241 #[must_use]
243 pub fn new() -> Self {
244 Self {
245 _backend: PhantomData,
246 }
247 }
248
249 #[allow(clippy::too_many_arguments)]
250 fn build_candidates(
251 targets: &[usize],
252 neighbors: &[usize],
253 dims: &[usize],
254 phi: &[f32],
255 colony: &Tensor<B, 2>,
256 pop_size: usize,
257 genome_dim: usize,
258 device: &<B as burn::tensor::backend::BackendTypes>::Device,
259 ) -> Tensor<B, 2> {
260 #[allow(clippy::cast_possible_wrap)]
262 let target_idx: Vec<i64> = targets.iter().map(|&i| i as i64).collect();
263 let _ = pop_size; let n_cand = targets.len();
265 let target_tensor =
266 Tensor::<B, 1, Int>::from_data(TensorData::new(target_idx, [n_cand]), device);
267 let base = colony.clone().select(0, target_tensor);
268
269 #[allow(clippy::cast_possible_wrap)]
271 let neighbor_idx: Vec<i64> = neighbors.iter().map(|&i| i as i64).collect();
272 let neighbor_tensor =
273 Tensor::<B, 1, Int>::from_data(TensorData::new(neighbor_idx, [n_cand]), device);
274 let neighbor_rows = colony.clone().select(0, neighbor_tensor);
275
276 let mut mask = vec![0i64; n_cand * genome_dim];
278 for (row, &j) in dims.iter().enumerate() {
279 mask[row * genome_dim + j] = 1;
280 }
281 let mask_bool =
282 Tensor::<B, 2, Int>::from_data(TensorData::new(mask, [n_cand, genome_dim]), device)
283 .equal_elem(1);
284
285 let phi_row = Tensor::<B, 1>::from_data(TensorData::new(phi.to_vec(), [n_cand]), device)
287 .unsqueeze_dim::<2>(1)
288 .expand([n_cand, genome_dim]);
289 let delta = phi_row.mul(base.clone() - neighbor_rows);
290 let perturbed = base.clone() + delta;
291 base.mask_where(mask_bool, perturbed)
292 }
293}
294
295impl<B: Backend> Strategy<B> for ArtificialBeeColony<B>
296where
297 B::Device: Clone,
298{
299 type Params = AbcConfig;
300 type State = AbcState<B>;
301 type Genome = Tensor<B, 2>;
302
303 fn init(
304 &self,
305 params: &AbcConfig,
306 rng: &mut dyn Rng,
307 device: &<B as burn::tensor::backend::BackendTypes>::Device,
308 ) -> AbcState<B> {
309 debug_assert!(
310 params.validate().is_ok(),
311 "invalid AbcConfig reached init: {params:?}"
312 );
313 let (lo, hi): (f32, f32) = params.bounds.into();
314 let pop = params.pop_size;
319 let genome_dim = params.genome_dim;
320 let mut stream = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
321 let mut colony_rows = Vec::with_capacity(pop * genome_dim);
322 for _ in 0..pop * genome_dim {
323 colony_rows.push(lo + (hi - lo) * stream.random::<f32>());
324 }
325 let colony =
326 Tensor::<B, 2>::from_data(TensorData::new(colony_rows, [pop, genome_dim]), device);
327 AbcState {
328 colony,
329 fitness: Vec::new(),
330 trial: vec![0; params.pop_size],
331 target_of_candidate: Vec::new(),
332 best_genome: None,
333 best_fitness: f32::NEG_INFINITY,
334 generation: 0,
335 }
336 }
337
338 fn ask(
339 &self,
340 params: &AbcConfig,
341 state: &AbcState<B>,
342 rng: &mut dyn Rng,
343 device: &<B as burn::tensor::backend::BackendTypes>::Device,
344 ) -> (Tensor<B, 2>, AbcState<B>) {
345 if state.fitness.is_empty() {
346 return (state.colony.clone(), state.clone());
347 }
348
349 let pop = params.pop_size;
350 let genome_dim = params.genome_dim;
351 let n_cand = 2 * pop;
352
353 let mut stream = seed_stream(rng.next_u64(), state.generation as u64, SeedPurpose::Other);
354
355 let mut targets = Vec::with_capacity(n_cand);
356 let mut neighbors = Vec::with_capacity(n_cand);
357 let mut dims = Vec::with_capacity(n_cand);
358 let mut phis = Vec::with_capacity(n_cand);
359
360 for i in 0..pop {
362 targets.push(i);
363 }
364 targets.extend(
367 tournament_indices_host(&state.fitness, params.tournament_size, pop, &mut stream)
368 .into_iter()
369 .map(|w| usize::try_from(w).expect("winner index is non-negative")),
370 );
371 for &t in &targets {
373 let mut k = stream.random_range(0..pop);
374 if k == t {
375 k = (k + 1) % pop;
376 }
377 neighbors.push(k);
378 dims.push(stream.random_range(0..genome_dim));
379 let phi = 2.0 * stream.random::<f32>() - 1.0;
380 phis.push(phi);
381 }
382
383 let candidates = Self::build_candidates(
384 &targets,
385 &neighbors,
386 &dims,
387 &phis,
388 &state.colony,
389 pop,
390 genome_dim,
391 device,
392 );
393 let (lo, hi): (f32, f32) = params.bounds.into();
394 let candidates = candidates.clamp(lo, hi);
395
396 let mut next = state.clone();
397 next.target_of_candidate = targets;
398 (candidates, next)
399 }
400
401 #[allow(clippy::too_many_lines)]
402 fn tell(
403 &self,
404 params: &AbcConfig,
405 candidates: Tensor<B, 2>,
406 fitness: Tensor<B, 1>,
407 mut state: AbcState<B>,
408 rng: &mut dyn Rng,
409 ) -> (AbcState<B>, StrategyMetrics) {
410 let fitness_host = fitness
411 .into_data()
412 .into_vec::<f32>()
413 .expect("fitness tensor must be readable as f32");
414 let device = candidates.device();
415 let pop = params.pop_size;
416 let genome_dim = params.genome_dim;
417
418 if state.fitness.is_empty() {
420 state.fitness.clone_from(&fitness_host);
421 let best_idx = argmax_host(&fitness_host);
422 state.best_fitness = fitness_host[best_idx];
423 #[allow(clippy::cast_possible_wrap)]
424 let idx = Tensor::<B, 1, Int>::from_data(
425 TensorData::new(vec![best_idx as i64], [1]),
426 &device,
427 );
428 state.best_genome = Some(candidates.clone().select(0, idx));
429 state.colony = candidates;
430 state.generation += 1;
431 let m = StrategyMetrics::from_host_fitness(
432 state.generation,
433 &fitness_host,
434 state.best_fitness,
435 );
436 state.best_fitness = m.best_fitness_ever();
437 return (state, m);
438 }
439
440 let mut best_per_target: Vec<Option<(usize, f32)>> = vec![None; pop];
443 for (cand_idx, &t) in state.target_of_candidate.iter().enumerate() {
444 let cand_fit = fitness_host[cand_idx];
445 if cand_fit >= state.fitness[t] {
446 match best_per_target[t] {
447 None => best_per_target[t] = Some((cand_idx, cand_fit)),
448 Some((_, prev)) if cand_fit > prev => {
449 best_per_target[t] = Some((cand_idx, cand_fit));
450 }
451 _ => {}
452 }
453 }
454 }
455
456 let stacked = Tensor::cat(vec![state.colony.clone(), candidates.clone()], 0);
461 #[allow(clippy::cast_possible_wrap)]
462 let mut rs: Vec<i64> = (0..pop).map(|i| i as i64).collect();
463 let mut new_fitness = state.fitness.clone();
464 for t in 0..pop {
465 match best_per_target[t] {
466 Some((cand_idx, cand_fit)) => {
467 #[allow(clippy::cast_possible_wrap)]
468 {
469 rs[t] = (pop + cand_idx) as i64;
470 }
471 new_fitness[t] = cand_fit;
472 state.trial[t] = 0;
473 }
474 None => {
475 state.trial[t] += 1;
476 }
477 }
478 }
479 let idx = Tensor::<B, 1, Int>::from_data(TensorData::new(rs, [pop]), &device);
480 state.colony = stacked.select(0, idx);
481 state.fitness = new_fitness;
482
483 let mut scouts: Vec<usize> = Vec::new();
485 for (i, trial) in state.trial.iter_mut().enumerate() {
486 if *trial > params.limit {
487 scouts.push(i);
488 *trial = 0;
489 }
490 }
491 if !scouts.is_empty() {
492 let (lo, hi): (f32, f32) = params.bounds.into();
493 let mut scout_stream = seed_stream(
497 rng.next_u64(),
498 state.generation as u64,
499 SeedPurpose::Replacement,
500 );
501 let mut fresh_rows = Vec::with_capacity(scouts.len() * genome_dim);
502 for _ in 0..scouts.len() * genome_dim {
503 fresh_rows.push(lo + (hi - lo) * scout_stream.random::<f32>());
504 }
505 let fresh = Tensor::<B, 2>::from_data(
506 TensorData::new(fresh_rows, [scouts.len(), genome_dim]),
507 &device,
508 );
509 #[allow(clippy::cast_possible_wrap)]
511 let mut rs2: Vec<i64> = (0..pop).map(|i| i as i64).collect();
512 for (k, &scout) in scouts.iter().enumerate() {
513 #[allow(clippy::cast_possible_wrap)]
514 {
515 rs2[scout] = (pop + k) as i64;
516 }
517 state.fitness[scout] = f32::NEG_INFINITY;
521 }
522 let stacked2 = Tensor::cat(vec![state.colony.clone(), fresh], 0);
523 let idx2 = Tensor::<B, 1, Int>::from_data(TensorData::new(rs2, [pop]), &device);
524 state.colony = stacked2.select(0, idx2);
525 }
526
527 let best_idx = argmax_host(&state.fitness);
530 if state.fitness[best_idx].is_finite() && state.fitness[best_idx] > state.best_fitness {
531 state.best_fitness = state.fitness[best_idx];
532 #[allow(clippy::cast_possible_wrap)]
533 let idx = Tensor::<B, 1, Int>::from_data(
534 TensorData::new(vec![best_idx as i64], [1]),
535 &device,
536 );
537 state.best_genome = Some(state.colony.clone().select(0, idx));
538 }
539
540 state.generation += 1;
541 let m =
542 StrategyMetrics::from_host_fitness(state.generation, &fitness_host, state.best_fitness);
543 state.best_fitness = m.best_fitness_ever();
544 (state, m)
545 }
546
547 fn best(&self, state: &AbcState<B>) -> Option<(Tensor<B, 2>, f32)> {
548 state
549 .best_genome
550 .as_ref()
551 .map(|g| (g.clone(), state.best_fitness))
552 }
553}
554
555#[cfg(test)]
556mod tests {
557 use super::*;
558 use crate::fitness::FromFitnessEvaluable;
559 use crate::strategy::EvolutionaryHarness;
560 use burn::backend::Flex;
561 use rand::SeedableRng;
562 use rand::rngs::StdRng;
563 use rlevo_core::fitness::FitnessEvaluable;
564
565 type TestBackend = Flex;
566
567 #[test]
568 fn try_new_checks_cache_lengths() {
569 let device = Default::default();
570 let colony = Tensor::<TestBackend, 2>::zeros([3, 2], &device);
571 assert!(
573 AbcState::try_new(
574 colony.clone(),
575 vec![],
576 vec![0; 3],
577 vec![],
578 None,
579 f32::MIN,
580 0
581 )
582 .is_ok()
583 );
584 assert!(
585 AbcState::try_new(
586 colony.clone(),
587 vec![1.0; 3],
588 vec![0; 3],
589 vec![7; 6],
590 None,
591 1.0,
592 1
593 )
594 .is_ok()
595 );
596 assert!(
598 AbcState::try_new(
599 colony.clone(),
600 vec![1.0; 2],
601 vec![0; 3],
602 vec![],
603 None,
604 1.0,
605 1
606 )
607 .is_err()
608 );
609 assert!(AbcState::try_new(colony, vec![], vec![], vec![0; 5], None, 1.0, 1).is_err());
610 let empty = Tensor::<TestBackend, 2>::zeros([0, 2], &device);
612 assert!(AbcState::try_new(empty, vec![], vec![], vec![], None, 1.0, 0).is_err());
613 }
614
615 #[test]
616 fn default_config_validates() {
617 assert!(AbcConfig::default_for(30, 10).validate().is_ok());
618 }
619
620 #[test]
621 fn rejects_pop_size_below_two() {
622 let mut cfg = AbcConfig::default_for(30, 10);
623 cfg.pop_size = 1;
624 assert_eq!(cfg.validate().unwrap_err().field, "pop_size");
625 }
626
627 struct Sphere;
628 struct SphereFit;
629 impl FitnessEvaluable for SphereFit {
630 type Individual = Vec<f64>;
631 type Landscape = Sphere;
632 fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
633 x.iter().map(|v| v * v).sum()
634 }
635 }
636
637 #[test]
641 fn onlooker_targets_prefer_best_source() {
642 let device = Default::default();
643 let strategy = ArtificialBeeColony::<TestBackend>::new();
644 let mut params = AbcConfig::default_for(16, 4);
645 params.tournament_size = 8;
646 let mut rng = StdRng::seed_from_u64(11);
647 let mut state = strategy.init(¶ms, &mut rng, &device);
648 state.fitness = vec![0.0; 16];
650 state.fitness[3] = 100.0;
651 let (_candidates, next) = strategy.ask(¶ms, &state, &mut rng, &device);
652 let onlooker_hits = next.target_of_candidate[16..]
654 .iter()
655 .filter(|&&t| t == 3)
656 .count();
657 assert!(
660 onlooker_hits >= 3,
661 "onlooker hits on the best bee = {onlooker_hits} (expected ~6)",
662 );
663 }
664
665 #[test]
666 fn abc_converges_on_sphere_d10() {
667 let device = Default::default();
668 let strategy = ArtificialBeeColony::<TestBackend>::new();
669 let params = AbcConfig::default_for(30, 10);
670 let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
671 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
672 strategy, params, fitness_fn, 13, device, 400,
673 )
674 .expect("valid params");
675 harness.reset();
676 while !harness.step(()).done {}
677 let best = harness.latest_metrics().unwrap().best_fitness_ever();
678 assert!(best < 1e-4, "ABC D10 best={best}");
679 }
680
681 struct PartialNanFitness;
685 impl<B: Backend> crate::fitness::BatchFitnessFn<B, Tensor<B, 2>> for PartialNanFitness {
686 fn evaluate_batch(
687 &mut self,
688 population: &Tensor<B, 2>,
689 device: &<B as burn::tensor::backend::BackendTypes>::Device,
690 ) -> Tensor<B, 1> {
691 let n = population.dims()[0];
692 #[allow(clippy::cast_precision_loss)]
693 let mut vals: Vec<f32> = (0..n).map(|i| -(i as f32)).collect();
694 vals[0] = f32::NAN;
695 Tensor::<B, 1>::from_data(TensorData::new(vals, [n]), device)
696 }
697 fn sense(&self) -> rlevo_core::objective::ObjectiveSense {
698 rlevo_core::objective::ObjectiveSense::Maximize
699 }
700 }
701
702 #[test]
706 #[allow(clippy::float_cmp)] fn ask_on_empty_fitness_returns_colony_unchanged() {
708 let device = Default::default();
709 let strategy = ArtificialBeeColony::<TestBackend>::new();
710 let params = AbcConfig::default_for(6, 4);
711 let mut rng = StdRng::seed_from_u64(3);
712 let state = strategy.init(¶ms, &mut rng, &device);
713 let (genome, next) = strategy.ask(¶ms, &state, &mut rng, &device);
714 let before = state
715 .colony()
716 .clone()
717 .into_data()
718 .into_vec::<f32>()
719 .expect("colony readable as f32");
720 let after = genome
721 .into_data()
722 .into_vec::<f32>()
723 .expect("genome readable as f32");
724 assert_eq!(before, after);
725 assert!(next.target_of_candidate().is_empty());
726 }
727
728 #[test]
732 fn pop_size_two_minimal_colony_runs() {
733 let device = Default::default();
734 let strategy = ArtificialBeeColony::<TestBackend>::new();
735 let params = AbcConfig::default_for(2, 3);
736 let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
737 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
738 strategy, params, fitness_fn, 5, device, 8,
739 )
740 .expect("valid params");
741 harness.reset();
742 while !harness.step(()).done {}
743 assert!(
744 harness
745 .latest_metrics()
746 .unwrap()
747 .best_fitness_ever()
748 .is_finite()
749 );
750 }
751
752 #[test]
756 fn genome_dim_one_degenerate_mask_runs() {
757 let device = Default::default();
758 let strategy = ArtificialBeeColony::<TestBackend>::new();
759 let params = AbcConfig::default_for(6, 1);
760 let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
761 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
762 strategy, params, fitness_fn, 8, device, 8,
763 )
764 .expect("valid params");
765 harness.reset();
766 while !harness.step(()).done {}
767 assert!(
768 harness
769 .latest_metrics()
770 .unwrap()
771 .best_fitness_ever()
772 .is_finite()
773 );
774 }
775
776 #[test]
781 #[allow(clippy::float_cmp)] fn scout_reinit_triggers_on_stagnation() {
783 let device = Default::default();
784 let strategy = ArtificialBeeColony::<TestBackend>::new();
785 let mut params = AbcConfig::default_for(4, 2);
786 params.limit = 1;
787 let colony = Tensor::<TestBackend, 2>::zeros([4, 2], &device);
789 let state = AbcState::try_new(
791 colony,
792 vec![0.0; 4], vec![1; 4], vec![0, 1, 2, 3, 0, 1, 2, 3], None,
796 f32::NEG_INFINITY,
797 5,
798 )
799 .expect("valid state");
800 let candidates = Tensor::<TestBackend, 2>::full([8, 2], 3.0, &device);
803 let fit =
804 Tensor::<TestBackend, 1>::from_data(TensorData::new(vec![-18.0_f32; 8], [8]), &device);
805 let mut rng = StdRng::seed_from_u64(9);
806 let (next, _m) = strategy.tell(¶ms, candidates, fit, state, &mut rng);
807 assert!(
809 next.trial().iter().all(|&t| t == 0),
810 "trials not reset: {:?}",
811 next.trial()
812 );
813 let colony_vals = next
815 .colony()
816 .clone()
817 .into_data()
818 .into_vec::<f32>()
819 .expect("colony readable as f32");
820 assert!(
821 colony_vals.iter().any(|&v| v != 0.0),
822 "no scout reinit happened; colony still all-zero"
823 );
824 }
825
826 #[test]
830 fn nan_fitness_survives_harness() {
831 let device = Default::default();
832 let strategy = ArtificialBeeColony::<TestBackend>::new();
833 let params = AbcConfig::default_for(6, 3);
834 let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
835 strategy,
836 params,
837 PartialNanFitness,
838 4,
839 device,
840 4,
841 )
842 .expect("valid params");
843 harness.reset();
844 while !harness.step(()).done {}
845 let m = harness.latest_metrics().unwrap();
846 assert!(
847 m.best_fitness_ever().is_finite(),
848 "best_fitness_ever not finite: {}",
849 m.best_fitness_ever()
850 );
851 assert!(m.broken_count() > 0, "expected a broken (NaN) member");
852 }
853}