From 432c74b927e01faf1cd13a4e80490eb19589b291 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 16:32:03 +0100 Subject: [PATCH] feat: add MedianPruner for statistics-based trial pruning Prunes trials whose intermediate values are worse than the median of completed trials at the same step. Supports configurable warmup steps and minimum trial count before pruning activates. --- src/lib.rs | 4 +- src/pruner/median.rs | 135 +++++++++++++++++++++ src/pruner/mod.rs | 2 + tests/median_pruner_tests.rs | 227 +++++++++++++++++++++++++++++++++++ 4 files changed, 366 insertions(+), 2 deletions(-) create mode 100644 src/pruner/median.rs create mode 100644 tests/median_pruner_tests.rs 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)); +}