feat: add Study::top_trials(n) for retrieving best N trials
Returns the top N completed trials sorted by objective value, respecting the study's optimization direction. Pruned and failed trials are excluded.
This commit is contained in:
@@ -591,6 +591,37 @@ where
|
|||||||
self.best_trial().map(|trial| trial.value)
|
self.best_trial().map(|trial| trial.value)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the top `n` trials sorted by objective value.
|
||||||
|
///
|
||||||
|
/// For `Direction::Minimize`, returns trials with the lowest values.
|
||||||
|
/// For `Direction::Maximize`, returns trials with the highest values.
|
||||||
|
/// Only includes completed trials (not failed or pruned).
|
||||||
|
///
|
||||||
|
/// If fewer than `n` completed trials exist, returns all of them.
|
||||||
|
pub fn top_trials(&self, n: usize) -> Vec<CompletedTrial<V>>
|
||||||
|
where
|
||||||
|
V: Clone,
|
||||||
|
{
|
||||||
|
let trials = self.completed_trials.read();
|
||||||
|
let mut completed: Vec<_> = trials
|
||||||
|
.iter()
|
||||||
|
.filter(|t| t.state == TrialState::Complete)
|
||||||
|
.cloned()
|
||||||
|
.collect();
|
||||||
|
completed.sort_by(|a, b| match self.direction {
|
||||||
|
Direction::Minimize => a
|
||||||
|
.value
|
||||||
|
.partial_cmp(&b.value)
|
||||||
|
.unwrap_or(core::cmp::Ordering::Equal),
|
||||||
|
Direction::Maximize => b
|
||||||
|
.value
|
||||||
|
.partial_cmp(&a.value)
|
||||||
|
.unwrap_or(core::cmp::Ordering::Equal),
|
||||||
|
});
|
||||||
|
completed.truncate(n);
|
||||||
|
completed
|
||||||
|
}
|
||||||
|
|
||||||
/// Runs optimization with the given objective function.
|
/// Runs optimization with the given objective function.
|
||||||
///
|
///
|
||||||
/// This method runs `n_trials` evaluations sequentially. For each trial:
|
/// This method runs `n_trials` evaluations sequentially. For each trial:
|
||||||
|
|||||||
@@ -1522,3 +1522,82 @@ fn test_optimize_until_with_non_f64_value_type() {
|
|||||||
let best = study.best_trial().unwrap();
|
let best = study.best_trial().unwrap();
|
||||||
assert!(best.value >= 0);
|
assert!(best.value >= 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// =============================================================================
|
||||||
|
// Tests for top_trials
|
||||||
|
// =============================================================================
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_top_trials_minimize() {
|
||||||
|
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||||
|
|
||||||
|
// Manually complete trials with known values
|
||||||
|
for &val in &[5.0, 1.0, 3.0, 2.0, 4.0] {
|
||||||
|
let trial = study.create_trial();
|
||||||
|
study.complete_trial(trial, val);
|
||||||
|
}
|
||||||
|
|
||||||
|
let top3 = study.top_trials(3);
|
||||||
|
assert_eq!(top3.len(), 3);
|
||||||
|
assert_eq!(top3[0].value, 1.0);
|
||||||
|
assert_eq!(top3[1].value, 2.0);
|
||||||
|
assert_eq!(top3[2].value, 3.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_top_trials_maximize() {
|
||||||
|
let study: Study<f64> = Study::new(Direction::Maximize);
|
||||||
|
|
||||||
|
for &val in &[5.0, 1.0, 3.0, 2.0, 4.0] {
|
||||||
|
let trial = study.create_trial();
|
||||||
|
study.complete_trial(trial, val);
|
||||||
|
}
|
||||||
|
|
||||||
|
let top3 = study.top_trials(3);
|
||||||
|
assert_eq!(top3.len(), 3);
|
||||||
|
assert_eq!(top3[0].value, 5.0);
|
||||||
|
assert_eq!(top3[1].value, 4.0);
|
||||||
|
assert_eq!(top3[2].value, 3.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_top_trials_n_greater_than_total() {
|
||||||
|
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||||
|
|
||||||
|
for &val in &[3.0, 1.0] {
|
||||||
|
let trial = study.create_trial();
|
||||||
|
study.complete_trial(trial, val);
|
||||||
|
}
|
||||||
|
|
||||||
|
let top = study.top_trials(10);
|
||||||
|
assert_eq!(top.len(), 2);
|
||||||
|
assert_eq!(top[0].value, 1.0);
|
||||||
|
assert_eq!(top[1].value, 3.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_top_trials_empty() {
|
||||||
|
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||||
|
|
||||||
|
let top = study.top_trials(5);
|
||||||
|
assert!(top.is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_top_trials_excludes_pruned() {
|
||||||
|
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||||
|
|
||||||
|
// Complete some trials
|
||||||
|
for &val in &[5.0, 1.0, 3.0] {
|
||||||
|
let trial = study.create_trial();
|
||||||
|
study.complete_trial(trial, val);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Prune a trial (it gets a default value of 0.0 but should be excluded)
|
||||||
|
let trial = study.create_trial();
|
||||||
|
study.prune_trial(trial);
|
||||||
|
|
||||||
|
let top = study.top_trials(5);
|
||||||
|
assert_eq!(top.len(), 3, "pruned trial should be excluded");
|
||||||
|
assert_eq!(top[0].value, 1.0);
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user