feat: add PercentilePruner for configurable percentile-based trial pruning

Generalizes MedianPruner with a configurable percentile threshold.
MedianPruner now reuses the shared compute_percentile function at 50%.

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::{MedianPruner, NopPruner, Pruner, ThresholdPruner};
+pub use pruner::{MedianPruner, NopPruner, PercentilePruner, 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::{MedianPruner, NopPruner, Pruner, ThresholdPruner};
+    pub use crate::pruner::{MedianPruner, NopPruner, PercentilePruner, Pruner, ThresholdPruner};
     pub use crate::sampler::CompletedTrial;
     pub use crate::sampler::grid::GridSearchSampler;
     pub use crate::sampler::random::RandomSampler;

src/pruner/median.rs

@@ -1,4 +1,5 @@
 use super::Pruner;
+use super::percentile::compute_percentile;
 use crate::sampler::CompletedTrial;
 use crate::types::{Direction, TrialState};
 
@@ -9,6 +10,8 @@ use crate::types::{Direction, TrialState};
 /// intermediate value at each step with the median of all completed trials'
 /// values at that same step.
 ///
+/// Equivalent to `PercentilePruner::new(50.0, direction)`.
+///
 /// # Examples
 ///
 /// ```
@@ -92,8 +95,8 @@ impl Pruner for MedianPruner {
             return false;
         }
 
-        // 4. Compute median
-        let median = compute_median(&mut values_at_step);
+        // 4. Compute median (50th percentile)
+        let median = compute_percentile(&mut values_at_step, 50.0);
 
         // 5. Compare against median based on direction
         match self.direction {
@@ -103,33 +106,22 @@ impl Pruner for MedianPruner {
     }
 }
 
-/// 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);
+        assert!((compute_percentile(&mut [3.0, 1.0, 2.0], 50.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);
+        assert!((compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 50.0) - 2.5).abs() < f64::EPSILON);
     }
 
     #[test]
     fn compute_median_single() {
-        assert!((compute_median(&mut [5.0]) - 5.0).abs() < f64::EPSILON);
+        assert!((compute_percentile(&mut [5.0], 50.0) - 5.0).abs() < f64::EPSILON);
     }
 }

src/pruner/mod.rs

@@ -6,10 +6,12 @@
 
 mod median;
 mod nop;
+pub(crate) mod percentile;
 mod threshold;
 
 pub use median::MedianPruner;
 pub use nop::NopPruner;
+pub use percentile::PercentilePruner;
 pub use threshold::ThresholdPruner;
 
 use crate::sampler::CompletedTrial;

src/pruner/percentile.rs

