use std::marker::PhantomData;
use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
use rand::{Rng, RngExt};
use rlevo_core::bounds::Bounds;
use rlevo_core::config::{self, ConfigError, Validate};
use crate::ops::selection::argmax_host;
use crate::rng::{SeedPurpose, seed_stream};
use crate::strategy::{Strategy, StrategyMetrics};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DeVariant {
Rand1Bin,
Best1Bin,
CurrentToBest1Bin,
Rand2Bin,
Rand1Exp,
}
impl DeVariant {
const fn random_indices(self) -> usize {
match self {
DeVariant::Rand1Bin | DeVariant::Rand1Exp => 3,
DeVariant::Best1Bin | DeVariant::CurrentToBest1Bin => 2,
DeVariant::Rand2Bin => 5,
}
}
const fn is_exponential(self) -> bool {
matches!(self, DeVariant::Rand1Exp)
}
}
#[derive(Debug, Clone)]
pub struct DeConfig {
pub pop_size: usize,
pub genome_dim: usize,
pub bounds: Bounds,
pub f: f32,
pub cr: f32,
pub variant: DeVariant,
}
impl DeConfig {
#[must_use]
pub fn default_for(pop_size: usize, genome_dim: usize) -> Self {
Self {
pop_size,
genome_dim,
bounds: Bounds::new(-5.12, 5.12),
f: 0.5,
cr: 0.9,
variant: DeVariant::Rand1Bin,
}
}
}
impl Validate for DeConfig {
fn validate(&self) -> Result<(), ConfigError> {
const C: &str = "DeConfig";
let min_pop = if self.variant == DeVariant::Rand2Bin {
5
} else {
4
};
config::at_least(C, "pop_size", self.pop_size, min_pop)?;
config::nonzero(C, "genome_dim", self.genome_dim)?;
config::in_range(C, "f", 0.0, 2.0, f64::from(self.f))?;
config::in_range(C, "cr", 0.0, 1.0, f64::from(self.cr))?;
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct DeState<B: Backend> {
pub population: Tensor<B, 2>,
pub fitness: Vec<f32>,
pub best_index: usize,
pub best_genome: Option<Tensor<B, 2>>,
pub best_fitness: f32,
pub generation: usize,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct DifferentialEvolution<B: Backend> {
_backend: PhantomData<fn() -> B>,
}
impl<B: Backend> DifferentialEvolution<B> {
#[must_use]
pub fn new() -> Self {
Self {
_backend: PhantomData,
}
}
fn sample_initial_population(
params: &DeConfig,
rng: &mut dyn Rng,
device: &<B as burn::tensor::backend::BackendTypes>::Device,
) -> Tensor<B, 2> {
let (lo, hi): (f32, f32) = params.bounds.into();
let pop = params.pop_size;
let genome_dim = params.genome_dim;
let mut stream = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
let mut rows = Vec::with_capacity(pop * genome_dim);
for _ in 0..pop * genome_dim {
rows.push(lo + (hi - lo) * stream.random::<f32>());
}
Tensor::<B, 2>::from_data(TensorData::new(rows, [pop, genome_dim]), device)
}
fn sample_distinct_excluding(
self_idx: usize,
pop_size: usize,
k: usize,
rng: &mut dyn Rng,
) -> Vec<usize> {
assert!(
pop_size > k,
"DE: pop_size must exceed the number of distinct indices required"
);
let mut chosen = Vec::with_capacity(k);
while chosen.len() < k {
let candidate = rng.random_range(0..pop_size);
if candidate != self_idx && !chosen.contains(&candidate) {
chosen.push(candidate);
}
}
chosen
}
}
impl<B: Backend> Strategy<B> for DifferentialEvolution<B>
where
B::Device: Clone,
{
type Params = DeConfig;
type State = DeState<B>;
type Genome = Tensor<B, 2>;
fn init(
&self,
params: &DeConfig,
rng: &mut dyn Rng,
device: &<B as burn::tensor::backend::BackendTypes>::Device,
) -> DeState<B> {
debug_assert!(
params.validate().is_ok(),
"invalid DeConfig reached init: {params:?}"
);
let population = Self::sample_initial_population(params, rng, device);
DeState {
population,
fitness: Vec::new(),
best_index: 0,
best_genome: None,
best_fitness: f32::NEG_INFINITY,
generation: 0,
}
}
#[allow(clippy::too_many_lines, clippy::many_single_char_names)]
fn ask(
&self,
params: &DeConfig,
state: &DeState<B>,
rng: &mut dyn Rng,
device: &<B as burn::tensor::backend::BackendTypes>::Device,
) -> (Tensor<B, 2>, DeState<B>) {
if state.fitness.is_empty() {
return (state.population.clone(), state.clone());
}
let DeConfig {
pop_size,
genome_dim,
f,
cr,
variant,
..
} = *params;
let mut trial_rng =
seed_stream(rng.next_u64(), state.generation as u64, SeedPurpose::Trial);
let k = variant.random_indices();
let mut rand_indices: Vec<Vec<usize>> =
(0..k).map(|_| Vec::with_capacity(pop_size)).collect();
for i in 0..pop_size {
let chosen = Self::sample_distinct_excluding(i, pop_size, k, &mut trial_rng);
for (j, idx) in chosen.into_iter().enumerate() {
rand_indices[j].push(idx);
}
}
let gather = |idxs: &[usize]| -> Tensor<B, 2> {
#[allow(clippy::cast_possible_wrap)]
let v: Vec<i64> = idxs.iter().map(|&i| i as i64).collect();
let t = Tensor::<B, 1, Int>::from_data(TensorData::new(v, [pop_size]), device);
state.population.clone().select(0, t)
};
let v = match variant {
DeVariant::Rand1Bin | DeVariant::Rand1Exp => {
let a = gather(&rand_indices[0]);
let b = gather(&rand_indices[1]);
let c = gather(&rand_indices[2]);
a + (b - c).mul_scalar(f)
}
DeVariant::Best1Bin => {
#[allow(clippy::single_range_in_vec_init)]
let best = state
.population
.clone()
.slice([state.best_index..state.best_index + 1])
.expand([pop_size, genome_dim]);
let b = gather(&rand_indices[0]);
let c = gather(&rand_indices[1]);
best + (b - c).mul_scalar(f)
}
DeVariant::CurrentToBest1Bin => {
#[allow(clippy::single_range_in_vec_init)]
let best = state
.population
.clone()
.slice([state.best_index..state.best_index + 1])
.expand([pop_size, genome_dim]);
let current = state.population.clone();
let a = gather(&rand_indices[0]);
let b = gather(&rand_indices[1]);
current.clone() + (best - current).mul_scalar(f) + (a - b).mul_scalar(f)
}
DeVariant::Rand2Bin => {
let a = gather(&rand_indices[0]);
let b = gather(&rand_indices[1]);
let c = gather(&rand_indices[2]);
let d = gather(&rand_indices[3]);
let e = gather(&rand_indices[4]);
a + (b - c).mul_scalar(f) + (d - e).mul_scalar(f)
}
};
let mut cross_rng = seed_stream(
rng.next_u64(),
state.generation as u64,
SeedPurpose::Crossover,
);
let mut cross_mask = vec![false; pop_size * genome_dim];
if variant.is_exponential() {
for row in 0..pop_size {
let start = cross_rng.random_range(0..genome_dim);
let mut len = 1;
while len < genome_dim && cross_rng.random::<f32>() < cr {
len += 1;
}
for k in 0..len {
let j = (start + k) % genome_dim;
cross_mask[row * genome_dim + j] = true;
}
}
} else {
for row in 0..pop_size {
let j_rand = cross_rng.random_range(0..genome_dim);
for j in 0..genome_dim {
if j == j_rand || cross_rng.random::<f32>() < cr {
cross_mask[row * genome_dim + j] = true;
}
}
}
}
#[allow(clippy::cast_possible_wrap)]
let mask_int: Vec<i64> = cross_mask.iter().map(|&b| i64::from(b)).collect();
let mask_tensor = Tensor::<B, 2, Int>::from_data(
TensorData::new(mask_int, [pop_size, genome_dim]),
device,
);
let mask_bool = mask_tensor.equal_elem(1);
let trial = state.population.clone().mask_where(mask_bool, v);
let (lo, hi): (f32, f32) = params.bounds.into();
let trial = trial.clamp(lo, hi);
(trial, state.clone())
}
fn tell(
&self,
_params: &DeConfig,
trial: Tensor<B, 2>,
fitness: Tensor<B, 1>,
mut state: DeState<B>,
_rng: &mut dyn Rng,
) -> (DeState<B>, StrategyMetrics) {
let fitness_host = fitness
.into_data()
.into_vec::<f32>()
.expect("fitness tensor must be readable as f32");
if state.fitness.is_empty() {
state.fitness.clone_from(&fitness_host);
state.best_index = argmax_host(&fitness_host);
state.generation += 1;
update_best(&mut state, &trial, &fitness_host);
let m = StrategyMetrics::from_host_fitness(
state.generation,
&fitness_host,
state.best_fitness,
);
state.best_fitness = m.best_fitness_ever();
state.population = trial;
return (state, m);
}
let device = trial.device();
let pop_size = state.fitness.len();
let mut replace_mask = vec![0i64; pop_size];
let mut new_fit = state.fitness.clone();
for i in 0..pop_size {
if fitness_host[i] >= state.fitness[i] {
replace_mask[i] = 1;
new_fit[i] = fitness_host[i];
}
}
let mask_int =
Tensor::<B, 1, Int>::from_data(TensorData::new(replace_mask, [pop_size]), &device);
let mask_bool_row = mask_int.equal_elem(1);
let genome_dim = state.population.dims()[1];
let mask_bool = mask_bool_row
.unsqueeze_dim::<2>(1)
.expand([pop_size, genome_dim]);
let next_pop = state
.population
.clone()
.mask_where(mask_bool, trial.clone());
state.population = next_pop;
state.fitness.clone_from(&new_fit);
state.best_index = argmax_host(&new_fit);
state.generation += 1;
update_best(&mut state, &trial, &fitness_host);
let m = StrategyMetrics::from_host_fitness(state.generation, &new_fit, state.best_fitness);
state.best_fitness = m.best_fitness_ever();
(state, m)
}
fn best(&self, state: &DeState<B>) -> Option<(Tensor<B, 2>, f32)> {
state
.best_genome
.as_ref()
.map(|g| (g.clone(), state.best_fitness))
}
}
fn update_best<B: Backend>(state: &mut DeState<B>, pop: &Tensor<B, 2>, fitness: &[f32]) {
if fitness.is_empty() {
return;
}
let best_idx = argmax_host(fitness);
let best_f = fitness[best_idx];
if best_f > state.best_fitness {
let device = pop.device();
#[allow(clippy::cast_possible_wrap)]
let idx =
Tensor::<B, 1, Int>::from_data(TensorData::new(vec![best_idx as i64], [1]), &device);
state.best_genome = Some(pop.clone().select(0, idx));
state.best_fitness = best_f;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::fitness::FromFitnessEvaluable;
use crate::strategy::EvolutionaryHarness;
use burn::backend::Flex;
use rlevo_core::fitness::FitnessEvaluable;
type TestBackend = Flex;
#[test]
fn default_config_validates() {
assert!(DeConfig::default_for(30, 10).validate().is_ok());
}
#[test]
fn rejects_pop_size_below_min() {
let mut cfg = DeConfig::default_for(3, 10);
cfg.pop_size = 3;
assert_eq!(cfg.validate().unwrap_err().field, "pop_size");
}
#[test]
fn sample_distinct_excluding_yields_valid_indices() {
use rand::SeedableRng;
use rand::rngs::StdRng;
let mut rng = StdRng::seed_from_u64(20_240_607);
for pop_size in [4usize, 5, 8, 20] {
for k in 1..pop_size.min(6) {
for self_idx in 0..pop_size {
for _ in 0..25 {
let chosen: Vec<usize> =
DifferentialEvolution::<TestBackend>::sample_distinct_excluding(
self_idx, pop_size, k, &mut rng,
);
assert_eq!(chosen.len(), k, "must return exactly k indices");
for (a, &x) in chosen.iter().enumerate() {
assert!(x < pop_size, "index {x} out of range for pop {pop_size}");
assert_ne!(x, self_idx, "index must differ from self_idx");
for &y in &chosen[a + 1..] {
assert_ne!(x, y, "indices must be pairwise distinct");
}
}
}
}
}
}
}
#[test]
#[should_panic(expected = "pop_size must exceed")]
fn sample_distinct_excluding_panics_when_pop_too_small() {
use rand::SeedableRng;
use rand::rngs::StdRng;
let mut rng = StdRng::seed_from_u64(1);
let _ = DifferentialEvolution::<TestBackend>::sample_distinct_excluding(0, 3, 3, &mut rng);
}
#[test]
fn trial_genes_stay_within_bounds() {
use rand::SeedableRng;
use rand::rngs::StdRng;
let device = Default::default();
let strategy = DifferentialEvolution::<TestBackend>::new();
let mut params = DeConfig::default_for(12, 4);
params.variant = DeVariant::Rand1Bin;
let (lo, hi): (f32, f32) = params.bounds.into();
let mut rng = StdRng::seed_from_u64(4242);
let state = strategy.init(¶ms, &mut rng, &device);
let (pop0, s) = strategy.ask(¶ms, &state, &mut rng, &device);
let n = pop0.dims()[0];
let fitness =
Tensor::<TestBackend, 1>::from_data(TensorData::new(vec![0.0_f32; n], [n]), &device);
let (s, _) = strategy.tell(¶ms, pop0, fitness, s, &mut rng);
let (trial, _) = strategy.ask(¶ms, &s, &mut rng, &device);
let genes: Vec<f32> = trial
.into_data()
.into_vec::<f32>()
.expect("trial host-read of a tensor this test just built");
for g in genes {
assert!(
g.is_finite() && g >= lo && g <= hi,
"trial gene {g} left [{lo}, {hi}]"
);
}
}
struct NanSphere;
struct NanSphereFit;
impl FitnessEvaluable for NanSphereFit {
type Individual = Vec<f64>;
type Landscape = NanSphere;
fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
let s: f64 = x.iter().map(|v| v * v).sum();
if x[0] > 0.0 { f64::NAN } else { s }
}
}
#[test]
fn nan_fitness_never_becomes_best() {
let device = Default::default();
let params = DeConfig::default_for(30, 4);
let fitness_fn = FromFitnessEvaluable::new(NanSphereFit, NanSphere);
let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
DifferentialEvolution::<TestBackend>::new(),
params,
fitness_fn,
99,
device,
40,
)
.expect("valid params");
harness.reset();
loop {
if harness.step(()).done {
break;
}
}
let best = harness.latest_metrics().unwrap().best_fitness_ever();
assert!(
best.is_finite(),
"NaN fitness poisoned best_fitness_ever: {best}"
);
}
struct Sphere;
struct SphereFit;
impl FitnessEvaluable for SphereFit {
type Individual = Vec<f64>;
type Landscape = Sphere;
fn evaluate(&self, x: &Self::Individual, _: &Self::Landscape) -> f64 {
x.iter().map(|v| v * v).sum()
}
}
fn run_de(variant: DeVariant, dim: usize, gens: usize) -> f32 {
let device = Default::default();
let mut params = DeConfig::default_for(30, dim);
params.variant = variant;
let fitness_fn = FromFitnessEvaluable::new(SphereFit, Sphere);
let mut harness = EvolutionaryHarness::<TestBackend, _, _>::new(
DifferentialEvolution::<TestBackend>::new(),
params,
fitness_fn,
11,
device,
gens,
)
.expect("valid params");
harness.reset();
loop {
if harness.step(()).done {
break;
}
}
harness.latest_metrics().unwrap().best_fitness_ever()
}
#[test]
fn all_variants_converge_on_sphere_d10() {
let rand1bin = run_de(DeVariant::Rand1Bin, 10, 500);
assert!(rand1bin < 1e-6, "DE/rand/1/bin best={rand1bin}");
let rand2bin = run_de(DeVariant::Rand2Bin, 10, 800);
assert!(rand2bin < 1e-6, "DE/rand/2/bin best={rand2bin}");
let rand1exp = run_de(DeVariant::Rand1Exp, 10, 500);
assert!(rand1exp < 1e-6, "DE/rand/1/exp best={rand1exp}");
let best1bin = run_de(DeVariant::Best1Bin, 10, 500);
assert!(best1bin < 1.0, "DE/best/1/bin best={best1bin}");
let c2b = run_de(DeVariant::CurrentToBest1Bin, 10, 500);
assert!(c2b < 2.0, "DE/current-to-best/1/bin best={c2b}");
}
}