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

@@ -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;

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<f64> = 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);
+    }
+}

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;
 

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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &completed));
+    // 2.4 < 2.5 => don't prune
+    let current = vec![(0, 2.4)];
+    assert!(!pruner.should_prune(4, 0, &current, &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, &current, &completed));
+    // 2.5 < 3.0 => don't prune
+    let current = vec![(0, 2.5)];
+    assert!(!pruner.should_prune(5, 0, &current, &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, &current, &completed));
+
+    let current = vec![(0, 1.0), (10, 2.0), (100, 3.5)];
+    assert!(!pruner.should_prune(3, 100, &current, &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, &current, &completed));
+
+    // 6.0 > 5.0 => prune
+    let current = vec![(0, 6.0)];
+    assert!(pruner.should_prune(2, 0, &current, &completed));
+}