rlevo_evolution/ops/
selection.rs

1//! Parent-selection operators.
2//!
3//! Selection operators turn a fitness tensor into winner indices that the
4//! caller uses (via `Tensor::select`) to gather parents out of the
5//! population. The baseline is host-side sampling + a single device
6//! gather; a fused [kernel](super::kernels) variant is tracked as
7//! follow-up work.
8//!
9//! # Fitness convention
10//!
11//! Fitness values are interpreted as "lower is better" (cost).
12//! Tournament selection picks the smaller of the two sampled fitnesses.
13
14use burn::tensor::{backend::Backend, Int, Tensor, TensorData};
15use rand::{Rng, RngExt};
16
17/// Binary tournament selection over a fitness slice.
18///
19/// Samples two indices uniformly at random and keeps the one with the
20/// smaller fitness. Repeats `n_winners` times, returning one winner
21/// index per draw.
22///
23/// # Panics
24///
25/// Panics if `fitness.is_empty()` or if `tournament_size < 2`.
26#[must_use]
27pub fn tournament_indices_host(
28    fitness: &[f32],
29    tournament_size: usize,
30    n_winners: usize,
31    rng: &mut dyn Rng,
32) -> Vec<i64> {
33    assert!(!fitness.is_empty(), "fitness must be non-empty");
34    assert!(tournament_size >= 2, "tournament size must be >= 2");
35    let pop_size = fitness.len();
36    let mut winners = Vec::with_capacity(n_winners);
37    for _ in 0..n_winners {
38        let mut best_idx = rng.random_range(0..pop_size);
39        let mut best_f = fitness[best_idx];
40        for _ in 1..tournament_size {
41            let idx = rng.random_range(0..pop_size);
42            if fitness[idx] < best_f {
43                best_f = fitness[idx];
44                best_idx = idx;
45            }
46        }
47        #[allow(clippy::cast_possible_wrap)]
48        winners.push(best_idx as i64);
49    }
50    winners
51}
52
53/// Gather `n_winners` rows out of the population by tournament.
54///
55/// Convenience wrapper that lifts [`tournament_indices_host`] into a
56/// device gather on the population tensor.
57///
58/// # Panics
59///
60/// Inherits the panic conditions of [`tournament_indices_host`]:
61/// `fitness.is_empty()` or `tournament_size < 2`.
62#[must_use]
63pub fn tournament_select<B: Backend>(
64    population: &Tensor<B, 2>,
65    fitness: &[f32],
66    tournament_size: usize,
67    n_winners: usize,
68    rng: &mut dyn Rng,
69    device: &B::Device,
70) -> Tensor<B, 2> {
71    let winners = tournament_indices_host(fitness, tournament_size, n_winners, rng);
72    let indices = Tensor::<B, 1, Int>::from_data(TensorData::new(winners, [n_winners]), device);
73    population.clone().select(0, indices)
74}
75
76/// Truncation selection: returns the indices of the `top_k` lowest-
77/// fitness members.
78///
79/// # Panics
80///
81/// Panics if `top_k > fitness.len()` or `fitness.is_empty()`.
82#[must_use]
83pub fn truncation_indices_host(fitness: &[f32], top_k: usize) -> Vec<i64> {
84    assert!(!fitness.is_empty(), "fitness must be non-empty");
85    assert!(top_k <= fitness.len(), "top_k must be <= population size");
86    let mut indexed: Vec<(usize, f32)> = fitness.iter().copied().enumerate().collect();
87    indexed
88        .sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
89    #[allow(clippy::cast_possible_wrap)]
90    indexed
91        .into_iter()
92        .take(top_k)
93        .map(|(i, _)| i as i64)
94        .collect()
95}
96
97/// Gather the `top_k` lowest-fitness rows out of the population.
98///
99/// # Panics
100///
101/// Inherits the panic conditions of [`truncation_indices_host`]:
102/// `fitness.is_empty()` or `top_k > fitness.len()`.
103#[must_use]
104pub fn truncation_select<B: Backend>(
105    population: &Tensor<B, 2>,
106    fitness: &[f32],
107    top_k: usize,
108    device: &B::Device,
109) -> Tensor<B, 2> {
110    let winners = truncation_indices_host(fitness, top_k);
111    let indices = Tensor::<B, 1, Int>::from_data(TensorData::new(winners, [top_k]), device);
112    population.clone().select(0, indices)
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use burn::backend::NdArray;
119    use rand::SeedableRng;
120    use rand::rngs::StdRng;
121
122    type TestBackend = NdArray;
123
124    #[test]
125    fn tournament_prefers_better_fitness_in_expectation() {
126        let mut rng = StdRng::seed_from_u64(1);
127        let fitness = [100.0f32, 0.0, 100.0, 100.0];
128        let winners = tournament_indices_host(&fitness, 2, 1000, &mut rng);
129        let wins_for_best = winners.iter().filter(|&&w| w == 1).count();
130        // For pop_size=4 and tournament_size=2, P(best wins) =
131        // 1 − (3/4)² = 7/16 ≈ 0.4375 → ~437 wins per 1000 trials.
132        // Use a generous band to stay stable across RNG versions.
133        assert!(
134            (350..=550).contains(&wins_for_best),
135            "wins_for_best={wins_for_best} (expected ~437)",
136        );
137    }
138
139    #[test]
140    fn truncation_returns_smallest_fitness_first() {
141        let fitness = [5.0f32, 1.0, 3.0, 2.0, 4.0];
142        let idx = truncation_indices_host(&fitness, 3);
143        // The three smallest fitnesses live at indices 1 (=1.0), 3 (=2.0), 2 (=3.0).
144        assert_eq!(idx.len(), 3);
145        assert!(idx.contains(&1));
146        assert!(idx.contains(&3));
147        assert!(idx.contains(&2));
148    }
149
150    #[test]
151    fn tournament_select_returns_shaped_tensor() {
152        let device = Default::default();
153        let data = TensorData::new(vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0], [3, 2]);
154        let pop = Tensor::<TestBackend, 2>::from_data(data, &device);
155        let fitness = [10.0_f32, 0.0, 10.0];
156        let mut rng = StdRng::seed_from_u64(2);
157        let parents = tournament_select(&pop, &fitness, 2, 4, &mut rng, &device);
158        assert_eq!(parents.shape().dims, vec![4, 2]);
159    }
160}
rlevo_evolution/ops/selection.rs

rlevo_evolution/ops/
selection.rs