@@ -0,0 +1,293 @@
+use super::Pruner;
+use crate::sampler::CompletedTrial;
+use crate::types::{Direction, TrialState};
+
+/// Prune trials that are not in the top `percentile`% of completed trials
+/// at the same training step.
+///
+/// `PercentilePruner::new(50.0, direction)` is equivalent to `MedianPruner`.
+/// `PercentilePruner::new(25.0, direction)` keeps only the top 25% of trials.
+///
+/// # Examples
+///
+/// ```
+/// use optimizer::Direction;
+/// use optimizer::pruner::PercentilePruner;
+///
+/// // Keep only the top 25% of trials (aggressive pruning)
+/// let pruner = PercentilePruner::new(25.0, Direction::Minimize)
+///     .n_warmup_steps(5)
+///     .n_min_trials(3);
+/// ```
+pub struct PercentilePruner {
+    /// Keep trials in the top `percentile`%. Range: (0.0, 100.0).
+    percentile: f64,
+    /// 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,
+    /// The optimization direction.
+    direction: Direction,
+}
+
+impl PercentilePruner {
+    /// Create a new `PercentilePruner` for the given percentile and direction.
+    ///
+    /// The `percentile` value must be in `(0.0, 100.0)`.
+    /// A percentile of 50.0 is equivalent to median pruning.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `percentile` is not in `(0.0, 100.0)`.
+    #[must_use]
+    pub fn new(percentile: f64, direction: Direction) -> Self {
+        assert!(
+            percentile > 0.0 && percentile < 100.0,
+            "percentile must be in (0.0, 100.0), got {percentile}"
+        );
+        Self {
+            percentile,
+            n_warmup_steps: 0,
+            n_min_trials: 1,
+            direction,
+        }
+    }
+
+    /// 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 PercentilePruner {
+    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 percentile threshold
+        let threshold = compute_percentile(&mut values_at_step, self.percentile);
+
+        // 5. Compare against threshold based on direction
+        match self.direction {
+            Direction::Minimize => current_value > threshold,
+            Direction::Maximize => current_value < threshold,
+        }
+    }
+}
+
+/// Compute the given percentile of a non-empty slice. Sorts the slice in place.
+///
+/// Uses linear interpolation between the two nearest ranks.
+#[allow(
+    clippy::cast_precision_loss,
+    clippy::cast_possible_truncation,
+    clippy::cast_sign_loss
+)]
+pub(crate) fn compute_percentile(values: &mut [f64], percentile: f64) -> f64 {
+    values.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
+    let len = values.len();
+    if len == 1 {
+        return values[0];
+    }
+    // Rank in [0, len-1] range
+    let rank = percentile / 100.0 * (len - 1) as f64;
+    let lower = rank.floor() as usize;
+    let upper = rank.ceil() as usize;
+    if lower == upper {
+        values[lower]
+    } else {
+        let frac = rank - lower as f64;
+        values[lower] * (1.0 - frac) + values[upper] * frac
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn compute_percentile_median_odd() {
+        // Percentile 50 on odd-length slice = median
+        let val = compute_percentile(&mut [3.0, 1.0, 2.0], 50.0);
+        assert!((val - 2.0).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn compute_percentile_median_even() {
+        // Percentile 50 on even-length slice = median (interpolated)
+        let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 50.0);
+        assert!((val - 2.5).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn compute_percentile_25() {
+        // [1.0, 2.0, 3.0, 4.0], rank = 0.25 * 3 = 0.75
+        // interpolate: 1.0 * 0.25 + 2.0 * 0.75 = 1.75
+        let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 25.0);
+        assert!((val - 1.75).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn compute_percentile_75() {
+        // [1.0, 2.0, 3.0, 4.0], rank = 0.75 * 3 = 2.25
+        // interpolate: 3.0 * 0.75 + 4.0 * 0.25 = 3.25
+        let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 75.0);
+        assert!((val - 3.25).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    fn compute_percentile_single() {
+        let val = compute_percentile(&mut [5.0], 50.0);
+        assert!((val - 5.0).abs() < f64::EPSILON);
+    }
+
+    #[test]
+    #[should_panic(expected = "percentile must be in (0.0, 100.0)")]
+    fn new_rejects_zero() {
+        let _ = PercentilePruner::new(0.0, Direction::Minimize);
+    }
+
+    #[test]
+    #[should_panic(expected = "percentile must be in (0.0, 100.0)")]
+    fn new_rejects_hundred() {
+        let _ = PercentilePruner::new(100.0, Direction::Minimize);
+    }
+
+    fn make_completed_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
+        use std::collections::HashMap;
+
+        use crate::parameter::ParamId;
+
+        CompletedTrial::with_intermediate_values(
+            id,
+            HashMap::<ParamId, crate::ParamValue>::new(),
+            HashMap::new(),
+            HashMap::new(),
+            0.0,
+            values.to_vec(),
+        )
+    }
+
+    #[test]
+    fn percentile_50_matches_median_behavior() {
+        let pruner = PercentilePruner::new(50.0, Direction::Minimize);
+        let completed = vec![
+            make_completed_trial(0, &[(0, 1.0), (1, 2.0)]),
+            make_completed_trial(1, &[(0, 3.0), (1, 4.0)]),
+            make_completed_trial(2, &[(0, 5.0), (1, 6.0)]),
+        ];
+        // Median at step 1 is 4.0
+        // Value 5.0 > 4.0 → prune
+        assert!(pruner.should_prune(3, 1, &[(0, 3.0), (1, 5.0)], &completed));
+        // Value 3.0 < 4.0 → keep
+        assert!(!pruner.should_prune(3, 1, &[(0, 3.0), (1, 3.0)], &completed));
+    }
+
+    #[test]
+    fn percentile_25_is_more_aggressive() {
+        let pruner_25 = PercentilePruner::new(25.0, Direction::Minimize);
+        let pruner_75 = PercentilePruner::new(75.0, Direction::Minimize);
+        let completed = vec![
+            make_completed_trial(0, &[(0, 1.0)]),
+            make_completed_trial(1, &[(0, 2.0)]),
+            make_completed_trial(2, &[(0, 3.0)]),
+            make_completed_trial(3, &[(0, 4.0)]),
+        ];
+        // 25th percentile at step 0: 1.75
+        // 75th percentile at step 0: 3.25
+        // Value 2.5: above 25th (prune), below 75th (keep)
+        assert!(pruner_25.should_prune(4, 0, &[(0, 2.5)], &completed));
+        assert!(!pruner_75.should_prune(4, 0, &[(0, 2.5)], &completed));
+    }
+
+    #[test]
+    fn warmup_prevents_pruning() {
+        let pruner = PercentilePruner::new(50.0, Direction::Minimize).n_warmup_steps(5);
+        let completed = vec![make_completed_trial(0, &[(0, 1.0)])];
+        // Step 3 < warmup 5 → no prune even with bad value
+        assert!(!pruner.should_prune(1, 3, &[(3, 100.0)], &completed));
+    }
+
+    #[test]
+    fn n_min_trials_prevents_pruning() {
+        let pruner = PercentilePruner::new(50.0, Direction::Minimize).n_min_trials(5);
+        let completed = vec![
+            make_completed_trial(0, &[(0, 1.0)]),
+            make_completed_trial(1, &[(0, 2.0)]),
+        ];
+        // Only 2 trials, need 5 → no prune
+        assert!(!pruner.should_prune(2, 0, &[(0, 100.0)], &completed));
+    }
+
+    #[test]
+    fn maximize_direction() {
+        let pruner = PercentilePruner::new(50.0, Direction::Maximize);
+        let completed = vec![
+            make_completed_trial(0, &[(0, 1.0)]),
+            make_completed_trial(1, &[(0, 3.0)]),
+            make_completed_trial(2, &[(0, 5.0)]),
+        ];
+        // Median at step 0 is 3.0
+        // Value 2.0 < 3.0 → prune (maximize wants higher)
+        assert!(pruner.should_prune(3, 0, &[(0, 2.0)], &completed));
+        // Value 4.0 > 3.0 → keep
+        assert!(!pruner.should_prune(3, 0, &[(0, 4.0)], &completed));
+    }
+
+    #[test]
+    fn near_boundary_percentiles() {
+        let pruner_low = PercentilePruner::new(1.0, Direction::Minimize);
+        let pruner_high = PercentilePruner::new(99.0, Direction::Minimize);
+        let completed = vec![
+            make_completed_trial(0, &[(0, 1.0)]),
+            make_completed_trial(1, &[(0, 2.0)]),
+            make_completed_trial(2, &[(0, 3.0)]),
+            make_completed_trial(3, &[(0, 100.0)]),
+        ];
+        // Percentile 1 is very aggressive (threshold near 1.0)
+        // Value 1.5 should be pruned
+        assert!(pruner_low.should_prune(4, 0, &[(0, 1.5)], &completed));
+        // Percentile 99 is very lenient (threshold near 100.0)
+        // Value 50.0 should not be pruned
+        assert!(!pruner_high.should_prune(4, 0, &[(0, 50.0)], &completed));
+    }
+}