diff --git a/src/lib.rs b/src/lib.rs index 84919af..485da35 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -201,7 +201,7 @@ pub use param::ParamValue; pub use parameter::{ BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, ParamId, Parameter, }; -pub use pruner::{NopPruner, Pruner, ThresholdPruner}; +pub use pruner::{MedianPruner, NopPruner, Pruner, ThresholdPruner}; pub use sampler::CompletedTrial; pub use sampler::grid::GridSearchSampler; pub use sampler::random::RandomSampler; @@ -224,7 +224,7 @@ pub mod prelude { pub use crate::parameter::{ BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter, }; - pub use crate::pruner::{NopPruner, Pruner, ThresholdPruner}; + pub use crate::pruner::{MedianPruner, NopPruner, Pruner, ThresholdPruner}; pub use crate::sampler::CompletedTrial; pub use crate::sampler::grid::GridSearchSampler; pub use crate::sampler::random::RandomSampler; diff --git a/src/pruner/median.rs b/src/pruner/median.rs new file mode 100644 index 0000000..b062d84 --- /dev/null +++ b/src/pruner/median.rs @@ -0,0 +1,135 @@ +use super::Pruner; +use crate::sampler::CompletedTrial; +use crate::types::{Direction, TrialState}; + +/// Prune trials that are performing worse than the median of completed trials +/// at the same step. +/// +/// This is the most commonly used pruner. It compares the current trial's +/// intermediate value at each step with the median of all completed trials' +/// values at that same step. +/// +/// # Examples +/// +/// ``` +/// use optimizer::Direction; +/// use optimizer::pruner::MedianPruner; +/// +/// // Prune trials worse than median when minimizing, after 5 warmup steps +/// let pruner = MedianPruner::new(Direction::Minimize) +/// .n_warmup_steps(5) +/// .n_min_trials(3); +/// ``` +pub struct MedianPruner { + /// The optimization direction. + direction: Direction, + /// Don't prune in the first N steps (let the trial warm up). + n_warmup_steps: u64, + /// Require at least N completed trials before pruning. + n_min_trials: usize, +} + +impl MedianPruner { + /// Create a new `MedianPruner` for the given optimization direction. + /// + /// By default, `n_warmup_steps` is 0 and `n_min_trials` is 1. + #[must_use] + pub fn new(direction: Direction) -> Self { + Self { + direction, + n_warmup_steps: 0, + n_min_trials: 1, + } + } + + /// Set the number of warmup steps. No pruning occurs before this step. + #[must_use] + pub fn n_warmup_steps(mut self, n: u64) -> Self { + self.n_warmup_steps = n; + self + } + + /// Set the minimum number of completed trials required before pruning. + #[must_use] + pub fn n_min_trials(mut self, n: usize) -> Self { + self.n_min_trials = n; + self + } +} + +impl Pruner for MedianPruner { + fn should_prune( + &self, + _trial_id: u64, + step: u64, + intermediate_values: &[(u64, f64)], + completed_trials: &[CompletedTrial], + ) -> bool { + // 1. Don't prune during warmup + if step < self.n_warmup_steps { + return false; + } + + // Get the current trial's latest value + let Some(&(_, current_value)) = intermediate_values.last() else { + return false; + }; + + // 2. Collect values at this step from completed (non-pruned) trials + let mut values_at_step: Vec = completed_trials + .iter() + .filter(|t| t.state == TrialState::Complete) + .filter_map(|t| { + t.intermediate_values + .iter() + .find(|(s, _)| *s == step) + .map(|(_, v)| *v) + }) + .collect(); + + // 3. Not enough trials + if values_at_step.len() < self.n_min_trials { + return false; + } + + // 4. Compute median + let median = compute_median(&mut values_at_step); + + // 5. Compare against median based on direction + match self.direction { + Direction::Minimize => current_value > median, + Direction::Maximize => current_value < median, + } + } +} + +/// Compute the median of a non-empty slice. Sorts the slice in place. +fn compute_median(values: &mut [f64]) -> f64 { + values.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal)); + let len = values.len(); + if len % 2 == 1 { + values[len / 2] + } else { + f64::midpoint(values[len / 2 - 1], values[len / 2]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compute_median_odd() { + assert!((compute_median(&mut [3.0, 1.0, 2.0]) - 2.0).abs() < f64::EPSILON); + } + + #[test] + fn compute_median_even() { + assert!((compute_median(&mut [4.0, 1.0, 3.0, 2.0]) - 2.5).abs() < f64::EPSILON); + } + + #[test] + fn compute_median_single() { + assert!((compute_median(&mut [5.0]) - 5.0).abs() < f64::EPSILON); + } +} diff --git a/src/pruner/mod.rs b/src/pruner/mod.rs index 4719157..6ae64fb 100644 --- a/src/pruner/mod.rs +++ b/src/pruner/mod.rs @@ -4,9 +4,11 @@ //! intermediate values compared to other trials. This is useful for //! discarding unpromising trials before they complete, saving compute. +mod median; mod nop; mod threshold; +pub use median::MedianPruner; pub use nop::NopPruner; pub use threshold::ThresholdPruner; diff --git a/tests/median_pruner_tests.rs b/tests/median_pruner_tests.rs new file mode 100644 index 0000000..980675f --- /dev/null +++ b/tests/median_pruner_tests.rs @@ -0,0 +1,227 @@ +use std::collections::HashMap; + +use optimizer::Direction; +use optimizer::pruner::{MedianPruner, Pruner}; +use optimizer::sampler::CompletedTrial; + +/// Helper to build a completed trial with given intermediate values. +fn trial_with_values(id: u64, intermediate_values: Vec<(u64, f64)>) -> CompletedTrial { + CompletedTrial::with_intermediate_values( + id, + HashMap::new(), + HashMap::new(), + HashMap::new(), + 0.0, + intermediate_values, + ) +} + +// --- Minimize direction --- + +#[test] +fn prune_when_worse_than_median_minimize() { + let pruner = MedianPruner::new(Direction::Minimize); + // 3 completed trials with values at step 2: [1.0, 2.0, 3.0] => median = 2.0 + let completed = vec![ + trial_with_values(0, vec![(0, 0.5), (1, 0.8), (2, 1.0)]), + trial_with_values(1, vec![(0, 0.6), (1, 1.5), (2, 2.0)]), + trial_with_values(2, vec![(0, 0.7), (1, 2.0), (2, 3.0)]), + ]; + // Current trial value at step 2 is 2.5 > median 2.0 => prune + let current = vec![(0, 0.5), (1, 1.0), (2, 2.5)]; + assert!(pruner.should_prune(3, 2, ¤t, &completed)); +} + +#[test] +fn no_prune_when_better_than_median_minimize() { + let pruner = MedianPruner::new(Direction::Minimize); + let completed = vec![ + trial_with_values(0, vec![(0, 0.5), (1, 0.8), (2, 1.0)]), + trial_with_values(1, vec![(0, 0.6), (1, 1.5), (2, 2.0)]), + trial_with_values(2, vec![(0, 0.7), (1, 2.0), (2, 3.0)]), + ]; + // Current trial value at step 2 is 1.5 < median 2.0 => don't prune + let current = vec![(0, 0.5), (1, 1.0), (2, 1.5)]; + assert!(!pruner.should_prune(3, 2, ¤t, &completed)); +} + +// --- Maximize direction --- + +#[test] +fn prune_when_worse_than_median_maximize() { + let pruner = MedianPruner::new(Direction::Maximize); + // Values at step 1: [5.0, 7.0, 9.0] => median = 7.0 + let completed = vec![ + trial_with_values(0, vec![(0, 3.0), (1, 5.0)]), + trial_with_values(1, vec![(0, 4.0), (1, 7.0)]), + trial_with_values(2, vec![(0, 5.0), (1, 9.0)]), + ]; + // Current value 6.0 < median 7.0 => prune (worse for maximize) + let current = vec![(0, 4.0), (1, 6.0)]; + assert!(pruner.should_prune(3, 1, ¤t, &completed)); +} + +#[test] +fn no_prune_when_better_than_median_maximize() { + let pruner = MedianPruner::new(Direction::Maximize); + let completed = vec![ + trial_with_values(0, vec![(0, 3.0), (1, 5.0)]), + trial_with_values(1, vec![(0, 4.0), (1, 7.0)]), + trial_with_values(2, vec![(0, 5.0), (1, 9.0)]), + ]; + // Current value 8.0 > median 7.0 => don't prune + let current = vec![(0, 4.0), (1, 8.0)]; + assert!(!pruner.should_prune(3, 1, ¤t, &completed)); +} + +// --- Warmup steps --- + +#[test] +fn no_prune_during_warmup() { + let pruner = MedianPruner::new(Direction::Minimize).n_warmup_steps(5); + let completed = vec![trial_with_values(0, vec![(0, 1.0), (1, 1.0), (2, 1.0)])]; + // Step 2 < warmup 5 => never prune, even if value is terrible + let current = vec![(0, 100.0), (1, 100.0), (2, 100.0)]; + assert!(!pruner.should_prune(1, 2, ¤t, &completed)); +} + +#[test] +fn prune_after_warmup() { + let pruner = MedianPruner::new(Direction::Minimize).n_warmup_steps(2); + let completed = vec![trial_with_values(0, vec![(0, 1.0), (1, 1.0), (2, 1.0)])]; + // Step 2 >= warmup 2 => pruning allowed; current 100.0 > median 1.0 + let current = vec![(0, 100.0), (1, 100.0), (2, 100.0)]; + assert!(pruner.should_prune(1, 2, ¤t, &completed)); +} + +// --- n_min_trials --- + +#[test] +fn no_prune_when_fewer_than_n_min_trials() { + let pruner = MedianPruner::new(Direction::Minimize).n_min_trials(3); + // Only 2 completed trials — below threshold of 3 + let completed = vec![ + trial_with_values(0, vec![(0, 1.0)]), + trial_with_values(1, vec![(0, 2.0)]), + ]; + let current = vec![(0, 100.0)]; + assert!(!pruner.should_prune(2, 0, ¤t, &completed)); +} + +#[test] +fn prune_when_at_least_n_min_trials() { + let pruner = MedianPruner::new(Direction::Minimize).n_min_trials(3); + // 3 completed trials with step 0: [1.0, 2.0, 3.0] => median 2.0 + let completed = vec![ + trial_with_values(0, vec![(0, 1.0)]), + trial_with_values(1, vec![(0, 2.0)]), + trial_with_values(2, vec![(0, 3.0)]), + ]; + // 5.0 > median 2.0 => prune + let current = vec![(0, 5.0)]; + assert!(pruner.should_prune(3, 0, ¤t, &completed)); +} + +// --- No completed trials with values at step --- + +#[test] +fn no_prune_when_no_completed_trials_at_step() { + let pruner = MedianPruner::new(Direction::Minimize); + // Completed trials only have values at step 0, not step 5 + let completed = vec![ + trial_with_values(0, vec![(0, 1.0)]), + trial_with_values(1, vec![(0, 2.0)]), + ]; + let current = vec![(0, 0.5), (5, 100.0)]; + assert!(!pruner.should_prune(2, 5, ¤t, &completed)); +} + +// --- Median calculation edge cases --- + +#[test] +fn correct_median_with_even_number_of_trials() { + let pruner = MedianPruner::new(Direction::Minimize); + // 4 trials at step 0: [1.0, 2.0, 3.0, 4.0] => median = 2.5 + let completed = vec![ + trial_with_values(0, vec![(0, 1.0)]), + trial_with_values(1, vec![(0, 2.0)]), + trial_with_values(2, vec![(0, 3.0)]), + trial_with_values(3, vec![(0, 4.0)]), + ]; + // 2.6 > 2.5 => prune + let current = vec![(0, 2.6)]; + assert!(pruner.should_prune(4, 0, ¤t, &completed)); + // 2.4 < 2.5 => don't prune + let current = vec![(0, 2.4)]; + assert!(!pruner.should_prune(4, 0, ¤t, &completed)); +} + +#[test] +fn correct_median_with_odd_number_of_trials() { + let pruner = MedianPruner::new(Direction::Minimize); + // 5 trials at step 0: [1.0, 2.0, 3.0, 4.0, 5.0] => median = 3.0 + let completed = vec![ + trial_with_values(0, vec![(0, 1.0)]), + trial_with_values(1, vec![(0, 2.0)]), + trial_with_values(2, vec![(0, 3.0)]), + trial_with_values(3, vec![(0, 4.0)]), + trial_with_values(4, vec![(0, 5.0)]), + ]; + // 3.5 > 3.0 => prune + let current = vec![(0, 3.5)]; + assert!(pruner.should_prune(5, 0, ¤t, &completed)); + // 2.5 < 3.0 => don't prune + let current = vec![(0, 2.5)]; + assert!(!pruner.should_prune(5, 0, ¤t, &completed)); +} + +// --- Non-contiguous step numbers --- + +#[test] +fn works_with_non_contiguous_steps() { + let pruner = MedianPruner::new(Direction::Minimize); + // Steps are 0, 10, 100 — non-contiguous + let completed = vec![ + trial_with_values(0, vec![(0, 1.0), (10, 2.0), (100, 3.0)]), + trial_with_values(1, vec![(0, 1.5), (10, 2.5), (100, 4.0)]), + trial_with_values(2, vec![(0, 2.0), (10, 3.0), (100, 5.0)]), + ]; + // At step 100: [3.0, 4.0, 5.0] => median = 4.0 + let current = vec![(0, 1.0), (10, 2.0), (100, 4.5)]; + assert!(pruner.should_prune(3, 100, ¤t, &completed)); + + let current = vec![(0, 1.0), (10, 2.0), (100, 3.5)]; + assert!(!pruner.should_prune(3, 100, ¤t, &completed)); +} + +// --- No intermediate values for current trial --- + +#[test] +fn no_prune_when_no_intermediate_values() { + let pruner = MedianPruner::new(Direction::Minimize); + let completed = vec![trial_with_values(0, vec![(0, 1.0)])]; + assert!(!pruner.should_prune(1, 0, &[], &completed)); +} + +// --- Pruned trials are excluded from median calculation --- + +#[test] +fn pruned_trials_excluded_from_median() { + use optimizer::TrialState; + + let pruner = MedianPruner::new(Direction::Minimize); + + let mut pruned = trial_with_values(0, vec![(0, 0.1)]); + pruned.state = TrialState::Pruned; + + // Only the completed trial (value 5.0) counts. Pruned trial (0.1) is excluded. + let completed = vec![pruned, trial_with_values(1, vec![(0, 5.0)])]; + + // 3.0 < 5.0 => don't prune (only 1 completed trial with median 5.0) + let current = vec![(0, 3.0)]; + assert!(!pruner.should_prune(2, 0, ¤t, &completed)); + + // 6.0 > 5.0 => prune + let current = vec![(0, 6.0)]; + assert!(pruner.should_prune(2, 0, ¤t, &completed)); +}