1use std::time::Instant;
11
12use rand::Rng;
13use rand_distr::{Bernoulli, Distribution, Normal};
14
15use crate::diagnostics::{EvolutionResult, EvolutionStats, GenerationStats, TimingStats};
16use crate::error::EvolutionError;
17use crate::fitness::traits::{Fitness, FitnessValue};
18use crate::genome::bit_string::BitString;
19use crate::genome::bounds::MultiBounds;
20use crate::genome::real_vector::RealVector;
21use crate::genome::traits::{BinaryGenome, EvolutionaryGenome, RealValuedGenome};
22use crate::population::individual::Individual;
23use crate::population::population::Population;
24use crate::termination::{EvolutionState, MaxGenerations, TerminationCriterion};
25
26#[derive(Clone, Debug)]
28pub struct UMDAConfig {
29 pub population_size: usize,
31 pub selection_ratio: f64,
33 pub min_variance: f64,
35 pub prob_bounds: (f64, f64),
37 pub learning_rate: f64,
39}
40
41impl Default for UMDAConfig {
42 fn default() -> Self {
43 Self {
44 population_size: 100,
45 selection_ratio: 0.5,
46 min_variance: 0.01,
47 prob_bounds: (0.01, 0.99),
48 learning_rate: 1.0,
49 }
50 }
51}
52
53pub struct UMDABuilder<G, F, Fit, Term>
55where
56 G: EvolutionaryGenome,
57 F: FitnessValue,
58{
59 config: UMDAConfig,
60 bounds: Option<MultiBounds>,
61 fitness: Option<Fit>,
62 termination: Option<Term>,
63 _phantom: std::marker::PhantomData<(G, F)>,
64}
65
66impl<G, F> UMDABuilder<G, F, (), ()>
67where
68 G: EvolutionaryGenome,
69 F: FitnessValue,
70{
71 pub fn new() -> Self {
73 Self {
74 config: UMDAConfig::default(),
75 bounds: None,
76 fitness: None,
77 termination: None,
78 _phantom: std::marker::PhantomData,
79 }
80 }
81}
82
83impl<G, F> Default for UMDABuilder<G, F, (), ()>
84where
85 G: EvolutionaryGenome,
86 F: FitnessValue,
87{
88 fn default() -> Self {
89 Self::new()
90 }
91}
92
93impl<G, F, Fit, Term> UMDABuilder<G, F, Fit, Term>
94where
95 G: EvolutionaryGenome,
96 F: FitnessValue,
97{
98 pub fn population_size(mut self, size: usize) -> Self {
100 self.config.population_size = size;
101 self
102 }
103
104 pub fn selection_ratio(mut self, ratio: f64) -> Self {
106 self.config.selection_ratio = ratio.clamp(0.1, 0.9);
107 self
108 }
109
110 pub fn min_variance(mut self, variance: f64) -> Self {
112 self.config.min_variance = variance;
113 self
114 }
115
116 pub fn prob_bounds(mut self, min: f64, max: f64) -> Self {
118 self.config.prob_bounds = (min.clamp(0.001, 0.5), max.clamp(0.5, 0.999));
119 self
120 }
121
122 pub fn learning_rate(mut self, rate: f64) -> Self {
124 self.config.learning_rate = rate.clamp(0.0, 1.0);
125 self
126 }
127
128 pub fn bounds(mut self, bounds: MultiBounds) -> Self {
130 self.bounds = Some(bounds);
131 self
132 }
133
134 pub fn fitness<NewFit>(self, fitness: NewFit) -> UMDABuilder<G, F, NewFit, Term>
136 where
137 NewFit: Fitness<Genome = G, Value = F>,
138 {
139 UMDABuilder {
140 config: self.config,
141 bounds: self.bounds,
142 fitness: Some(fitness),
143 termination: self.termination,
144 _phantom: std::marker::PhantomData,
145 }
146 }
147
148 pub fn termination<NewTerm>(self, termination: NewTerm) -> UMDABuilder<G, F, Fit, NewTerm>
150 where
151 NewTerm: TerminationCriterion<G, F>,
152 {
153 UMDABuilder {
154 config: self.config,
155 bounds: self.bounds,
156 fitness: self.fitness,
157 termination: Some(termination),
158 _phantom: std::marker::PhantomData,
159 }
160 }
161
162 pub fn max_generations(self, max: usize) -> UMDABuilder<G, F, Fit, MaxGenerations> {
164 UMDABuilder {
165 config: self.config,
166 bounds: self.bounds,
167 fitness: self.fitness,
168 termination: Some(MaxGenerations::new(max)),
169 _phantom: std::marker::PhantomData,
170 }
171 }
172}
173
174#[derive(Clone, Debug)]
180pub struct ContinuousUnivariateModel {
181 pub means: Vec<f64>,
183 pub variances: Vec<f64>,
185}
186
187impl ContinuousUnivariateModel {
188 pub fn from_bounds(bounds: &MultiBounds) -> Self {
190 let means: Vec<f64> = bounds.bounds.iter().map(|b| b.center()).collect();
191 let variances: Vec<f64> = bounds
192 .bounds
193 .iter()
194 .map(|b| (b.range() / 4.0).powi(2))
195 .collect();
196 Self { means, variances }
197 }
198
199 pub fn update(&mut self, selected: &[&RealVector], config: &UMDAConfig) {
201 let n = selected.len() as f64;
202 if n == 0.0 {
203 return;
204 }
205
206 for i in 0..self.means.len() {
207 let mean: f64 = selected.iter().map(|g| g.genes()[i]).sum::<f64>() / n;
209
210 let variance: f64 = selected
212 .iter()
213 .map(|g| (g.genes()[i] - mean).powi(2))
214 .sum::<f64>()
215 / n;
216
217 self.means[i] =
219 config.learning_rate * mean + (1.0 - config.learning_rate) * self.means[i];
220 self.variances[i] = config.learning_rate * variance.max(config.min_variance)
221 + (1.0 - config.learning_rate) * self.variances[i];
222 }
223 }
224
225 pub fn sample<R: Rng>(&self, bounds: &MultiBounds, rng: &mut R) -> RealVector {
227 let genes: Vec<f64> = self
228 .means
229 .iter()
230 .zip(self.variances.iter())
231 .zip(bounds.bounds.iter())
232 .map(|((mean, var), bound)| {
233 let normal =
234 Normal::new(*mean, var.sqrt()).unwrap_or(Normal::new(*mean, 0.1).unwrap());
235 let value = normal.sample(rng);
236 value.clamp(bound.min, bound.max)
237 })
238 .collect();
239
240 RealVector::new(genes)
241 }
242}
243
244impl<F, Fit, Term> UMDABuilder<RealVector, F, Fit, Term>
245where
246 F: FitnessValue + Send,
247 Fit: Fitness<Genome = RealVector, Value = F> + Sync,
248 Term: TerminationCriterion<RealVector, F>,
249{
250 pub fn build(self) -> Result<ContinuousUMDA<F, Fit, Term>, EvolutionError> {
252 let bounds = self
253 .bounds
254 .ok_or_else(|| EvolutionError::Configuration("Bounds must be specified".to_string()))?;
255
256 let fitness = self.fitness.ok_or_else(|| {
257 EvolutionError::Configuration("Fitness function must be specified".to_string())
258 })?;
259
260 let termination = self.termination.ok_or_else(|| {
261 EvolutionError::Configuration("Termination criterion must be specified".to_string())
262 })?;
263
264 Ok(ContinuousUMDA {
265 config: self.config,
266 bounds,
267 fitness,
268 termination,
269 _phantom: std::marker::PhantomData,
270 })
271 }
272}
273
274pub struct ContinuousUMDA<F, Fit, Term>
276where
277 F: FitnessValue,
278{
279 config: UMDAConfig,
280 bounds: MultiBounds,
281 fitness: Fit,
282 termination: Term,
283 _phantom: std::marker::PhantomData<F>,
284}
285
286impl<F, Fit, Term> ContinuousUMDA<F, Fit, Term>
287where
288 F: FitnessValue + Send,
289 Fit: Fitness<Genome = RealVector, Value = F> + Sync,
290 Term: TerminationCriterion<RealVector, F>,
291{
292 pub fn builder() -> UMDABuilder<RealVector, F, (), ()> {
294 UMDABuilder::new()
295 }
296
297 pub fn run<R: Rng>(
299 &self,
300 rng: &mut R,
301 ) -> Result<EvolutionResult<RealVector, F>, EvolutionError> {
302 let start_time = Instant::now();
303
304 let mut model = ContinuousUnivariateModel::from_bounds(&self.bounds);
306
307 let mut population: Population<RealVector, F> =
309 Population::random(self.config.population_size, &self.bounds, rng);
310 population.evaluate(&self.fitness);
311
312 let mut stats = EvolutionStats::new();
313 let mut evaluations = self.config.population_size;
314 let mut fitness_history: Vec<f64> = Vec::new();
315 let mut generation = 0usize;
316
317 let mut best = population
319 .best()
320 .ok_or(EvolutionError::EmptyPopulation)?
321 .clone();
322
323 let gen_stats = GenerationStats::from_population(&population, 0, evaluations);
325 fitness_history.push(gen_stats.best_fitness);
326 stats.record(gen_stats);
327
328 loop {
330 let state = EvolutionState {
332 generation,
333 evaluations,
334 best_fitness: best.fitness_value().to_f64(),
335 population: &population,
336 fitness_history: &fitness_history,
337 };
338
339 if self.termination.should_terminate(&state) {
340 stats.set_termination_reason(self.termination.reason());
341 break;
342 }
343
344 let gen_start = Instant::now();
345
346 population.sort_by_fitness();
348 let select_count =
349 (self.config.population_size as f64 * self.config.selection_ratio).ceil() as usize;
350 let selected: Vec<&RealVector> = population
351 .iter()
352 .take(select_count)
353 .map(|ind| &ind.genome)
354 .collect();
355
356 model.update(&selected, &self.config);
358
359 population = Population::with_capacity(self.config.population_size);
361 for _ in 0..self.config.population_size {
362 let genome = model.sample(&self.bounds, rng);
363 population.push(Individual::new(genome));
364 }
365
366 population.evaluate(&self.fitness);
368 evaluations += self.config.population_size;
369
370 if let Some(pop_best) = population.best() {
372 if pop_best.is_better_than(&best) {
373 best = pop_best.clone();
374 }
375 }
376
377 generation += 1;
378 population.set_generation(generation);
379
380 let timing = TimingStats::new().with_total(gen_start.elapsed());
382 let gen_stats = GenerationStats::from_population(&population, generation, evaluations)
383 .with_timing(timing);
384 fitness_history.push(gen_stats.best_fitness);
385 stats.record(gen_stats);
386 }
387
388 stats.set_runtime(start_time.elapsed());
389
390 Ok(
391 EvolutionResult::new(best.genome, best.fitness.unwrap(), generation, evaluations)
392 .with_stats(stats),
393 )
394 }
395}
396
397#[derive(Clone, Debug)]
403pub struct BinaryUnivariateModel {
404 pub probabilities: Vec<f64>,
406}
407
408impl BinaryUnivariateModel {
409 pub fn uniform(dimension: usize) -> Self {
411 Self {
412 probabilities: vec![0.5; dimension],
413 }
414 }
415
416 pub fn update(&mut self, selected: &[&BitString], config: &UMDAConfig) {
418 let n = selected.len() as f64;
419 if n == 0.0 {
420 return;
421 }
422
423 for i in 0..self.probabilities.len() {
424 let ones: f64 = selected
426 .iter()
427 .filter(|g| g.bits().get(i).copied().unwrap_or(false))
428 .count() as f64;
429
430 let new_prob = ones / n;
432 let bounded_prob = new_prob.clamp(config.prob_bounds.0, config.prob_bounds.1);
433 self.probabilities[i] = config.learning_rate * bounded_prob
434 + (1.0 - config.learning_rate) * self.probabilities[i];
435 }
436 }
437
438 pub fn sample<R: Rng>(&self, rng: &mut R) -> BitString {
440 let bits: Vec<bool> = self
441 .probabilities
442 .iter()
443 .map(|p| {
444 let dist = Bernoulli::new(*p).unwrap_or(Bernoulli::new(0.5).unwrap());
445 dist.sample(rng)
446 })
447 .collect();
448
449 BitString::new(bits)
450 }
451}
452
453impl<F, Fit, Term> UMDABuilder<BitString, F, Fit, Term>
454where
455 F: FitnessValue + Send,
456 Fit: Fitness<Genome = BitString, Value = F> + Sync,
457 Term: TerminationCriterion<BitString, F>,
458{
459 pub fn build(self) -> Result<BinaryUMDA<F, Fit, Term>, EvolutionError> {
461 let bounds = self
462 .bounds
463 .ok_or_else(|| EvolutionError::Configuration("Bounds must be specified".to_string()))?;
464
465 let fitness = self.fitness.ok_or_else(|| {
466 EvolutionError::Configuration("Fitness function must be specified".to_string())
467 })?;
468
469 let termination = self.termination.ok_or_else(|| {
470 EvolutionError::Configuration("Termination criterion must be specified".to_string())
471 })?;
472
473 Ok(BinaryUMDA {
474 config: self.config,
475 bounds,
476 fitness,
477 termination,
478 _phantom: std::marker::PhantomData,
479 })
480 }
481}
482
483pub struct BinaryUMDA<F, Fit, Term>
485where
486 F: FitnessValue,
487{
488 config: UMDAConfig,
489 bounds: MultiBounds,
490 fitness: Fit,
491 termination: Term,
492 _phantom: std::marker::PhantomData<F>,
493}
494
495impl<F, Fit, Term> BinaryUMDA<F, Fit, Term>
496where
497 F: FitnessValue + Send,
498 Fit: Fitness<Genome = BitString, Value = F> + Sync,
499 Term: TerminationCriterion<BitString, F>,
500{
501 pub fn builder() -> UMDABuilder<BitString, F, (), ()> {
503 UMDABuilder::new()
504 }
505
506 pub fn run<R: Rng>(
508 &self,
509 rng: &mut R,
510 ) -> Result<EvolutionResult<BitString, F>, EvolutionError> {
511 let start_time = Instant::now();
512
513 let dimension = self.bounds.dimension();
515
516 let mut model = BinaryUnivariateModel::uniform(dimension);
518
519 let mut population: Population<BitString, F> =
521 Population::random(self.config.population_size, &self.bounds, rng);
522 population.evaluate(&self.fitness);
523
524 let mut stats = EvolutionStats::new();
525 let mut evaluations = self.config.population_size;
526 let mut fitness_history: Vec<f64> = Vec::new();
527 let mut generation = 0usize;
528
529 let mut best = population
531 .best()
532 .ok_or(EvolutionError::EmptyPopulation)?
533 .clone();
534
535 let gen_stats = GenerationStats::from_population(&population, 0, evaluations);
537 fitness_history.push(gen_stats.best_fitness);
538 stats.record(gen_stats);
539
540 loop {
542 let state = EvolutionState {
544 generation,
545 evaluations,
546 best_fitness: best.fitness_value().to_f64(),
547 population: &population,
548 fitness_history: &fitness_history,
549 };
550
551 if self.termination.should_terminate(&state) {
552 stats.set_termination_reason(self.termination.reason());
553 break;
554 }
555
556 let gen_start = Instant::now();
557
558 population.sort_by_fitness();
560 let select_count =
561 (self.config.population_size as f64 * self.config.selection_ratio).ceil() as usize;
562 let selected: Vec<&BitString> = population
563 .iter()
564 .take(select_count)
565 .map(|ind| &ind.genome)
566 .collect();
567
568 model.update(&selected, &self.config);
570
571 population = Population::with_capacity(self.config.population_size);
573 for _ in 0..self.config.population_size {
574 let genome: BitString = model.sample(rng);
575 population.push(Individual::new(genome));
576 }
577
578 population.evaluate(&self.fitness);
580 evaluations += self.config.population_size;
581
582 if let Some(pop_best) = population.best() {
584 if pop_best.is_better_than(&best) {
585 best = pop_best.clone();
586 }
587 }
588
589 generation += 1;
590 population.set_generation(generation);
591
592 let timing = TimingStats::new().with_total(gen_start.elapsed());
594 let gen_stats = GenerationStats::from_population(&population, generation, evaluations)
595 .with_timing(timing);
596 fitness_history.push(gen_stats.best_fitness);
597 stats.record(gen_stats);
598 }
599
600 stats.set_runtime(start_time.elapsed());
601
602 Ok(
603 EvolutionResult::new(best.genome, best.fitness.unwrap(), generation, evaluations)
604 .with_stats(stats),
605 )
606 }
607}
608
609#[cfg(test)]
610mod tests {
611 use super::*;
612 use crate::fitness::benchmarks::{OneMax, Sphere};
613 use crate::termination::MaxEvaluations;
614
615 #[test]
616 fn test_continuous_umda_builder() {
617 let bounds = MultiBounds::symmetric(5.0, 10);
618 let umda: Result<ContinuousUMDA<f64, _, _>, _> = UMDABuilder::new()
619 .population_size(50)
620 .selection_ratio(0.3)
621 .bounds(bounds)
622 .fitness(Sphere::new(10))
623 .max_generations(10)
624 .build();
625
626 assert!(umda.is_ok());
627 }
628
629 #[test]
630 fn test_continuous_umda_sphere() {
631 let mut rng = rand::thread_rng();
632 let bounds = MultiBounds::symmetric(5.12, 10);
633
634 let umda: ContinuousUMDA<f64, _, _> = UMDABuilder::new()
635 .population_size(100)
636 .selection_ratio(0.3)
637 .min_variance(0.01)
638 .bounds(bounds)
639 .fitness(Sphere::new(10))
640 .termination(MaxEvaluations::new(5000))
641 .build()
642 .unwrap();
643
644 let result = umda.run(&mut rng).unwrap();
645
646 assert!(
648 result.best_fitness > -50.0,
649 "Expected fitness > -50, got {}",
650 result.best_fitness
651 );
652 }
653
654 #[test]
655 fn test_binary_umda_onemax() {
656 let mut rng = rand::thread_rng();
657 let bounds = MultiBounds::uniform(crate::genome::bounds::Bounds::unit(), 20);
658
659 let umda: BinaryUMDA<usize, _, _> = UMDABuilder::new()
660 .population_size(100)
661 .selection_ratio(0.3)
662 .prob_bounds(0.05, 0.95)
663 .bounds(bounds)
664 .fitness(OneMax::new(20))
665 .termination(MaxEvaluations::new(3000))
666 .build()
667 .unwrap();
668
669 let result = umda.run(&mut rng).unwrap();
670
671 assert!(
673 result.best_fitness >= 15,
674 "Expected fitness >= 15, got {}",
675 result.best_fitness
676 );
677 }
678
679 #[test]
680 fn test_continuous_model_update() {
681 let bounds = MultiBounds::symmetric(5.0, 3);
682 let mut model = ContinuousUnivariateModel::from_bounds(&bounds);
683 let config = UMDAConfig::default();
684
685 let g1 = RealVector::new(vec![1.0, 2.0, 3.0]);
687 let g2 = RealVector::new(vec![2.0, 3.0, 4.0]);
688 let g3 = RealVector::new(vec![3.0, 4.0, 5.0]);
689 let selected = vec![&g1, &g2, &g3];
690
691 model.update(&selected, &config);
692
693 assert!((model.means[0] - 2.0).abs() < 0.01);
695 assert!((model.means[1] - 3.0).abs() < 0.01);
696 assert!((model.means[2] - 4.0).abs() < 0.01);
697 }
698
699 #[test]
700 fn test_binary_model_update() {
701 let mut model = BinaryUnivariateModel::uniform(4);
702 let config = UMDAConfig::default();
703
704 let g1 = BitString::new(vec![true, true, false, false]);
706 let g2 = BitString::new(vec![true, false, true, false]);
707 let g3 = BitString::new(vec![true, false, false, true]);
708 let selected = vec![&g1, &g2, &g3];
709
710 model.update(&selected, &config);
711
712 assert!(model.probabilities[0] > 0.9);
717 assert!(model.probabilities[1] > 0.2 && model.probabilities[1] < 0.5);
718 }
719
720 #[test]
721 fn test_learning_rate() {
722 let bounds = MultiBounds::symmetric(5.0, 2);
723 let mut model = ContinuousUnivariateModel::from_bounds(&bounds);
724 let config = UMDAConfig {
725 learning_rate: 0.5,
726 ..Default::default()
727 };
728
729 let initial_mean = model.means[0];
731
732 let g1 = RealVector::new(vec![2.0, 2.0]);
734 let selected = vec![&g1];
735
736 model.update(&selected, &config);
737
738 assert!((model.means[0] - (0.5 * 2.0 + 0.5 * initial_mean)).abs() < 0.01);
740 }
741}