Manuel Raimann
52f3c074dc
Add AttrValue enum (Float, Int, String, Bool) with From impls and set_user_attr/user_attr/user_attrs methods on both Trial and CompletedTrial. Attributes set during optimization are propagated through complete_trial and prune_trial. AttrValue is re-exported at crate root and in the prelude.
295 lines
9.7 KiB
Rust
295 lines
9.7 KiB
Rust
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(),
|
|
HashMap::new(),
|
|
)
|
|
}
|
|
|
|
#[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));
|
|
}
|
|
}
|