Skip to main content

fugue_evo/algorithms/eda/
umda.rs

1//! Univariate Marginal Distribution Algorithm (UMDA)
2//!
3//! UMDA is a simple but effective EDA that assumes independence between variables.
4//! It estimates univariate marginal distributions for each variable and samples
5//! new solutions from the product of these distributions.
6//!
7//! For continuous problems: estimates mean and variance per dimension
8//! For binary problems: estimates probability of 1 per bit
9
10use 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/// Configuration for UMDA
27#[derive(Clone, Debug)]
28pub struct UMDAConfig {
29    /// Population size
30    pub population_size: usize,
31    /// Selection ratio (top proportion selected for model learning)
32    pub selection_ratio: f64,
33    /// Minimum variance to prevent collapse (for continuous)
34    pub min_variance: f64,
35    /// Probability bounds for binary (to prevent determinism)
36    pub prob_bounds: (f64, f64),
37    /// Learning rate for model update (1.0 = replace, <1.0 = blend with previous)
38    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
53/// Builder for UMDA
54pub 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    /// Create a new builder with default configuration
72    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    /// Set the population size
99    pub fn population_size(mut self, size: usize) -> Self {
100        self.config.population_size = size;
101        self
102    }
103
104    /// Set the selection ratio
105    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    /// Set the minimum variance
111    pub fn min_variance(mut self, variance: f64) -> Self {
112        self.config.min_variance = variance;
113        self
114    }
115
116    /// Set probability bounds for binary UMDA
117    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    /// Set the learning rate for model update
123    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    /// Set the search space bounds
129    pub fn bounds(mut self, bounds: MultiBounds) -> Self {
130        self.bounds = Some(bounds);
131        self
132    }
133
134    /// Set the fitness function
135    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    /// Set the termination criterion
149    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    /// Set max generations (convenience method)
163    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// ============================================================================
175// Continuous UMDA (RealVector)
176// ============================================================================
177
178/// Univariate model for continuous variables
179#[derive(Clone, Debug)]
180pub struct ContinuousUnivariateModel {
181    /// Mean for each dimension
182    pub means: Vec<f64>,
183    /// Variance for each dimension
184    pub variances: Vec<f64>,
185}
186
187impl ContinuousUnivariateModel {
188    /// Create from bounds (initial uniform-ish distribution)
189    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    /// Update model from selected individuals
200    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            // Compute sample mean
208            let mean: f64 = selected.iter().map(|g| g.genes()[i]).sum::<f64>() / n;
209
210            // Compute sample variance
211            let variance: f64 = selected
212                .iter()
213                .map(|g| (g.genes()[i] - mean).powi(2))
214                .sum::<f64>()
215                / n;
216
217            // Apply learning rate and bounds
218            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    /// Sample a new individual from the model
226    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    /// Build the continuous UMDA instance
251    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
274/// UMDA for continuous optimization
275pub 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    /// Create a builder for continuous UMDA
293    pub fn builder() -> UMDABuilder<RealVector, F, (), ()> {
294        UMDABuilder::new()
295    }
296
297    /// Run the UMDA algorithm
298    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        // Initialize model
305        let mut model = ContinuousUnivariateModel::from_bounds(&self.bounds);
306
307        // Initialize population
308        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        // Track best individual
318        let mut best = population
319            .best()
320            .ok_or(EvolutionError::EmptyPopulation)?
321            .clone();
322
323        // Record initial statistics
324        let gen_stats = GenerationStats::from_population(&population, 0, evaluations);
325        fitness_history.push(gen_stats.best_fitness);
326        stats.record(gen_stats);
327
328        // Main loop
329        loop {
330            // Check termination
331            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            // Sort and select top individuals
347            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            // Update model from selected individuals
357            model.update(&selected, &self.config);
358
359            // Sample new population from model
360            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            // Evaluate
367            population.evaluate(&self.fitness);
368            evaluations += self.config.population_size;
369
370            // Update best
371            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            // Record statistics
381            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// ============================================================================
398// Binary UMDA (BitString)
399// ============================================================================
400
401/// Univariate model for binary variables
402#[derive(Clone, Debug)]
403pub struct BinaryUnivariateModel {
404    /// Probability of 1 for each bit position
405    pub probabilities: Vec<f64>,
406}
407
408impl BinaryUnivariateModel {
409    /// Create with uniform 0.5 probability
410    pub fn uniform(dimension: usize) -> Self {
411        Self {
412            probabilities: vec![0.5; dimension],
413        }
414    }
415
416    /// Update model from selected individuals
417    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            // Count ones at position i
425            let ones: f64 = selected
426                .iter()
427                .filter(|g| g.bits().get(i).copied().unwrap_or(false))
428                .count() as f64;
429
430            // Compute probability with learning rate and bounds
431            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    /// Sample a new individual from the model
439    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    /// Build the binary UMDA instance
460    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
483/// UMDA for binary optimization
484pub 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    /// Create a builder for binary UMDA
502    pub fn builder() -> UMDABuilder<BitString, F, (), ()> {
503        UMDABuilder::new()
504    }
505
506    /// Run the UMDA algorithm
507    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        // Get dimension from bounds
514        let dimension = self.bounds.dimension();
515
516        // Initialize model
517        let mut model = BinaryUnivariateModel::uniform(dimension);
518
519        // Initialize population
520        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        // Track best individual
530        let mut best = population
531            .best()
532            .ok_or(EvolutionError::EmptyPopulation)?
533            .clone();
534
535        // Record initial statistics
536        let gen_stats = GenerationStats::from_population(&population, 0, evaluations);
537        fitness_history.push(gen_stats.best_fitness);
538        stats.record(gen_stats);
539
540        // Main loop
541        loop {
542            // Check termination
543            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            // Sort and select top individuals
559            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            // Update model from selected individuals
569            model.update(&selected, &self.config);
570
571            // Sample new population from model
572            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            // Evaluate
579            population.evaluate(&self.fitness);
580            evaluations += self.config.population_size;
581
582            // Update best
583            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            // Record statistics
593            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        // Should find improvement
647        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        // Should find good solution
672        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        // Create some test individuals
686        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        // Mean should be average: (1+2+3)/3=2, (2+3+4)/3=3, (3+4+5)/3=4
694        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        // Create test individuals
705        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        // Position 0: all true -> p=0.99 (bounded)
713        // Position 1: 1/3 true -> p=0.33
714        // Position 2: 1/3 true -> p=0.33
715        // Position 3: 1/3 true -> p=0.33
716        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        // Initial means should be 0.0 (center of [-5, 5])
730        let initial_mean = model.means[0];
731
732        // Create individuals at 2.0
733        let g1 = RealVector::new(vec![2.0, 2.0]);
734        let selected = vec![&g1];
735
736        model.update(&selected, &config);
737
738        // With learning rate 0.5, new mean = 0.5 * 2.0 + 0.5 * 0.0 = 1.0
739        assert!((model.means[0] - (0.5 * 2.0 + 0.5 * initial_mean)).abs() < 0.01);
740    }
741}