pub mod bayesian_network;
pub mod compact_genetic;
pub mod dependency_chain;
pub mod univariate_bernoulli;
pub mod univariate_gaussian;
pub use bayesian_network::{BayesianNetwork, BayesianNetworkParams, BayesianNetworkState};
pub use compact_genetic::{CompactGenetic, CompactGeneticParams, CompactGeneticState};
pub use dependency_chain::{DependencyChain, DependencyChainParams, DependencyChainState};
pub use univariate_bernoulli::{
UnivariateBernoulli, UnivariateBernoulliParams, UnivariateBernoulliState,
};
pub use univariate_gaussian::{
UnivariateGaussian, UnivariateGaussianParams, UnivariateGaussianState,
};
use std::fmt::Debug;
use std::marker::PhantomData;
use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
use rand::Rng;
use rlevo_core::bounds::Bounds;
use rlevo_core::config::{self, ConfigError, Validate};
use crate::probability_model::ProbabilityModel;
use crate::rng::{SeedPurpose, seed_stream};
use crate::strategy::{Strategy, StrategyMetrics};
#[derive(Debug, Clone)]
pub struct EdaParams<MP> {
pub pop_size: usize,
pub selection_ratio: f32,
pub bounds: Option<Bounds>,
pub model: MP,
}
impl<MP> Validate for EdaParams<MP> {
fn validate(&self) -> Result<(), ConfigError> {
const C: &str = "EdaParams";
config::at_least(C, "pop_size", self.pop_size, 2)?;
config::positive(C, "selection_ratio", f64::from(self.selection_ratio))?;
config::ordered(C, "selection_ratio", f64::from(self.selection_ratio), 1.0)?;
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct EdaState<B: Backend, MS> {
pub model_state: MS,
pub best_genome: Option<Tensor<B, 1>>,
pub best_fitness_ever: f32,
pub generation: usize,
}
pub struct EdaStrategy<B: Backend, M> {
model: M,
_backend: PhantomData<fn() -> B>,
}
impl<B: Backend, M: Debug> Debug for EdaStrategy<B, M> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EdaStrategy")
.field("model", &self.model)
.finish_non_exhaustive()
}
}
impl<B: Backend, M> EdaStrategy<B, M> {
#[must_use]
pub fn new(model: M) -> Self {
Self {
model,
_backend: PhantomData,
}
}
}
impl<B: Backend, M: ProbabilityModel<B>> Strategy<B> for EdaStrategy<B, M> {
type Params = EdaParams<M::Params>;
type State = EdaState<B, M::State>;
type Genome = Tensor<B, 2>;
fn init(
&self,
params: &Self::Params,
rng: &mut dyn Rng,
device: &<B as burn::tensor::backend::BackendTypes>::Device,
) -> Self::State {
debug_assert!(
params.validate().is_ok(),
"invalid EdaParams reached init: {params:?}"
);
let _ = rng;
let model_state = self.model.fit(
¶ms.model,
None,
Tensor::empty([0, 0], device),
Tensor::empty([0], device),
device,
);
EdaState {
model_state,
best_genome: None,
best_fitness_ever: f32::NEG_INFINITY,
generation: 0,
}
}
fn ask(
&self,
params: &Self::Params,
state: &Self::State,
rng: &mut dyn Rng,
device: &<B as burn::tensor::backend::BackendTypes>::Device,
) -> (Self::Genome, Self::State) {
let draw = rng.next_u64();
let mut stream = seed_stream(draw, state.generation as u64, SeedPurpose::EdaSampling);
let mut pop = self
.model
.sample(&state.model_state, params.pop_size, &mut stream, device);
if let Some(b) = params.bounds {
pop = pop.clamp(b.lo(), b.hi());
}
(pop, state.clone())
}
fn tell(
&self,
params: &Self::Params,
population: Self::Genome,
fitness: Tensor<B, 1>,
mut state: Self::State,
_rng: &mut dyn Rng,
) -> (Self::State, StrategyMetrics) {
let raw = fitness
.into_data()
.into_vec::<f32>()
.expect("fitness tensor must be readable as f32");
let sanitized: Vec<f32> = raw
.iter()
.map(|&f| crate::fitness::sanitize_fitness(f))
.collect();
let n = sanitized.len();
let device = population.device();
let mut best_idx = 0_usize;
let mut best_f = f32::NEG_INFINITY;
for (i, &f) in sanitized.iter().enumerate() {
if f.total_cmp(&best_f) == std::cmp::Ordering::Greater {
best_f = f;
best_idx = i;
}
}
if best_f > state.best_fitness_ever {
#[allow(clippy::cast_possible_wrap)]
let idx = Tensor::<B, 1, Int>::from_data(
TensorData::new(vec![best_idx as i64], [1]),
&device,
);
let row = population.clone().select(0, idx).squeeze_dim::<1>(0);
state.best_genome = Some(row);
}
#[allow(clippy::cast_precision_loss)]
let target = (params.selection_ratio * params.pop_size as f32).ceil();
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
let k = (target as usize).max(2).min(n.max(1));
let mut order: Vec<usize> = (0..n).collect();
order.sort_by(|&a, &b| sanitized[b].total_cmp(&sanitized[a]).then(a.cmp(&b)));
order.truncate(k);
#[allow(clippy::cast_possible_wrap)]
let idx_vec: Vec<i64> = order.iter().map(|&i| i as i64).collect();
let idx = Tensor::<B, 1, Int>::from_data(TensorData::new(idx_vec, [k]), &device);
let selected = population.clone().select(0, idx);
let nan_mask = selected.clone().is_nan();
let selected = selected
.mask_fill(nan_mask, 0.0_f32)
.clamp(-f32::MAX, f32::MAX);
let selected_fitness_host: Vec<f32> = order.iter().map(|&i| sanitized[i]).collect();
let selected_fitness =
Tensor::<B, 1>::from_data(TensorData::new(selected_fitness_host, [k]), &device);
let model_state = self.model.fit(
¶ms.model,
Some(&state.model_state),
selected,
selected_fitness,
&device,
);
state.model_state = model_state;
state.generation += 1;
let metrics = StrategyMetrics::from_host_fitness(
state.generation,
&sanitized,
state.best_fitness_ever,
);
state.best_fitness_ever = metrics.best_fitness_ever();
(state, metrics)
}
fn best(&self, state: &Self::State) -> Option<(Self::Genome, f32)> {
state
.best_genome
.as_ref()
.map(|g| (g.clone().unsqueeze::<2>(), state.best_fitness_ever))
}
}
#[cfg(test)]
mod tests {
use super::*;
use burn::backend::Flex;
use rand::SeedableRng;
use rand::rngs::StdRng;
type TestBackend = Flex;
fn make_pop(rows: &[f32], n: usize, d: usize) -> Tensor<TestBackend, 2> {
let device = Default::default();
Tensor::<TestBackend, 2>::from_data(TensorData::new(rows.to_vec(), [n, d]), &device)
}
fn make_fitness(values: &[f32]) -> Tensor<TestBackend, 1> {
let device = Default::default();
let n = values.len();
Tensor::<TestBackend, 1>::from_data(TensorData::new(values.to_vec(), [n]), &device)
}
fn params(pop_size: usize, ratio: f32, dim: usize) -> EdaParams<UnivariateGaussianParams> {
EdaParams {
pop_size,
selection_ratio: ratio,
bounds: None,
model: UnivariateGaussianParams::default_for(dim),
}
}
#[test]
fn best_is_none_before_tell_some_after() {
let device = Default::default();
let strategy = EdaStrategy::<TestBackend, _>::new(UnivariateGaussian);
let p = params(4, 0.5, 2);
let mut rng = StdRng::seed_from_u64(0);
let state = strategy.init(&p, &mut rng, &device);
assert!(strategy.best(&state).is_none());
let pop = make_pop(&[0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0], 4, 2);
let fitness = make_fitness(&[5.0, 1.0, 9.0, 7.0]);
let (state, _m) = strategy.tell(&p, pop, fitness, state, &mut rng);
let best = strategy.best(&state);
assert!(best.is_some());
let (genome, f) = best.unwrap();
assert_eq!(genome.dims(), [1, 2]);
approx::assert_relative_eq!(f, 9.0, epsilon = 1e-6);
}
#[test]
fn k_is_clamped_to_at_least_two_and_at_most_n() {
let device = Default::default();
let strategy = EdaStrategy::<TestBackend, _>::new(UnivariateGaussian);
let mut rng = StdRng::seed_from_u64(0);
let p = params(5, 0.1, 1);
let state = strategy.init(&p, &mut rng, &device);
let pop = make_pop(&[5.0, 4.0, 3.0, 2.0, 1.0], 5, 1);
let fitness = make_fitness(&[5.0, 4.0, 3.0, 2.0, 1.0]);
let (state, _m) = strategy.tell(&p, pop, fitness, state, &mut rng);
assert_eq!(state.generation, 1);
let p2 = params(3, 0.99, 1);
let state2 = strategy.init(&p2, &mut rng, &device);
let pop2 = make_pop(&[1.0, 2.0, 3.0], 3, 1);
let fitness2 = make_fitness(&[1.0, 2.0, 3.0]);
let (state2, _m) = strategy.tell(&p2, pop2, fitness2, state2, &mut rng);
assert_eq!(state2.generation, 1);
}
#[test]
fn tie_breaking_prefers_lowest_index() {
let device = Default::default();
let strategy = EdaStrategy::<TestBackend, _>::new(UnivariateGaussian);
let mut rng = StdRng::seed_from_u64(0);
let p = params(4, 0.5, 1);
let state = strategy.init(&p, &mut rng, &device);
let pop = make_pop(&[0.0, 10.0, 20.0, 30.0], 4, 1);
let fitness = make_fitness(&[1.0, 5.0, 1.0, 5.0]);
let (state, _m) = strategy.tell(&p, pop, fitness, state, &mut rng);
let (genome, _f) = strategy.best(&state).unwrap();
let v = genome
.into_data()
.into_vec::<f32>()
.expect("genome host-read of a tensor this test just built");
approx::assert_relative_eq!(v[0], 10.0, epsilon = 1e-6);
}
#[test]
fn nan_fitness_never_becomes_best_and_does_not_break_ordering() {
let device = Default::default();
let strategy = EdaStrategy::<TestBackend, _>::new(UnivariateGaussian);
let mut rng = StdRng::seed_from_u64(0);
let p = params(4, 0.5, 1);
let state = strategy.init(&p, &mut rng, &device);
let pop = make_pop(&[0.0, 1.0, 2.0, 3.0], 4, 1);
let fitness = make_fitness(&[f32::NAN, 9.0, 2.0, 7.0]);
let (state, m) = strategy.tell(&p, pop, fitness, state, &mut rng);
let (genome, f) = strategy.best(&state).unwrap();
let v = genome
.into_data()
.into_vec::<f32>()
.expect("genome host-read of a tensor this test just built");
approx::assert_relative_eq!(v[0], 1.0, epsilon = 1e-6);
approx::assert_relative_eq!(f, 9.0, epsilon = 1e-6);
assert!(m.best_fitness().is_finite());
}
#[test]
fn nonfinite_genome_sanitized_before_fit_yields_finite_samples() {
let device = Default::default();
let strategy = EdaStrategy::<TestBackend, _>::new(UnivariateGaussian);
let mut rng = StdRng::seed_from_u64(1);
let p = params(4, 0.75, 2);
let state = strategy.init(&p, &mut rng, &device);
let pop = make_pop(
&[
f32::NAN,
0.0, f32::INFINITY,
1.0, f32::NEG_INFINITY,
2.0, 0.5,
3.0,
],
4,
2,
);
let fitness = make_fitness(&[1.0, 2.0, 3.0, 4.0]);
let (state, _m) = strategy.tell(&p, pop, fitness, state, &mut rng);
for &m in state.model_state.mean() {
assert!(m.is_finite(), "fitted mean must be finite, got {m}");
}
for &v in state.model_state.variance() {
assert!(
v.is_finite() && v > 0.0,
"fitted variance must be finite/positive, got {v}"
);
}
let (next_pop, _s) = strategy.ask(&p, &state, &mut rng, &device);
for x in next_pop
.into_data()
.into_vec::<f32>()
.expect("population host-read of a tensor this test just built")
{
assert!(x.is_finite(), "sampled genome must be finite, got {x}");
}
}
#[test]
fn bounds_clamp_is_applied_in_ask() {
let device = Default::default();
let strategy = EdaStrategy::<TestBackend, _>::new(UnivariateGaussian);
let mut rng = StdRng::seed_from_u64(7);
let p = EdaParams {
pop_size: 64,
selection_ratio: 0.5,
bounds: Some(Bounds::new(-0.5, 0.5)),
model: UnivariateGaussianParams {
genome_dim: 3,
init_mean: 0.0,
init_std: 5.0,
min_variance: 1e-6,
},
};
let state = strategy.init(&p, &mut rng, &device);
let (pop, _s) = strategy.ask(&p, &state, &mut rng, &device);
let values = pop
.into_data()
.into_vec::<f32>()
.expect("population host-read of a tensor this test just built");
for v in values {
assert!((-0.5..=0.5).contains(&v), "value {v} escaped the clamp");
}
}
#[cfg(debug_assertions)]
#[test]
#[should_panic(expected = "selection_ratio")]
fn selection_ratio_zero_panics() {
let device = Default::default();
let strategy = EdaStrategy::<TestBackend, _>::new(UnivariateGaussian);
let mut rng = StdRng::seed_from_u64(0);
let p = params(4, 0.0, 1);
let _ = strategy.init(&p, &mut rng, &device);
}
#[cfg(debug_assertions)]
#[test]
#[should_panic(expected = "selection_ratio")]
fn selection_ratio_one_panics() {
let device = Default::default();
let strategy = EdaStrategy::<TestBackend, _>::new(UnivariateGaussian);
let mut rng = StdRng::seed_from_u64(0);
let p = params(4, 1.0, 1);
let _ = strategy.init(&p, &mut rng, &device);
}
}