From bcc4549e66ba76227410fa3375e03c105d209fcf Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 19:35:18 +0100 Subject: [PATCH] feat: add multi-objective optimization with NSGA-II Add MultiObjectiveStudy for optimizing multiple objectives simultaneously, backed by NSGA-II (Non-dominated Sorting Genetic Algorithm II) with SBX crossover, polynomial mutation, and constraint-aware dominance. New public API: - MultiObjectiveStudy with optimize(), pareto_front(), ask()/tell() - MultiObjectiveTrial with get(), is_feasible(), user attributes - MultiObjectiveSampler trait for custom MO samplers - Nsga2Sampler with builder for population size, crossover/mutation params - ObjectiveDimensionMismatch error variant --- src/error.rs | 9 + src/lib.rs | 7 + src/multi_objective.rs | 412 ++++++++++++++++++ src/pareto.rs | 247 +++++++++++ src/sampler/mod.rs | 1 + src/sampler/nsga2.rs | 763 +++++++++++++++++++++++++++++++++ tests/multi_objective_tests.rs | 393 +++++++++++++++++ 7 files changed, 1832 insertions(+) create mode 100644 src/multi_objective.rs create mode 100644 src/pareto.rs create mode 100644 src/sampler/nsga2.rs create mode 100644 tests/multi_objective_tests.rs diff --git a/src/error.rs b/src/error.rs index 44801aa..fe22296 100644 --- a/src/error.rs +++ b/src/error.rs @@ -76,6 +76,15 @@ pub enum Error { #[error("trial was pruned")] TrialPruned, + /// Returned when the objective returns the wrong number of values. + #[error("objective dimension mismatch: expected {expected} values, got {got}")] + ObjectiveDimensionMismatch { + /// The expected number of objective values. + expected: usize, + /// The actual number of objective values returned. + got: usize, + }, + /// Returned when an internal invariant is violated. #[error("internal error: {0}")] Internal(&'static str), diff --git a/src/lib.rs b/src/lib.rs index 4079fff..5a4de76 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,6 +20,7 @@ //! - **Sobol (QMC)** - Quasi-random sampling for better space coverage (requires `sobol` feature) //! - **CMA-ES** - Covariance Matrix Adaptation Evolution Strategy for continuous optimization (requires `cma-es` feature) //! - **BOHB** - Bayesian Optimization + `HyperBand` for budget-aware TPE sampling +//! - **NSGA-II** - Non-dominated Sorting Genetic Algorithm II for multi-objective optimization //! //! Additional features include: //! @@ -218,8 +219,10 @@ mod distribution; mod error; mod importance; mod kde; +pub mod multi_objective; mod param; pub mod parameter; +mod pareto; pub mod pruner; pub mod sampler; mod study; @@ -227,6 +230,7 @@ mod trial; mod types; pub use error::{Error, Result, TrialPruned}; +pub use multi_objective::{MultiObjectiveSampler, MultiObjectiveStudy, MultiObjectiveTrial}; #[cfg(feature = "derive")] pub use optimizer_derive::Categorical; pub use param::ParamValue; @@ -242,6 +246,7 @@ pub use sampler::bohb::BohbSampler; #[cfg(feature = "cma-es")] pub use sampler::cma_es::CmaEsSampler; pub use sampler::grid::GridSearchSampler; +pub use sampler::nsga2::Nsga2Sampler; pub use sampler::random::RandomSampler; #[cfg(feature = "sobol")] pub use sampler::sobol::SobolSampler; @@ -262,6 +267,7 @@ pub mod prelude { pub use optimizer_derive::Categorical as DeriveCategory; pub use crate::error::{Error, Result, TrialPruned}; + pub use crate::multi_objective::{MultiObjectiveStudy, MultiObjectiveTrial}; pub use crate::param::ParamValue; pub use crate::parameter::{ BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter, @@ -275,6 +281,7 @@ pub mod prelude { #[cfg(feature = "cma-es")] pub use crate::sampler::cma_es::CmaEsSampler; pub use crate::sampler::grid::GridSearchSampler; + pub use crate::sampler::nsga2::Nsga2Sampler; pub use crate::sampler::random::RandomSampler; #[cfg(feature = "sobol")] pub use crate::sampler::sobol::SobolSampler; diff --git a/src/multi_objective.rs b/src/multi_objective.rs new file mode 100644 index 0000000..5a4b3b8 --- /dev/null +++ b/src/multi_objective.rs @@ -0,0 +1,412 @@ +//! Multi-objective optimization via a dedicated study type. +//! +//! [`MultiObjectiveStudy`] manages trials that return multiple objective +//! values. It supports arbitrary numbers of objectives with per-objective +//! directions (minimize or maximize). Use [`pareto_front()`](MultiObjectiveStudy::pareto_front) +//! to retrieve the Pareto-optimal solutions. +//! +//! # Examples +//! +//! ``` +//! use optimizer::Direction; +//! use optimizer::multi_objective::MultiObjectiveStudy; +//! use optimizer::parameter::{FloatParam, Parameter}; +//! +//! let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); +//! let x = FloatParam::new(0.0, 1.0); +//! +//! study +//! .optimize(20, |trial| { +//! let xv = x.suggest(trial)?; +//! Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) +//! }) +//! .unwrap(); +//! +//! let front = study.pareto_front(); +//! assert!(!front.is_empty()); +//! ``` + +use core::sync::atomic::{AtomicU64, Ordering}; +use std::collections::HashMap; +use std::sync::Arc; + +use parking_lot::RwLock; + +use crate::distribution::Distribution; +use crate::param::ParamValue; +use crate::parameter::{ParamId, Parameter}; +use crate::pruner::NopPruner; +use crate::sampler::random::RandomSampler; +use crate::sampler::{CompletedTrial, Sampler}; +use crate::trial::{AttrValue, Trial}; +use crate::types::{Direction, TrialState}; + +// --------------------------------------------------------------------------- +// MultiObjectiveTrial +// --------------------------------------------------------------------------- + +/// A completed trial with multiple objective values. +#[derive(Clone, Debug)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub struct MultiObjectiveTrial { + /// The unique identifier for this trial. + pub id: u64, + /// The sampled parameter values, keyed by parameter id. + pub params: HashMap, + /// The parameter distributions used, keyed by parameter id. + pub distributions: HashMap, + /// Human-readable labels for parameters, keyed by parameter id. + pub param_labels: HashMap, + /// The objective values (one per objective). + pub values: Vec, + /// The state of the trial. + pub state: TrialState, + /// User-defined attributes stored during the trial. + pub user_attrs: HashMap, + /// Constraint values for this trial (<=0.0 means feasible). + #[cfg_attr(feature = "serde", serde(default))] + pub constraints: Vec, +} + +impl MultiObjectiveTrial { + /// Returns the typed value for the given parameter. + /// + /// Returns `None` if the parameter was not used in this trial. + /// + /// # Panics + /// + /// Panics if the stored value is incompatible with the parameter type. + pub fn get(&self, param: &P) -> Option { + self.params.get(¶m.id()).map(|v| { + param + .cast_param_value(v) + .expect("parameter type mismatch: stored value incompatible with parameter") + }) + } + + /// Returns `true` if all constraints are satisfied (values <= 0.0). + /// + /// A trial with no constraints is considered feasible. + #[must_use] + pub fn is_feasible(&self) -> bool { + self.constraints.iter().all(|&c| c <= 0.0) + } + + /// Gets a user attribute by key. + #[must_use] + pub fn user_attr(&self, key: &str) -> Option<&AttrValue> { + self.user_attrs.get(key) + } + + /// Returns all user attributes. + #[must_use] + pub fn user_attrs(&self) -> &HashMap { + &self.user_attrs + } +} + +// --------------------------------------------------------------------------- +// MultiObjectiveSampler trait +// --------------------------------------------------------------------------- + +/// Trait for samplers aware of multi-objective history. +/// +/// Separate from [`Sampler`] because NSGA-II needs access to +/// `&[MultiObjectiveTrial]` (with vector-valued objectives) and +/// `&[Direction]` (one direction per objective). +pub trait MultiObjectiveSampler: Send + Sync { + /// Samples a parameter value from the given distribution. + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + history: &[MultiObjectiveTrial], + directions: &[Direction], + ) -> ParamValue; +} + +// --------------------------------------------------------------------------- +// RandomMultiObjectiveSampler +// --------------------------------------------------------------------------- + +/// Default MO sampler that delegates to [`RandomSampler`]. +pub(crate) struct RandomMultiObjectiveSampler(RandomSampler); + +impl RandomMultiObjectiveSampler { + pub(crate) fn new() -> Self { + Self(RandomSampler::new()) + } +} + +impl MultiObjectiveSampler for RandomMultiObjectiveSampler { + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + _history: &[MultiObjectiveTrial], + _directions: &[Direction], + ) -> ParamValue { + self.0.sample(distribution, trial_id, &[]) + } +} + +// --------------------------------------------------------------------------- +// MoSamplerBridge — bridges MultiObjectiveSampler to Sampler trait +// --------------------------------------------------------------------------- + +/// Bridges a [`MultiObjectiveSampler`] to the [`Sampler`] trait so that +/// `Trial::with_sampler()` can use it. +struct MoSamplerBridge { + inner: Arc, + history: Arc>>, + directions: Vec, +} + +impl Sampler for MoSamplerBridge { + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + _history: &[CompletedTrial], + ) -> ParamValue { + let mo_history = self.history.read(); + self.inner + .sample(distribution, trial_id, &mo_history, &self.directions) + } +} + +// --------------------------------------------------------------------------- +// MultiObjectiveStudy +// --------------------------------------------------------------------------- + +/// A study for multi-objective optimization. +/// +/// Manages trials that return multiple objective values. Supports +/// arbitrary numbers of objectives with independent minimize/maximize +/// directions. +/// +/// # Examples +/// +/// ``` +/// use optimizer::Direction; +/// use optimizer::multi_objective::MultiObjectiveStudy; +/// use optimizer::parameter::{FloatParam, Parameter}; +/// +/// // Bi-objective: minimize both +/// let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); +/// 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(); +/// +/// let front = study.pareto_front(); +/// assert!(!front.is_empty()); +/// ``` +pub struct MultiObjectiveStudy { + directions: Vec, + sampler: Arc, + completed_trials: Arc>>, + next_trial_id: AtomicU64, +} + +impl MultiObjectiveStudy { + /// Creates a new multi-objective study with the given directions. + /// + /// Uses a random sampler by default. + /// + /// # Arguments + /// + /// * `directions` - One direction per objective (minimize or maximize). + #[must_use] + pub fn new(directions: Vec) -> Self { + Self { + directions, + sampler: Arc::new(RandomMultiObjectiveSampler::new()), + completed_trials: Arc::new(RwLock::new(Vec::new())), + next_trial_id: AtomicU64::new(0), + } + } + + /// Creates a new study with a custom multi-objective sampler. + #[must_use] + pub fn with_sampler( + directions: Vec, + sampler: impl MultiObjectiveSampler + 'static, + ) -> Self { + Self { + directions, + sampler: Arc::new(sampler), + completed_trials: Arc::new(RwLock::new(Vec::new())), + next_trial_id: AtomicU64::new(0), + } + } + + /// Returns the optimization directions. + #[must_use] + pub fn directions(&self) -> &[Direction] { + &self.directions + } + + /// Returns the number of objectives. + #[must_use] + pub fn n_objectives(&self) -> usize { + self.directions.len() + } + + /// Returns the number of completed trials. + #[must_use] + pub fn n_trials(&self) -> usize { + self.completed_trials.read().len() + } + + /// Returns all completed trials. + #[must_use] + pub fn trials(&self) -> Vec { + self.completed_trials.read().clone() + } + + /// Returns the Pareto-optimal trials (front 0). + #[must_use] + pub fn pareto_front(&self) -> Vec { + let trials = self.completed_trials.read(); + let complete: Vec<_> = trials + .iter() + .filter(|t| t.state == TrialState::Complete) + .collect(); + + if complete.is_empty() { + return Vec::new(); + } + + let values: Vec> = complete.iter().map(|t| t.values.clone()).collect(); + let fronts = crate::pareto::fast_non_dominated_sort(&values, &self.directions); + + if fronts.is_empty() { + return Vec::new(); + } + + fronts[0].iter().map(|&i| complete[i].clone()).collect() + } + + /// Creates a new trial wired to the study's MO sampler. + fn create_trial(&self) -> Trial { + let id = self.next_trial_id.fetch_add(1, Ordering::SeqCst); + + let bridge: Arc = Arc::new(MoSamplerBridge { + inner: Arc::clone(&self.sampler), + history: Arc::clone(&self.completed_trials), + directions: self.directions.clone(), + }); + + // Dummy f64 history — the bridge ignores it. + let dummy_history: Arc>>> = + Arc::new(RwLock::new(Vec::new())); + + Trial::with_sampler(id, bridge, dummy_history, Arc::new(NopPruner)) + } + + /// Records a completed trial. + fn complete_trial(&self, mut trial: Trial, values: Vec) { + trial.set_complete(); + let mo_trial = MultiObjectiveTrial { + id: trial.id(), + params: trial.params().clone(), + distributions: trial.distributions().clone(), + param_labels: trial.param_labels().clone(), + values, + state: TrialState::Complete, + user_attrs: trial.user_attrs().clone(), + constraints: trial.constraint_values().to_vec(), + }; + self.completed_trials.write().push(mo_trial); + } + + /// Records a failed trial (not stored in history). + fn fail_trial(trial: &mut Trial) { + trial.set_failed(); + } + + /// Request a new trial for the ask/tell interface. + /// + /// After creating the trial, suggest parameters on it, evaluate your + /// objective externally, then pass the trial back to [`tell()`](Self::tell). + pub fn ask(&self) -> Trial { + self.create_trial() + } + + /// Report the result of a trial obtained from [`ask()`](Self::ask). + /// + /// Pass `Ok(values)` for a successful evaluation or `Err(reason)` for a failure. + /// + /// # Errors + /// + /// Returns `ObjectiveDimensionMismatch` if the number of values doesn't + /// match the number of directions. + pub fn tell( + &self, + mut trial: Trial, + result: core::result::Result, impl ToString>, + ) -> crate::Result<()> { + if let Ok(values) = result { + if values.len() != self.directions.len() { + return Err(crate::Error::ObjectiveDimensionMismatch { + expected: self.directions.len(), + got: values.len(), + }); + } + self.complete_trial(trial, values); + } else { + Self::fail_trial(&mut trial); + } + Ok(()) + } + + /// Runs multi-objective optimization for `n_trials` trials. + /// + /// The objective function must return a `Vec` with one value per + /// objective. + /// + /// # Errors + /// + /// Returns `ObjectiveDimensionMismatch` if the objective returns the wrong + /// number of values. Returns `NoCompletedTrials` if all trials fail. + pub fn optimize(&self, n_trials: usize, mut objective: F) -> crate::Result<()> + where + F: FnMut(&mut Trial) -> core::result::Result, E>, + E: ToString, + { + for _ in 0..n_trials { + let mut trial = self.create_trial(); + + match objective(&mut trial) { + Ok(values) => { + if values.len() != self.directions.len() { + return Err(crate::Error::ObjectiveDimensionMismatch { + expected: self.directions.len(), + got: values.len(), + }); + } + self.complete_trial(trial, values); + } + Err(_) => { + Self::fail_trial(&mut trial); + } + } + } + + let has_complete = self + .completed_trials + .read() + .iter() + .any(|t| t.state == TrialState::Complete); + if !has_complete { + return Err(crate::Error::NoCompletedTrials); + } + + Ok(()) + } +} diff --git a/src/pareto.rs b/src/pareto.rs new file mode 100644 index 0000000..be4f718 --- /dev/null +++ b/src/pareto.rs @@ -0,0 +1,247 @@ +//! Pareto dominance utilities for multi-objective optimization. +//! +//! Provides fast non-dominated sorting (Deb et al., 2002) and crowding +//! distance computation used by both `MultiObjectiveStudy::pareto_front()` +//! and `Nsga2Sampler`. + +use crate::types::Direction; + +/// Returns `true` if solution `a` Pareto-dominates solution `b`. +/// +/// A solution dominates another if it is at least as good in all objectives +/// and strictly better in at least one, respecting the given directions. +#[allow(clippy::module_name_repetitions)] +pub(crate) fn dominates(a: &[f64], b: &[f64], directions: &[Direction]) -> bool { + debug_assert_eq!(a.len(), b.len()); + debug_assert_eq!(a.len(), directions.len()); + + let mut strictly_better = false; + for ((&av, &bv), dir) in a.iter().zip(b.iter()).zip(directions.iter()) { + let better = match dir { + Direction::Minimize => av < bv, + Direction::Maximize => av > bv, + }; + let worse = match dir { + Direction::Minimize => av > bv, + Direction::Maximize => av < bv, + }; + if worse { + return false; + } + if better { + strictly_better = true; + } + } + strictly_better +} + +/// Constrained dominance: feasible beats infeasible, among infeasible +/// prefer lower total constraint violation, among feasible use Pareto dominance. +pub(crate) fn constrained_dominates( + a_values: &[f64], + b_values: &[f64], + a_constraints: &[f64], + b_constraints: &[f64], + directions: &[Direction], +) -> bool { + let a_feasible = a_constraints.iter().all(|&c| c <= 0.0); + let b_feasible = b_constraints.iter().all(|&c| c <= 0.0); + + match (a_feasible, b_feasible) { + (true, false) => true, + (false, true) => false, + (false, false) => { + let a_violation: f64 = a_constraints.iter().map(|c| c.max(0.0)).sum(); + let b_violation: f64 = b_constraints.iter().map(|c| c.max(0.0)).sum(); + a_violation < b_violation + } + (true, true) => dominates(a_values, b_values, directions), + } +} + +/// Fast non-dominated sorting (Deb et al., 2002). +/// +/// Returns `Vec>` where `fronts[0]` is the Pareto front, +/// each inner vec contains indices into `values`. +/// +/// Complexity: O(M * N^2) where M = objectives, N = solutions. +#[allow(clippy::cast_possible_truncation)] +pub(crate) fn fast_non_dominated_sort( + values: &[Vec], + directions: &[Direction], +) -> Vec> { + fast_non_dominated_sort_constrained(values, directions, &[]) +} + +/// Fast non-dominated sorting with constraint support. +/// +/// `constraints` is either empty (no constraints) or has the same length +/// as `values`, where each entry is the constraint vector for that solution. +#[allow(clippy::cast_possible_truncation)] +pub(crate) fn fast_non_dominated_sort_constrained( + values: &[Vec], + directions: &[Direction], + constraints: &[Vec], +) -> Vec> { + let n = values.len(); + if n == 0 { + return Vec::new(); + } + + let has_constraints = !constraints.is_empty(); + let empty_constraints: Vec = Vec::new(); + + // S_p: set of solutions dominated by p + let mut dominated_by: Vec> = vec![Vec::new(); n]; + // n_p: domination count for p + let mut domination_count: Vec = vec![0; n]; + + for i in 0..n { + for j in (i + 1)..n { + let (a_c, b_c) = if has_constraints { + (&constraints[i], &constraints[j]) + } else { + (&empty_constraints, &empty_constraints) + }; + + let i_dom_j = if has_constraints { + constrained_dominates(&values[i], &values[j], a_c, b_c, directions) + } else { + dominates(&values[i], &values[j], directions) + }; + let j_dom_i = if has_constraints { + constrained_dominates(&values[j], &values[i], b_c, a_c, directions) + } else { + dominates(&values[j], &values[i], directions) + }; + + if i_dom_j { + dominated_by[i].push(j); + domination_count[j] += 1; + } else if j_dom_i { + dominated_by[j].push(i); + domination_count[i] += 1; + } + } + } + + let mut fronts: Vec> = Vec::new(); + let mut current_front: Vec = (0..n).filter(|&i| domination_count[i] == 0).collect(); + + while !current_front.is_empty() { + let mut next_front: Vec = Vec::new(); + for &p in ¤t_front { + for &q in &dominated_by[p] { + domination_count[q] -= 1; + if domination_count[q] == 0 { + next_front.push(q); + } + } + } + fronts.push(current_front); + current_front = next_front; + } + + fronts +} + +/// Crowding distance for one front. +/// +/// Boundary solutions get `f64::INFINITY`. Returns one distance value per +/// solution in the front, in the same order as `front_indices`. +#[allow(clippy::cast_precision_loss)] +pub(crate) fn crowding_distance(front_indices: &[usize], values: &[Vec]) -> Vec { + let n = front_indices.len(); + if n <= 2 { + return vec![f64::INFINITY; n]; + } + + let m = values[front_indices[0]].len(); // number of objectives + let mut distances = vec![0.0_f64; n]; + + // Helper to look up objective value for a front member. + let val = |front_pos: usize, obj: usize| -> f64 { values[front_indices[front_pos]][obj] }; + + for obj in 0..m { + // Sort front positions by this objective + let mut sorted: Vec = (0..n).collect(); + sorted.sort_by(|&a, &b| { + val(a, obj) + .partial_cmp(&val(b, obj)) + .unwrap_or(core::cmp::Ordering::Equal) + }); + + // Boundary solutions get infinity + distances[sorted[0]] = f64::INFINITY; + distances[sorted[n - 1]] = f64::INFINITY; + + let range = val(sorted[n - 1], obj) - val(sorted[0], obj); + if range > 0.0 { + for i in 1..(n - 1) { + distances[sorted[i]] += (val(sorted[i + 1], obj) - val(sorted[i - 1], obj)) / range; + } + } + } + + distances +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_dominates_basic() { + let dirs = [Direction::Minimize, Direction::Minimize]; + assert!(dominates(&[1.0, 1.0], &[2.0, 2.0], &dirs)); + assert!(!dominates(&[2.0, 2.0], &[1.0, 1.0], &dirs)); + // Equal does not dominate + assert!(!dominates(&[1.0, 1.0], &[1.0, 1.0], &dirs)); + } + + #[test] + fn test_dominates_incomparable() { + let dirs = [Direction::Minimize, Direction::Minimize]; + assert!(!dominates(&[1.0, 3.0], &[3.0, 1.0], &dirs)); + assert!(!dominates(&[3.0, 1.0], &[1.0, 3.0], &dirs)); + } + + #[test] + fn test_dominates_maximize() { + let dirs = [Direction::Maximize, Direction::Minimize]; + // a = (5, 1) vs b = (3, 2): a is better in both + assert!(dominates(&[5.0, 1.0], &[3.0, 2.0], &dirs)); + assert!(!dominates(&[3.0, 2.0], &[5.0, 1.0], &dirs)); + } + + #[test] + fn test_nds_known() { + let values = vec![ + vec![1.0, 5.0], // front 0 + vec![5.0, 1.0], // front 0 + vec![3.0, 3.0], // front 0 (non-dominated) + vec![4.0, 4.0], // front 1 (dominated by #2) + vec![6.0, 6.0], // front 2 + ]; + let dirs = [Direction::Minimize, Direction::Minimize]; + let fronts = fast_non_dominated_sort(&values, &dirs); + + assert_eq!(fronts.len(), 3); + let mut f0 = fronts[0].clone(); + f0.sort_unstable(); + assert_eq!(f0, vec![0, 1, 2]); + assert_eq!(fronts[1], vec![3]); + assert_eq!(fronts[2], vec![4]); + } + + #[test] + fn test_crowding_boundaries() { + let values = vec![vec![1.0, 5.0], vec![3.0, 3.0], vec![5.0, 1.0]]; + let front = vec![0, 1, 2]; + let cd = crowding_distance(&front, &values); + assert!(cd[0].is_infinite()); + assert!(cd[2].is_infinite()); + assert!(cd[1].is_finite()); + assert!(cd[1] > 0.0); + } +} diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 883d834..b9bdf12 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -4,6 +4,7 @@ pub mod bohb; #[cfg(feature = "cma-es")] pub mod cma_es; pub mod grid; +pub mod nsga2; pub mod random; #[cfg(feature = "sobol")] pub mod sobol; diff --git a/src/sampler/nsga2.rs b/src/sampler/nsga2.rs new file mode 100644 index 0000000..6046ad8 --- /dev/null +++ b/src/sampler/nsga2.rs @@ -0,0 +1,763 @@ +//! NSGA-II (Non-dominated Sorting Genetic Algorithm II) sampler. +//! +//! Implements multi-objective optimization using non-dominated sorting, +//! crowding distance, SBX crossover, and polynomial mutation. +//! +//! # Examples +//! +//! ``` +//! use optimizer::Direction; +//! use optimizer::multi_objective::MultiObjectiveStudy; +//! use optimizer::parameter::{FloatParam, Parameter}; +//! use optimizer::sampler::nsga2::Nsga2Sampler; +//! +//! let sampler = Nsga2Sampler::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)?; +//! Ok::<_, optimizer::Error>(vec![xv * xv, (xv - 1.0).powi(2)]) +//! }) +//! .unwrap(); +//! ``` + +use std::collections::HashMap; + +use parking_lot::Mutex; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +use crate::distribution::Distribution; +use crate::multi_objective::MultiObjectiveTrial; +use crate::param::ParamValue; +use crate::pareto; +use crate::types::Direction; + +/// NSGA-II sampler for multi-objective optimization. +/// +/// Provides non-dominated sorting, crowding distance selection, +/// SBX crossover, and polynomial mutation. +pub struct Nsga2Sampler { + state: Mutex, +} + +impl Nsga2Sampler { + /// Creates a new NSGA-II sampler with a random seed. + #[must_use] + pub fn new() -> Self { + Self { + state: Mutex::new(Nsga2State::new(Nsga2Config::default(), None)), + } + } + + /// Creates a new NSGA-II sampler with a fixed seed. + #[must_use] + pub fn with_seed(seed: u64) -> Self { + Self { + state: Mutex::new(Nsga2State::new(Nsga2Config::default(), Some(seed))), + } + } + + /// Creates a builder for configuring an `Nsga2Sampler`. + #[must_use] + pub fn builder() -> Nsga2SamplerBuilder { + Nsga2SamplerBuilder::default() + } +} + +impl Default for Nsga2Sampler { + fn default() -> Self { + Self::new() + } +} + +/// Builder for [`Nsga2Sampler`]. +#[derive(Debug, Clone, Default)] +pub struct Nsga2SamplerBuilder { + population_size: Option, + crossover_prob: Option, + crossover_eta: Option, + mutation_eta: Option, + seed: Option, +} + +impl Nsga2SamplerBuilder { + /// Sets the population size. Default: `4 + floor(3 * ln(n_params))`, minimum 4. + #[must_use] + pub fn population_size(mut self, size: usize) -> Self { + self.population_size = Some(size); + self + } + + /// Sets the crossover probability. Default: 0.9. + #[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 [`Nsga2Sampler`]. + #[must_use] + pub fn build(self) -> Nsga2Sampler { + let config = Nsga2Config { + user_population_size: self.population_size, + crossover_prob: self.crossover_prob.unwrap_or(0.9), + crossover_eta: self.crossover_eta.unwrap_or(20.0), + mutation_eta: self.mutation_eta.unwrap_or(20.0), + }; + Nsga2Sampler { + state: Mutex::new(Nsga2State::new(config, self.seed)), + } + } +} + +// --------------------------------------------------------------------------- +// Internal types +// --------------------------------------------------------------------------- + +#[derive(Clone, Debug)] +struct Nsga2Config { + user_population_size: Option, + crossover_prob: f64, + crossover_eta: f64, + mutation_eta: f64, +} + +impl Default for Nsga2Config { + fn default() -> Self { + Self { + user_population_size: None, + crossover_prob: 0.9, + crossover_eta: 20.0, + mutation_eta: 20.0, + } + } +} + +/// 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: StdRng, + 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(rand::make_rng, StdRng::seed_from_u64); + Self { + rng, + 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, + } + } +} + +// --------------------------------------------------------------------------- +// MultiObjectiveSampler implementation +// --------------------------------------------------------------------------- + +impl crate::multi_objective::MultiObjectiveSampler for Nsga2Sampler { + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + history: &[MultiObjectiveTrial], + directions: &[Direction], + ) -> 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), + } + } +} + +/// 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( + state: &mut Nsga2State, + 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); + } + 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; + } + + // 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; +} + +// --------------------------------------------------------------------------- +// NSGA-II generation algorithm +// --------------------------------------------------------------------------- + +/// Performs NSGA-II selection: non-dominated sort + crowding distance, +/// then selects `pop_size` parents from the population. +fn nsga2_select( + state: &mut Nsga2State, + population: &[&MultiObjectiveTrial], + directions: &[Direction], +) -> (Vec>, Vec, Vec) { + let pop_size = state.population_size; + + let values: Vec> = population.iter().map(|t| t.values.clone()).collect(); + 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(&values, directions, &constraints) + } else { + pareto::fast_non_dominated_sort(&values, directions) + }; + + let n = population.len(); + let mut rank = vec![0_usize; n]; + let mut crowding = vec![0.0_f64; n]; + + for (front_rank, front) in fronts.iter().enumerate() { + let cd = pareto::crowding_distance(front, &values); + for (i, &idx) in front.iter().enumerate() { + rank[idx] = front_rank; + crowding[idx] = cd[i]; + } + } + + let mut selected: Vec = Vec::with_capacity(pop_size); + for front in &fronts { + if selected.len() + front.len() <= pop_size { + selected.extend_from_slice(front); + } else { + let remaining = pop_size - selected.len(); + let mut front_sorted: Vec = front.clone(); + front_sorted.sort_by(|&a, &b| { + crowding[b] + .partial_cmp(&crowding[a]) + .unwrap_or(core::cmp::Ordering::Equal) + }); + selected.extend_from_slice(&front_sorted[..remaining]); + break; + } + } + + while selected.len() < pop_size { + selected.push(state.rng.random_range(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)) + .collect(); + + let sel_rank: Vec = selected.iter().map(|&i| rank[i]).collect(); + let sel_crowding: Vec = selected.iter().map(|&i| crowding[i]).collect(); + + (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 StdRng, +) -> 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; + + if population.len() < 2 { + return (0..pop_size) + .map(|_| { + let params = state + .dimensions + .iter() + .map(|d| sample_random(&mut state.rng, &d.distribution)) + .collect(); + Candidate { params } + }) + .collect(); + } + + let (parents, sel_rank, sel_crowding) = nsga2_select(state, population, directions); + + 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 (mut child1, mut child2) = crossover( + &mut state.rng, + &parents[p1], + &parents[p2], + &state.dimensions, + state.config.crossover_prob, + state.config.crossover_eta, + ); + + mutate( + &mut state.rng, + &mut child1, + &state.dimensions, + state.config.mutation_eta, + ); + mutate( + &mut state.rng, + &mut child2, + &state.dimensions, + state.config.mutation_eta, + ); + + offspring.push(Candidate { params: child1 }); + if offspring.len() < pop_size { + offspring.push(Candidate { params: child2 }); + } + } + + 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(rng: &mut StdRng, ranks: &[usize], crowding: &[f64], n: usize) -> usize { + let a = rng.random_range(0..n); + let b = rng.random_range(0..n); + + if ranks[a] < ranks[b] { + a + } else if ranks[b] < ranks[a] { + b + } else if crowding[a] >= crowding[b] { + a + } else { + b + } +} + +/// SBX crossover for continuous params, uniform crossover for categorical. +fn crossover( + rng: &mut StdRng, + 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.random_range(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.random_range(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 StdRng, + p1: f64, + p2: f64, + low: f64, + high: f64, + eta: f64, +) -> (f64, f64) { + let u: f64 = rng.random_range(0.0_f64..1.0_f64); + + 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 StdRng, 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.random_range(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.random_range(0..d.n_choices)); + } + _ => {} + } + } +} + +/// Polynomial mutation for a single float value. +fn polynomial_mutation_f64(rng: &mut StdRng, x: f64, low: f64, high: f64, eta: f64) -> f64 { + let u: f64 = rng.random_range(0.0_f64..1.0_f64); + 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 StdRng, 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.random_range(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.random_range(0..=n_steps); + d.low + (k as f64) * step + } else { + rng.random_range(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.random_range(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.random_range(0..=n_steps); + d.low + k * step + } else { + rng.random_range(d.low..=d.high) + }; + ParamValue::Int(value) + } + Distribution::Categorical(d) => ParamValue::Categorical(rng.random_range(0..d.n_choices)), + } +} diff --git a/tests/multi_objective_tests.rs b/tests/multi_objective_tests.rs new file mode 100644 index 0000000..34ed767 --- /dev/null +++ b/tests/multi_objective_tests.rs @@ -0,0 +1,393 @@ +//! Integration tests for multi-objective optimization. + +use optimizer::Direction; +use optimizer::multi_objective::MultiObjectiveStudy; +use optimizer::parameter::{CategoricalParam, FloatParam, Parameter}; +use optimizer::sampler::nsga2::Nsga2Sampler; + +// --------------------------------------------------------------------------- +// Pareto utility tests (via public MultiObjectiveStudy) +// --------------------------------------------------------------------------- + +#[test] +fn test_basic_two_objective_random() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + 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(); + + let front = study.pareto_front(); + assert!(!front.is_empty(), "Pareto front should be non-empty"); + + // Verify no solution in the front dominates another + 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_dimension_mismatch_error() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + let x = FloatParam::new(0.0, 1.0); + + let result = study.optimize(1, |trial| { + let xv = x.suggest(trial)?; + // Return wrong number of values + Ok::<_, optimizer::Error>(vec![xv]) + }); + + assert!(result.is_err()); + let err = result.unwrap_err(); + assert!( + matches!( + err, + optimizer::Error::ObjectiveDimensionMismatch { + expected: 2, + got: 1 + } + ), + "Expected ObjectiveDimensionMismatch, got: {err}" + ); +} + +#[test] +fn test_ask_tell() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Maximize]); + let x = FloatParam::new(0.0, 10.0); + + for _ in 0..10 { + let mut trial = study.ask(); + let xv = x.suggest(&mut trial).unwrap(); + study + .tell(trial, Ok::<_, &str>(vec![xv, 10.0 - xv])) + .unwrap(); + } + + assert_eq!(study.n_trials(), 10); + let front = study.pareto_front(); + assert!(!front.is_empty()); +} + +#[test] +fn test_ask_tell_dimension_mismatch() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + let trial = study.ask(); + let result = study.tell(trial, Ok::<_, &str>(vec![1.0, 2.0, 3.0])); + assert!(result.is_err()); +} + +#[test] +fn test_n_trials_counting() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + assert_eq!(study.n_trials(), 0); + + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(5, |trial| { + let xv = x.suggest(trial)?; + Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 5); +} + +#[test] +fn test_three_objectives() { + let study = MultiObjectiveStudy::new(vec![ + Direction::Minimize, + Direction::Minimize, + Direction::Maximize, + ]); + let x = FloatParam::new(0.0, 1.0); + let y = FloatParam::new(0.0, 1.0); + + study + .optimize(30, |trial| { + let xv = x.suggest(trial)?; + let yv = y.suggest(trial)?; + Ok::<_, optimizer::Error>(vec![xv, yv, 1.0 - xv - yv]) + }) + .unwrap(); + + let front = study.pareto_front(); + assert!(!front.is_empty()); + assert_eq!(study.n_objectives(), 3); +} + +#[test] +fn test_directions_accessor() { + let dirs = vec![Direction::Minimize, Direction::Maximize]; + let study = MultiObjectiveStudy::new(dirs.clone()); + assert_eq!(study.directions(), &dirs); + assert_eq!(study.n_objectives(), 2); +} + +#[test] +fn test_trials_accessor() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(3, |trial| { + let xv = x.suggest(trial)?; + Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) + }) + .unwrap(); + + let trials = study.trials(); + assert_eq!(trials.len(), 3); + for t in &trials { + assert_eq!(t.values.len(), 2); + } +} + +// --------------------------------------------------------------------------- +// NSGA-II sampler tests +// --------------------------------------------------------------------------- + +#[test] +fn test_nsga2_zdt1() { + // ZDT1 benchmark: minimize both objectives + let n_vars = 5; + let params: Vec = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect(); + + let sampler = Nsga2Sampler::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(), "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_nsga2_with_seed_reproducible() { + let x = FloatParam::new(0.0, 1.0); + let y = FloatParam::new(0.0, 1.0); + + let run = |seed: u64| -> Vec> { + let sampler = Nsga2Sampler::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_nsga2_builder() { + let sampler = Nsga2Sampler::builder() + .population_size(10) + .crossover_prob(0.8) + .crossover_eta(15.0) + .mutation_eta(25.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_nsga2_categorical_params() { + let sampler = Nsga2Sampler::with_seed(42); + let study = + MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler); + + let x = FloatParam::new(0.0, 1.0); + let cat = CategoricalParam::new(vec!["a", "b", "c"]); + + study + .optimize(30, |trial| { + let xv = x.suggest(trial)?; + let cv = cat.suggest(trial)?; + let bonus = match cv { + "a" => 0.0, + "b" => 0.5, + _ => 1.0, + }; + Ok::<_, optimizer::Error>(vec![xv + bonus, 1.0 - xv]) + }) + .unwrap(); + + assert_eq!(study.n_trials(), 30); + let front = study.pareto_front(); + assert!(!front.is_empty()); +} + +#[test] +fn test_nsga2_constraints() { + let sampler = Nsga2Sampler::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)?; + // Constraint: x >= 0.3 (i.e. 0.3 - x <= 0) + 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()); + + // Check that feasible solutions exist on the front + let feasible_count = front.iter().filter(|t| t.is_feasible()).count(); + assert!( + feasible_count > 0, + "Should have feasible solutions on front" + ); +} + +#[test] +fn test_multi_objective_trial_get() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + let x = FloatParam::new(0.0, 10.0).name("x"); + + study + .optimize(5, |trial| { + let xv = x.suggest(trial)?; + Ok::<_, optimizer::Error>(vec![xv, 10.0 - xv]) + }) + .unwrap(); + + let front = study.pareto_front(); + for t in &front { + let xv: f64 = t.get(&x).unwrap(); + assert!((0.0..=10.0).contains(&xv)); + } +} + +#[test] +fn test_multi_objective_trial_is_feasible() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(10, |trial| { + let xv = x.suggest(trial)?; + trial.set_constraints(vec![0.5 - xv]); // feasible if x >= 0.5 + Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) + }) + .unwrap(); + + let trials = study.trials(); + for t in &trials { + let xv = t.values[0]; + if xv >= 0.5 { + assert!(t.is_feasible()); + } else { + assert!(!t.is_feasible()); + } + } +} + +#[test] +fn test_multi_objective_trial_user_attrs() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + let x = FloatParam::new(0.0, 1.0); + + study + .optimize(3, |trial| { + let xv = x.suggest(trial)?; + trial.set_user_attr("iteration", 42_i64); + Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv]) + }) + .unwrap(); + + let trials = study.trials(); + for t in &trials { + assert!(t.user_attr("iteration").is_some()); + } +} + +#[test] +fn test_tell_with_failure() { + let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]); + + let trial = study.ask(); + study + .tell(trial, Err::, _>("evaluation failed")) + .unwrap(); + + // Failed trial not counted + assert_eq!(study.n_trials(), 0); +}