From 8199ba6b439871faf22fc8fa24382c32fd41c0f7 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 16:28:22 +0100 Subject: [PATCH] feat: add ThresholdPruner for fixed-bound trial pruning --- src/lib.rs | 4 +- src/pruner/mod.rs | 2 + src/pruner/threshold.rs | 83 +++++++++++++++++++++++++++++++++ tests/threshold_pruner_tests.rs | 64 +++++++++++++++++++++++++ 4 files changed, 151 insertions(+), 2 deletions(-) create mode 100644 src/pruner/threshold.rs create mode 100644 tests/threshold_pruner_tests.rs diff --git a/src/lib.rs b/src/lib.rs index d5c53a1..84919af 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}; +pub use pruner::{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}; + pub use crate::pruner::{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/mod.rs b/src/pruner/mod.rs index 06ea9f4..4719157 100644 --- a/src/pruner/mod.rs +++ b/src/pruner/mod.rs @@ -5,8 +5,10 @@ //! discarding unpromising trials before they complete, saving compute. mod nop; +mod threshold; pub use nop::NopPruner; +pub use threshold::ThresholdPruner; use crate::sampler::CompletedTrial; diff --git a/src/pruner/threshold.rs b/src/pruner/threshold.rs new file mode 100644 index 0000000..7058a71 --- /dev/null +++ b/src/pruner/threshold.rs @@ -0,0 +1,83 @@ +use super::Pruner; +use crate::sampler::CompletedTrial; + +/// Prune trials whose intermediate values exceed fixed thresholds. +/// +/// Useful for cutting off trials that are clearly diverging or stuck +/// at bad values early in training. +/// +/// # Examples +/// +/// ``` +/// use optimizer::pruner::ThresholdPruner; +/// +/// // Prune if the intermediate value exceeds 100.0 or falls below 0.0 +/// let pruner = ThresholdPruner::new().upper(100.0).lower(0.0); +/// ``` +pub struct ThresholdPruner { + /// Prune if intermediate value is greater than this. `None` = no upper bound. + upper: Option, + /// Prune if intermediate value is less than this. `None` = no lower bound. + lower: Option, +} + +impl ThresholdPruner { + /// Create a new `ThresholdPruner` with no thresholds set. + /// + /// By default, no pruning occurs. Use [`upper`](Self::upper) and + /// [`lower`](Self::lower) to set bounds. + #[must_use] + pub fn new() -> Self { + Self { + upper: None, + lower: None, + } + } + + /// Set the upper threshold. Trials with intermediate values above this + /// will be pruned. + #[must_use] + pub fn upper(mut self, threshold: f64) -> Self { + self.upper = Some(threshold); + self + } + + /// Set the lower threshold. Trials with intermediate values below this + /// will be pruned. + #[must_use] + pub fn lower(mut self, threshold: f64) -> Self { + self.lower = Some(threshold); + self + } +} + +impl Default for ThresholdPruner { + fn default() -> Self { + Self::new() + } +} + +impl Pruner for ThresholdPruner { + fn should_prune( + &self, + _trial_id: u64, + _step: u64, + intermediate_values: &[(u64, f64)], + _completed_trials: &[CompletedTrial], + ) -> bool { + let Some(&(_, latest_value)) = intermediate_values.last() else { + return false; + }; + if let Some(upper) = self.upper + && latest_value > upper + { + return true; + } + if let Some(lower) = self.lower + && latest_value < lower + { + return true; + } + false + } +} diff --git a/tests/threshold_pruner_tests.rs b/tests/threshold_pruner_tests.rs new file mode 100644 index 0000000..1583e2d --- /dev/null +++ b/tests/threshold_pruner_tests.rs @@ -0,0 +1,64 @@ +use optimizer::pruner::{Pruner, ThresholdPruner}; + +#[test] +fn prune_when_value_exceeds_upper_threshold() { + let pruner = ThresholdPruner::new().upper(10.0); + let values = vec![(0, 5.0), (1, 8.0), (2, 11.0)]; + assert!(pruner.should_prune(0, 2, &values, &[])); +} + +#[test] +fn prune_when_value_falls_below_lower_threshold() { + let pruner = ThresholdPruner::new().lower(0.0); + let values = vec![(0, 5.0), (1, 2.0), (2, -1.0)]; + assert!(pruner.should_prune(0, 2, &values, &[])); +} + +#[test] +fn no_prune_when_value_within_bounds() { + let pruner = ThresholdPruner::new().upper(10.0).lower(0.0); + let values = vec![(0, 3.0), (1, 5.0), (2, 7.0)]; + assert!(!pruner.should_prune(0, 2, &values, &[])); +} + +#[test] +fn no_prune_when_no_intermediate_values() { + let pruner = ThresholdPruner::new().upper(10.0).lower(0.0); + assert!(!pruner.should_prune(0, 0, &[], &[])); +} + +#[test] +fn works_with_only_upper_set() { + let pruner = ThresholdPruner::new().upper(5.0); + let below = vec![(0, 3.0)]; + let above = vec![(0, 6.0)]; + assert!(!pruner.should_prune(0, 0, &below, &[])); + assert!(pruner.should_prune(0, 0, &above, &[])); +} + +#[test] +fn works_with_only_lower_set() { + let pruner = ThresholdPruner::new().lower(2.0); + let above = vec![(0, 5.0)]; + let below = vec![(0, 1.0)]; + assert!(!pruner.should_prune(0, 0, &above, &[])); + assert!(pruner.should_prune(0, 0, &below, &[])); +} + +#[test] +fn works_with_both_thresholds_set() { + let pruner = ThresholdPruner::new().upper(10.0).lower(0.0); + + // Within bounds + assert!(!pruner.should_prune(0, 0, &[(0, 5.0)], &[])); + + // Exceeds upper + assert!(pruner.should_prune(0, 0, &[(0, 15.0)], &[])); + + // Below lower + assert!(pruner.should_prune(0, 0, &[(0, -3.0)], &[])); + + // At exact boundary (not pruned — strictly greater/less) + assert!(!pruner.should_prune(0, 0, &[(0, 10.0)], &[])); + assert!(!pruner.should_prune(0, 0, &[(0, 0.0)], &[])); +}