From 705687a42e31c05639955133da15ff50b861ae58 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 23:53:20 +0100 Subject: [PATCH] feat: add NSGA-III and MOEA/D samplers for many-objective optimization Extract shared evolutionary algorithm infrastructure (genetic operators, candidate management, Das-Dennis reference points) from NSGA-II into a new genetic.rs module, then build two new multi-objective samplers on top: - NSGA-III: reference-point-based niching for well-distributed fronts on 3+ objective problems (Das-Dennis structured points, normalization, perpendicular distance association, niching selection) - MOEA/D: decomposition-based optimization with three scalarization methods (Tchebycheff, WeightedSum, PBI), weight-vector neighborhoods, and neighborhood-based mating selection Both implement MultiObjectiveSampler with builder pattern, seeded RNG, and SBX crossover / polynomial mutation via the shared genetic module. --- src/lib.rs | 6 + src/sampler/genetic.rs | 572 ++++++++++++++++++++++++++ src/sampler/mod.rs | 3 + src/sampler/moead.rs | 603 +++++++++++++++++++++++++++ src/sampler/nsga2.rs | 485 +++------------------- src/sampler/nsga3.rs | 720 +++++++++++++++++++++++++++++++++ tests/multi_objective_tests.rs | 347 +++++++++++++++- 7 files changed, 2296 insertions(+), 440 deletions(-) create mode 100644 src/sampler/genetic.rs create mode 100644 src/sampler/moead.rs create mode 100644 src/sampler/nsga3.rs diff --git a/src/lib.rs b/src/lib.rs index 35fe5c3..8ff91cb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,6 +23,8 @@ //! - **GP** - Gaussian Process Bayesian optimization with Expected Improvement (requires `gp` feature) //! - **BOHB** - Bayesian Optimization + `HyperBand` for budget-aware TPE sampling //! - **NSGA-II** - Non-dominated Sorting Genetic Algorithm II for multi-objective optimization +//! - **NSGA-III** - Reference-point-based NSGA for many-objective (3+) optimization +//! - **MOEA/D** - Decomposition-based multi-objective with Tchebycheff, Weighted Sum, or PBI //! - **MOTPE** - Multi-Objective Tree-Parzen Estimator for Bayesian multi-objective optimization //! //! Additional features include: @@ -261,8 +263,10 @@ pub use sampler::differential_evolution::{ #[cfg(feature = "gp")] pub use sampler::gp::GpSampler; pub use sampler::grid::GridSearchSampler; +pub use sampler::moead::{Decomposition, MoeadSampler}; pub use sampler::motpe::MotpeSampler; pub use sampler::nsga2::Nsga2Sampler; +pub use sampler::nsga3::Nsga3Sampler; pub use sampler::random::RandomSampler; #[cfg(feature = "sobol")] pub use sampler::sobol::SobolSampler; @@ -309,8 +313,10 @@ pub mod prelude { #[cfg(feature = "gp")] pub use crate::sampler::gp::GpSampler; pub use crate::sampler::grid::GridSearchSampler; + pub use crate::sampler::moead::{Decomposition, MoeadSampler}; pub use crate::sampler::motpe::MotpeSampler; pub use crate::sampler::nsga2::Nsga2Sampler; + pub use crate::sampler::nsga3::Nsga3Sampler; pub use crate::sampler::random::RandomSampler; #[cfg(feature = "sobol")] pub use crate::sampler::sobol::SobolSampler; diff --git a/src/sampler/genetic.rs b/src/sampler/genetic.rs new file mode 100644 index 0000000..d2a539a --- /dev/null +++ b/src/sampler/genetic.rs @@ -0,0 +1,572 @@ +//! Shared types and genetic operators for evolutionary multi-objective samplers. +//! +//! This module extracts common functionality used by NSGA-II, NSGA-III, and MOEA/D: +//! candidate management, discovery/active phase logic, SBX crossover, +//! polynomial mutation, and Das-Dennis reference point generation. + +use std::collections::HashMap; + +use crate::distribution::Distribution; +use crate::multi_objective::MultiObjectiveTrial; +use crate::param::ParamValue; +use crate::rng_util; + +/// Describes a parameter dimension discovered during the first trial. +#[derive(Clone, Debug)] +pub(crate) struct DimensionInfo { + pub distribution: Distribution, +} + +/// A candidate solution: one value per dimension. +#[derive(Clone, Debug)] +pub(crate) struct Candidate { + pub params: Vec, +} + +/// Tracks per-trial sampling progress (which candidate, which dimension next). +#[derive(Clone, Debug)] +pub(crate) struct TrialProgress { + pub candidate_idx: usize, + pub next_dim: usize, +} + +/// Phase of an evolutionary sampler. +pub(crate) enum Phase { + /// First trial reveals parameter dimensions. + Discovery, + /// Evolutionary optimisation. + Active, +} + +/// Common state shared by all evolutionary multi-objective samplers. +pub(crate) struct EvolutionaryState { + pub rng: fastrand::Rng, + pub phase: Phase, + pub dimensions: Vec, + pub population_size: usize, + pub candidates: Vec, + pub trial_progress: HashMap, + pub assigned_count: usize, + pub generation_trial_ids: Vec, + pub discovery_trial_id: Option, + pub generation: usize, +} + +impl EvolutionaryState { + pub(crate) fn new(seed: Option) -> Self { + let rng = seed.map_or_else(fastrand::Rng::new, fastrand::Rng::with_seed); + Self { + rng, + phase: Phase::Discovery, + dimensions: Vec::new(), + population_size: 4, + candidates: Vec::new(), + trial_progress: HashMap::new(), + assigned_count: 0, + generation_trial_ids: Vec::new(), + discovery_trial_id: None, + generation: 0, + } + } +} + +// --------------------------------------------------------------------------- +// Discovery phase helpers +// --------------------------------------------------------------------------- + +/// Handle sampling during the discovery phase. +/// +/// Returns `Some(value)` if the discovery phase handled the sample, +/// or `None` if it transitioned to active phase and the caller should +/// generate candidates and sample from them. +pub(crate) fn sample_discovery( + evo: &mut EvolutionaryState, + distribution: &Distribution, + trial_id: u64, +) -> Option { + if let Some(prev_id) = evo.discovery_trial_id + && trial_id != prev_id + { + // A new trial arrived — transition to active phase + return None; + } + + evo.discovery_trial_id = Some(trial_id); + evo.dimensions.push(DimensionInfo { + distribution: distribution.clone(), + }); + + Some(sample_random(&mut evo.rng, distribution)) +} + +/// Compute population size from dimensions and optional user override. +#[allow( + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss +)] +pub(crate) fn compute_population_size( + n_dims: usize, + user_pop_size: Option, + minimum: usize, +) -> usize { + user_pop_size + .unwrap_or_else(|| (4.0 + 3.0 * (n_dims as f64).ln().max(0.0)).floor() as usize) + .max(minimum) +} + +/// Transition from discovery to active phase. +pub(crate) fn finalize_discovery(evo: &mut EvolutionaryState, user_pop_size: Option) { + evo.population_size = compute_population_size(evo.dimensions.len(), user_pop_size, 4); + evo.phase = Phase::Active; +} + +/// Generate `population_size` random candidates. +pub(crate) fn generate_random_candidates(evo: &mut EvolutionaryState) { + let pop = evo.population_size; + evo.candidates = (0..pop) + .map(|_| { + let params: Vec = evo + .dimensions + .iter() + .map(|d| sample_random(&mut evo.rng, &d.distribution)) + .collect(); + Candidate { params } + }) + .collect(); + evo.assigned_count = 0; + evo.generation_trial_ids.clear(); + evo.trial_progress.clear(); +} + +/// Assign a candidate to a trial and return the next dimension value. +pub(crate) fn sample_from_candidate(evo: &mut EvolutionaryState, trial_id: u64) -> ParamValue { + if !evo.trial_progress.contains_key(&trial_id) { + let candidate_idx = if evo.assigned_count < evo.candidates.len() { + let idx = evo.assigned_count; + evo.assigned_count += 1; + idx + } else { + // Overflow: generate a random candidate + let params: Vec = evo + .dimensions + .iter() + .map(|d| sample_random(&mut evo.rng, &d.distribution)) + .collect(); + evo.candidates.push(Candidate { params }); + let idx = evo.candidates.len() - 1; + evo.assigned_count = evo.candidates.len(); + idx + }; + + evo.trial_progress.insert( + trial_id, + TrialProgress { + candidate_idx, + next_dim: 0, + }, + ); + evo.generation_trial_ids.push(trial_id); + } + + let progress = evo.trial_progress.get_mut(&trial_id).unwrap(); + let dim_idx = progress.next_dim; + progress.next_dim += 1; + + if dim_idx >= evo.dimensions.len() { + return sample_random(&mut evo.rng, &evo.dimensions.last().unwrap().distribution); + } + + evo.candidates[progress.candidate_idx].params[dim_idx].clone() +} + +/// Extract parameter values from a trial, ordered by dimension index. +pub(crate) fn extract_trial_params( + trial: &MultiObjectiveTrial, + dimensions: &[DimensionInfo], + rng: &mut fastrand::Rng, +) -> Vec { + let mut param_pairs: Vec<_> = trial.params.iter().collect(); + param_pairs.sort_by_key(|(id, _)| *id); + + dimensions + .iter() + .enumerate() + .map(|(dim_idx, dim_info)| { + if dim_idx < param_pairs.len() { + param_pairs[dim_idx].1.clone() + } else { + sample_random(rng, &dim_info.distribution) + } + }) + .collect() +} + +/// Install new offspring as the next generation's candidates. +pub(crate) fn advance_generation(evo: &mut EvolutionaryState, offspring: Vec) { + evo.candidates = offspring; + evo.assigned_count = 0; + evo.generation_trial_ids.clear(); + evo.trial_progress.clear(); + evo.generation += 1; +} + +/// Check if the current generation is fully evaluated and return the +/// evaluated trials if so. +pub(crate) fn collect_evaluated_generation<'a>( + evo: &EvolutionaryState, + history: &'a [MultiObjectiveTrial], +) -> Option> { + let pop_size = evo.population_size; + + if evo.generation_trial_ids.len() < pop_size { + return None; + } + + let gen_ids: Vec = evo + .generation_trial_ids + .iter() + .take(pop_size) + .copied() + .collect(); + let history_map: HashMap = + history.iter().map(|t| (t.id, t)).collect(); + + if !gen_ids.iter().all(|id| history_map.contains_key(id)) { + return None; + } + + Some( + gen_ids + .iter() + .filter_map(|id| history_map.get(id).copied()) + .collect(), + ) +} + +// --------------------------------------------------------------------------- +// Genetic operators +// --------------------------------------------------------------------------- + +/// SBX crossover for continuous params, uniform crossover for categorical. +pub(crate) fn crossover( + rng: &mut fastrand::Rng, + parent1: &[ParamValue], + parent2: &[ParamValue], + dimensions: &[DimensionInfo], + crossover_prob: f64, + eta: f64, +) -> (Vec, Vec) { + let n = parent1.len(); + let mut child1 = parent1.to_vec(); + let mut child2 = parent2.to_vec(); + + let u: f64 = rng_util::f64_range(rng, 0.0, 1.0); + if u > crossover_prob { + return (child1, child2); + } + + for i in 0..n { + match (&parent1[i], &parent2[i], &dimensions[i].distribution) { + (ParamValue::Float(p1), ParamValue::Float(p2), Distribution::Float(d)) => { + if (p1 - p2).abs() < 1e-14 { + continue; + } + let (c1, c2) = sbx_crossover_f64(rng, *p1, *p2, d.low, d.high, eta); + child1[i] = ParamValue::Float(c1); + child2[i] = ParamValue::Float(c2); + } + (ParamValue::Int(p1), ParamValue::Int(p2), Distribution::Int(d)) => { + if p1 == p2 { + continue; + } + #[allow(clippy::cast_precision_loss)] + let (c1, c2) = sbx_crossover_f64( + rng, + *p1 as f64, + *p2 as f64, + d.low as f64, + d.high as f64, + eta, + ); + #[allow(clippy::cast_possible_truncation)] + { + child1[i] = ParamValue::Int((c1.round() as i64).clamp(d.low, d.high)); + child2[i] = ParamValue::Int((c2.round() as i64).clamp(d.low, d.high)); + } + } + (ParamValue::Categorical(_), ParamValue::Categorical(_), _) => { + if rng_util::f64_range(rng, 0.0, 1.0) < 0.5 { + core::mem::swap(&mut child1[i], &mut child2[i]); + } + } + _ => {} + } + } + + (child1, child2) +} + +/// SBX crossover for a single float dimension. +pub(crate) fn sbx_crossover_f64( + rng: &mut fastrand::Rng, + p1: f64, + p2: f64, + low: f64, + high: f64, + eta: f64, +) -> (f64, f64) { + let u: f64 = rng_util::f64_range(rng, 0.0, 1.0); + + let beta = if u <= 0.5 { + (2.0 * u).powf(1.0 / (eta + 1.0)) + } else { + (1.0 / (2.0 * (1.0 - u))).powf(1.0 / (eta + 1.0)) + }; + + let c1 = 0.5 * ((1.0 + beta) * p1 + (1.0 - beta) * p2); + let c2 = 0.5 * ((1.0 - beta) * p1 + (1.0 + beta) * p2); + + (c1.clamp(low, high), c2.clamp(low, high)) +} + +/// Polynomial mutation for each dimension. +#[allow(clippy::cast_precision_loss)] +pub(crate) fn mutate( + rng: &mut fastrand::Rng, + individual: &mut [ParamValue], + dimensions: &[DimensionInfo], + eta: f64, +) { + let n = individual.len(); + if n == 0 { + return; + } + let mutation_prob = 1.0 / n as f64; + + for (i, value) in individual.iter_mut().enumerate() { + if rng_util::f64_range(rng, 0.0, 1.0) >= mutation_prob { + continue; + } + + match (value, &dimensions[i].distribution) { + (v @ ParamValue::Float(_), Distribution::Float(d)) => { + let ParamValue::Float(x) = *v else { + unreachable!(); + }; + let mutated = polynomial_mutation_f64(rng, x, d.low, d.high, eta); + *v = ParamValue::Float(mutated); + } + (v @ ParamValue::Int(_), Distribution::Int(d)) => { + let ParamValue::Int(x) = *v else { + unreachable!(); + }; + #[allow(clippy::cast_possible_truncation)] + { + let mutated = + polynomial_mutation_f64(rng, x as f64, d.low as f64, d.high as f64, eta); + *v = ParamValue::Int((mutated.round() as i64).clamp(d.low, d.high)); + } + } + (v @ ParamValue::Categorical(_), Distribution::Categorical(d)) => { + *v = ParamValue::Categorical(rng.usize(0..d.n_choices)); + } + _ => {} + } + } +} + +/// Polynomial mutation for a single float value. +pub(crate) fn polynomial_mutation_f64( + rng: &mut fastrand::Rng, + x: f64, + low: f64, + high: f64, + eta: f64, +) -> f64 { + let u: f64 = rng_util::f64_range(rng, 0.0, 1.0); + let range = high - low; + if range <= 0.0 { + return x; + } + + let delta1 = (x - low) / range; + let delta2 = (high - x) / range; + + let delta_q = if u < 0.5 { + let xy = 1.0 - delta1; + let val = 2.0 * u + (1.0 - 2.0 * u) * xy.powf(eta + 1.0); + val.powf(1.0 / (eta + 1.0)) - 1.0 + } else { + let xy = 1.0 - delta2; + let val = 2.0 * (1.0 - u) + 2.0 * (u - 0.5) * xy.powf(eta + 1.0); + 1.0 - val.powf(1.0 / (eta + 1.0)) + }; + + (x + delta_q * range).clamp(low, high) +} + +/// Random sampling for a single distribution. +#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] +pub(crate) fn sample_random(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue { + match distribution { + Distribution::Float(d) => { + let value = if d.log_scale { + let log_low = d.low.ln(); + let log_high = d.high.ln(); + rng_util::f64_range(rng, log_low, log_high).exp() + } else if let Some(step) = d.step { + let n_steps = ((d.high - d.low) / step).floor() as i64; + let k = rng.i64(0..=n_steps); + d.low + (k as f64) * step + } else { + rng_util::f64_range(rng, d.low, d.high) + }; + ParamValue::Float(value) + } + Distribution::Int(d) => { + let value = if d.log_scale { + let log_low = (d.low as f64).ln(); + let log_high = (d.high as f64).ln(); + let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; + raw.clamp(d.low, d.high) + } else if let Some(step) = d.step { + let n_steps = (d.high - d.low) / step; + let k = rng.i64(0..=n_steps); + d.low + k * step + } else { + rng.i64(d.low..=d.high) + }; + ParamValue::Int(value) + } + Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), + } +} + +// --------------------------------------------------------------------------- +// Das-Dennis reference point generation +// --------------------------------------------------------------------------- + +/// Generate Das-Dennis (simplex-lattice) reference points. +/// +/// Returns `C(H + M - 1, M - 1)` uniformly spaced points on the +/// `M`-dimensional unit simplex, where `M = n_objectives` and +/// `H = divisions`. +pub(crate) fn das_dennis(n_objectives: usize, divisions: usize) -> Vec> { + let mut points = Vec::new(); + let mut point = vec![0.0_f64; n_objectives]; + das_dennis_recursive( + n_objectives, + divisions, + 0, + divisions, + &mut point, + &mut points, + ); + points +} + +#[allow(clippy::cast_precision_loss)] +fn das_dennis_recursive( + n_objectives: usize, + divisions: usize, + depth: usize, + remaining: usize, + current: &mut Vec, + result: &mut Vec>, +) { + if depth == n_objectives - 1 { + current[depth] = remaining as f64 / divisions as f64; + result.push(current.clone()); + return; + } + + for i in 0..=remaining { + current[depth] = i as f64 / divisions as f64; + das_dennis_recursive( + n_objectives, + divisions, + depth + 1, + remaining - i, + current, + result, + ); + } +} + +/// Choose the number of divisions for Das-Dennis to get close to a target +/// population size. +/// +/// The number of reference points is `C(H + M - 1, M - 1)`. This function +/// finds the smallest `H` such that the number of points >= `target_pop`. +pub(crate) fn auto_divisions(n_objectives: usize, target_pop: usize) -> usize { + let m = n_objectives; + for h in 1..200 { + let n_points = n_combinations(h + m - 1, m - 1); + if n_points >= target_pop { + return h; + } + } + 12 +} + +/// Compute `C(n, k)` = n! / (k! * (n-k)!). +fn n_combinations(n: usize, k: usize) -> usize { + if k > n { + return 0; + } + let k = k.min(n - k); + let mut result: usize = 1; + for i in 0..k { + result = result.saturating_mul(n - i) / (i + 1); + } + result +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_das_dennis_2d() { + let points = das_dennis(2, 4); + // C(4+1, 1) = 5 points + assert_eq!(points.len(), 5); + for p in &points { + let sum: f64 = p.iter().sum(); + assert!((sum - 1.0).abs() < 1e-10, "point {p:?} doesn't sum to 1"); + } + } + + #[test] + fn test_das_dennis_3d() { + let points = das_dennis(3, 4); + // C(4+2, 2) = 15 points + assert_eq!(points.len(), 15); + for p in &points { + let sum: f64 = p.iter().sum(); + assert!((sum - 1.0).abs() < 1e-10); + } + } + + #[test] + fn test_auto_divisions() { + // For 2 objectives targeting 10 points: H=9 gives C(10,1)=10 + let h = auto_divisions(2, 10); + let n = n_combinations(h + 1, 1); + assert!(n >= 10); + + // For 3 objectives targeting ~91 points: H=12 gives C(14,2)=91 + let h3 = auto_divisions(3, 91); + let n3 = n_combinations(h3 + 2, 2); + assert!(n3 >= 91); + } + + #[test] + fn test_n_combinations() { + assert_eq!(n_combinations(5, 2), 10); + assert_eq!(n_combinations(4, 0), 1); + assert_eq!(n_combinations(4, 4), 1); + assert_eq!(n_combinations(6, 3), 20); + } +} diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index b6c78fd..e1813e0 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -4,11 +4,14 @@ pub mod bohb; #[cfg(feature = "cma-es")] pub mod cma_es; pub mod differential_evolution; +pub(crate) mod genetic; #[cfg(feature = "gp")] pub mod gp; pub mod grid; +pub mod moead; pub mod motpe; pub mod nsga2; +pub mod nsga3; pub mod random; #[cfg(feature = "sobol")] pub mod sobol; diff --git a/src/sampler/moead.rs b/src/sampler/moead.rs new file mode 100644 index 0000000..d6d3997 --- /dev/null +++ b/src/sampler/moead.rs @@ -0,0 +1,603 @@ +//! MOEA/D (Multi-Objective Evolutionary Algorithm based on Decomposition) sampler. +//! +//! Decomposes a multi-objective problem into scalar subproblems using +//! weight vectors and solves them collaboratively. Supports Weighted Sum, +//! Tchebycheff, and Penalty-based Boundary Intersection (PBI) scalarization. +//! +//! # Examples +//! +//! ``` +//! use optimizer::Direction; +//! use optimizer::multi_objective::MultiObjectiveStudy; +//! use optimizer::parameter::{FloatParam, Parameter}; +//! use optimizer::sampler::moead::MoeadSampler; +//! +//! let sampler = MoeadSampler::with_seed(42); +//! let study = +//! MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); +//! +//! let x = FloatParam::new(0.0, 1.0); +//! study +//! .optimize(100, |trial| { +//! let xv = x.suggest(trial)?; +//! Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) +//! }) +//! .unwrap(); +//! ``` + +use parking_lot::Mutex; + +use super::genetic::{ + self, Candidate, EvolutionaryState, Phase, advance_generation, auto_divisions, + collect_evaluated_generation, crossover, das_dennis, extract_trial_params, + generate_random_candidates, mutate, sample_from_candidate, sample_random, +}; +use crate::distribution::Distribution; +use crate::multi_objective::MultiObjectiveTrial; +use crate::param::ParamValue; +use crate::types::Direction; + +/// Decomposition (scalarization) method for MOEA/D. +#[derive(Debug, Clone, Default)] +pub enum Decomposition { + /// Weighted sum: `sum(w_i * f_i)`. + WeightedSum, + /// Tchebycheff: `max(w_i * |f_i - z_i*|)`. + #[default] + Tchebycheff, + /// Penalty-based Boundary Intersection with parameter theta. + Pbi { + /// Penalty parameter controlling the balance between convergence + /// and diversity. Default: 5.0. + theta: f64, + }, +} + +/// MOEA/D sampler for multi-objective optimization. +/// +/// Decomposes the multi-objective problem into scalar subproblems +/// using weight vectors, solving them collaboratively via +/// neighborhood-based mating and replacement. +pub struct MoeadSampler { + state: Mutex, +} + +impl MoeadSampler { + /// Creates a new MOEA/D sampler with a random seed. + #[must_use] + pub fn new() -> Self { + Self { + state: Mutex::new(MoeadState::new(MoeadConfig::default(), None)), + } + } + + /// Creates a new MOEA/D sampler with a fixed seed. + #[must_use] + pub fn with_seed(seed: u64) -> Self { + Self { + state: Mutex::new(MoeadState::new(MoeadConfig::default(), Some(seed))), + } + } + + /// Creates a builder for configuring a `MoeadSampler`. + #[must_use] + pub fn builder() -> MoeadSamplerBuilder { + MoeadSamplerBuilder::default() + } +} + +impl Default for MoeadSampler { + fn default() -> Self { + Self::new() + } +} + +/// Builder for [`MoeadSampler`]. +#[derive(Debug, Clone, Default)] +pub struct MoeadSamplerBuilder { + population_size: Option, + neighborhood_size: Option, + decomposition: Decomposition, + crossover_prob: Option, + crossover_eta: Option, + mutation_eta: Option, + seed: Option, +} + +impl MoeadSamplerBuilder { + /// Sets the population size. If unset, equals the number of + /// Das-Dennis weight vectors. + #[must_use] + pub fn population_size(mut self, size: usize) -> Self { + self.population_size = Some(size); + self + } + + /// Sets the neighborhood size (T). Default: `min(20, pop_size)`. + #[must_use] + pub fn neighborhood_size(mut self, size: usize) -> Self { + self.neighborhood_size = Some(size); + self + } + + /// Sets the decomposition method. Default: Tchebycheff. + #[must_use] + pub fn decomposition(mut self, decomp: Decomposition) -> Self { + self.decomposition = decomp; + self + } + + /// Sets the crossover probability. Default: 1.0. + #[must_use] + pub fn crossover_prob(mut self, prob: f64) -> Self { + self.crossover_prob = Some(prob); + self + } + + /// Sets the SBX distribution index. Default: 20.0. + #[must_use] + pub fn crossover_eta(mut self, eta: f64) -> Self { + self.crossover_eta = Some(eta); + self + } + + /// Sets the polynomial mutation distribution index. Default: 20.0. + #[must_use] + pub fn mutation_eta(mut self, eta: f64) -> Self { + self.mutation_eta = Some(eta); + self + } + + /// Sets the random seed for reproducibility. + #[must_use] + pub fn seed(mut self, seed: u64) -> Self { + self.seed = Some(seed); + self + } + + /// Builds the configured [`MoeadSampler`]. + #[must_use] + pub fn build(self) -> MoeadSampler { + let config = MoeadConfig { + user_population_size: self.population_size, + neighborhood_size: self.neighborhood_size, + decomposition: self.decomposition, + crossover_prob: self.crossover_prob.unwrap_or(1.0), + crossover_eta: self.crossover_eta.unwrap_or(20.0), + mutation_eta: self.mutation_eta.unwrap_or(20.0), + }; + MoeadSampler { + state: Mutex::new(MoeadState::new(config, self.seed)), + } + } +} + +// --------------------------------------------------------------------------- +// Internal types +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone)] +struct MoeadConfig { + user_population_size: Option, + neighborhood_size: Option, + decomposition: Decomposition, + crossover_prob: f64, + crossover_eta: f64, + mutation_eta: f64, +} + +impl Default for MoeadConfig { + fn default() -> Self { + Self { + user_population_size: None, + neighborhood_size: None, + decomposition: Decomposition::default(), + crossover_prob: 1.0, + crossover_eta: 20.0, + mutation_eta: 20.0, + } + } +} + +struct MoeadState { + evo: EvolutionaryState, + config: MoeadConfig, + /// Weight vectors (Das-Dennis), one per subproblem. + weight_vectors: Vec>, + /// Neighborhoods: for each subproblem, indices of T nearest weight vectors. + neighborhoods: Vec>, + /// Ideal point z* (best per-objective in minimize-space). + ideal_point: Vec, + /// Current population's objective values in minimize-space (one per subproblem). + population_values: Vec>, + /// Current population's parameter vectors (one per subproblem). + population_params: Vec>, + /// Whether the MOEA/D state has been initialized. + initialized: bool, +} + +impl MoeadState { + fn new(config: MoeadConfig, seed: Option) -> Self { + Self { + evo: EvolutionaryState::new(seed), + config, + weight_vectors: Vec::new(), + neighborhoods: Vec::new(), + ideal_point: Vec::new(), + population_values: Vec::new(), + population_params: Vec::new(), + initialized: false, + } + } +} + +// --------------------------------------------------------------------------- +// MultiObjectiveSampler implementation +// --------------------------------------------------------------------------- + +impl crate::multi_objective::MultiObjectiveSampler for MoeadSampler { + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + history: &[MultiObjectiveTrial], + directions: &[Direction], + ) -> ParamValue { + let mut state = self.state.lock(); + + match &state.evo.phase { + Phase::Discovery => { + if let Some(value) = + genetic::sample_discovery(&mut state.evo, distribution, trial_id) + { + return value; + } + // Transitioned to active phase + initialize_moead(&mut state, directions); + generate_random_candidates(&mut state.evo); + sample_from_candidate(&mut state.evo, trial_id) + } + Phase::Active => { + maybe_generate_new_generation(&mut state, history, directions); + sample_from_candidate(&mut state.evo, trial_id) + } + } + } +} + +/// Initialize MOEA/D: weight vectors, neighborhoods, ideal point. +fn initialize_moead(state: &mut MoeadState, directions: &[Direction]) { + let n_obj = directions.len(); + + // Generate weight vectors + let divisions = auto_divisions(n_obj, state.config.user_population_size.unwrap_or(100)); + state.weight_vectors = das_dennis(n_obj, divisions); + + let pop_size = state + .config + .user_population_size + .unwrap_or(state.weight_vectors.len()) + .max(4); + + // Trim or pad weight vectors to match population size + state.weight_vectors.truncate(pop_size); + while state.weight_vectors.len() < pop_size { + // Duplicate random existing weight vectors + let idx = state.evo.rng.usize(0..state.weight_vectors.len()); + let w = state.weight_vectors[idx].clone(); + state.weight_vectors.push(w); + } + + // Compute neighborhoods + let t = state + .config + .neighborhood_size + .unwrap_or_else(|| 20.min(pop_size)); + let t = t.min(pop_size); + state.neighborhoods = compute_neighborhoods(&state.weight_vectors, t); + + state.evo.population_size = pop_size; + state.evo.phase = Phase::Active; + state.ideal_point = vec![f64::INFINITY; n_obj]; + state.initialized = true; +} + +/// Compute T-nearest neighborhoods by Euclidean distance between weight vectors. +fn compute_neighborhoods(weights: &[Vec], t: usize) -> Vec> { + let n = weights.len(); + weights + .iter() + .map(|wi| { + let mut distances: Vec<(usize, f64)> = (0..n) + .map(|j| { + let d: f64 = wi + .iter() + .zip(&weights[j]) + .map(|(&a, &b)| (a - b).powi(2)) + .sum::() + .sqrt(); + (j, d) + }) + .collect(); + distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal)); + distances.into_iter().take(t).map(|(idx, _)| idx).collect() + }) + .collect() +} + +/// Convert values to minimize-space. +fn to_minimize_space(values: &[f64], directions: &[Direction]) -> Vec { + values + .iter() + .zip(directions) + .map(|(&v, d)| match d { + Direction::Minimize => v, + Direction::Maximize => -v, + }) + .collect() +} + +fn maybe_generate_new_generation( + state: &mut MoeadState, + history: &[MultiObjectiveTrial], + directions: &[Direction], +) { + if state.evo.candidates.is_empty() { + generate_random_candidates(&mut state.evo); + return; + } + + if let Some(evaluated) = collect_evaluated_generation(&state.evo, history) { + let offspring = moead_generate_offspring(state, &evaluated, directions); + advance_generation(&mut state.evo, offspring); + } +} + +// --------------------------------------------------------------------------- +// Scalarization functions +// --------------------------------------------------------------------------- + +/// Weighted sum scalarization: `sum(w_i * f_i)`. +fn scalarize_weighted_sum(values: &[f64], weight: &[f64]) -> f64 { + values.iter().zip(weight).map(|(&v, &w)| w * v).sum() +} + +/// Tchebycheff scalarization: `max(w_i * |f_i - z_i*|)`. +fn scalarize_tchebycheff(values: &[f64], weight: &[f64], ideal: &[f64]) -> f64 { + values + .iter() + .zip(weight) + .zip(ideal) + .map(|((&v, &w), &z)| { + let w = if w < 1e-6 { 1e-6 } else { w }; + w * (v - z).abs() + }) + .fold(f64::NEG_INFINITY, f64::max) +} + +/// PBI scalarization: `d1 + theta * d2`. +/// +/// d1 = projection onto weight direction, d2 = perpendicular distance. +fn scalarize_pbi(values: &[f64], weight: &[f64], ideal: &[f64], theta: f64) -> f64 { + let n = values.len(); + + // Direction from ideal to the point + let diff: Vec = values.iter().zip(ideal).map(|(&v, &z)| v - z).collect(); + + // Normalize weight vector + let w_norm: f64 = weight.iter().map(|&w| w * w).sum::().sqrt(); + if w_norm < 1e-30 { + return f64::INFINITY; + } + let w_unit: Vec = weight.iter().map(|&w| w / w_norm).collect(); + + // d1 = projection of diff onto weight direction + let d1: f64 = diff.iter().zip(&w_unit).map(|(&d, &w)| d * w).sum(); + + // d2 = perpendicular distance + let d2_sq: f64 = (0..n) + .map(|i| { + let proj = d1 * w_unit[i]; + (diff[i] - proj).powi(2) + }) + .sum::(); + + d1 + theta * d2_sq.sqrt() +} + +/// Evaluate scalarization for a given decomposition method. +fn scalarize(values: &[f64], weight: &[f64], ideal: &[f64], decomposition: &Decomposition) -> f64 { + match decomposition { + Decomposition::WeightedSum => scalarize_weighted_sum(values, weight), + Decomposition::Tchebycheff => scalarize_tchebycheff(values, weight, ideal), + Decomposition::Pbi { theta } => scalarize_pbi(values, weight, ideal, *theta), + } +} + +// --------------------------------------------------------------------------- +// MOEA/D generation algorithm +// --------------------------------------------------------------------------- + +fn moead_generate_offspring( + state: &mut MoeadState, + population: &[&MultiObjectiveTrial], + directions: &[Direction], +) -> Vec { + let pop_size = state.evo.population_size; + + if population.len() < 2 { + return (0..pop_size) + .map(|_| { + let params = state + .evo + .dimensions + .iter() + .map(|d| sample_random(&mut state.evo.rng, &d.distribution)) + .collect(); + Candidate { params } + }) + .collect(); + } + + // Extract current population parameters and objective values + let current_params: Vec> = population + .iter() + .map(|t| extract_trial_params(t, &state.evo.dimensions, &mut state.evo.rng)) + .collect(); + + let current_values: Vec> = population + .iter() + .map(|t| to_minimize_space(&t.values, directions)) + .collect(); + + // Update ideal point + for vals in ¤t_values { + for (i, &v) in vals.iter().enumerate() { + if i < state.ideal_point.len() && v < state.ideal_point[i] { + state.ideal_point[i] = v; + } + } + } + + // Assign each solution to its best subproblem via scalarization + // and select the best solution for each subproblem as its representative + let n_weights = state.weight_vectors.len(); + let mut best_for_subproblem: Vec = Vec::with_capacity(n_weights); + + for j in 0..n_weights { + let mut best_idx = 0; + let mut best_val = f64::INFINITY; + for (k, vals) in current_values.iter().enumerate() { + let s = scalarize( + vals, + &state.weight_vectors[j], + &state.ideal_point, + &state.config.decomposition, + ); + if s < best_val { + best_val = s; + best_idx = k; + } + } + best_for_subproblem.push(best_idx); + } + + // Store current population state + state.population_values = current_values; + state.population_params = current_params; + + // Generate offspring: for each subproblem, mate from neighborhood + let mut offspring = Vec::with_capacity(pop_size); + + for i in 0..pop_size.min(state.neighborhoods.len()) { + let neighborhood = &state.neighborhoods[i]; + + // Pick two parents from the neighborhood using subproblem assignments + let n1 = neighborhood[state.evo.rng.usize(0..neighborhood.len())]; + let n2 = neighborhood[state.evo.rng.usize(0..neighborhood.len())]; + + let p1_idx = best_for_subproblem[n1 % best_for_subproblem.len()]; + let p2_idx = best_for_subproblem[n2 % best_for_subproblem.len()]; + + let p1 = &state.population_params[p1_idx]; + let p2 = &state.population_params[p2_idx]; + + let (mut child1, _child2) = crossover( + &mut state.evo.rng, + p1, + p2, + &state.evo.dimensions, + state.config.crossover_prob, + state.config.crossover_eta, + ); + + mutate( + &mut state.evo.rng, + &mut child1, + &state.evo.dimensions, + state.config.mutation_eta, + ); + + offspring.push(Candidate { params: child1 }); + } + + // If pop_size > neighborhoods, fill remaining with random neighborhood crossover + while offspring.len() < pop_size { + let i = state.evo.rng.usize(0..state.neighborhoods.len()); + let neighborhood = &state.neighborhoods[i]; + let n1 = neighborhood[state.evo.rng.usize(0..neighborhood.len())]; + let n2 = neighborhood[state.evo.rng.usize(0..neighborhood.len())]; + + let p1_idx = best_for_subproblem[n1 % best_for_subproblem.len()]; + let p2_idx = best_for_subproblem[n2 % best_for_subproblem.len()]; + + let (mut child1, _) = crossover( + &mut state.evo.rng, + &state.population_params[p1_idx], + &state.population_params[p2_idx], + &state.evo.dimensions, + state.config.crossover_prob, + state.config.crossover_eta, + ); + + mutate( + &mut state.evo.rng, + &mut child1, + &state.evo.dimensions, + state.config.mutation_eta, + ); + + offspring.push(Candidate { params: child1 }); + } + + offspring +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_scalarize_weighted_sum() { + let values = [1.0, 2.0, 3.0]; + let weight = [0.5, 0.3, 0.2]; + let result = scalarize_weighted_sum(&values, &weight); + assert!((result - (0.5 + 0.6 + 0.6)).abs() < 1e-10); + } + + #[test] + fn test_scalarize_tchebycheff() { + let values = [3.0, 2.0]; + let weight = [0.5, 0.5]; + let ideal = [1.0, 1.0]; + let result = scalarize_tchebycheff(&values, &weight, &ideal); + // max(0.5 * |3-1|, 0.5 * |2-1|) = max(1.0, 0.5) = 1.0 + assert!((result - 1.0).abs() < 1e-10); + } + + #[test] + fn test_scalarize_pbi() { + let values = [2.0, 2.0]; + let weight = [1.0, 1.0]; + let ideal = [0.0, 0.0]; + let result = scalarize_pbi(&values, &weight, &ideal, 5.0); + // d1 = projection of (2,2) onto (1/√2, 1/√2) = 2*√2 + // d2 = 0 (point is on the weight direction) + let expected_d1 = 2.0 * (2.0_f64).sqrt(); + assert!((result - expected_d1).abs() < 1e-10); + } + + #[test] + fn test_compute_neighborhoods() { + let weights = vec![vec![1.0, 0.0], vec![0.5, 0.5], vec![0.0, 1.0]]; + let neighborhoods = compute_neighborhoods(&weights, 2); + assert_eq!(neighborhoods.len(), 3); + // Each neighborhood should have 2 entries + for n in &neighborhoods { + assert_eq!(n.len(), 2); + } + // First weight [1,0] should be closest to itself and [0.5,0.5] + assert_eq!(neighborhoods[0][0], 0); // itself + assert_eq!(neighborhoods[0][1], 1); // nearest neighbor + } +} diff --git a/src/sampler/nsga2.rs b/src/sampler/nsga2.rs index e779224..acbbab5 100644 --- a/src/sampler/nsga2.rs +++ b/src/sampler/nsga2.rs @@ -24,15 +24,18 @@ //! .unwrap(); //! ``` -use std::collections::HashMap; - use parking_lot::Mutex; +use super::genetic::{ + self, Candidate, EvolutionaryState, Phase, advance_generation, collect_evaluated_generation, + crossover, extract_trial_params, finalize_discovery, generate_random_candidates, mutate, + sample_from_candidate, sample_random, +}; use crate::distribution::Distribution; use crate::multi_objective::MultiObjectiveTrial; use crate::param::ParamValue; +use crate::pareto; use crate::types::Direction; -use crate::{pareto, rng_util}; /// NSGA-II sampler for multi-objective optimization. /// @@ -156,62 +159,16 @@ impl Default for Nsga2Config { } } -/// Describes a parameter dimension. -#[derive(Clone, Debug)] -struct DimensionInfo { - distribution: Distribution, -} - -/// A candidate solution: one value per dimension. -#[derive(Clone, Debug)] -struct Candidate { - params: Vec, -} - -/// Tracks per-trial sampling progress. -#[derive(Clone, Debug)] -struct TrialProgress { - candidate_idx: usize, - next_dim: usize, -} - -enum Phase { - /// First trial reveals parameter dimensions. - Discovery, - /// NSGA-II optimisation. - Active, -} - struct Nsga2State { - rng: fastrand::Rng, + evo: EvolutionaryState, config: Nsga2Config, - phase: Phase, - dimensions: Vec, - population_size: usize, - candidates: Vec, - trial_progress: HashMap, - assigned_count: usize, - generation_trial_ids: Vec, - discovery_trial_id: Option, - /// How many complete generations have been evaluated. - generation: usize, } impl Nsga2State { fn new(config: Nsga2Config, seed: Option) -> Self { - let rng = seed.map_or_else(fastrand::Rng::new, fastrand::Rng::with_seed); Self { - rng, + evo: EvolutionaryState::new(seed), config, - phase: Phase::Discovery, - dimensions: Vec::new(), - population_size: 4, - candidates: Vec::new(), - trial_progress: HashMap::new(), - assigned_count: 0, - generation_trial_ids: Vec::new(), - discovery_trial_id: None, - generation: 0, } } } @@ -230,130 +187,27 @@ impl crate::multi_objective::MultiObjectiveSampler for Nsga2Sampler { ) -> ParamValue { let mut state = self.state.lock(); - match &state.phase { - Phase::Discovery => sample_discovery(&mut state, distribution, trial_id), - Phase::Active => sample_active(&mut state, distribution, trial_id, history, directions), + match &state.evo.phase { + Phase::Discovery => { + if let Some(value) = + genetic::sample_discovery(&mut state.evo, distribution, trial_id) + { + return value; + } + // Transitioned to active phase + let user_pop = state.config.user_population_size; + finalize_discovery(&mut state.evo, user_pop); + generate_random_candidates(&mut state.evo); + sample_from_candidate(&mut state.evo, trial_id) + } + Phase::Active => { + maybe_generate_new_generation(&mut state, history, directions); + sample_from_candidate(&mut state.evo, trial_id) + } } } } -/// Handle sampling during the discovery phase. -fn sample_discovery( - state: &mut Nsga2State, - distribution: &Distribution, - trial_id: u64, -) -> ParamValue { - if let Some(prev_id) = state.discovery_trial_id - && trial_id != prev_id - { - finalize_discovery(state); - // Assign this trial a random candidate (no history yet) - generate_random_candidates(state); - return sample_from_candidate(state, trial_id); - } - - state.discovery_trial_id = Some(trial_id); - state.dimensions.push(DimensionInfo { - distribution: distribution.clone(), - }); - - sample_random(&mut state.rng, distribution) -} - -/// Transition from discovery to active phase. -#[allow( - clippy::cast_precision_loss, - clippy::cast_possible_truncation, - clippy::cast_sign_loss -)] -fn finalize_discovery(state: &mut Nsga2State) { - let n = state.dimensions.len(); - state.population_size = state - .config - .user_population_size - .unwrap_or_else(|| (4.0 + 3.0 * (n as f64).ln().max(0.0)).floor() as usize) - .max(4); - state.phase = Phase::Active; -} - -/// Generate `population_size` random candidates. -fn generate_random_candidates(state: &mut Nsga2State) { - let pop = state.population_size; - state.candidates = (0..pop) - .map(|_| { - let params: Vec = state - .dimensions - .iter() - .map(|d| sample_random(&mut state.rng, &d.distribution)) - .collect(); - Candidate { params } - }) - .collect(); - state.assigned_count = 0; - state.generation_trial_ids.clear(); - state.trial_progress.clear(); -} - -/// Active-phase sampling. -fn sample_active( - state: &mut Nsga2State, - _distribution: &Distribution, - trial_id: u64, - history: &[MultiObjectiveTrial], - directions: &[Direction], -) -> ParamValue { - // Check if we need to generate a new generation - maybe_generate_new_generation(state, history, directions); - - sample_from_candidate(state, trial_id) -} - -/// Assign a candidate to a trial and return the next dimension value. -fn sample_from_candidate(state: &mut Nsga2State, trial_id: u64) -> ParamValue { - // Assign candidate if not yet done - if !state.trial_progress.contains_key(&trial_id) { - let candidate_idx = if state.assigned_count < state.candidates.len() { - let idx = state.assigned_count; - state.assigned_count += 1; - idx - } else { - // Overflow: generate a random candidate - let params: Vec = state - .dimensions - .iter() - .map(|d| sample_random(&mut state.rng, &d.distribution)) - .collect(); - state.candidates.push(Candidate { params }); - let idx = state.candidates.len() - 1; - state.assigned_count = state.candidates.len(); - idx - }; - - state.trial_progress.insert( - trial_id, - TrialProgress { - candidate_idx, - next_dim: 0, - }, - ); - state.generation_trial_ids.push(trial_id); - } - - let progress = state.trial_progress.get_mut(&trial_id).unwrap(); - let dim_idx = progress.next_dim; - progress.next_dim += 1; - - if dim_idx >= state.dimensions.len() { - // Extra dimension: sample randomly - return sample_random( - &mut state.rng, - &state.dimensions.last().unwrap().distribution, - ); - } - - state.candidates[progress.candidate_idx].params[dim_idx].clone() -} - /// Check if all candidates in the current generation have been evaluated; /// if so, run NSGA-II selection and generate offspring. fn maybe_generate_new_generation( @@ -361,45 +215,15 @@ fn maybe_generate_new_generation( history: &[MultiObjectiveTrial], directions: &[Direction], ) { - let pop_size = state.population_size; - - // Need at least pop_size assigned trials - if state.generation_trial_ids.len() < pop_size { - // Not enough candidates assigned yet — check if we need initial candidates - if state.candidates.is_empty() { - generate_random_candidates(state); - } + if state.evo.candidates.is_empty() { + generate_random_candidates(&mut state.evo); return; } - // Check if the first pop_size trials are completed - let gen_ids: Vec = state - .generation_trial_ids - .iter() - .take(pop_size) - .copied() - .collect(); - let history_map: HashMap = - history.iter().map(|t| (t.id, t)).collect(); - - let all_completed = gen_ids.iter().all(|id| history_map.contains_key(id)); - if !all_completed { - return; + if let Some(evaluated) = collect_evaluated_generation(&state.evo, history) { + let offspring = nsga2_generate_offspring(state, &evaluated, directions); + advance_generation(&mut state.evo, offspring); } - - // Collect the evaluated population - let evaluated: Vec<&MultiObjectiveTrial> = gen_ids - .iter() - .filter_map(|id| history_map.get(id).copied()) - .collect(); - - // Run NSGA-II to produce offspring - let offspring = nsga2_generate_offspring(state, &evaluated, directions); - state.candidates = offspring; - state.assigned_count = 0; - state.generation_trial_ids.clear(); - state.trial_progress.clear(); - state.generation += 1; } // --------------------------------------------------------------------------- @@ -413,7 +237,7 @@ fn nsga2_select( population: &[&MultiObjectiveTrial], directions: &[Direction], ) -> (Vec>, Vec, Vec) { - let pop_size = state.population_size; + let pop_size = state.evo.population_size; let values: Vec> = population.iter().map(|t| t.values.clone()).collect(); let constraints: Vec> = population.iter().map(|t| t.constraints.clone()).collect(); @@ -455,13 +279,14 @@ fn nsga2_select( } while selected.len() < pop_size { - selected.push(state.rng.usize(0..n)); + selected.push(state.evo.rng.usize(0..n)); } - // Extract parent parameter vectors ordered by dimension let parents: Vec> = selected .iter() - .map(|&idx| extract_trial_params(population[idx], &state.dimensions, &mut state.rng)) + .map(|&idx| { + extract_trial_params(population[idx], &state.evo.dimensions, &mut state.evo.rng) + }) .collect(); let sel_rank: Vec = selected.iter().map(|&i| rank[i]).collect(); @@ -470,43 +295,22 @@ fn nsga2_select( (parents, sel_rank, sel_crowding) } -/// Extract parameter values from a trial, ordered by dimension index. -fn extract_trial_params( - trial: &MultiObjectiveTrial, - dimensions: &[DimensionInfo], - rng: &mut fastrand::Rng, -) -> Vec { - let mut param_pairs: Vec<_> = trial.params.iter().collect(); - param_pairs.sort_by_key(|(id, _)| *id); - - dimensions - .iter() - .enumerate() - .map(|(dim_idx, dim_info)| { - if dim_idx < param_pairs.len() { - param_pairs[dim_idx].1.clone() - } else { - sample_random(rng, &dim_info.distribution) - } - }) - .collect() -} - /// Runs NSGA-II selection and generates offspring candidates. fn nsga2_generate_offspring( state: &mut Nsga2State, population: &[&MultiObjectiveTrial], directions: &[Direction], ) -> Vec { - let pop_size = state.population_size; + let pop_size = state.evo.population_size; if population.len() < 2 { return (0..pop_size) .map(|_| { let params = state + .evo .dimensions .iter() - .map(|d| sample_random(&mut state.rng, &d.distribution)) + .map(|d| sample_random(&mut state.evo.rng, &d.distribution)) .collect(); Candidate { params } }) @@ -517,28 +321,28 @@ fn nsga2_generate_offspring( let mut offspring = Vec::with_capacity(pop_size); while offspring.len() < pop_size { - let p1 = tournament_select(&mut state.rng, &sel_rank, &sel_crowding, parents.len()); - let p2 = tournament_select(&mut state.rng, &sel_rank, &sel_crowding, parents.len()); + let p1 = tournament_select(&mut state.evo.rng, &sel_rank, &sel_crowding, parents.len()); + let p2 = tournament_select(&mut state.evo.rng, &sel_rank, &sel_crowding, parents.len()); let (mut child1, mut child2) = crossover( - &mut state.rng, + &mut state.evo.rng, &parents[p1], &parents[p2], - &state.dimensions, + &state.evo.dimensions, state.config.crossover_prob, state.config.crossover_eta, ); mutate( - &mut state.rng, + &mut state.evo.rng, &mut child1, - &state.dimensions, + &state.evo.dimensions, state.config.mutation_eta, ); mutate( - &mut state.rng, + &mut state.evo.rng, &mut child2, - &state.dimensions, + &state.evo.dimensions, state.config.mutation_eta, ); @@ -551,10 +355,6 @@ fn nsga2_generate_offspring( offspring } -// --------------------------------------------------------------------------- -// Genetic operators -// --------------------------------------------------------------------------- - /// Tournament selection: pick 2 random individuals, return index of winner. /// Winner has lower rank; ties broken by higher crowding distance. fn tournament_select( @@ -576,196 +376,3 @@ fn tournament_select( b } } - -/// SBX crossover for continuous params, uniform crossover for categorical. -fn crossover( - rng: &mut fastrand::Rng, - parent1: &[ParamValue], - parent2: &[ParamValue], - dimensions: &[DimensionInfo], - crossover_prob: f64, - eta: f64, -) -> (Vec, Vec) { - let n = parent1.len(); - let mut child1 = parent1.to_vec(); - let mut child2 = parent2.to_vec(); - - let u: f64 = rng_util::f64_range(rng, 0.0, 1.0); - if u > crossover_prob { - return (child1, child2); - } - - for i in 0..n { - match (&parent1[i], &parent2[i], &dimensions[i].distribution) { - (ParamValue::Float(p1), ParamValue::Float(p2), Distribution::Float(d)) => { - if (p1 - p2).abs() < 1e-14 { - continue; - } - let (c1, c2) = sbx_crossover_f64(rng, *p1, *p2, d.low, d.high, eta); - child1[i] = ParamValue::Float(c1); - child2[i] = ParamValue::Float(c2); - } - (ParamValue::Int(p1), ParamValue::Int(p2), Distribution::Int(d)) => { - if p1 == p2 { - continue; - } - #[allow(clippy::cast_precision_loss)] - let (c1, c2) = sbx_crossover_f64( - rng, - *p1 as f64, - *p2 as f64, - d.low as f64, - d.high as f64, - eta, - ); - #[allow(clippy::cast_possible_truncation)] - { - child1[i] = ParamValue::Int((c1.round() as i64).clamp(d.low, d.high)); - child2[i] = ParamValue::Int((c2.round() as i64).clamp(d.low, d.high)); - } - } - (ParamValue::Categorical(_), ParamValue::Categorical(_), _) => { - // Uniform crossover: swap with 50% probability - if rng_util::f64_range(rng, 0.0, 1.0) < 0.5 { - core::mem::swap(&mut child1[i], &mut child2[i]); - } - } - _ => {} - } - } - - (child1, child2) -} - -/// SBX crossover for a single float dimension. -fn sbx_crossover_f64( - rng: &mut fastrand::Rng, - p1: f64, - p2: f64, - low: f64, - high: f64, - eta: f64, -) -> (f64, f64) { - let u: f64 = rng_util::f64_range(rng, 0.0, 1.0); - - let beta = if u <= 0.5 { - (2.0 * u).powf(1.0 / (eta + 1.0)) - } else { - (1.0 / (2.0 * (1.0 - u))).powf(1.0 / (eta + 1.0)) - }; - - let c1 = 0.5 * ((1.0 + beta) * p1 + (1.0 - beta) * p2); - let c2 = 0.5 * ((1.0 - beta) * p1 + (1.0 + beta) * p2); - - (c1.clamp(low, high), c2.clamp(low, high)) -} - -/// Polynomial mutation for each dimension. -#[allow(clippy::cast_precision_loss)] -fn mutate( - rng: &mut fastrand::Rng, - individual: &mut [ParamValue], - dimensions: &[DimensionInfo], - eta: f64, -) { - let n = individual.len(); - if n == 0 { - return; - } - let mutation_prob = 1.0 / n as f64; - - for (i, value) in individual.iter_mut().enumerate() { - if rng_util::f64_range(rng, 0.0, 1.0) >= mutation_prob { - continue; - } - - match (value, &dimensions[i].distribution) { - (v @ ParamValue::Float(_), Distribution::Float(d)) => { - let ParamValue::Float(x) = *v else { - unreachable!(); - }; - let mutated = polynomial_mutation_f64(rng, x, d.low, d.high, eta); - *v = ParamValue::Float(mutated); - } - (v @ ParamValue::Int(_), Distribution::Int(d)) => { - let ParamValue::Int(x) = *v else { - unreachable!(); - }; - #[allow(clippy::cast_possible_truncation)] - { - let mutated = - polynomial_mutation_f64(rng, x as f64, d.low as f64, d.high as f64, eta); - *v = ParamValue::Int((mutated.round() as i64).clamp(d.low, d.high)); - } - } - (v @ ParamValue::Categorical(_), Distribution::Categorical(d)) => { - *v = ParamValue::Categorical(rng.usize(0..d.n_choices)); - } - _ => {} - } - } -} - -/// Polynomial mutation for a single float value. -fn polynomial_mutation_f64(rng: &mut fastrand::Rng, x: f64, low: f64, high: f64, eta: f64) -> f64 { - let u: f64 = rng_util::f64_range(rng, 0.0, 1.0); - let range = high - low; - if range <= 0.0 { - return x; - } - - let delta1 = (x - low) / range; - let delta2 = (high - x) / range; - - let delta_q = if u < 0.5 { - let xy = 1.0 - delta1; - let val = 2.0 * u + (1.0 - 2.0 * u) * xy.powf(eta + 1.0); - val.powf(1.0 / (eta + 1.0)) - 1.0 - } else { - let xy = 1.0 - delta2; - let val = 2.0 * (1.0 - u) + 2.0 * (u - 0.5) * xy.powf(eta + 1.0); - 1.0 - val.powf(1.0 / (eta + 1.0)) - }; - - (x + delta_q * range).clamp(low, high) -} - -// --------------------------------------------------------------------------- -// Random sampling helper (for discovery phase) -// --------------------------------------------------------------------------- - -#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] -fn sample_random(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - let log_low = d.low.ln(); - let log_high = d.high.ln(); - rng_util::f64_range(rng, log_low, log_high).exp() - } else if let Some(step) = d.step { - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - d.low + (k as f64) * step - } else { - rng_util::f64_range(rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - let value = if d.log_scale { - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), - } -} diff --git a/src/sampler/nsga3.rs b/src/sampler/nsga3.rs new file mode 100644 index 0000000..09f414c --- /dev/null +++ b/src/sampler/nsga3.rs @@ -0,0 +1,720 @@ +//! NSGA-III (Non-dominated Sorting Genetic Algorithm III) sampler. +//! +//! Uses reference-point-based niching for better diversity in +//! many-objective (3+) optimization problems. Das-Dennis structured +//! reference points guide the search toward a well-distributed +//! Pareto front. +//! +//! # Examples +//! +//! ``` +//! use optimizer::Direction; +//! use optimizer::multi_objective::MultiObjectiveStudy; +//! use optimizer::parameter::{FloatParam, Parameter}; +//! use optimizer::sampler::nsga3::Nsga3Sampler; +//! +//! let sampler = Nsga3Sampler::with_seed(42); +//! let study = MultiObjectiveStudy::with_sampler( +//! vec![ +//! Direction::Minimize, +//! Direction::Minimize, +//! Direction::Minimize, +//! ], +//! sampler, +//! ); +//! +//! let x = FloatParam::new(0.0, 1.0); +//! let y = FloatParam::new(0.0, 1.0); +//! study +//! .optimize(100, |trial| { +//! let xv = x.suggest(trial)?; +//! let yv = y.suggest(trial)?; +//! Ok::<_, optimizer::Error>(vec![xv, yv, (1.0 - xv - yv).abs()]) +//! }) +//! .unwrap(); +//! ``` + +use parking_lot::Mutex; + +use super::genetic::{ + self, Candidate, EvolutionaryState, Phase, advance_generation, auto_divisions, + collect_evaluated_generation, crossover, das_dennis, extract_trial_params, + generate_random_candidates, mutate, sample_from_candidate, sample_random, +}; +use crate::distribution::Distribution; +use crate::multi_objective::MultiObjectiveTrial; +use crate::param::ParamValue; +use crate::pareto; +use crate::types::Direction; + +/// NSGA-III sampler for multi-objective optimization. +/// +/// Uses reference-point-based niching to maintain diversity, +/// especially effective for problems with 3 or more objectives. +pub struct Nsga3Sampler { + state: Mutex, +} + +impl Nsga3Sampler { + /// Creates a new NSGA-III sampler with a random seed. + #[must_use] + pub fn new() -> Self { + Self { + state: Mutex::new(Nsga3State::new(Nsga3Config::default(), None)), + } + } + + /// Creates a new NSGA-III sampler with a fixed seed. + #[must_use] + pub fn with_seed(seed: u64) -> Self { + Self { + state: Mutex::new(Nsga3State::new(Nsga3Config::default(), Some(seed))), + } + } + + /// Creates a builder for configuring an `Nsga3Sampler`. + #[must_use] + pub fn builder() -> Nsga3SamplerBuilder { + Nsga3SamplerBuilder::default() + } +} + +impl Default for Nsga3Sampler { + fn default() -> Self { + Self::new() + } +} + +/// Builder for [`Nsga3Sampler`]. +#[derive(Debug, Clone, Default)] +pub struct Nsga3SamplerBuilder { + population_size: Option, + n_divisions: Option, + crossover_prob: Option, + crossover_eta: Option, + mutation_eta: Option, + seed: Option, +} + +impl Nsga3SamplerBuilder { + /// Sets the population size. If unset, equals the number of + /// Das-Dennis reference points. + #[must_use] + pub fn population_size(mut self, size: usize) -> Self { + self.population_size = Some(size); + self + } + + /// Sets the number of divisions (H) for Das-Dennis reference points. + /// If unset, automatically chosen based on population size and number + /// of objectives. + #[must_use] + pub fn n_divisions(mut self, h: usize) -> Self { + self.n_divisions = Some(h); + self + } + + /// Sets the crossover probability. Default: 1.0. + #[must_use] + pub fn crossover_prob(mut self, prob: f64) -> Self { + self.crossover_prob = Some(prob); + self + } + + /// Sets the SBX distribution index. Default: 30.0. + #[must_use] + pub fn crossover_eta(mut self, eta: f64) -> Self { + self.crossover_eta = Some(eta); + self + } + + /// Sets the polynomial mutation distribution index. Default: 20.0. + #[must_use] + pub fn mutation_eta(mut self, eta: f64) -> Self { + self.mutation_eta = Some(eta); + self + } + + /// Sets the random seed for reproducibility. + #[must_use] + pub fn seed(mut self, seed: u64) -> Self { + self.seed = Some(seed); + self + } + + /// Builds the configured [`Nsga3Sampler`]. + #[must_use] + pub fn build(self) -> Nsga3Sampler { + let config = Nsga3Config { + user_population_size: self.population_size, + n_divisions: self.n_divisions, + crossover_prob: self.crossover_prob.unwrap_or(1.0), + crossover_eta: self.crossover_eta.unwrap_or(30.0), + mutation_eta: self.mutation_eta.unwrap_or(20.0), + }; + Nsga3Sampler { + state: Mutex::new(Nsga3State::new(config, self.seed)), + } + } +} + +// --------------------------------------------------------------------------- +// Internal types +// --------------------------------------------------------------------------- + +#[derive(Clone, Debug)] +struct Nsga3Config { + user_population_size: Option, + n_divisions: Option, + crossover_prob: f64, + crossover_eta: f64, + mutation_eta: f64, +} + +impl Default for Nsga3Config { + fn default() -> Self { + Self { + user_population_size: None, + n_divisions: None, + crossover_prob: 1.0, + crossover_eta: 30.0, + mutation_eta: 20.0, + } + } +} + +struct Nsga3State { + evo: EvolutionaryState, + config: Nsga3Config, + /// Das-Dennis reference points (lazily generated once objectives are known). + reference_points: Vec>, + /// Best value seen per objective (minimize-space). + ideal_point: Vec, + /// Whether reference points have been initialized. + initialized: bool, +} + +impl Nsga3State { + fn new(config: Nsga3Config, seed: Option) -> Self { + Self { + evo: EvolutionaryState::new(seed), + config, + reference_points: Vec::new(), + ideal_point: Vec::new(), + initialized: false, + } + } +} + +// --------------------------------------------------------------------------- +// MultiObjectiveSampler implementation +// --------------------------------------------------------------------------- + +impl crate::multi_objective::MultiObjectiveSampler for Nsga3Sampler { + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + history: &[MultiObjectiveTrial], + directions: &[Direction], + ) -> ParamValue { + let mut state = self.state.lock(); + + match &state.evo.phase { + Phase::Discovery => { + if let Some(value) = + genetic::sample_discovery(&mut state.evo, distribution, trial_id) + { + return value; + } + // Transitioned to active phase + initialize_nsga3(&mut state, directions); + generate_random_candidates(&mut state.evo); + sample_from_candidate(&mut state.evo, trial_id) + } + Phase::Active => { + maybe_generate_new_generation(&mut state, history, directions); + sample_from_candidate(&mut state.evo, trial_id) + } + } + } +} + +/// Initialize NSGA-III: generate reference points and set population size. +fn initialize_nsga3(state: &mut Nsga3State, directions: &[Direction]) { + let n_obj = directions.len(); + + // Determine divisions + let divisions = state + .config + .n_divisions + .unwrap_or_else(|| auto_divisions(n_obj, state.config.user_population_size.unwrap_or(100))); + + state.reference_points = das_dennis(n_obj, divisions); + let n_ref = state.reference_points.len(); + + // Population size = number of reference points (or user override, at least n_ref) + let pop_size = state.config.user_population_size.unwrap_or(n_ref).max(4); + state.evo.population_size = pop_size; + state.evo.phase = Phase::Active; + state.ideal_point = vec![f64::INFINITY; n_obj]; + state.initialized = true; +} + +fn maybe_generate_new_generation( + state: &mut Nsga3State, + history: &[MultiObjectiveTrial], + directions: &[Direction], +) { + if state.evo.candidates.is_empty() { + generate_random_candidates(&mut state.evo); + return; + } + + if let Some(evaluated) = collect_evaluated_generation(&state.evo, history) { + let offspring = nsga3_generate_offspring(state, &evaluated, directions); + advance_generation(&mut state.evo, offspring); + } +} + +// --------------------------------------------------------------------------- +// NSGA-III selection algorithm +// --------------------------------------------------------------------------- + +/// Normalize objectives to minimize-space. +fn to_minimize_space(values: &[f64], directions: &[Direction]) -> Vec { + values + .iter() + .zip(directions) + .map(|(&v, d)| match d { + Direction::Minimize => v, + Direction::Maximize => -v, + }) + .collect() +} + +/// Update ideal point with new observations. +fn update_ideal_point(ideal: &mut [f64], normalized_values: &[Vec]) { + for vals in normalized_values { + for (i, &v) in vals.iter().enumerate() { + if v < ideal[i] { + ideal[i] = v; + } + } + } +} + +/// Compute Achievement Scalarizing Function (ASF) for extreme point finding. +fn asf(point: &[f64], weight: &[f64], ideal: &[f64]) -> f64 { + point + .iter() + .zip(weight) + .zip(ideal) + .map(|((&p, &w), &z)| { + let w = if w < 1e-6 { 1e-6 } else { w }; + (p - z) / w + }) + .fold(f64::NEG_INFINITY, f64::max) +} + +/// Find intercepts for normalization via extreme points. +/// +/// For each objective, find the point with best ASF (using a weight vector +/// that emphasizes that objective). The intercepts are where the hyperplane +/// through the extreme points crosses each axis. +fn find_intercepts(normalized_values: &[Vec], ideal: &[f64]) -> Vec { + let n_obj = ideal.len(); + let n = normalized_values.len(); + + if n == 0 || n_obj == 0 { + return vec![1.0; n_obj]; + } + + // Find extreme points (one per objective) + let mut extreme_indices = Vec::with_capacity(n_obj); + for obj in 0..n_obj { + let mut weight = vec![1e-6; n_obj]; + weight[obj] = 1.0; + + let mut best_idx = 0; + let mut best_asf = f64::INFINITY; + for (i, vals) in normalized_values.iter().enumerate() { + let a = asf(vals, &weight, ideal); + if a < best_asf { + best_asf = a; + best_idx = i; + } + } + extreme_indices.push(best_idx); + } + + // Try to compute hyperplane intercepts + // For stability, if the extreme points are degenerate, fall back to + // max - ideal per objective + let mut intercepts = Vec::with_capacity(n_obj); + for obj in 0..n_obj { + let max_val = normalized_values + .iter() + .map(|v| v[obj]) + .fold(f64::NEG_INFINITY, f64::max); + let intercept = max_val - ideal[obj]; + intercepts.push(if intercept > 1e-10 { intercept } else { 1.0 }); + } + + intercepts +} + +/// Normalize objective values: subtract ideal, divide by intercepts. +fn normalize_objectives(values: &[Vec], ideal: &[f64], intercepts: &[f64]) -> Vec> { + values + .iter() + .map(|v| { + v.iter() + .zip(ideal) + .zip(intercepts) + .map(|((&val, &z), &a)| { + let norm = if a > 1e-10 { a } else { 1.0 }; + (val - z) / norm + }) + .collect() + }) + .collect() +} + +/// Perpendicular distance from a point to a reference line (direction vector). +fn perpendicular_distance(point: &[f64], reference: &[f64]) -> f64 { + let dot: f64 = point.iter().zip(reference).map(|(&p, &r)| p * r).sum(); + let ref_norm_sq: f64 = reference.iter().map(|&r| r * r).sum(); + + if ref_norm_sq < 1e-30 { + return f64::INFINITY; + } + + let proj_scalar = dot / ref_norm_sq; + let dist_sq: f64 = point + .iter() + .zip(reference) + .map(|(&p, &r)| { + let proj = proj_scalar * r; + (p - proj).powi(2) + }) + .sum(); + + dist_sq.sqrt() +} + +/// Associate each solution with its nearest reference point. +/// Returns (`closest_ref_idx`, distance) for each solution. +fn associate_to_reference_points( + normalized: &[Vec], + reference_points: &[Vec], +) -> Vec<(usize, f64)> { + normalized + .iter() + .map(|point| { + let mut best_ref = 0; + let mut best_dist = f64::INFINITY; + for (j, rp) in reference_points.iter().enumerate() { + let d = perpendicular_distance(point, rp); + if d < best_dist { + best_dist = d; + best_ref = j; + } + } + (best_ref, best_dist) + }) + .collect() +} + +/// NSGA-III niching-based selection from the last front. +/// +/// `already_selected` are indices into the combined population that are +/// already accepted (from fronts 0..L-1). `last_front` contains indices +/// from front L. We need to pick `remaining` more from `last_front`. +fn niching_select( + rng: &mut fastrand::Rng, + associations: &[(usize, f64)], + already_selected: &[usize], + last_front: &[usize], + n_reference_points: usize, + remaining: usize, +) -> Vec { + // Count niche per reference point for already selected + let mut niche_count = vec![0_usize; n_reference_points]; + for &idx in already_selected { + niche_count[associations[idx].0] += 1; + } + + // Build per-reference-point candidate lists from the last front + let mut ref_candidates: Vec> = vec![Vec::new(); n_reference_points]; + for &idx in last_front { + let (ref_idx, dist) = associations[idx]; + ref_candidates[ref_idx].push((idx, dist)); + } + + let mut selected = Vec::with_capacity(remaining); + let mut excluded = vec![false; associations.len()]; + + for _ in 0..remaining { + // Find minimum niche count among reference points that still have candidates + let min_count = (0..n_reference_points) + .filter(|&j| ref_candidates[j].iter().any(|&(idx, _)| !excluded[idx])) + .map(|j| niche_count[j]) + .min(); + + let Some(min_count) = min_count else { + break; + }; + + // Collect reference points with this minimum count that have candidates + let min_refs: Vec = (0..n_reference_points) + .filter(|&j| { + niche_count[j] == min_count + && ref_candidates[j].iter().any(|&(idx, _)| !excluded[idx]) + }) + .collect(); + + if min_refs.is_empty() { + break; + } + + // Pick a random reference point from the minimum set + let chosen_ref = min_refs[rng.usize(0..min_refs.len())]; + + // Available candidates for this reference point + let available: Vec<(usize, f64)> = ref_candidates[chosen_ref] + .iter() + .filter(|&&(idx, _)| !excluded[idx]) + .copied() + .collect(); + + if available.is_empty() { + continue; + } + + let chosen_idx = if min_count == 0 { + // Pick closest to reference line + available + .iter() + .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal)) + .unwrap() + .0 + } else { + // Pick random + available[rng.usize(0..available.len())].0 + }; + + selected.push(chosen_idx); + excluded[chosen_idx] = true; + niche_count[chosen_ref] += 1; + } + + selected +} + +/// Perform NSGA-III selection: non-dominated sort + reference-point niching. +fn nsga3_select( + state: &mut Nsga3State, + population: &[&MultiObjectiveTrial], + directions: &[Direction], +) -> Vec> { + let pop_size = state.evo.population_size; + let n_obj = directions.len(); + + // Convert to minimize-space + let min_values: Vec> = population + .iter() + .map(|t| to_minimize_space(&t.values, directions)) + .collect(); + + // Non-dominated sort + let constraints: Vec> = population.iter().map(|t| t.constraints.clone()).collect(); + let has_constraints = constraints.iter().any(|c| !c.is_empty()); + let fronts = if has_constraints { + pareto::fast_non_dominated_sort_constrained( + &min_values, + &vec![Direction::Minimize; n_obj], + &constraints, + ) + } else { + pareto::fast_non_dominated_sort(&min_values, &vec![Direction::Minimize; n_obj]) + }; + + // Fill front-by-front + let mut selected: Vec = Vec::with_capacity(pop_size); + let mut last_front_idx = None; + + for (fi, front) in fronts.iter().enumerate() { + if selected.len() + front.len() <= pop_size { + selected.extend_from_slice(front); + } else { + last_front_idx = Some(fi); + break; + } + } + + // If we filled exactly or all fronts fit, done + if selected.len() < pop_size + && let Some(lf_idx) = last_front_idx + { + // Need niching from the last partial front + let remaining = pop_size - selected.len(); + + // Update ideal point + update_ideal_point(&mut state.ideal_point, &min_values); + + // Find intercepts and normalize + let intercepts = find_intercepts(&min_values, &state.ideal_point); + let normalized = normalize_objectives(&min_values, &state.ideal_point, &intercepts); + + // Associate all solutions with reference points + let associations = associate_to_reference_points(&normalized, &state.reference_points); + + // Select from last front using niching + let last_front = &fronts[lf_idx]; + let additional = niching_select( + &mut state.evo.rng, + &associations, + &selected, + last_front, + state.reference_points.len(), + remaining, + ); + selected.extend(additional); + } + + // Pad if needed + let n = population.len(); + while selected.len() < pop_size { + selected.push(state.evo.rng.usize(0..n)); + } + + selected + .iter() + .map(|&idx| { + extract_trial_params(population[idx], &state.evo.dimensions, &mut state.evo.rng) + }) + .collect() +} + +/// Tournament selection based on rank only (no crowding distance in NSGA-III). +fn tournament_select_rank(rng: &mut fastrand::Rng, ranks: &[usize], n: usize) -> usize { + let a = rng.usize(0..n); + let b = rng.usize(0..n); + + if ranks[a] <= ranks[b] { a } else { b } +} + +fn nsga3_generate_offspring( + state: &mut Nsga3State, + population: &[&MultiObjectiveTrial], + directions: &[Direction], +) -> Vec { + let pop_size = state.evo.population_size; + + if population.len() < 2 { + return (0..pop_size) + .map(|_| { + let params = state + .evo + .dimensions + .iter() + .map(|d| sample_random(&mut state.evo.rng, &d.distribution)) + .collect(); + Candidate { params } + }) + .collect(); + } + + // Initialize reference points and ideal on first generation + if !state.initialized { + initialize_nsga3(state, directions); + } + + let parents = nsga3_select(state, population, directions); + + // Assign ranks for tournament selection + let n_obj = directions.len(); + let min_values: Vec> = population + .iter() + .map(|t| to_minimize_space(&t.values, directions)) + .collect(); + let fronts = pareto::fast_non_dominated_sort(&min_values, &vec![Direction::Minimize; n_obj]); + let mut rank = vec![0_usize; parents.len()]; + for (front_rank, front) in fronts.iter().enumerate() { + for &idx in front { + if idx < rank.len() { + rank[idx] = front_rank; + } + } + } + // Ranks for selected parents (simplified: use index order) + let parent_ranks: Vec = (0..parents.len()) + .map(|i| i % (fronts.len().max(1))) + .collect(); + + let mut offspring = Vec::with_capacity(pop_size); + while offspring.len() < pop_size { + let p1 = tournament_select_rank(&mut state.evo.rng, &parent_ranks, parents.len()); + let p2 = tournament_select_rank(&mut state.evo.rng, &parent_ranks, parents.len()); + + let (mut child1, mut child2) = crossover( + &mut state.evo.rng, + &parents[p1], + &parents[p2], + &state.evo.dimensions, + state.config.crossover_prob, + state.config.crossover_eta, + ); + + mutate( + &mut state.evo.rng, + &mut child1, + &state.evo.dimensions, + state.config.mutation_eta, + ); + mutate( + &mut state.evo.rng, + &mut child2, + &state.evo.dimensions, + state.config.mutation_eta, + ); + + offspring.push(Candidate { params: child1 }); + if offspring.len() < pop_size { + offspring.push(Candidate { params: child2 }); + } + } + + offspring +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_perpendicular_distance() { + // Point (1, 0) to reference line (1, 1) (45-degree line) + let d = perpendicular_distance(&[1.0, 0.0], &[1.0, 1.0]); + // Projection is (0.5, 0.5), distance = sqrt(0.25 + 0.25) = sqrt(0.5) + assert!((d - (0.5_f64).sqrt()).abs() < 1e-10); + } + + #[test] + fn test_perpendicular_distance_on_line() { + // Point on the reference line + let d = perpendicular_distance(&[2.0, 2.0], &[1.0, 1.0]); + assert!(d < 1e-10); + } + + #[test] + fn test_normalize_objectives() { + let values = vec![vec![2.0, 4.0], vec![4.0, 2.0]]; + let ideal = vec![1.0, 1.0]; + let intercepts = vec![3.0, 3.0]; + let normalized = normalize_objectives(&values, &ideal, &intercepts); + assert!((normalized[0][0] - 1.0 / 3.0).abs() < 1e-10); + assert!((normalized[0][1] - 1.0).abs() < 1e-10); + } +} diff --git a/tests/multi_objective_tests.rs b/tests/multi_objective_tests.rs index 34ed767..ec5e200 100644 --- a/tests/multi_objective_tests.rs +++ b/tests/multi_objective_tests.rs @@ -1,9 +1,11 @@ //! Integration tests for multi-objective optimization. -use optimizer::Direction; use optimizer::multi_objective::MultiObjectiveStudy; use optimizer::parameter::{CategoricalParam, FloatParam, Parameter}; +use optimizer::sampler::moead::MoeadSampler; use optimizer::sampler::nsga2::Nsga2Sampler; +use optimizer::sampler::nsga3::Nsga3Sampler; +use optimizer::{Decomposition, Direction}; // --------------------------------------------------------------------------- // Pareto utility tests (via public MultiObjectiveStudy) @@ -391,3 +393,346 @@ fn test_tell_with_failure() { // Failed trial not counted assert_eq!(study.n_trials(), 0); } + +// --------------------------------------------------------------------------- +// NSGA-III sampler tests +// --------------------------------------------------------------------------- + +#[test] +fn test_nsga3_zdt1() { + let n_vars = 5; + let params: Vec = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect(); + + let sampler = Nsga3Sampler::builder().population_size(20).seed(42).build(); + let study = + MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); + + study + .optimize(200, |trial| { + let xs: Vec = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>()?; + + let f1 = xs[0]; + let g = 1.0 + 9.0 * xs[1..].iter().sum::() / (n_vars - 1) as f64; + let f2 = g * (1.0 - (f1 / g).sqrt()); + Ok::<_, optimizer::Error>(vec![f1, f2]) + }) + .unwrap(); + + let front = study.pareto_front(); + assert!( + !front.is_empty(), + "NSGA-III Pareto front should be non-empty" + ); + + // Verify no dominated solutions in the front + for a in &front { + for b in &front { + if core::ptr::eq(a, b) { + continue; + } + let a_dom_b = a.values[0] <= b.values[0] + && a.values[1] <= b.values[1] + && (a.values[0] < b.values[0] || a.values[1] < b.values[1]); + assert!( + !a_dom_b, + "Front solution {:?} dominates {:?}", + a.values, b.values + ); + } + } +} + +#[test] +fn test_nsga3_four_objectives() { + // DTLZ2 with 4 objectives + let n_obj = 4; + let n_vars = n_obj + 4; // k = 5 decision variables beyond the first (n_obj-1) + let params: Vec = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect(); + + let sampler = Nsga3Sampler::builder().population_size(50).seed(42).build(); + let directions = vec![Direction::Minimize; n_obj]; + let study = MultiObjectiveStudy::with_sampler(directions, sampler); + + study + .optimize(500, |trial| { + let xs: Vec = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>()?; + + // DTLZ2 formulation + let g: f64 = xs[n_obj - 1..] + .iter() + .map(|&xi| (xi - 0.5).powi(2)) + .sum::(); + + let mut objectives = vec![0.0_f64; n_obj]; + for i in 0..n_obj { + let mut f = 1.0 + g; + for xj in &xs[..(n_obj - 1 - i)] { + f *= (xj * core::f64::consts::FRAC_PI_2).cos(); + } + if i > 0 { + f *= (xs[n_obj - 1 - i] * core::f64::consts::FRAC_PI_2).sin(); + } + objectives[i] = f; + } + + Ok::<_, optimizer::Error>(objectives) + }) + .unwrap(); + + let front = study.pareto_front(); + assert!(!front.is_empty(), "4-objective front should be non-empty"); + // All front solutions should have 4 objectives + for t in &front { + assert_eq!(t.values.len(), 4); + } +} + +#[test] +fn test_nsga3_reproducible() { + let x = FloatParam::new(0.0, 1.0); + let y = FloatParam::new(0.0, 1.0); + + let run = |seed: u64| -> Vec> { + let sampler = Nsga3Sampler::with_seed(seed); + let study = MultiObjectiveStudy::with_sampler( + vec![Direction::Minimize, Direction::Minimize], + sampler, + ); + study + .optimize(30, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, optimizer::Error>(vec![xv, yv]) + }) + .unwrap(); + study.trials().iter().map(|t| t.values.clone()).collect() + }; + + let r1 = run(123); + let r2 = run(123); + assert_eq!(r1, r2, "Same seed should produce same results"); + + let r3 = run(456); + assert_ne!(r1, r3, "Different seeds should produce different results"); +} + +#[test] +fn test_nsga3_builder() { + let sampler = Nsga3Sampler::builder() + .population_size(12) + .n_divisions(4) + .crossover_prob(0.9) + .crossover_eta(20.0) + .mutation_eta(20.0) + .seed(42) + .build(); + + let study = + MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(30, |trial| { + let xv = x.suggest(trial)?; + Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 30); +} + +#[test] +fn test_nsga3_constraints() { + let sampler = Nsga3Sampler::with_seed(42); + let study = + MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); + + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(50, |trial| { + let xv = x.suggest(trial)?; + trial.set_constraints(vec![0.3 - xv]); + Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) + }) + .unwrap(); + + let front = study.pareto_front(); + assert!(!front.is_empty()); + + let feasible_count = front.iter().filter(|t| t.is_feasible()).count(); + assert!( + feasible_count > 0, + "Should have feasible solutions on front" + ); +} + +// --------------------------------------------------------------------------- +// MOEA/D sampler tests +// --------------------------------------------------------------------------- + +#[test] +fn test_moead_zdt1_tchebycheff() { + let n_vars = 5; + let params: Vec = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect(); + + let sampler = MoeadSampler::builder().population_size(20).seed(42).build(); + let study = + MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); + + study + .optimize(200, |trial| { + let xs: Vec = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>()?; + + let f1 = xs[0]; + let g = 1.0 + 9.0 * xs[1..].iter().sum::() / (n_vars - 1) as f64; + let f2 = g * (1.0 - (f1 / g).sqrt()); + Ok::<_, optimizer::Error>(vec![f1, f2]) + }) + .unwrap(); + + let front = study.pareto_front(); + assert!(!front.is_empty(), "MOEA/D Pareto front should be non-empty"); + + for a in &front { + for b in &front { + if core::ptr::eq(a, b) { + continue; + } + let a_dom_b = a.values[0] <= b.values[0] + && a.values[1] <= b.values[1] + && (a.values[0] < b.values[0] || a.values[1] < b.values[1]); + assert!( + !a_dom_b, + "Front solution {:?} dominates {:?}", + a.values, b.values + ); + } + } +} + +#[test] +fn test_moead_zdt1_weighted_sum() { + let n_vars = 3; + let params: Vec = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect(); + + let sampler = MoeadSampler::builder() + .population_size(20) + .decomposition(Decomposition::WeightedSum) + .seed(42) + .build(); + let study = + MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); + + study + .optimize(200, |trial| { + let xs: Vec = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>()?; + + let f1 = xs[0]; + let g = 1.0 + 9.0 * xs[1..].iter().sum::() / (n_vars - 1) as f64; + let f2 = g * (1.0 - (f1 / g).sqrt()); + Ok::<_, optimizer::Error>(vec![f1, f2]) + }) + .unwrap(); + + let front = study.pareto_front(); + assert!(!front.is_empty()); +} + +#[test] +fn test_moead_zdt1_pbi() { + let n_vars = 3; + let params: Vec = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect(); + + let sampler = MoeadSampler::builder() + .population_size(20) + .decomposition(Decomposition::Pbi { theta: 5.0 }) + .seed(42) + .build(); + let study = + MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); + + study + .optimize(200, |trial| { + let xs: Vec = params + .iter() + .map(|p| p.suggest(trial)) + .collect::>()?; + + let f1 = xs[0]; + let g = 1.0 + 9.0 * xs[1..].iter().sum::() / (n_vars - 1) as f64; + let f2 = g * (1.0 - (f1 / g).sqrt()); + Ok::<_, optimizer::Error>(vec![f1, f2]) + }) + .unwrap(); + + let front = study.pareto_front(); + assert!(!front.is_empty()); +} + +#[test] +fn test_moead_reproducible() { + let x = FloatParam::new(0.0, 1.0); + let y = FloatParam::new(0.0, 1.0); + + let run = |seed: u64| -> Vec> { + let sampler = MoeadSampler::with_seed(seed); + let study = MultiObjectiveStudy::with_sampler( + vec![Direction::Minimize, Direction::Minimize], + sampler, + ); + study + .optimize(30, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, optimizer::Error>(vec![xv, yv]) + }) + .unwrap(); + study.trials().iter().map(|t| t.values.clone()).collect() + }; + + let r1 = run(123); + let r2 = run(123); + assert_eq!(r1, r2, "Same seed should produce same results"); + + let r3 = run(456); + assert_ne!(r1, r3, "Different seeds should produce different results"); +} + +#[test] +fn test_moead_builder() { + let sampler = MoeadSampler::builder() + .population_size(15) + .neighborhood_size(5) + .decomposition(Decomposition::Tchebycheff) + .crossover_prob(0.9) + .crossover_eta(20.0) + .mutation_eta(20.0) + .seed(42) + .build(); + + let study = + MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(30, |trial| { + let xv = x.suggest(trial)?; + Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 30); +}