Compare commits
18 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bb79d98027 | |||
| f4b0631178 | |||
| 52f3c074dc | |||
| f67188e3a7 | |||
| 72da883add | |||
| bc278d83a3 | |||
| c586650df4 | |||
| cc78a92ed6 | |||
| 12c35e7cb4 | |||
| a885d10436 | |||
| 41150057d6 | |||
| a522fb2f34 | |||
| ec04757740 | |||
| 432c74b927 | |||
| 8199ba6b43 | |||
| abe15bdd7b | |||
| 4d8af3242b | |||
| 2651e61b2c |
+1
-1
@@ -3,7 +3,7 @@ members = ["optimizer-derive"]
|
||||
|
||||
[package]
|
||||
name = "optimizer"
|
||||
version = "0.5.1"
|
||||
version = "0.6.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.88"
|
||||
license = "MIT"
|
||||
|
||||
@@ -72,6 +72,10 @@ pub enum Error {
|
||||
got: usize,
|
||||
},
|
||||
|
||||
/// Returned when a trial is pruned (stopped early by the objective function).
|
||||
#[error("trial was pruned")]
|
||||
TrialPruned,
|
||||
|
||||
/// Returned when an internal invariant is violated.
|
||||
#[error("internal error: {0}")]
|
||||
Internal(&'static str),
|
||||
@@ -83,3 +87,33 @@ pub enum Error {
|
||||
}
|
||||
|
||||
pub type Result<T> = core::result::Result<T, Error>;
|
||||
|
||||
/// Convenience type for signalling a pruned trial from an objective function.
|
||||
///
|
||||
/// Implements `Into<Error>` so it can be used with `?` in objectives that
|
||||
/// return `Result<V, Error>`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::{Error, TrialPruned};
|
||||
///
|
||||
/// fn objective_that_prunes() -> Result<f64, Error> {
|
||||
/// // ... some computation ...
|
||||
/// Err(TrialPruned)?
|
||||
/// }
|
||||
/// ```
|
||||
#[derive(Debug)]
|
||||
pub struct TrialPruned;
|
||||
|
||||
impl core::fmt::Display for TrialPruned {
|
||||
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
||||
write!(f, "trial was pruned")
|
||||
}
|
||||
}
|
||||
|
||||
impl From<TrialPruned> for Error {
|
||||
fn from(_: TrialPruned) -> Self {
|
||||
Error::TrialPruned
|
||||
}
|
||||
}
|
||||
|
||||
+13
-4
@@ -188,24 +188,29 @@ mod error;
|
||||
mod kde;
|
||||
mod param;
|
||||
pub mod parameter;
|
||||
pub mod pruner;
|
||||
pub mod sampler;
|
||||
mod study;
|
||||
mod trial;
|
||||
mod types;
|
||||
|
||||
pub use error::{Error, Result};
|
||||
pub use error::{Error, Result, TrialPruned};
|
||||
#[cfg(feature = "derive")]
|
||||
pub use optimizer_derive::Categorical;
|
||||
pub use param::ParamValue;
|
||||
pub use parameter::{
|
||||
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, ParamId, Parameter,
|
||||
};
|
||||
pub use pruner::{
|
||||
HyperbandPruner, MedianPruner, NopPruner, PatientPruner, PercentilePruner, Pruner,
|
||||
SuccessiveHalvingPruner, ThresholdPruner, WilcoxonPruner,
|
||||
};
|
||||
pub use sampler::CompletedTrial;
|
||||
pub use sampler::grid::GridSearchSampler;
|
||||
pub use sampler::random::RandomSampler;
|
||||
pub use sampler::tpe::TpeSampler;
|
||||
pub use study::Study;
|
||||
pub use trial::Trial;
|
||||
pub use trial::{AttrValue, Trial};
|
||||
pub use types::{Direction, TrialState};
|
||||
|
||||
/// Convenient wildcard import for the most common types.
|
||||
@@ -217,16 +222,20 @@ pub mod prelude {
|
||||
#[cfg(feature = "derive")]
|
||||
pub use optimizer_derive::Categorical as DeriveCategory;
|
||||
|
||||
pub use crate::error::{Error, Result};
|
||||
pub use crate::error::{Error, Result, TrialPruned};
|
||||
pub use crate::param::ParamValue;
|
||||
pub use crate::parameter::{
|
||||
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter,
|
||||
};
|
||||
pub use crate::pruner::{
|
||||
HyperbandPruner, MedianPruner, NopPruner, PatientPruner, PercentilePruner, Pruner,
|
||||
SuccessiveHalvingPruner, ThresholdPruner,
|
||||
};
|
||||
pub use crate::sampler::CompletedTrial;
|
||||
pub use crate::sampler::grid::GridSearchSampler;
|
||||
pub use crate::sampler::random::RandomSampler;
|
||||
pub use crate::sampler::tpe::TpeSampler;
|
||||
pub use crate::study::Study;
|
||||
pub use crate::trial::Trial;
|
||||
pub use crate::trial::{AttrValue, Trial};
|
||||
pub use crate::types::Direction;
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@
|
||||
//! ```
|
||||
|
||||
use core::fmt::Debug;
|
||||
use core::ops::RangeInclusive;
|
||||
use core::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use crate::distribution::{
|
||||
@@ -187,6 +188,12 @@ impl FloatParam {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<RangeInclusive<f64>> for FloatParam {
|
||||
fn from(range: RangeInclusive<f64>) -> Self {
|
||||
FloatParam::new(*range.start(), *range.end())
|
||||
}
|
||||
}
|
||||
|
||||
impl Parameter for FloatParam {
|
||||
type Value = f64;
|
||||
|
||||
@@ -306,6 +313,12 @@ impl IntParam {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<RangeInclusive<i64>> for IntParam {
|
||||
fn from(range: RangeInclusive<i64>) -> Self {
|
||||
IntParam::new(*range.start(), *range.end())
|
||||
}
|
||||
}
|
||||
|
||||
impl Parameter for IntParam {
|
||||
type Value = i64;
|
||||
|
||||
@@ -992,4 +1005,34 @@ mod tests {
|
||||
let cloned = param.clone();
|
||||
assert_eq!(param.id(), cloned.id());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn float_param_from_range() {
|
||||
let param = FloatParam::from(0.0..=1.0);
|
||||
assert_eq!(
|
||||
param.distribution(),
|
||||
Distribution::Float(FloatDistribution {
|
||||
low: 0.0,
|
||||
high: 1.0,
|
||||
log_scale: false,
|
||||
step: None,
|
||||
})
|
||||
);
|
||||
assert_eq!(param.label(), format!("{param:?}"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn int_param_from_range() {
|
||||
let param = IntParam::from(1..=10);
|
||||
assert_eq!(
|
||||
param.distribution(),
|
||||
Distribution::Int(IntDistribution {
|
||||
low: 1,
|
||||
high: 10,
|
||||
log_scale: false,
|
||||
step: None,
|
||||
})
|
||||
);
|
||||
assert_eq!(param.label(), format!("{param:?}"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,559 @@
|
||||
use core::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// Hyperband pruner that manages multiple Successive Halving brackets.
|
||||
///
|
||||
/// Hyperband addresses SHA's sensitivity to the `min_resource` choice by
|
||||
/// running multiple brackets, each with a different tradeoff between the
|
||||
/// number of configurations and the starting budget:
|
||||
///
|
||||
/// - Bracket 0: many trials, very small starting budget (aggressive pruning)
|
||||
/// - Bracket 1: fewer trials, larger starting budget (moderate pruning)
|
||||
/// - ...
|
||||
/// - Bracket `s_max`: few trials, full budget (no pruning)
|
||||
///
|
||||
/// Trials are assigned to brackets in round-robin fashion. Each bracket
|
||||
/// runs SHA with its own `min_resource` and rung schedule.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Direction;
|
||||
/// use optimizer::pruner::HyperbandPruner;
|
||||
///
|
||||
/// let pruner = HyperbandPruner::new()
|
||||
/// .min_resource(1)
|
||||
/// .max_resource(81)
|
||||
/// .reduction_factor(3)
|
||||
/// .direction(Direction::Minimize);
|
||||
/// ```
|
||||
pub struct HyperbandPruner {
|
||||
min_resource: u64,
|
||||
max_resource: u64,
|
||||
reduction_factor: u64,
|
||||
direction: Direction,
|
||||
/// Tracks which bracket each trial belongs to.
|
||||
trial_brackets: Mutex<HashMap<u64, usize>>,
|
||||
/// Counter for round-robin bracket assignment.
|
||||
next_bracket: AtomicU64,
|
||||
}
|
||||
|
||||
impl HyperbandPruner {
|
||||
/// Create a new `HyperbandPruner` with default parameters.
|
||||
///
|
||||
/// Defaults: `min_resource=1`, `max_resource=81`, `reduction_factor=3`,
|
||||
/// `direction=Minimize`.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
min_resource: 1,
|
||||
max_resource: 81,
|
||||
reduction_factor: 3,
|
||||
direction: Direction::Minimize,
|
||||
trial_brackets: Mutex::new(HashMap::new()),
|
||||
next_bracket: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the minimum resource (budget) per trial.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `r` is 0.
|
||||
#[must_use]
|
||||
pub fn min_resource(mut self, r: u64) -> Self {
|
||||
assert!(r > 0, "min_resource must be > 0, got {r}");
|
||||
self.min_resource = r;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the maximum resource (budget) per trial.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `r` is 0.
|
||||
#[must_use]
|
||||
pub fn max_resource(mut self, r: u64) -> Self {
|
||||
assert!(r > 0, "max_resource must be > 0, got {r}");
|
||||
self.max_resource = r;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the reduction factor (eta). At each rung, the top 1/eta trials survive.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `eta` is less than 2.
|
||||
#[must_use]
|
||||
pub fn reduction_factor(mut self, eta: u64) -> Self {
|
||||
assert!(eta >= 2, "reduction_factor must be >= 2, got {eta}");
|
||||
self.reduction_factor = eta;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the optimization direction.
|
||||
#[must_use]
|
||||
pub fn direction(mut self, d: Direction) -> Self {
|
||||
self.direction = d;
|
||||
self
|
||||
}
|
||||
|
||||
/// Compute `s_max = floor(log(max_resource / min_resource) / log(eta))`.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn s_max(&self) -> u64 {
|
||||
let eta = self.reduction_factor as f64;
|
||||
let ratio = self.max_resource as f64 / self.min_resource as f64;
|
||||
(ratio.ln() / eta.ln()).floor() as u64
|
||||
}
|
||||
|
||||
/// Compute the rung steps for a given bracket `s`.
|
||||
///
|
||||
/// For bracket `s`, the starting resource is `max_resource / eta^(s_max - s)`,
|
||||
/// and rungs are spaced at powers of eta from there up to `max_resource`.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn rung_steps_for_bracket(&self, bracket: usize) -> Vec<u64> {
|
||||
let s_max = self.s_max();
|
||||
let eta = self.reduction_factor as f64;
|
||||
|
||||
// Starting resource for this bracket
|
||||
let exponent = s_max.saturating_sub(bracket as u64);
|
||||
let min_resource_bracket =
|
||||
(self.max_resource as f64 / eta.powi(exponent as i32)).ceil() as u64;
|
||||
|
||||
let mut steps = Vec::new();
|
||||
let mut rung: u32 = 0;
|
||||
while let Some(power) = self.reduction_factor.checked_pow(rung) {
|
||||
let step = min_resource_bracket.saturating_mul(power);
|
||||
if step > self.max_resource {
|
||||
break;
|
||||
}
|
||||
steps.push(step);
|
||||
rung += 1;
|
||||
}
|
||||
steps
|
||||
}
|
||||
|
||||
/// Assign a trial to a bracket (round-robin) and return the bracket index.
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
fn assign_bracket(&self, trial_id: u64) -> usize {
|
||||
let n_brackets = (self.s_max() + 1) as usize;
|
||||
let mut map = self.trial_brackets.lock().expect("lock poisoned");
|
||||
*map.entry(trial_id).or_insert_with(|| {
|
||||
let idx = self.next_bracket.fetch_add(1, Ordering::Relaxed);
|
||||
(idx as usize) % n_brackets
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for HyperbandPruner {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
impl Pruner for HyperbandPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let bracket = self.assign_bracket(trial_id);
|
||||
let rungs = self.rung_steps_for_bracket(bracket);
|
||||
|
||||
// Find the highest rung step <= current step
|
||||
let Some(&rung_step) = rungs.iter().rev().find(|&&r| r <= step) else {
|
||||
return false;
|
||||
};
|
||||
|
||||
// Never prune at the last rung (full budget)
|
||||
if rung_step >= self.max_resource {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Get the current trial's value at this rung step
|
||||
let current_value =
|
||||
if let Some(&(_, v)) = intermediate_values.iter().find(|(s, _)| *s == rung_step) {
|
||||
v
|
||||
} else if let Some(&(_, v)) = intermediate_values
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|(s, _)| *s <= rung_step)
|
||||
{
|
||||
v
|
||||
} else {
|
||||
return false;
|
||||
};
|
||||
|
||||
self.is_pruned_at_rung(current_value, rung_step, bracket, completed_trials)
|
||||
}
|
||||
}
|
||||
|
||||
impl HyperbandPruner {
|
||||
/// Determine whether a trial should be pruned at the given rung within its bracket.
|
||||
///
|
||||
/// Only compares against other trials in the same bracket.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn is_pruned_at_rung(
|
||||
&self,
|
||||
current_value: f64,
|
||||
rung_step: u64,
|
||||
bracket: usize,
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let eta = self.reduction_factor as usize;
|
||||
|
||||
// Collect values at this rung step from trials in the same bracket
|
||||
let map = self.trial_brackets.lock().expect("lock poisoned");
|
||||
let mut values_at_rung: Vec<f64> = completed_trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete || t.state == TrialState::Pruned)
|
||||
.filter(|t| map.get(&t.id).copied() == Some(bracket))
|
||||
.filter_map(|t| {
|
||||
t.intermediate_values
|
||||
.iter()
|
||||
.find(|(s, _)| *s == rung_step)
|
||||
.map(|(_, v)| *v)
|
||||
})
|
||||
.collect();
|
||||
drop(map);
|
||||
|
||||
// Need at least eta trials to make a meaningful comparison
|
||||
if values_at_rung.len() < eta {
|
||||
return false;
|
||||
}
|
||||
|
||||
values_at_rung.push(current_value);
|
||||
|
||||
values_at_rung
|
||||
.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
|
||||
if self.direction == Direction::Maximize {
|
||||
values_at_rung.reverse();
|
||||
}
|
||||
|
||||
let n_keep = (values_at_rung.len() as f64 / eta as f64).ceil() as usize;
|
||||
let threshold_idx = n_keep.max(1) - 1;
|
||||
let threshold = values_at_rung[threshold_idx];
|
||||
|
||||
match self.direction {
|
||||
Direction::Minimize => current_value > threshold,
|
||||
Direction::Maximize => current_value < threshold,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_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(),
|
||||
)
|
||||
}
|
||||
|
||||
fn make_pruned_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
|
||||
let mut t = make_trial(id, values);
|
||||
t.state = TrialState::Pruned;
|
||||
t
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn s_max_default() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// s_max = floor(ln(81/1) / ln(3)) = floor(4.0) = 4
|
||||
assert_eq!(pruner.s_max(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn s_max_custom() {
|
||||
let pruner = HyperbandPruner::new()
|
||||
.min_resource(1)
|
||||
.max_resource(16)
|
||||
.reduction_factor(2);
|
||||
// s_max = floor(ln(16) / ln(2)) = floor(4.0) = 4
|
||||
assert_eq!(pruner.s_max(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bracket_count() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// s_max=4, so brackets 0..=4 → 5 brackets
|
||||
assert_eq!(pruner.s_max() + 1, 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_bracket_0_default() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// Bracket 0: min_resource_bracket = ceil(81 / 3^4) = ceil(81/81) = 1
|
||||
// Rungs: 1, 3, 9, 27, 81
|
||||
assert_eq!(pruner.rung_steps_for_bracket(0), vec![1, 3, 9, 27, 81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_bracket_2_default() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// Bracket 2: min_resource_bracket = ceil(81 / 3^(4-2)) = ceil(81/9) = 9
|
||||
// Rungs: 9, 27, 81
|
||||
assert_eq!(pruner.rung_steps_for_bracket(2), vec![9, 27, 81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_bracket_4_default() {
|
||||
let pruner = HyperbandPruner::new();
|
||||
// Bracket 4 (s_max): min_resource_bracket = ceil(81 / 3^0) = 81
|
||||
// Rungs: 81 only (no pruning, full budget)
|
||||
assert_eq!(pruner.rung_steps_for_bracket(4), vec![81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_eta2() {
|
||||
let pruner = HyperbandPruner::new()
|
||||
.min_resource(1)
|
||||
.max_resource(16)
|
||||
.reduction_factor(2);
|
||||
// s_max = 4
|
||||
// Bracket 0: min=ceil(16/2^4)=1, rungs: 1,2,4,8,16
|
||||
assert_eq!(pruner.rung_steps_for_bracket(0), vec![1, 2, 4, 8, 16]);
|
||||
// Bracket 2: min=ceil(16/2^2)=4, rungs: 4,8,16
|
||||
assert_eq!(pruner.rung_steps_for_bracket(2), vec![4, 8, 16]);
|
||||
// Bracket 4: min=16, rungs: 16
|
||||
assert_eq!(pruner.rung_steps_for_bracket(4), vec![16]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_robin_bracket_assignment() {
|
||||
let pruner = HyperbandPruner::new(); // 5 brackets (0..=4)
|
||||
// Trials get assigned in round-robin: 0→0, 1→1, 2→2, 3→3, 4→4, 5→0, ...
|
||||
assert_eq!(pruner.assign_bracket(100), 0);
|
||||
assert_eq!(pruner.assign_bracket(101), 1);
|
||||
assert_eq!(pruner.assign_bracket(102), 2);
|
||||
assert_eq!(pruner.assign_bracket(103), 3);
|
||||
assert_eq!(pruner.assign_bracket(104), 4);
|
||||
assert_eq!(pruner.assign_bracket(105), 0); // wraps around
|
||||
|
||||
// Repeated calls for same trial return same bracket
|
||||
assert_eq!(pruner.assign_bracket(100), 0);
|
||||
assert_eq!(pruner.assign_bracket(103), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_before_first_rung() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
// Assign trial 0 to bracket 0 (rungs: 1, 3, 9, 27, 81)
|
||||
pruner.assign_bracket(0);
|
||||
|
||||
// Register completed trials in bracket 0
|
||||
let mut completed = Vec::new();
|
||||
for i in 1..=9 {
|
||||
pruner.assign_bracket(i);
|
||||
completed.push(make_trial(i, &[(1, i as f64)]));
|
||||
}
|
||||
|
||||
// Trial at step 0 (before rung 1) → don't prune
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_at_max_resource() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Put all trials in bracket 0
|
||||
let mut completed = Vec::new();
|
||||
for i in 0..9 {
|
||||
pruner.assign_bracket(i);
|
||||
completed.push(make_trial(i, &[(81, (i + 1) as f64)]));
|
||||
}
|
||||
|
||||
let trial_id = 9;
|
||||
pruner.assign_bracket(trial_id);
|
||||
// At max_resource (81), never prune
|
||||
assert!(!pruner.should_prune(trial_id, 81, &[(81, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_worst_in_bracket_minimize() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Force all trials into bracket 0 by assigning sequentially
|
||||
// With 5 brackets, trials 0,5,10,... go to bracket 0
|
||||
let bracket_0_ids: Vec<u64> = (0..5).map(|i| i * 5).collect();
|
||||
// Assign all 25 trial IDs to fill brackets
|
||||
for i in 0..25 {
|
||||
pruner.assign_bracket(i);
|
||||
}
|
||||
|
||||
// Create 9 completed trials in bracket 0 at rung step=1
|
||||
let completed: Vec<_> = bracket_0_ids
|
||||
.iter()
|
||||
.take(3)
|
||||
.enumerate()
|
||||
.map(|(idx, &id)| make_trial(id, &[(1, (idx + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// Trial 25 → bracket 0 (25 % 5 == 0)
|
||||
let test_id = 25;
|
||||
pruner.assign_bracket(test_id);
|
||||
assert_eq!(pruner.assign_bracket(test_id), 0);
|
||||
|
||||
// 3 completed + 1 current = 4. eta=3. ceil(4/3)=2. Threshold = 2.0
|
||||
// Value 2.0 → keep
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 2.0)], &completed));
|
||||
// Value 3.0 → prune
|
||||
assert!(pruner.should_prune(test_id, 1, &[(1, 3.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_worst_in_bracket_maximize() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Maximize);
|
||||
|
||||
// Assign trials so they end up in bracket 0
|
||||
for i in 0..25 {
|
||||
pruner.assign_bracket(i);
|
||||
}
|
||||
|
||||
let completed: Vec<_> = [0u64, 5, 10]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, &id)| make_trial(id, &[(1, (idx + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
let test_id = 25;
|
||||
pruner.assign_bracket(test_id);
|
||||
|
||||
// For maximize, best = highest. Values: 1,2,3 + current
|
||||
// Value 2.0 → keep (threshold = 2.0 when sorted desc: 3,2,current,1)
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 2.0)], &completed));
|
||||
// Value 1.0 → prune
|
||||
assert!(pruner.should_prune(test_id, 1, &[(1, 0.5)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_brackets_have_different_aggressiveness() {
|
||||
let pruner = HyperbandPruner::new()
|
||||
.min_resource(1)
|
||||
.max_resource(81)
|
||||
.reduction_factor(3)
|
||||
.direction(Direction::Minimize);
|
||||
|
||||
let rungs_0 = pruner.rung_steps_for_bracket(0);
|
||||
let rungs_2 = pruner.rung_steps_for_bracket(2);
|
||||
let rungs_4 = pruner.rung_steps_for_bracket(4);
|
||||
|
||||
// Bracket 0 has the most rungs (most aggressive)
|
||||
assert!(rungs_0.len() > rungs_2.len());
|
||||
// Bracket 4 has just 1 rung (no pruning)
|
||||
assert_eq!(rungs_4.len(), 1);
|
||||
// Bracket 0 starts earliest
|
||||
assert!(rungs_0[0] < rungs_2[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trials_in_different_brackets_independent() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Assign trials: bracket 0 gets IDs 0,5,10,15,20
|
||||
for i in 0..25 {
|
||||
pruner.assign_bracket(i);
|
||||
}
|
||||
|
||||
// Bracket 0 trials: bad values at rung step=1
|
||||
let bracket_0_trials: Vec<_> = [0u64, 5, 10]
|
||||
.iter()
|
||||
.map(|&id| make_trial(id, &[(1, 100.0)]))
|
||||
.collect();
|
||||
|
||||
// Bracket 1 trials: good values at rung step=1
|
||||
let bracket_1_trials: Vec<_> = [1u64, 6, 11]
|
||||
.iter()
|
||||
.map(|&id| make_trial(id, &[(1, 1.0)]))
|
||||
.collect();
|
||||
|
||||
let mut all_trials = bracket_0_trials;
|
||||
all_trials.extend(bracket_1_trials);
|
||||
|
||||
// A new bracket-0 trial with value 50 should be compared against
|
||||
// bracket-0 peers (100,100,100), not bracket-1 peers (1,1,1)
|
||||
let test_id = 25; // bracket 0
|
||||
pruner.assign_bracket(test_id);
|
||||
// 3 peers at 100.0 + current at 50.0. ceil(4/3)=2. Sorted: 50,100,100,100. Threshold=100.0
|
||||
// Value 50.0 < 100.0 → keep
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 50.0)], &all_trials));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn includes_pruned_trials() {
|
||||
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
|
||||
|
||||
for i in 0..25 {
|
||||
pruner.assign_bracket(i);
|
||||
}
|
||||
|
||||
let completed = vec![
|
||||
make_trial(0, &[(1, 1.0)]),
|
||||
make_pruned_trial(5, &[(1, 8.0)]),
|
||||
make_pruned_trial(10, &[(1, 9.0)]),
|
||||
];
|
||||
|
||||
let test_id = 25;
|
||||
pruner.assign_bracket(test_id);
|
||||
|
||||
// Values: 1.0, 8.0, 9.0 + current. eta=3.
|
||||
// Value 1.0 → keep
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 1.0)], &completed));
|
||||
// Value 5.0 → prune (sorted: 1,5,8,9 → keep ceil(4/3)=2 → threshold=5.0, 5.0 not > 5.0 → keep)
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 5.0)], &completed));
|
||||
// Value 6.0 → prune (sorted: 1,6,8,9 → threshold=6.0, 6.0 not > 6.0 → keep)
|
||||
assert!(!pruner.should_prune(test_id, 1, &[(1, 6.0)], &completed));
|
||||
// Value 9.5 → prune
|
||||
assert!(pruner.should_prune(test_id, 1, &[(1, 9.5)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "min_resource must be > 0")]
|
||||
fn rejects_zero_min_resource() {
|
||||
let _ = HyperbandPruner::new().min_resource(0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "max_resource must be > 0")]
|
||||
fn rejects_zero_max_resource() {
|
||||
let _ = HyperbandPruner::new().max_resource(0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "reduction_factor must be >= 2")]
|
||||
fn rejects_reduction_factor_one() {
|
||||
let _ = HyperbandPruner::new().reduction_factor(1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
use super::Pruner;
|
||||
use super::percentile::compute_percentile;
|
||||
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.
|
||||
///
|
||||
/// Equivalent to `PercentilePruner::new(50.0, direction)`.
|
||||
///
|
||||
/// # 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 (50th percentile)
|
||||
let median = compute_percentile(&mut values_at_step, 50.0);
|
||||
|
||||
// 5. Compare against median based on direction
|
||||
match self.direction {
|
||||
Direction::Minimize => current_value > median,
|
||||
Direction::Maximize => current_value < median,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn compute_median_odd() {
|
||||
assert!((compute_percentile(&mut [3.0, 1.0, 2.0], 50.0) - 2.0).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compute_median_even() {
|
||||
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_percentile(&mut [5.0], 50.0) - 5.0).abs() < f64::EPSILON);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
//! Pruner trait and implementations for trial pruning.
|
||||
//!
|
||||
//! Pruners decide whether to stop (prune) a trial early based on its
|
||||
//! intermediate values compared to other trials. This is useful for
|
||||
//! discarding unpromising trials before they complete, saving compute.
|
||||
|
||||
mod hyperband;
|
||||
mod median;
|
||||
mod nop;
|
||||
mod patient;
|
||||
pub(crate) mod percentile;
|
||||
mod successive_halving;
|
||||
mod threshold;
|
||||
mod wilcoxon;
|
||||
|
||||
pub use hyperband::HyperbandPruner;
|
||||
pub use median::MedianPruner;
|
||||
pub use nop::NopPruner;
|
||||
pub use patient::PatientPruner;
|
||||
pub use percentile::PercentilePruner;
|
||||
pub use successive_halving::SuccessiveHalvingPruner;
|
||||
pub use threshold::ThresholdPruner;
|
||||
pub use wilcoxon::WilcoxonPruner;
|
||||
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
/// Trait for pluggable trial pruning strategies.
|
||||
///
|
||||
/// Pruners are consulted after each intermediate value is reported to
|
||||
/// decide whether the trial should be stopped early. The trait requires
|
||||
/// `Send + Sync` to support concurrent and async optimization.
|
||||
///
|
||||
/// # Implementing a custom pruner
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::pruner::Pruner;
|
||||
/// use optimizer::sampler::CompletedTrial;
|
||||
///
|
||||
/// struct MyPruner {
|
||||
/// threshold: f64,
|
||||
/// }
|
||||
///
|
||||
/// impl Pruner for MyPruner {
|
||||
/// fn should_prune(
|
||||
/// &self,
|
||||
/// _trial_id: u64,
|
||||
/// _step: u64,
|
||||
/// intermediate_values: &[(u64, f64)],
|
||||
/// _completed_trials: &[CompletedTrial],
|
||||
/// ) -> bool {
|
||||
/// // Prune if the latest value exceeds the threshold
|
||||
/// intermediate_values
|
||||
/// .last()
|
||||
/// .is_some_and(|&(_, v)| v > self.threshold)
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub trait Pruner: Send + Sync {
|
||||
/// Decide whether to prune a trial at the given step.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `trial_id` - The current trial's ID.
|
||||
/// * `step` - The step at which the intermediate value was reported.
|
||||
/// * `intermediate_values` - All `(step, value)` pairs reported so far for this trial.
|
||||
/// * `completed_trials` - History of all completed trials (for comparison).
|
||||
fn should_prune(
|
||||
&self,
|
||||
trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool;
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
/// A pruner that never prunes. This is the default when no pruner is configured.
|
||||
pub struct NopPruner;
|
||||
|
||||
impl Pruner for NopPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
_step: u64,
|
||||
_intermediate_values: &[(u64, f64)],
|
||||
_completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
/// Wraps another pruner and adds a patience window.
|
||||
///
|
||||
/// The inner pruner must recommend pruning for `patience` consecutive
|
||||
/// steps before this pruner actually prunes the trial. This is useful
|
||||
/// to prevent premature pruning when intermediate values are noisy.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::pruner::{PatientPruner, ThresholdPruner};
|
||||
///
|
||||
/// // Only prune after the threshold pruner recommends pruning 3 times in a row
|
||||
/// let inner = ThresholdPruner::new().upper(100.0);
|
||||
/// let pruner = PatientPruner::new(inner, 3);
|
||||
/// ```
|
||||
pub struct PatientPruner {
|
||||
inner: Box<dyn Pruner>,
|
||||
patience: u64,
|
||||
/// Track consecutive prune recommendations per trial.
|
||||
consecutive_counts: Mutex<HashMap<u64, u64>>,
|
||||
}
|
||||
|
||||
impl PatientPruner {
|
||||
/// Create a new `PatientPruner` wrapping the given inner pruner.
|
||||
///
|
||||
/// The inner pruner must recommend pruning for `patience` consecutive
|
||||
/// calls before this pruner returns `true`.
|
||||
pub fn new(inner: impl Pruner + 'static, patience: u64) -> Self {
|
||||
Self {
|
||||
inner: Box::new(inner),
|
||||
patience,
|
||||
consecutive_counts: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Pruner for PatientPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let inner_says_prune =
|
||||
self.inner
|
||||
.should_prune(trial_id, step, intermediate_values, completed_trials);
|
||||
let mut counts = self.consecutive_counts.lock().expect("lock poisoned");
|
||||
let count = counts.entry(trial_id).or_insert(0);
|
||||
if inner_says_prune {
|
||||
*count += 1;
|
||||
*count >= self.patience
|
||||
} else {
|
||||
*count = 0;
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::pruner::ThresholdPruner;
|
||||
|
||||
/// A test pruner that always returns the given value.
|
||||
struct ConstPruner(bool);
|
||||
|
||||
impl Pruner for ConstPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
_step: u64,
|
||||
_intermediate_values: &[(u64, f64)],
|
||||
_completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
/// A pruner that returns values from a sequence.
|
||||
struct SequencePruner(Mutex<Vec<bool>>);
|
||||
|
||||
impl Pruner for SequencePruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
_step: u64,
|
||||
_intermediate_values: &[(u64, f64)],
|
||||
_completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
self.0.lock().expect("lock poisoned").remove(0)
|
||||
}
|
||||
}
|
||||
|
||||
fn call(pruner: &PatientPruner, trial_id: u64, step: u64) -> bool {
|
||||
pruner.should_prune(trial_id, step, &[(step, 0.0)], &[])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patience_1_behaves_like_inner() {
|
||||
let pruner = PatientPruner::new(ConstPruner(true), 1);
|
||||
assert!(call(&pruner, 0, 0));
|
||||
assert!(call(&pruner, 0, 1));
|
||||
|
||||
let pruner = PatientPruner::new(ConstPruner(false), 1);
|
||||
assert!(!call(&pruner, 0, 0));
|
||||
assert!(!call(&pruner, 0, 1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patience_3_requires_consecutive_recommendations() {
|
||||
let pruner = PatientPruner::new(ConstPruner(true), 3);
|
||||
assert!(!call(&pruner, 0, 0)); // count=1
|
||||
assert!(!call(&pruner, 0, 1)); // count=2
|
||||
assert!(call(&pruner, 0, 2)); // count=3 → prune
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn counter_resets_on_no_prune() {
|
||||
// Sequence: prune, prune, no-prune, prune, prune, prune
|
||||
let seq = vec![true, true, false, true, true, true];
|
||||
let pruner = PatientPruner::new(SequencePruner(Mutex::new(seq)), 3);
|
||||
|
||||
assert!(!call(&pruner, 0, 0)); // count=1
|
||||
assert!(!call(&pruner, 0, 1)); // count=2
|
||||
assert!(!call(&pruner, 0, 2)); // reset → count=0
|
||||
assert!(!call(&pruner, 0, 3)); // count=1
|
||||
assert!(!call(&pruner, 0, 4)); // count=2
|
||||
assert!(call(&pruner, 0, 5)); // count=3 → prune
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn independent_per_trial() {
|
||||
let pruner = PatientPruner::new(ConstPruner(true), 2);
|
||||
assert!(!call(&pruner, 0, 0)); // trial 0: count=1
|
||||
assert!(!call(&pruner, 1, 0)); // trial 1: count=1
|
||||
assert!(call(&pruner, 0, 1)); // trial 0: count=2 → prune
|
||||
assert!(!call(&pruner, 2, 0)); // trial 2: count=1
|
||||
assert!(call(&pruner, 1, 1)); // trial 1: count=2 → prune
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn works_with_threshold_pruner() {
|
||||
let inner = ThresholdPruner::new().upper(10.0);
|
||||
let pruner = PatientPruner::new(inner, 2);
|
||||
|
||||
// Value below threshold → inner says no
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 5.0)], &[]));
|
||||
// Value above threshold → inner says yes, count=1
|
||||
assert!(!pruner.should_prune(0, 1, &[(0, 5.0), (1, 15.0)], &[]));
|
||||
// Value above threshold again → count=2 → prune
|
||||
assert!(pruner.should_prune(0, 2, &[(0, 5.0), (1, 15.0), (2, 20.0)], &[]));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
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));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,466 @@
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// Successive Halving pruner based on the SHA algorithm.
|
||||
///
|
||||
/// Trials are evaluated at exponentially-spaced "rungs". At each rung,
|
||||
/// only the top 1/eta fraction of trials survive to the next rung.
|
||||
///
|
||||
/// For example, with `min_resource=1`, `max_resource=81`, `reduction_factor=3`:
|
||||
/// - Rung 0: evaluate at step 1, keep top 1/3
|
||||
/// - Rung 1: evaluate at step 3, keep top 1/3
|
||||
/// - Rung 2: evaluate at step 9, keep top 1/3
|
||||
/// - Rung 3: evaluate at step 27, keep top 1/3
|
||||
/// - Rung 4: evaluate at step 81 (full budget)
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Direction;
|
||||
/// use optimizer::pruner::SuccessiveHalvingPruner;
|
||||
///
|
||||
/// let pruner = SuccessiveHalvingPruner::new()
|
||||
/// .min_resource(1)
|
||||
/// .max_resource(81)
|
||||
/// .reduction_factor(3)
|
||||
/// .direction(Direction::Minimize);
|
||||
/// ```
|
||||
pub struct SuccessiveHalvingPruner {
|
||||
min_resource: u64,
|
||||
max_resource: u64,
|
||||
reduction_factor: u64,
|
||||
min_early_stopping_rate: u64,
|
||||
direction: Direction,
|
||||
}
|
||||
|
||||
impl SuccessiveHalvingPruner {
|
||||
/// Create a new `SuccessiveHalvingPruner` with default parameters.
|
||||
///
|
||||
/// Defaults: `min_resource=1`, `max_resource=81`, `reduction_factor=3`,
|
||||
/// `min_early_stopping_rate=0`, `direction=Minimize`.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
min_resource: 1,
|
||||
max_resource: 81,
|
||||
reduction_factor: 3,
|
||||
min_early_stopping_rate: 0,
|
||||
direction: Direction::Minimize,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the minimum resource (budget) per trial.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `r` is 0.
|
||||
#[must_use]
|
||||
pub fn min_resource(mut self, r: u64) -> Self {
|
||||
assert!(r > 0, "min_resource must be > 0, got {r}");
|
||||
self.min_resource = r;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the maximum resource (budget) per trial.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `r` is 0.
|
||||
#[must_use]
|
||||
pub fn max_resource(mut self, r: u64) -> Self {
|
||||
assert!(r > 0, "max_resource must be > 0, got {r}");
|
||||
self.max_resource = r;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the reduction factor (eta). At each rung, the top 1/eta trials survive.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `eta` is less than 2.
|
||||
#[must_use]
|
||||
pub fn reduction_factor(mut self, eta: u64) -> Self {
|
||||
assert!(eta >= 2, "reduction_factor must be >= 2, got {eta}");
|
||||
self.reduction_factor = eta;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the minimum early stopping rate. Skips the first N rungs.
|
||||
#[must_use]
|
||||
pub fn min_early_stopping_rate(mut self, n: u64) -> Self {
|
||||
self.min_early_stopping_rate = n;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the optimization direction.
|
||||
#[must_use]
|
||||
pub fn direction(mut self, d: Direction) -> Self {
|
||||
self.direction = d;
|
||||
self
|
||||
}
|
||||
|
||||
/// Compute the rung steps: `[min_resource * eta^(s), ...]` up to `max_resource`,
|
||||
/// skipping the first `min_early_stopping_rate` rungs.
|
||||
fn rung_steps(&self) -> Vec<u64> {
|
||||
let eta = self.reduction_factor;
|
||||
let mut steps = Vec::new();
|
||||
let mut rung: u32 = 0;
|
||||
while let Some(power) = eta.checked_pow(rung) {
|
||||
let step = self.min_resource.saturating_mul(power);
|
||||
if step > self.max_resource {
|
||||
break;
|
||||
}
|
||||
if u64::from(rung) >= self.min_early_stopping_rate {
|
||||
steps.push(step);
|
||||
}
|
||||
rung += 1;
|
||||
}
|
||||
steps
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SuccessiveHalvingPruner {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
impl Pruner for SuccessiveHalvingPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let rungs = self.rung_steps();
|
||||
|
||||
// Find the highest rung step <= current step
|
||||
let Some(&rung_step) = rungs.iter().rev().find(|&&r| r <= step) else {
|
||||
// No rung matches (before the first rung) → don't prune
|
||||
return false;
|
||||
};
|
||||
|
||||
// If this is the last rung (full budget), don't prune
|
||||
if rung_step >= self.max_resource {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Get the current trial's value at this rung step
|
||||
let Some(&(_, current_value)) = intermediate_values.iter().find(|(s, _)| *s == rung_step)
|
||||
else {
|
||||
// Trial hasn't reported a value at this exact rung step.
|
||||
// Use the latest intermediate value at or before the rung step instead.
|
||||
let Some(&(_, current_value)) = intermediate_values
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|(s, _)| *s <= rung_step)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
return self.is_pruned_at_rung(current_value, rung_step, completed_trials);
|
||||
};
|
||||
|
||||
self.is_pruned_at_rung(current_value, rung_step, completed_trials)
|
||||
}
|
||||
}
|
||||
|
||||
impl SuccessiveHalvingPruner {
|
||||
/// Determine whether a trial with `current_value` should be pruned at the given rung.
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn is_pruned_at_rung(
|
||||
&self,
|
||||
current_value: f64,
|
||||
rung_step: u64,
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let eta = self.reduction_factor as usize;
|
||||
|
||||
// Collect values at this rung step from all trials that reached it
|
||||
let mut values_at_rung: Vec<f64> = completed_trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete || t.state == TrialState::Pruned)
|
||||
.filter_map(|t| {
|
||||
t.intermediate_values
|
||||
.iter()
|
||||
.find(|(s, _)| *s == rung_step)
|
||||
.map(|(_, v)| *v)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Need at least eta trials to make a meaningful comparison
|
||||
// (with fewer trials, we can't determine the top 1/eta fraction)
|
||||
if values_at_rung.len() < eta {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Include the current trial's value for ranking
|
||||
values_at_rung.push(current_value);
|
||||
|
||||
// Sort based on direction: best values first
|
||||
values_at_rung
|
||||
.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
|
||||
if self.direction == Direction::Maximize {
|
||||
values_at_rung.reverse();
|
||||
}
|
||||
|
||||
// Keep top 1/eta fraction
|
||||
let n_keep = (values_at_rung.len() as f64 / eta as f64).ceil() as usize;
|
||||
let threshold_idx = n_keep.max(1) - 1;
|
||||
let threshold = values_at_rung[threshold_idx];
|
||||
|
||||
// Prune if current value is worse than the threshold
|
||||
match self.direction {
|
||||
Direction::Minimize => current_value > threshold,
|
||||
Direction::Maximize => current_value < threshold,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_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(),
|
||||
)
|
||||
}
|
||||
|
||||
fn make_pruned_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
|
||||
let mut t = make_trial(id, values);
|
||||
t.state = TrialState::Pruned;
|
||||
t
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_default() {
|
||||
let pruner = SuccessiveHalvingPruner::new();
|
||||
let rungs = pruner.rung_steps();
|
||||
// min=1, max=81, eta=3 → 1, 3, 9, 27, 81
|
||||
assert_eq!(rungs, vec![1, 3, 9, 27, 81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_custom() {
|
||||
let pruner = SuccessiveHalvingPruner::new()
|
||||
.min_resource(2)
|
||||
.max_resource(32)
|
||||
.reduction_factor(2);
|
||||
let rungs = pruner.rung_steps();
|
||||
// 2, 4, 8, 16, 32
|
||||
assert_eq!(rungs, vec![2, 4, 8, 16, 32]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rung_steps_with_early_stopping_rate() {
|
||||
let pruner = SuccessiveHalvingPruner::new().min_early_stopping_rate(2);
|
||||
let rungs = pruner.rung_steps();
|
||||
// Skip rung 0 (step=1) and rung 1 (step=3), keep rung 2+ (9, 27, 81)
|
||||
assert_eq!(rungs, vec![9, 27, 81]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_before_first_rung() {
|
||||
let pruner = SuccessiveHalvingPruner::new()
|
||||
.min_resource(10)
|
||||
.max_resource(100)
|
||||
.reduction_factor(3);
|
||||
let completed = vec![
|
||||
make_trial(0, &[(5, 1.0)]),
|
||||
make_trial(1, &[(5, 2.0)]),
|
||||
make_trial(2, &[(5, 3.0)]),
|
||||
];
|
||||
// Step 5 is before the first rung (10)
|
||||
assert!(!pruner.should_prune(3, 5, &[(5, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_with_single_trial() {
|
||||
let pruner = SuccessiveHalvingPruner::new();
|
||||
let completed = vec![make_trial(0, &[(1, 5.0)]), make_trial(1, &[(1, 3.0)])];
|
||||
// Only 2 completed trials at rung + 1 current = 3 total, threshold = ceil(3/3) = 1
|
||||
// With eta=3, we need at least 3 completed trials
|
||||
assert!(!pruner.should_prune(2, 1, &[(1, 10.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_worst_trials_at_rung() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// 9 completed trials at rung step=1, with values 1..=9
|
||||
let completed: Vec<_> = (0..9)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// With eta=3, keep top 1/3. 10 total values → ceil(10/3) = 4 kept
|
||||
// Best 4 values: 1, 2, 3, 4. Threshold = 4.0
|
||||
// Value 3.0 → keep (in top 1/3)
|
||||
assert!(!pruner.should_prune(9, 1, &[(1, 3.0)], &completed));
|
||||
// Value 5.0 → prune (not in top 1/3)
|
||||
assert!(pruner.should_prune(9, 1, &[(1, 5.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn top_fraction_survives() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// 6 completed trials at step=1
|
||||
let completed: Vec<_> = (0..6)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// 7 total (6 + current). ceil(7/3) = 3 keep. Threshold = 3.0
|
||||
// Value 2.0 → keep
|
||||
assert!(!pruner.should_prune(6, 1, &[(1, 2.0)], &completed));
|
||||
// Value 3.0 → keep (at threshold)
|
||||
assert!(!pruner.should_prune(6, 1, &[(1, 3.0)], &completed));
|
||||
// Value 4.0 → prune
|
||||
assert!(pruner.should_prune(6, 1, &[(1, 4.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maximize_direction() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Maximize);
|
||||
|
||||
let completed: Vec<_> = (0..6)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// 7 total. For maximize, best = highest. ceil(7/3)=3. Top 3: 6,5,4. Threshold=4.0
|
||||
// Value 5.0 → keep
|
||||
assert!(!pruner.should_prune(6, 1, &[(1, 5.0)], &completed));
|
||||
// Value 4.0 → keep (at threshold)
|
||||
assert!(!pruner.should_prune(6, 1, &[(1, 4.0)], &completed));
|
||||
// Value 3.0 → prune
|
||||
assert!(pruner.should_prune(6, 1, &[(1, 3.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reduction_factor_2() {
|
||||
let pruner = SuccessiveHalvingPruner::new()
|
||||
.reduction_factor(2)
|
||||
.min_resource(1)
|
||||
.max_resource(16)
|
||||
.direction(Direction::Minimize);
|
||||
|
||||
// Rungs: 1, 2, 4, 8, 16
|
||||
assert_eq!(pruner.rung_steps(), vec![1, 2, 4, 8, 16]);
|
||||
|
||||
// 4 completed trials at rung step=1
|
||||
let completed: Vec<_> = (0..4)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// With eta=2, 5 total. ceil(5/2) = 3 keep. Threshold = 3.0
|
||||
// Value 3.0 → keep
|
||||
assert!(!pruner.should_prune(4, 1, &[(1, 3.0)], &completed));
|
||||
// Value 4.0 → prune
|
||||
assert!(pruner.should_prune(4, 1, &[(1, 4.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reduction_factor_4() {
|
||||
let pruner = SuccessiveHalvingPruner::new()
|
||||
.reduction_factor(4)
|
||||
.min_resource(1)
|
||||
.max_resource(64)
|
||||
.direction(Direction::Minimize);
|
||||
|
||||
// Rungs: 1, 4, 16, 64
|
||||
assert_eq!(pruner.rung_steps(), vec![1, 4, 16, 64]);
|
||||
|
||||
// 12 completed trials at rung step=1
|
||||
let completed: Vec<_> = (0..12)
|
||||
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// With eta=4, 13 total. ceil(13/4) = 4 keep. Threshold = 4.0
|
||||
// Value 4.0 → keep
|
||||
assert!(!pruner.should_prune(12, 1, &[(1, 4.0)], &completed));
|
||||
// Value 5.0 → prune
|
||||
assert!(pruner.should_prune(12, 1, &[(1, 5.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_contiguous_steps() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Trials reporting at rung step=3 (not step=1)
|
||||
let completed: Vec<_> = (0..6)
|
||||
.map(|i| make_trial(i, &[(3, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// Current trial reports at step 5 (between rung 3 and rung 9)
|
||||
// Highest rung <= 5 is 3. Use value at rung step 3.
|
||||
// Trial has value at step 3 → use it
|
||||
assert!(!pruner.should_prune(6, 5, &[(3, 2.0)], &completed));
|
||||
assert!(pruner.should_prune(6, 5, &[(3, 5.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_at_max_resource() {
|
||||
let pruner = SuccessiveHalvingPruner::new();
|
||||
let completed: Vec<_> = (0..9)
|
||||
.map(|i| make_trial(i, &[(81, (i + 1) as f64)]))
|
||||
.collect();
|
||||
|
||||
// At the max resource rung, never prune (trial should complete)
|
||||
assert!(!pruner.should_prune(9, 81, &[(81, 100.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn includes_pruned_trials_in_comparison() {
|
||||
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
|
||||
|
||||
// Mix of completed and pruned trials at rung step=1
|
||||
let completed = vec![
|
||||
make_trial(0, &[(1, 1.0)]),
|
||||
make_trial(1, &[(1, 2.0)]),
|
||||
make_pruned_trial(2, &[(1, 8.0)]),
|
||||
make_pruned_trial(3, &[(1, 9.0)]),
|
||||
make_pruned_trial(4, &[(1, 10.0)]),
|
||||
];
|
||||
|
||||
// 6 total. ceil(6/3) = 2 keep. Threshold = 2.0
|
||||
// Value 2.0 → keep
|
||||
assert!(!pruner.should_prune(5, 1, &[(1, 2.0)], &completed));
|
||||
// Value 3.0 → prune
|
||||
assert!(pruner.should_prune(5, 1, &[(1, 3.0)], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "min_resource must be > 0")]
|
||||
fn rejects_zero_min_resource() {
|
||||
let _ = SuccessiveHalvingPruner::new().min_resource(0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "max_resource must be > 0")]
|
||||
fn rejects_zero_max_resource() {
|
||||
let _ = SuccessiveHalvingPruner::new().max_resource(0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "reduction_factor must be >= 2")]
|
||||
fn rejects_reduction_factor_one() {
|
||||
let _ = SuccessiveHalvingPruner::new().reduction_factor(1);
|
||||
}
|
||||
}
|
||||
@@ -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<f64>,
|
||||
/// Prune if intermediate value is less than this. `None` = no lower bound.
|
||||
lower: Option<f64>,
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,531 @@
|
||||
use core::cmp::Ordering;
|
||||
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// Prune trials using a Wilcoxon signed-rank test comparing intermediate
|
||||
/// values against the best completed trial.
|
||||
///
|
||||
/// More principled than `MedianPruner` for noisy objectives — it accounts
|
||||
/// for the paired nature of step-aligned comparisons and doesn't prune
|
||||
/// on random fluctuations.
|
||||
///
|
||||
/// The test compares intermediate values at matching steps between the
|
||||
/// current trial and the best completed trial. If the current trial is
|
||||
/// statistically significantly worse (p < threshold), it is pruned.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Direction;
|
||||
/// use optimizer::pruner::WilcoxonPruner;
|
||||
///
|
||||
/// let pruner = WilcoxonPruner::new(Direction::Minimize)
|
||||
/// .p_value_threshold(0.05)
|
||||
/// .n_warmup_steps(5)
|
||||
/// .n_min_trials(1);
|
||||
/// ```
|
||||
pub struct WilcoxonPruner {
|
||||
/// Significance level (default 0.05). Lower = more conservative.
|
||||
p_value_threshold: 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 WilcoxonPruner {
|
||||
/// Create a new `WilcoxonPruner` for the given optimization direction.
|
||||
///
|
||||
/// By default, `p_value_threshold` is 0.05, `n_warmup_steps` is 0,
|
||||
/// and `n_min_trials` is 1.
|
||||
#[must_use]
|
||||
pub fn new(direction: Direction) -> Self {
|
||||
Self {
|
||||
p_value_threshold: 0.05,
|
||||
n_warmup_steps: 0,
|
||||
n_min_trials: 1,
|
||||
direction,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the p-value threshold for significance.
|
||||
///
|
||||
/// Must be in (0.0, 1.0). Lower values are more conservative (harder to prune).
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if `p` is not in the open interval (0.0, 1.0).
|
||||
#[must_use]
|
||||
pub fn p_value_threshold(mut self, p: f64) -> Self {
|
||||
assert!(
|
||||
p > 0.0 && p < 1.0,
|
||||
"p_value_threshold must be in (0.0, 1.0)"
|
||||
);
|
||||
self.p_value_threshold = p;
|
||||
self
|
||||
}
|
||||
|
||||
/// 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 WilcoxonPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
if step < self.n_warmup_steps {
|
||||
return false;
|
||||
}
|
||||
|
||||
let completed: Vec<&CompletedTrial> = completed_trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete)
|
||||
.collect();
|
||||
|
||||
if completed.len() < self.n_min_trials {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Find the best completed trial by final objective value.
|
||||
let best = match self.direction {
|
||||
Direction::Minimize => completed
|
||||
.iter()
|
||||
.min_by(|a, b| a.value.partial_cmp(&b.value).unwrap_or(Ordering::Equal)),
|
||||
Direction::Maximize => completed
|
||||
.iter()
|
||||
.max_by(|a, b| a.value.partial_cmp(&b.value).unwrap_or(Ordering::Equal)),
|
||||
};
|
||||
|
||||
let Some(best) = best else {
|
||||
return false;
|
||||
};
|
||||
|
||||
// Pair intermediate values at matching steps.
|
||||
let pairs: Vec<(f64, f64)> = intermediate_values
|
||||
.iter()
|
||||
.filter_map(|&(s, current_v)| {
|
||||
best.intermediate_values
|
||||
.iter()
|
||||
.find(|(bs, _)| *bs == s)
|
||||
.map(|&(_, best_v)| (current_v, best_v))
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Need at least 6 pairs for a meaningful test.
|
||||
if pairs.len() < 6 {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Compute signed differences: current - best.
|
||||
// For minimization: positive diff means current is worse.
|
||||
// For maximization: negative diff means current is worse.
|
||||
let differences: Vec<f64> = pairs
|
||||
.iter()
|
||||
.map(|&(current, best_v)| current - best_v)
|
||||
.collect();
|
||||
|
||||
// Run the Wilcoxon signed-rank test.
|
||||
let p_value = wilcoxon_signed_rank_test(&differences, self.direction);
|
||||
|
||||
p_value < self.p_value_threshold
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a one-sided Wilcoxon signed-rank test.
|
||||
///
|
||||
/// Tests whether the values tend to be worse than zero (positive for
|
||||
/// minimization, negative for maximization).
|
||||
///
|
||||
/// Returns a p-value. Small p-values indicate the current trial is
|
||||
/// significantly worse.
|
||||
fn wilcoxon_signed_rank_test(differences: &[f64], direction: Direction) -> f64 {
|
||||
// 1. Remove zero differences.
|
||||
let nonzero: Vec<f64> = differences.iter().copied().filter(|d| *d != 0.0).collect();
|
||||
let n = nonzero.len();
|
||||
|
||||
if n < 6 {
|
||||
return 1.0; // Not enough data
|
||||
}
|
||||
|
||||
// 2. Rank by absolute value.
|
||||
let mut abs_ranked: Vec<(usize, f64, f64)> = nonzero
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &d)| (i, d.abs(), d))
|
||||
.collect();
|
||||
abs_ranked.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
|
||||
|
||||
// 3. Assign ranks with tie correction.
|
||||
let ranks = assign_ranks(&abs_ranked);
|
||||
|
||||
// 4. Compute W+ (sum of ranks for positive differences) and
|
||||
// W- (sum of ranks for negative differences).
|
||||
let mut w_plus = 0.0;
|
||||
let mut w_minus = 0.0;
|
||||
for (i, &(_, _, orig)) in abs_ranked.iter().enumerate() {
|
||||
if orig > 0.0 {
|
||||
w_plus += ranks[i];
|
||||
} else {
|
||||
w_minus += ranks[i];
|
||||
}
|
||||
}
|
||||
|
||||
// For one-sided test:
|
||||
// - Minimization: we want to detect positive diffs (current worse).
|
||||
// Large W+ means significantly worse. Test statistic = W-.
|
||||
// - Maximization: we want to detect negative diffs (current worse).
|
||||
// Large W- means significantly worse. Test statistic = W+.
|
||||
let w = match direction {
|
||||
Direction::Minimize => w_minus,
|
||||
Direction::Maximize => w_plus,
|
||||
};
|
||||
|
||||
// 5. Normal approximation for the p-value.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let n_f = n as f64;
|
||||
let mean = n_f * (n_f + 1.0) / 4.0;
|
||||
let variance = n_f * (n_f + 1.0) * (2.0 * n_f + 1.0) / 24.0;
|
||||
|
||||
// Tie correction for variance.
|
||||
let tie_correction = compute_tie_correction(&ranks);
|
||||
let adjusted_variance = variance - tie_correction;
|
||||
|
||||
if adjusted_variance <= 0.0 {
|
||||
return 1.0;
|
||||
}
|
||||
|
||||
let std_dev = adjusted_variance.sqrt();
|
||||
|
||||
// Continuity correction: shift W by 0.5 towards mean.
|
||||
let continuity = if w < mean { 0.5 } else { -0.5 };
|
||||
let z = (w + continuity - mean) / std_dev;
|
||||
|
||||
// One-sided p-value (lower tail): probability that the test statistic
|
||||
// is this small or smaller under H0.
|
||||
normal_cdf(z)
|
||||
}
|
||||
|
||||
/// Assign average ranks, handling ties.
|
||||
fn assign_ranks(sorted: &[(usize, f64, f64)]) -> Vec<f64> {
|
||||
let n = sorted.len();
|
||||
let mut ranks = vec![0.0; n];
|
||||
let mut i = 0;
|
||||
|
||||
while i < n {
|
||||
let mut j = i;
|
||||
// Find all items tied with sorted[i].
|
||||
while j < n
|
||||
&& (sorted[j].1 - sorted[i].1).abs() < f64::EPSILON * sorted[i].1.max(1.0) * 100.0
|
||||
{
|
||||
j += 1;
|
||||
}
|
||||
// Average rank for the tie group. Ranks are 1-based.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let avg_rank = (i + 1 + j) as f64 / 2.0;
|
||||
for rank in ranks.iter_mut().take(j).skip(i) {
|
||||
*rank = avg_rank;
|
||||
}
|
||||
i = j;
|
||||
}
|
||||
|
||||
ranks
|
||||
}
|
||||
|
||||
/// Compute the tie correction term for the variance.
|
||||
/// For each tie group of size t, subtract t^3 - t from the sum,
|
||||
/// then divide by 48.
|
||||
fn compute_tie_correction(ranks: &[f64]) -> f64 {
|
||||
let mut correction = 0.0;
|
||||
let mut i = 0;
|
||||
|
||||
while i < ranks.len() {
|
||||
let mut j = i;
|
||||
while j < ranks.len() && (ranks[j] - ranks[i]).abs() < f64::EPSILON {
|
||||
j += 1;
|
||||
}
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let t = (j - i) as f64;
|
||||
if t > 1.0 {
|
||||
correction += t * t * t - t;
|
||||
}
|
||||
i = j;
|
||||
}
|
||||
|
||||
correction / 48.0
|
||||
}
|
||||
|
||||
/// Standard normal CDF using an approximation (Abramowitz & Stegun).
|
||||
fn normal_cdf(x: f64) -> f64 {
|
||||
// Use the complementary error function relationship:
|
||||
// Φ(x) = 0.5 * erfc(-x / √2)
|
||||
0.5 * erfc(-x / core::f64::consts::SQRT_2)
|
||||
}
|
||||
|
||||
/// Complementary error function approximation.
|
||||
/// Maximum error: 1.5 × 10⁻⁷ (Abramowitz & Stegun formula 7.1.26).
|
||||
fn erfc(x: f64) -> f64 {
|
||||
let t = 1.0 / (1.0 + 0.327_591_1 * x.abs());
|
||||
let poly = t
|
||||
* (0.254_829_592
|
||||
+ t * (-0.284_496_736
|
||||
+ t * (1.421_413_741 + t * (-1.453_152_027 + t * 1.061_405_429))));
|
||||
let result = poly * (-x * x).exp();
|
||||
|
||||
if x >= 0.0 { result } else { 2.0 - result }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn trial_with_values(
|
||||
id: u64,
|
||||
value: f64,
|
||||
intermediate_values: Vec<(u64, f64)>,
|
||||
) -> CompletedTrial {
|
||||
CompletedTrial::with_intermediate_values(
|
||||
id,
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
HashMap::new(),
|
||||
value,
|
||||
intermediate_values,
|
||||
HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_during_warmup() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize).n_warmup_steps(10);
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..20).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
let current: Vec<(u64, f64)> = (0..8).map(|s| (s, 100.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 7, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_with_insufficient_trials() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize).n_min_trials(5);
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..20).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
let current: Vec<(u64, f64)> = (0..10).map(|s| (s, 100.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 9, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_with_fewer_than_6_pairs() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..5).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
// Only 5 matching steps
|
||||
let current: Vec<(u64, f64)> = (0..5).map(|s| (s, 100.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 4, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_when_consistently_worse_minimize() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
|
||||
// Best trial has low values.
|
||||
let best_values: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
|
||||
let completed = vec![trial_with_values(0, 0.1, best_values)];
|
||||
|
||||
// Current trial is consistently much worse.
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 10.0)).collect();
|
||||
assert!(pruner.should_prune(1, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_when_consistently_worse_maximize() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Maximize);
|
||||
|
||||
// Best trial has high values.
|
||||
let best_values: Vec<(u64, f64)> = (0..20).map(|s| (s, 10.0)).collect();
|
||||
let completed = vec![trial_with_values(0, 10.0, best_values)];
|
||||
|
||||
// Current trial is consistently much worse.
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
|
||||
assert!(pruner.should_prune(1, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_statistically_similar() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
|
||||
// Best trial and current trial have very similar values with noise.
|
||||
let best_values: Vec<(u64, f64)> = (0..20_u64)
|
||||
.map(|s| {
|
||||
let noise = if s.is_multiple_of(2) { 0.01 } else { -0.01 };
|
||||
(s, 1.0 + noise)
|
||||
})
|
||||
.collect();
|
||||
let completed = vec![trial_with_values(0, 1.0, best_values)];
|
||||
|
||||
// Current trial is similar — alternating above/below.
|
||||
let current: Vec<(u64, f64)> = (0..20_u64)
|
||||
.map(|s| {
|
||||
let noise = if s.is_multiple_of(2) { -0.01 } else { 0.01 };
|
||||
(s, 1.0 + noise)
|
||||
})
|
||||
.collect();
|
||||
assert!(!pruner.should_prune(1, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_best_trial_minimize() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
|
||||
// Two completed trials: trial 0 is better (lower).
|
||||
let completed = vec![
|
||||
trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect()),
|
||||
trial_with_values(1, 5.0, (0..20).map(|s| (s, 5.0)).collect()),
|
||||
];
|
||||
|
||||
// Current trial is worse than the best but similar to the second.
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 5.0)).collect();
|
||||
assert!(pruner.should_prune(2, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_best_trial_maximize() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Maximize);
|
||||
|
||||
// Two completed trials: trial 1 is better (higher).
|
||||
let completed = vec![
|
||||
trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect()),
|
||||
trial_with_values(1, 10.0, (0..20).map(|s| (s, 10.0)).collect()),
|
||||
];
|
||||
|
||||
// Current trial is worse than the best.
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
|
||||
assert!(pruner.should_prune(2, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_pruned_trials() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
|
||||
// Only a pruned trial — no complete trials.
|
||||
let mut trial = trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect());
|
||||
trial.state = TrialState::Pruned;
|
||||
let completed = vec![trial];
|
||||
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 100.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 19, ¤t, &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lower_p_value_is_more_conservative() {
|
||||
let strict = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.001);
|
||||
let lenient = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.1);
|
||||
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..20).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
|
||||
// Moderately worse — should pass lenient but maybe not strict.
|
||||
let current: Vec<(u64, f64)> = (0..20)
|
||||
.map(|s| if s < 15 { (s, 0.2) } else { (s, 0.15) })
|
||||
.collect();
|
||||
|
||||
let lenient_prunes = lenient.should_prune(1, 19, ¤t, &completed);
|
||||
let strict_prunes = strict.should_prune(1, 19, ¤t, &completed);
|
||||
|
||||
// A stricter threshold should never prune when a lenient one doesn't.
|
||||
if !lenient_prunes {
|
||||
assert!(!strict_prunes);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "p_value_threshold must be in (0.0, 1.0)")]
|
||||
fn panics_on_zero_p_value() {
|
||||
let _ = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "p_value_threshold must be in (0.0, 1.0)")]
|
||||
fn panics_on_one_p_value() {
|
||||
let _ = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn correct_signed_rank_statistic() {
|
||||
// Known example: differences [1, 2, 3, 4, 5, 6] (all positive).
|
||||
// Ranks: 1, 2, 3, 4, 5, 6. W+ = 21, W- = 0.
|
||||
// For minimization (testing if positive = worse), W- = 0.
|
||||
// This should give a very small p-value.
|
||||
let diffs = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
|
||||
let p = wilcoxon_signed_rank_test(&diffs, Direction::Minimize);
|
||||
assert!(
|
||||
p < 0.05,
|
||||
"p-value {p} should be < 0.05 for all-positive diffs"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn symmetric_differences_not_significant() {
|
||||
// Balanced differences: half positive, half negative.
|
||||
let diffs = vec![1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0, -4.0];
|
||||
let p = wilcoxon_signed_rank_test(&diffs, Direction::Minimize);
|
||||
assert!(p > 0.05, "p-value {p} should be > 0.05 for symmetric diffs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normal_cdf_known_values() {
|
||||
assert!((normal_cdf(0.0) - 0.5).abs() < 1e-6);
|
||||
assert!(normal_cdf(-10.0) < 1e-6);
|
||||
assert!((normal_cdf(10.0) - 1.0).abs() < 1e-6);
|
||||
assert!((normal_cdf(-1.96) - 0.025).abs() < 0.001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_intermediate_values() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
let completed = vec![trial_with_values(
|
||||
0,
|
||||
0.1,
|
||||
(0..20).map(|s| (s, 0.1)).collect(),
|
||||
)];
|
||||
assert!(!pruner.should_prune(1, 0, &[], &completed));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_completed_trials() {
|
||||
let pruner = WilcoxonPruner::new(Direction::Minimize);
|
||||
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 1.0)).collect();
|
||||
assert!(!pruner.should_prune(1, 19, ¤t, &[]));
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,8 @@ use std::collections::HashMap;
|
||||
use crate::distribution::Distribution;
|
||||
use crate::param::ParamValue;
|
||||
use crate::parameter::{ParamId, Parameter};
|
||||
use crate::trial::AttrValue;
|
||||
use crate::types::TrialState;
|
||||
|
||||
/// A completed trial with its parameters, distributions, and objective value.
|
||||
///
|
||||
@@ -27,6 +29,12 @@ pub struct CompletedTrial<V = f64> {
|
||||
pub param_labels: HashMap<ParamId, String>,
|
||||
/// The objective value returned by the objective function.
|
||||
pub value: V,
|
||||
/// Intermediate objective values reported during the trial.
|
||||
pub intermediate_values: Vec<(u64, f64)>,
|
||||
/// The state of the trial (Complete, Pruned, or Failed).
|
||||
pub state: TrialState,
|
||||
/// User-defined attributes stored during the trial.
|
||||
pub user_attrs: HashMap<String, AttrValue>,
|
||||
}
|
||||
|
||||
impl<V> CompletedTrial<V> {
|
||||
@@ -44,6 +52,31 @@ impl<V> CompletedTrial<V> {
|
||||
distributions,
|
||||
param_labels,
|
||||
value,
|
||||
intermediate_values: Vec::new(),
|
||||
state: TrialState::Complete,
|
||||
user_attrs: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a new completed trial with intermediate values and user attributes.
|
||||
pub fn with_intermediate_values(
|
||||
id: u64,
|
||||
params: HashMap<ParamId, ParamValue>,
|
||||
distributions: HashMap<ParamId, Distribution>,
|
||||
param_labels: HashMap<ParamId, String>,
|
||||
value: V,
|
||||
intermediate_values: Vec<(u64, f64)>,
|
||||
user_attrs: HashMap<String, AttrValue>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id,
|
||||
params,
|
||||
distributions,
|
||||
param_labels,
|
||||
value,
|
||||
intermediate_values,
|
||||
state: TrialState::Complete,
|
||||
user_attrs,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -87,6 +120,18 @@ impl<V> CompletedTrial<V> {
|
||||
.expect("parameter type mismatch: stored value incompatible with parameter")
|
||||
})
|
||||
}
|
||||
|
||||
/// Gets a user attribute by key.
|
||||
#[must_use]
|
||||
pub fn user_attr(&self, key: &str) -> Option<&AttrValue> {
|
||||
self.user_attrs.get(key)
|
||||
}
|
||||
|
||||
/// Returns all user attributes.
|
||||
#[must_use]
|
||||
pub fn user_attrs(&self) -> &HashMap<String, AttrValue> {
|
||||
&self.user_attrs
|
||||
}
|
||||
}
|
||||
|
||||
/// A pending (running) trial with its parameters and distributions, but no objective value yet.
|
||||
|
||||
+691
-26
@@ -5,14 +5,20 @@ use core::any::Any;
|
||||
use core::future::Future;
|
||||
use core::ops::ControlFlow;
|
||||
use core::sync::atomic::{AtomicU64, Ordering};
|
||||
use core::time::Duration;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use parking_lot::RwLock;
|
||||
use parking_lot::{Mutex, RwLock};
|
||||
|
||||
use crate::param::ParamValue;
|
||||
use crate::parameter::ParamId;
|
||||
use crate::pruner::{NopPruner, Pruner};
|
||||
use crate::sampler::random::RandomSampler;
|
||||
use crate::sampler::{CompletedTrial, Sampler};
|
||||
use crate::trial::Trial;
|
||||
use crate::types::Direction;
|
||||
use crate::types::{Direction, TrialState};
|
||||
|
||||
/// A study manages the optimization process, tracking trials and their results.
|
||||
///
|
||||
@@ -40,6 +46,8 @@ where
|
||||
direction: Direction,
|
||||
/// The sampler used to generate parameter values.
|
||||
sampler: Arc<dyn Sampler>,
|
||||
/// The pruner used to decide whether to stop trials early.
|
||||
pruner: Arc<dyn Pruner>,
|
||||
/// Completed trials (wrapped in Arc for sharing with Trial).
|
||||
completed_trials: Arc<RwLock<Vec<CompletedTrial<V>>>>,
|
||||
/// Counter for generating unique trial IDs.
|
||||
@@ -48,6 +56,8 @@ where
|
||||
/// Set automatically for `Study<f64>` so that `create_trial()` and all
|
||||
/// optimization methods use the sampler without requiring `_with_sampler` suffixes.
|
||||
trial_factory: Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>,
|
||||
/// Queue of parameter configurations to evaluate next.
|
||||
enqueued_params: Arc<Mutex<VecDeque<HashMap<ParamId, ParamValue>>>>,
|
||||
}
|
||||
|
||||
impl<V> Study<V>
|
||||
@@ -78,6 +88,56 @@ where
|
||||
Self::with_sampler(direction, RandomSampler::new())
|
||||
}
|
||||
|
||||
/// Creates a study that minimizes the objective value.
|
||||
///
|
||||
/// This is a shorthand for `Study::with_sampler(Direction::Minimize, sampler)`.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `sampler` - The sampler to use for parameter sampling.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Study;
|
||||
/// use optimizer::sampler::tpe::TpeSampler;
|
||||
///
|
||||
/// let study: Study<f64> = Study::minimize(TpeSampler::new());
|
||||
/// assert_eq!(study.direction(), optimizer::Direction::Minimize);
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn minimize(sampler: impl Sampler + 'static) -> Self
|
||||
where
|
||||
V: 'static,
|
||||
{
|
||||
Self::with_sampler(Direction::Minimize, sampler)
|
||||
}
|
||||
|
||||
/// Creates a study that maximizes the objective value.
|
||||
///
|
||||
/// This is a shorthand for `Study::with_sampler(Direction::Maximize, sampler)`.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `sampler` - The sampler to use for parameter sampling.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::Study;
|
||||
/// use optimizer::sampler::tpe::TpeSampler;
|
||||
///
|
||||
/// let study: Study<f64> = Study::maximize(TpeSampler::new());
|
||||
/// assert_eq!(study.direction(), optimizer::Direction::Maximize);
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn maximize(sampler: impl Sampler + 'static) -> Self
|
||||
where
|
||||
V: 'static,
|
||||
{
|
||||
Self::with_sampler(Direction::Maximize, sampler)
|
||||
}
|
||||
|
||||
/// Creates a new study with a custom sampler.
|
||||
///
|
||||
/// # Arguments
|
||||
@@ -102,16 +162,20 @@ where
|
||||
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
||||
let completed_trials = Arc::new(RwLock::new(Vec::new()));
|
||||
|
||||
let pruner: Arc<dyn Pruner> = Arc::new(NopPruner);
|
||||
|
||||
// For Study<f64>, set up a trial factory that provides sampler integration.
|
||||
// This uses Any downcasting to check at runtime whether V = f64.
|
||||
let trial_factory = Self::make_trial_factory(&sampler, &completed_trials);
|
||||
let trial_factory = Self::make_trial_factory(&sampler, &completed_trials, &pruner);
|
||||
|
||||
Self {
|
||||
direction,
|
||||
sampler,
|
||||
pruner,
|
||||
completed_trials,
|
||||
next_trial_id: AtomicU64::new(0),
|
||||
trial_factory,
|
||||
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -119,6 +183,7 @@ where
|
||||
fn make_trial_factory(
|
||||
sampler: &Arc<dyn Sampler>,
|
||||
completed_trials: &Arc<RwLock<Vec<CompletedTrial<V>>>>,
|
||||
pruner: &Arc<dyn Pruner>,
|
||||
) -> Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>
|
||||
where
|
||||
V: 'static,
|
||||
@@ -131,8 +196,14 @@ where
|
||||
f64_trials.map(|trials| {
|
||||
let sampler = Arc::clone(sampler);
|
||||
let trials = Arc::clone(trials);
|
||||
let pruner = Arc::clone(pruner);
|
||||
let factory: Arc<dyn Fn(u64) -> Trial + Send + Sync> = Arc::new(move |id| {
|
||||
Trial::with_sampler(id, Arc::clone(&sampler), Arc::clone(&trials))
|
||||
Trial::with_sampler(
|
||||
id,
|
||||
Arc::clone(&sampler),
|
||||
Arc::clone(&trials),
|
||||
Arc::clone(&pruner),
|
||||
)
|
||||
});
|
||||
factory
|
||||
})
|
||||
@@ -158,12 +229,120 @@ where
|
||||
/// let mut study: Study<f64> = Study::new(Direction::Minimize);
|
||||
/// study.set_sampler(TpeSampler::new());
|
||||
/// ```
|
||||
/// Creates a new study with a custom sampler and pruner.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `direction` - Whether to minimize or maximize the objective function.
|
||||
/// * `sampler` - The sampler to use for parameter sampling.
|
||||
/// * `pruner` - The pruner to use for trial pruning.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::pruner::NopPruner;
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
///
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
/// let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, sampler, NopPruner);
|
||||
/// ```
|
||||
pub fn with_sampler_and_pruner(
|
||||
direction: Direction,
|
||||
sampler: impl Sampler + 'static,
|
||||
pruner: impl Pruner + 'static,
|
||||
) -> Self
|
||||
where
|
||||
V: 'static,
|
||||
{
|
||||
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
||||
let pruner: Arc<dyn Pruner> = Arc::new(pruner);
|
||||
let completed_trials = Arc::new(RwLock::new(Vec::new()));
|
||||
let trial_factory = Self::make_trial_factory(&sampler, &completed_trials, &pruner);
|
||||
|
||||
Self {
|
||||
direction,
|
||||
sampler,
|
||||
pruner,
|
||||
completed_trials,
|
||||
next_trial_id: AtomicU64::new(0),
|
||||
trial_factory,
|
||||
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_sampler(&mut self, sampler: impl Sampler + 'static)
|
||||
where
|
||||
V: 'static,
|
||||
{
|
||||
self.sampler = Arc::new(sampler);
|
||||
self.trial_factory = Self::make_trial_factory(&self.sampler, &self.completed_trials);
|
||||
self.trial_factory =
|
||||
Self::make_trial_factory(&self.sampler, &self.completed_trials, &self.pruner);
|
||||
}
|
||||
|
||||
/// Sets a new pruner for the study.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `pruner` - The pruner to use for trial pruning.
|
||||
pub fn set_pruner(&mut self, pruner: impl Pruner + 'static)
|
||||
where
|
||||
V: 'static,
|
||||
{
|
||||
self.pruner = Arc::new(pruner);
|
||||
self.trial_factory =
|
||||
Self::make_trial_factory(&self.sampler, &self.completed_trials, &self.pruner);
|
||||
}
|
||||
|
||||
/// Returns a reference to the study's pruner.
|
||||
pub fn pruner(&self) -> &dyn Pruner {
|
||||
&*self.pruner
|
||||
}
|
||||
|
||||
/// Enqueues a specific parameter configuration to be evaluated next.
|
||||
///
|
||||
/// The next call to [`ask()`](Self::ask) or the next trial in [`optimize()`](Self::optimize)
|
||||
/// will use these exact parameters instead of sampling from the sampler.
|
||||
///
|
||||
/// Multiple configurations can be enqueued; they are evaluated in FIFO order.
|
||||
/// If an enqueued configuration is missing a parameter that the objective calls
|
||||
/// `suggest()` on, that parameter falls back to normal sampling.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `params` - A map from parameter IDs to the values to use.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use std::collections::HashMap;
|
||||
///
|
||||
/// use optimizer::parameter::{FloatParam, IntParam, Parameter};
|
||||
/// use optimizer::{Direction, ParamValue, Study};
|
||||
///
|
||||
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
/// let x = FloatParam::new(0.0, 10.0);
|
||||
/// let y = IntParam::new(1, 100);
|
||||
///
|
||||
/// // Evaluate these specific configurations first
|
||||
/// study.enqueue(HashMap::from([
|
||||
/// (x.id(), ParamValue::Float(0.001)),
|
||||
/// (y.id(), ParamValue::Int(3)),
|
||||
/// ]));
|
||||
///
|
||||
/// // Next trial will use x=0.001, y=3
|
||||
/// let mut trial = study.ask();
|
||||
/// assert_eq!(x.suggest(&mut trial).unwrap(), 0.001);
|
||||
/// assert_eq!(y.suggest(&mut trial).unwrap(), 3);
|
||||
/// ```
|
||||
pub fn enqueue(&self, params: HashMap<ParamId, ParamValue>) {
|
||||
self.enqueued_params.lock().push_back(params);
|
||||
}
|
||||
|
||||
/// Returns the number of enqueued parameter configurations.
|
||||
#[must_use]
|
||||
pub fn n_enqueued(&self) -> usize {
|
||||
self.enqueued_params.lock().len()
|
||||
}
|
||||
|
||||
/// Generates the next unique trial ID.
|
||||
@@ -195,11 +374,18 @@ where
|
||||
/// ```
|
||||
pub fn create_trial(&self) -> Trial {
|
||||
let id = self.next_trial_id();
|
||||
if let Some(factory) = &self.trial_factory {
|
||||
let mut trial = if let Some(factory) = &self.trial_factory {
|
||||
factory(id)
|
||||
} else {
|
||||
Trial::new(id)
|
||||
};
|
||||
|
||||
// If there are enqueued params, inject them into this trial
|
||||
if let Some(fixed_params) = self.enqueued_params.lock().pop_front() {
|
||||
trial.set_fixed_params(fixed_params);
|
||||
}
|
||||
|
||||
trial
|
||||
}
|
||||
|
||||
/// Records a completed trial with its objective value.
|
||||
@@ -230,13 +416,16 @@ where
|
||||
/// ```
|
||||
pub fn complete_trial(&self, mut trial: Trial, value: V) {
|
||||
trial.set_complete();
|
||||
let completed = CompletedTrial::new(
|
||||
let mut completed = CompletedTrial::with_intermediate_values(
|
||||
trial.id(),
|
||||
trial.params().clone(),
|
||||
trial.distributions().clone(),
|
||||
trial.param_labels().clone(),
|
||||
value,
|
||||
trial.intermediate_values().to_vec(),
|
||||
trial.user_attrs().clone(),
|
||||
);
|
||||
completed.state = TrialState::Complete;
|
||||
self.completed_trials.write().push(completed);
|
||||
}
|
||||
|
||||
@@ -270,6 +459,85 @@ where
|
||||
// They could be stored in a separate list for debugging if needed
|
||||
}
|
||||
|
||||
/// Request a new trial with suggested parameters.
|
||||
///
|
||||
/// This is the first half of the ask-and-tell interface. After calling
|
||||
/// `ask()`, use parameter types to suggest values on the returned trial,
|
||||
/// evaluate your objective externally, then pass the trial back to
|
||||
/// [`tell()`](Self::tell) with the result.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::parameter::{FloatParam, Parameter};
|
||||
/// use optimizer::{Direction, Study};
|
||||
///
|
||||
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
/// let x = FloatParam::new(0.0, 10.0);
|
||||
///
|
||||
/// let mut trial = study.ask();
|
||||
/// let x_val = x.suggest(&mut trial).unwrap();
|
||||
/// let value = x_val * x_val;
|
||||
/// study.tell(trial, Ok::<_, &str>(value));
|
||||
/// ```
|
||||
pub fn ask(&self) -> Trial {
|
||||
self.create_trial()
|
||||
}
|
||||
|
||||
/// Report the result of a trial obtained from [`ask()`](Self::ask).
|
||||
///
|
||||
/// Pass `Ok(value)` for a successful evaluation or `Err(reason)` for a
|
||||
/// failure. Failed trials are not stored in the study's history.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::{Direction, Study};
|
||||
///
|
||||
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
///
|
||||
/// let trial = study.ask();
|
||||
/// study.tell(trial, Ok::<_, &str>(42.0));
|
||||
/// assert_eq!(study.n_trials(), 1);
|
||||
///
|
||||
/// let trial = study.ask();
|
||||
/// study.tell(trial, Err::<f64, _>("evaluation failed"));
|
||||
/// assert_eq!(study.n_trials(), 1); // failed trials not counted
|
||||
/// ```
|
||||
pub fn tell(&self, trial: Trial, value: core::result::Result<V, impl ToString>) {
|
||||
match value {
|
||||
Ok(v) => self.complete_trial(trial, v),
|
||||
Err(e) => self.fail_trial(trial, e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Records a pruned trial, preserving its intermediate values.
|
||||
///
|
||||
/// Pruned trials are stored alongside completed trials so that samplers
|
||||
/// can optionally learn from partial evaluations. The trial's state is
|
||||
/// set to `Pruned`.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `trial` - The trial that was pruned.
|
||||
pub fn prune_trial(&self, mut trial: Trial)
|
||||
where
|
||||
V: Default,
|
||||
{
|
||||
trial.set_pruned();
|
||||
let mut completed = CompletedTrial::with_intermediate_values(
|
||||
trial.id(),
|
||||
trial.params().clone(),
|
||||
trial.distributions().clone(),
|
||||
trial.param_labels().clone(),
|
||||
V::default(),
|
||||
trial.intermediate_values().to_vec(),
|
||||
trial.user_attrs().clone(),
|
||||
);
|
||||
completed.state = TrialState::Pruned;
|
||||
self.completed_trials.write().push(completed);
|
||||
}
|
||||
|
||||
/// Returns an iterator over all completed trials.
|
||||
///
|
||||
/// The iterator yields references to `CompletedTrial` values, which contain
|
||||
@@ -324,6 +592,15 @@ where
|
||||
self.completed_trials.read().len()
|
||||
}
|
||||
|
||||
/// Returns the number of pruned trials.
|
||||
pub fn n_pruned_trials(&self) -> usize {
|
||||
self.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Pruned)
|
||||
.count()
|
||||
}
|
||||
|
||||
/// Returns the trial with the best objective value.
|
||||
///
|
||||
/// The "best" trial depends on the optimization direction:
|
||||
@@ -364,12 +641,9 @@ where
|
||||
{
|
||||
let trials = self.completed_trials.read();
|
||||
|
||||
if trials.is_empty() {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
let best = trials
|
||||
.iter()
|
||||
.filter(|t| t.state == TrialState::Complete)
|
||||
.max_by(|a, b| {
|
||||
// For Minimize, we want the smallest value to be "max" in ordering
|
||||
// For Maximize, we want the largest value to be "max" in ordering
|
||||
@@ -431,6 +705,37 @@ where
|
||||
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.
|
||||
///
|
||||
/// This method runs `n_trials` evaluations sequentially. For each trial:
|
||||
@@ -480,7 +785,8 @@ where
|
||||
pub fn optimize<F, E>(&self, n_trials: usize, mut objective: F) -> crate::Result<()>
|
||||
where
|
||||
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
|
||||
E: ToString,
|
||||
E: ToString + 'static,
|
||||
V: Default,
|
||||
{
|
||||
for _ in 0..n_trials {
|
||||
let mut trial = self.create_trial();
|
||||
@@ -490,13 +796,22 @@ where
|
||||
self.complete_trial(trial, value);
|
||||
}
|
||||
Err(e) => {
|
||||
self.fail_trial(trial, e.to_string());
|
||||
if is_trial_pruned(&e) {
|
||||
self.prune_trial(trial);
|
||||
} else {
|
||||
self.fail_trial(trial, e.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
// Return error if no trials completed successfully
|
||||
let has_complete = self
|
||||
.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.any(|t| t.state == TrialState::Complete);
|
||||
if !has_complete {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
@@ -581,8 +896,13 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
// Return error if no trials completed successfully
|
||||
let has_complete = self
|
||||
.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.any(|t| t.state == TrialState::Complete);
|
||||
if !has_complete {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
@@ -696,8 +1016,13 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
// Return error if no trials completed successfully
|
||||
let has_complete = self
|
||||
.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.any(|t| t.state == TrialState::Complete);
|
||||
if !has_complete {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
@@ -768,10 +1093,10 @@ where
|
||||
mut callback: C,
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
V: Clone,
|
||||
V: Clone + Default,
|
||||
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
|
||||
C: FnMut(&Study<V>, &CompletedTrial<V>) -> ControlFlow<()>,
|
||||
E: ToString,
|
||||
E: ToString + 'static,
|
||||
{
|
||||
for _ in 0..n_trials {
|
||||
let mut trial = self.create_trial();
|
||||
@@ -799,13 +1124,340 @@ where
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
self.fail_trial(trial, e.to_string());
|
||||
if is_trial_pruned(&e) {
|
||||
self.prune_trial(trial);
|
||||
} else {
|
||||
self.fail_trial(trial, e.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Return error if no trials succeeded
|
||||
if self.n_trials() == 0 {
|
||||
// Return error if no trials completed successfully
|
||||
let has_complete = self
|
||||
.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.any(|t| t.state == TrialState::Complete);
|
||||
if !has_complete {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
/// Runs optimization until the given duration has elapsed.
|
||||
///
|
||||
/// Trials that are already running when the timeout is reached will
|
||||
/// complete — we never interrupt mid-trial. The actual elapsed time
|
||||
/// may therefore slightly exceed the specified duration.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `duration` - The maximum wall-clock time to spend on optimization.
|
||||
/// * `objective` - A closure that takes a mutable reference to a `Trial` and
|
||||
/// returns the objective value or an error.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::NoCompletedTrials` if no trials completed successfully
|
||||
/// before the timeout.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use std::time::Duration;
|
||||
///
|
||||
/// use optimizer::parameter::{FloatParam, Parameter};
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
///
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
///
|
||||
/// let x_param = FloatParam::new(-10.0, 10.0);
|
||||
///
|
||||
/// study
|
||||
/// .optimize_until(Duration::from_millis(100), |trial| {
|
||||
/// let x = x_param.suggest(trial)?;
|
||||
/// Ok::<_, optimizer::Error>(x * x)
|
||||
/// })
|
||||
/// .unwrap();
|
||||
///
|
||||
/// assert!(study.n_trials() > 0);
|
||||
/// ```
|
||||
pub fn optimize_until<F, E>(&self, duration: Duration, mut objective: F) -> crate::Result<()>
|
||||
where
|
||||
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
|
||||
E: ToString + 'static,
|
||||
V: Default,
|
||||
{
|
||||
let deadline = Instant::now() + duration;
|
||||
while Instant::now() < deadline {
|
||||
let mut trial = self.create_trial();
|
||||
|
||||
match objective(&mut trial) {
|
||||
Ok(value) => {
|
||||
self.complete_trial(trial, value);
|
||||
}
|
||||
Err(e) => {
|
||||
if is_trial_pruned(&e) {
|
||||
self.prune_trial(trial);
|
||||
} else {
|
||||
self.fail_trial(trial, e.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let has_complete = self
|
||||
.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.any(|t| t.state == TrialState::Complete);
|
||||
if !has_complete {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Runs optimization until the given duration has elapsed, with a callback.
|
||||
///
|
||||
/// Like [`optimize_until`](Self::optimize_until), but calls a callback after
|
||||
/// each completed trial. The callback can stop optimization early by returning
|
||||
/// `ControlFlow::Break(())`.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `duration` - The maximum wall-clock time to spend on optimization.
|
||||
/// * `objective` - A closure that takes a mutable reference to a `Trial` and
|
||||
/// returns the objective value or an error.
|
||||
/// * `callback` - A closure called after each successful trial. Returns
|
||||
/// `ControlFlow::Continue(())` to proceed or `ControlFlow::Break(())` to stop.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
|
||||
/// Returns `Error::Internal` if a completed trial is not found after adding.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use std::ops::ControlFlow;
|
||||
/// use std::time::Duration;
|
||||
///
|
||||
/// use optimizer::parameter::{FloatParam, Parameter};
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
///
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
///
|
||||
/// let x_param = FloatParam::new(-10.0, 10.0);
|
||||
///
|
||||
/// study
|
||||
/// .optimize_until_with_callback(
|
||||
/// Duration::from_secs(1),
|
||||
/// |trial| {
|
||||
/// let x = x_param.suggest(trial)?;
|
||||
/// Ok::<_, optimizer::Error>(x * x)
|
||||
/// },
|
||||
/// |_study, completed_trial| {
|
||||
/// if completed_trial.value < 1.0 {
|
||||
/// ControlFlow::Break(())
|
||||
/// } else {
|
||||
/// ControlFlow::Continue(())
|
||||
/// }
|
||||
/// },
|
||||
/// )
|
||||
/// .unwrap();
|
||||
///
|
||||
/// assert!(study.n_trials() > 0);
|
||||
/// ```
|
||||
pub fn optimize_until_with_callback<F, C, E>(
|
||||
&self,
|
||||
duration: Duration,
|
||||
mut objective: F,
|
||||
mut callback: C,
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
V: Clone + Default,
|
||||
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
|
||||
C: FnMut(&Study<V>, &CompletedTrial<V>) -> ControlFlow<()>,
|
||||
E: ToString + 'static,
|
||||
{
|
||||
let deadline = Instant::now() + duration;
|
||||
while Instant::now() < deadline {
|
||||
let mut trial = self.create_trial();
|
||||
|
||||
match objective(&mut trial) {
|
||||
Ok(value) => {
|
||||
self.complete_trial(trial, value);
|
||||
|
||||
let trials = self.completed_trials.read();
|
||||
let Some(completed) = trials.last() else {
|
||||
return Err(crate::Error::Internal(
|
||||
"completed trial not found after adding",
|
||||
));
|
||||
};
|
||||
|
||||
let completed_clone = completed.clone();
|
||||
drop(trials);
|
||||
|
||||
if let ControlFlow::Break(()) = callback(self, &completed_clone) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
if is_trial_pruned(&e) {
|
||||
self.prune_trial(trial);
|
||||
} else {
|
||||
self.fail_trial(trial, e.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let has_complete = self
|
||||
.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.any(|t| t.state == TrialState::Complete);
|
||||
if !has_complete {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Runs optimization asynchronously until the given duration has elapsed.
|
||||
///
|
||||
/// The async variant of [`optimize_until`](Self::optimize_until). Trials are
|
||||
/// run sequentially, but the objective function can be async.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `duration` - The maximum wall-clock time to spend on optimization.
|
||||
/// * `objective` - A function that takes a `Trial` and returns a `Future`
|
||||
/// that resolves to a tuple of `(Trial, Result<V, E>)`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
|
||||
#[cfg(feature = "async")]
|
||||
pub async fn optimize_until_async<F, Fut, E>(
|
||||
&self,
|
||||
duration: Duration,
|
||||
objective: F,
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
F: Fn(Trial) -> Fut,
|
||||
Fut: Future<Output = core::result::Result<(Trial, V), E>>,
|
||||
E: ToString,
|
||||
{
|
||||
let deadline = Instant::now() + duration;
|
||||
while Instant::now() < deadline {
|
||||
let trial = self.create_trial();
|
||||
|
||||
match objective(trial).await {
|
||||
Ok((trial, value)) => {
|
||||
self.complete_trial(trial, value);
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = e.to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let has_complete = self
|
||||
.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.any(|t| t.state == TrialState::Complete);
|
||||
if !has_complete {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Runs optimization with bounded parallelism until the given duration has elapsed.
|
||||
///
|
||||
/// The parallel variant of [`optimize_until`](Self::optimize_until). Runs up to
|
||||
/// `concurrency` trials simultaneously using async tasks. New trials are spawned
|
||||
/// as long as the deadline has not been reached; trials already running when the
|
||||
/// deadline passes will complete.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `duration` - The maximum wall-clock time to spend spawning new trials.
|
||||
/// * `concurrency` - The maximum number of trials to run simultaneously.
|
||||
/// * `objective` - A function that takes a `Trial` and returns a `Future`
|
||||
/// that resolves to a tuple of `(Trial, V)` or an error.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
|
||||
/// Returns `Error::TaskError` if the semaphore is closed or a spawned task panics.
|
||||
#[cfg(feature = "async")]
|
||||
pub async fn optimize_until_parallel<F, Fut, E>(
|
||||
&self,
|
||||
duration: Duration,
|
||||
concurrency: usize,
|
||||
objective: F,
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
F: Fn(Trial) -> Fut + Send + Sync + 'static,
|
||||
Fut: Future<Output = core::result::Result<(Trial, V), E>> + Send,
|
||||
E: ToString + Send + 'static,
|
||||
V: Send + 'static,
|
||||
{
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
let deadline = Instant::now() + duration;
|
||||
let semaphore = Arc::new(Semaphore::new(concurrency));
|
||||
let objective = Arc::new(objective);
|
||||
|
||||
let mut handles = Vec::new();
|
||||
|
||||
while Instant::now() < deadline {
|
||||
let permit = semaphore
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|e| crate::Error::TaskError(e.to_string()))?;
|
||||
let trial = self.create_trial();
|
||||
let objective = Arc::clone(&objective);
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
let result = objective(trial).await;
|
||||
drop(permit);
|
||||
result
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
match handle
|
||||
.await
|
||||
.map_err(|e| crate::Error::TaskError(e.to_string()))?
|
||||
{
|
||||
Ok((trial, value)) => {
|
||||
self.complete_trial(trial, value);
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = e.to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let has_complete = self
|
||||
.completed_trials
|
||||
.read()
|
||||
.iter()
|
||||
.any(|t| t.state == TrialState::Complete);
|
||||
if !has_complete {
|
||||
return Err(crate::Error::NoCompletedTrials);
|
||||
}
|
||||
|
||||
@@ -843,7 +1495,7 @@ impl Study<f64> {
|
||||
pub fn optimize_with_sampler<F, E>(&self, n_trials: usize, objective: F) -> crate::Result<()>
|
||||
where
|
||||
F: FnMut(&mut Trial) -> core::result::Result<f64, E>,
|
||||
E: ToString,
|
||||
E: ToString + 'static,
|
||||
{
|
||||
self.optimize(n_trials, objective)
|
||||
}
|
||||
@@ -865,7 +1517,7 @@ impl Study<f64> {
|
||||
where
|
||||
F: FnMut(&mut Trial) -> core::result::Result<f64, E>,
|
||||
C: FnMut(&Study<f64>, &CompletedTrial<f64>) -> ControlFlow<()>,
|
||||
E: ToString,
|
||||
E: ToString + 'static,
|
||||
{
|
||||
self.optimize_with_callback(n_trials, objective, callback)
|
||||
}
|
||||
@@ -916,3 +1568,16 @@ impl Study<f64> {
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns `true` if the error represents a pruned trial.
|
||||
///
|
||||
/// Checks via `Any` downcasting whether `e` is `Error::TrialPruned` or
|
||||
/// the standalone `TrialPruned` struct.
|
||||
fn is_trial_pruned<E: 'static>(e: &E) -> bool {
|
||||
let any: &dyn Any = e;
|
||||
if let Some(err) = any.downcast_ref::<crate::Error>() {
|
||||
matches!(err, crate::Error::TrialPruned)
|
||||
} else {
|
||||
any.downcast_ref::<crate::error::TrialPruned>().is_some()
|
||||
}
|
||||
}
|
||||
|
||||
+141
-2
@@ -9,9 +9,53 @@ use crate::distribution::Distribution;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::param::ParamValue;
|
||||
use crate::parameter::{ParamId, Parameter};
|
||||
use crate::pruner::Pruner;
|
||||
use crate::sampler::{CompletedTrial, Sampler};
|
||||
use crate::types::TrialState;
|
||||
|
||||
/// A user attribute value that can be stored on a trial.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum AttrValue {
|
||||
/// A floating-point attribute.
|
||||
Float(f64),
|
||||
/// An integer attribute.
|
||||
Int(i64),
|
||||
/// A string attribute.
|
||||
String(String),
|
||||
/// A boolean attribute.
|
||||
Bool(bool),
|
||||
}
|
||||
|
||||
impl From<f64> for AttrValue {
|
||||
fn from(v: f64) -> Self {
|
||||
Self::Float(v)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<i64> for AttrValue {
|
||||
fn from(v: i64) -> Self {
|
||||
Self::Int(v)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for AttrValue {
|
||||
fn from(v: String) -> Self {
|
||||
Self::String(v)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for AttrValue {
|
||||
fn from(v: &str) -> Self {
|
||||
Self::String(v.to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<bool> for AttrValue {
|
||||
fn from(v: bool) -> Self {
|
||||
Self::Bool(v)
|
||||
}
|
||||
}
|
||||
|
||||
/// A trial represents a single evaluation of the objective function.
|
||||
///
|
||||
/// Each trial has a unique ID and stores the sampled parameters along with
|
||||
@@ -36,6 +80,14 @@ pub struct Trial {
|
||||
sampler: Option<Arc<dyn Sampler>>,
|
||||
/// Access to the history of completed trials (shared with Study).
|
||||
history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>,
|
||||
/// Intermediate objective values reported at each step.
|
||||
intermediate_values: Vec<(u64, f64)>,
|
||||
/// The pruner used to decide whether to stop this trial early.
|
||||
pruner: Option<Arc<dyn Pruner>>,
|
||||
/// User-defined attributes for logging, debugging, and analysis.
|
||||
user_attrs: HashMap<String, AttrValue>,
|
||||
/// Pre-filled parameter values from enqueue (used instead of sampling).
|
||||
fixed_params: HashMap<ParamId, ParamValue>,
|
||||
}
|
||||
|
||||
impl core::fmt::Debug for Trial {
|
||||
@@ -48,6 +100,10 @@ impl core::fmt::Debug for Trial {
|
||||
.field("param_labels", &self.param_labels)
|
||||
.field("has_sampler", &self.sampler.is_some())
|
||||
.field("has_history", &self.history.is_some())
|
||||
.field("intermediate_values", &self.intermediate_values)
|
||||
.field("has_pruner", &self.pruner.is_some())
|
||||
.field("user_attrs", &self.user_attrs)
|
||||
.field("fixed_params", &self.fixed_params)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -83,6 +139,10 @@ impl Trial {
|
||||
param_labels: HashMap::new(),
|
||||
sampler: None,
|
||||
history: None,
|
||||
intermediate_values: Vec::new(),
|
||||
pruner: None,
|
||||
user_attrs: HashMap::new(),
|
||||
fixed_params: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,6 +160,7 @@ impl Trial {
|
||||
id: u64,
|
||||
sampler: Arc<dyn Sampler>,
|
||||
history: Arc<RwLock<Vec<CompletedTrial<f64>>>>,
|
||||
pruner: Arc<dyn Pruner>,
|
||||
) -> Self {
|
||||
Self {
|
||||
id,
|
||||
@@ -109,9 +170,21 @@ impl Trial {
|
||||
param_labels: HashMap::new(),
|
||||
sampler: Some(sampler),
|
||||
history: Some(history),
|
||||
intermediate_values: Vec::new(),
|
||||
pruner: Some(pruner),
|
||||
user_attrs: HashMap::new(),
|
||||
fixed_params: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets pre-filled parameters on this trial.
|
||||
///
|
||||
/// When `suggest_param` is called for a parameter that has a fixed value,
|
||||
/// the fixed value is used instead of sampling.
|
||||
pub(crate) fn set_fixed_params(&mut self, params: HashMap<ParamId, ParamValue>) {
|
||||
self.fixed_params = params;
|
||||
}
|
||||
|
||||
/// Samples a value from the given distribution using the sampler.
|
||||
///
|
||||
/// If the trial has a sampler, it delegates to the sampler's sample method
|
||||
@@ -159,6 +232,61 @@ impl Trial {
|
||||
&self.param_labels
|
||||
}
|
||||
|
||||
/// Reports an intermediate objective value at a given step.
|
||||
///
|
||||
/// Steps should be monotonically increasing (e.g., epoch number).
|
||||
/// Duplicate steps overwrite the previous value.
|
||||
pub fn report(&mut self, step: u64, value: f64) {
|
||||
if let Some(entry) = self
|
||||
.intermediate_values
|
||||
.iter_mut()
|
||||
.find(|(s, _)| *s == step)
|
||||
{
|
||||
entry.1 = value;
|
||||
} else {
|
||||
self.intermediate_values.push((step, value));
|
||||
}
|
||||
}
|
||||
|
||||
/// Ask whether this trial should be pruned at the current step.
|
||||
///
|
||||
/// Returns `true` if the pruner recommends stopping this trial.
|
||||
/// The caller should return `Err(TrialPruned)` from the objective.
|
||||
#[must_use]
|
||||
pub fn should_prune(&self) -> bool {
|
||||
let (Some(pruner), Some(history)) = (&self.pruner, &self.history) else {
|
||||
return false;
|
||||
};
|
||||
let Some(&(step, _)) = self.intermediate_values.last() else {
|
||||
return false;
|
||||
};
|
||||
let history_guard = history.read();
|
||||
pruner.should_prune(self.id, step, &self.intermediate_values, &history_guard)
|
||||
}
|
||||
|
||||
/// Returns all intermediate values reported so far.
|
||||
#[must_use]
|
||||
pub fn intermediate_values(&self) -> &[(u64, f64)] {
|
||||
&self.intermediate_values
|
||||
}
|
||||
|
||||
/// Sets a user attribute on this trial.
|
||||
pub fn set_user_attr(&mut self, key: impl Into<String>, value: impl Into<AttrValue>) {
|
||||
self.user_attrs.insert(key.into(), value.into());
|
||||
}
|
||||
|
||||
/// Gets a user attribute by key.
|
||||
#[must_use]
|
||||
pub fn user_attr(&self, key: &str) -> Option<&AttrValue> {
|
||||
self.user_attrs.get(key)
|
||||
}
|
||||
|
||||
/// Returns all user attributes.
|
||||
#[must_use]
|
||||
pub fn user_attrs(&self) -> &HashMap<String, AttrValue> {
|
||||
&self.user_attrs
|
||||
}
|
||||
|
||||
/// Sets the trial state to Complete.
|
||||
pub(crate) fn set_complete(&mut self) {
|
||||
self.state = TrialState::Complete;
|
||||
@@ -169,6 +297,11 @@ impl Trial {
|
||||
self.state = TrialState::Failed;
|
||||
}
|
||||
|
||||
/// Sets the trial state to Pruned.
|
||||
pub(crate) fn set_pruned(&mut self) {
|
||||
self.state = TrialState::Pruned;
|
||||
}
|
||||
|
||||
/// Suggests a parameter value using a [`Parameter`] definition.
|
||||
///
|
||||
/// This is the primary entry point for sampling parameters. It handles
|
||||
@@ -223,8 +356,14 @@ impl Trial {
|
||||
});
|
||||
}
|
||||
|
||||
// Sample using the sampler
|
||||
let value = self.sample_value(&distribution);
|
||||
// Check for a pre-filled (enqueued) value for this parameter
|
||||
let value = if let Some(fixed_value) = self.fixed_params.remove(¶m_id) {
|
||||
fixed_value
|
||||
} else {
|
||||
// Sample using the sampler
|
||||
self.sample_value(&distribution)
|
||||
};
|
||||
|
||||
let result = param.cast_param_value(&value)?;
|
||||
|
||||
// Store distribution, value, and label
|
||||
|
||||
@@ -18,4 +18,6 @@ pub enum TrialState {
|
||||
Complete,
|
||||
/// The trial failed with an error.
|
||||
Failed,
|
||||
/// The trial was pruned (stopped early).
|
||||
Pruned,
|
||||
}
|
||||
|
||||
@@ -1363,3 +1363,532 @@ fn test_completed_trial_get() {
|
||||
assert!((-10.0..=10.0).contains(&x_val));
|
||||
assert!((1..=10).contains(&n_val));
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests for timeout-based optimization
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_optimize_until_runs_for_approximately_specified_duration() {
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x_param = FloatParam::new(-10.0, 10.0);
|
||||
|
||||
let duration = Duration::from_millis(200);
|
||||
let start = Instant::now();
|
||||
|
||||
study
|
||||
.optimize_until(duration, |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
Ok::<_, Error>(x * x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
assert!(
|
||||
elapsed >= duration,
|
||||
"should run for at least the specified duration, elapsed: {elapsed:?}"
|
||||
);
|
||||
// Allow generous upper bound — the last trial may overshoot
|
||||
assert!(
|
||||
elapsed < duration + Duration::from_millis(200),
|
||||
"should not overshoot excessively, elapsed: {elapsed:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optimize_until_completes_at_least_one_trial() {
|
||||
use std::time::Duration;
|
||||
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x_param = FloatParam::new(-10.0, 10.0);
|
||||
|
||||
study
|
||||
.optimize_until(Duration::from_millis(100), |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
Ok::<_, Error>(x * x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
study.n_trials() >= 1,
|
||||
"should complete at least one trial, got {}",
|
||||
study.n_trials()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optimize_until_works_with_minimize() {
|
||||
use std::time::Duration;
|
||||
|
||||
let sampler = RandomSampler::with_seed(42);
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
let x_param = FloatParam::new(-10.0, 10.0);
|
||||
|
||||
study
|
||||
.optimize_until(Duration::from_millis(100), |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
Ok::<_, Error>(x * x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_value().unwrap();
|
||||
assert!(best >= 0.0, "x^2 should be non-negative");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optimize_until_works_with_maximize() {
|
||||
use std::time::Duration;
|
||||
|
||||
let sampler = RandomSampler::with_seed(42);
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
|
||||
let x_param = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study
|
||||
.optimize_until(Duration::from_millis(100), |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
Ok::<_, Error>(x)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_value().unwrap();
|
||||
assert!(best >= 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optimize_until_with_callback_early_stopping() {
|
||||
use std::ops::ControlFlow;
|
||||
use std::time::Duration;
|
||||
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x_param = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study
|
||||
.optimize_until_with_callback(
|
||||
Duration::from_secs(10), // long timeout — callback should stop early
|
||||
|trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
Ok::<_, Error>(x)
|
||||
},
|
||||
|study, _trial| {
|
||||
if study.n_trials() >= 5 {
|
||||
ControlFlow::Break(())
|
||||
} else {
|
||||
ControlFlow::Continue(())
|
||||
}
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
study.n_trials(),
|
||||
5,
|
||||
"callback should have stopped after 5 trials"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optimize_until_all_trials_fail() {
|
||||
use std::time::Duration;
|
||||
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
|
||||
let result = study.optimize_until(Duration::from_millis(50), |_trial| {
|
||||
Err::<f64, &str>("always fails")
|
||||
});
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(Error::NoCompletedTrials)),
|
||||
"should return NoCompletedTrials when all trials fail"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optimize_until_with_non_f64_value_type() {
|
||||
use std::time::Duration;
|
||||
|
||||
let study: Study<i32> = Study::new(Direction::Minimize);
|
||||
let x_param = IntParam::new(-10, 10);
|
||||
|
||||
study
|
||||
.optimize_until(Duration::from_millis(100), |trial| {
|
||||
let x = x_param.suggest(trial)?;
|
||||
Ok::<_, Error>(x.abs() as i32)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
assert!(study.n_trials() >= 1);
|
||||
let best = study.best_trial().unwrap();
|
||||
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);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Test: ask-and-tell interface
|
||||
// =============================================================================
|
||||
|
||||
#[test]
|
||||
fn test_ask_and_tell_basic() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x_param = FloatParam::new(0.0, 10.0);
|
||||
|
||||
for _ in 0..10 {
|
||||
let mut trial = study.ask();
|
||||
let x = x_param.suggest(&mut trial).unwrap();
|
||||
let value = x * x;
|
||||
study.tell(trial, Ok::<_, &str>(value));
|
||||
}
|
||||
|
||||
assert_eq!(study.n_trials(), 10);
|
||||
assert!(study.best_value().unwrap() >= 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ask_and_tell_with_failures() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x_param = FloatParam::new(-5.0, 5.0);
|
||||
|
||||
// Alternate success and failure
|
||||
for i in 0..10 {
|
||||
let mut trial = study.ask();
|
||||
let x = x_param.suggest(&mut trial).unwrap();
|
||||
if i % 2 == 0 {
|
||||
study.tell(trial, Ok::<_, &str>(x * x));
|
||||
} else {
|
||||
study.tell(trial, Err::<f64, _>("simulated failure"));
|
||||
}
|
||||
}
|
||||
|
||||
// Only successful trials are counted
|
||||
assert_eq!(study.n_trials(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ask_and_tell_with_tpe_sampler() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let study: Study<f64> = Study::minimize(sampler);
|
||||
let x_param = FloatParam::new(-10.0, 10.0);
|
||||
|
||||
for _ in 0..30 {
|
||||
let mut trial = study.ask();
|
||||
let x = x_param.suggest(&mut trial).unwrap();
|
||||
study.tell(trial, Ok::<_, &str>((x - 3.0).powi(2)));
|
||||
}
|
||||
|
||||
assert_eq!(study.n_trials(), 30);
|
||||
assert!(
|
||||
study.best_value().unwrap() < 5.0,
|
||||
"TPE ask-and-tell should find a reasonable value"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ask_and_tell_batch() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x_param = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Ask a batch of trials
|
||||
let batch: Vec<_> = (0..5)
|
||||
.map(|_| {
|
||||
let mut t = study.ask();
|
||||
let x = x_param.suggest(&mut t).unwrap();
|
||||
(t, x)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Tell results for the batch
|
||||
for (trial, x) in batch {
|
||||
study.tell(trial, Ok::<_, &str>(x * x));
|
||||
}
|
||||
|
||||
assert_eq!(study.n_trials(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ask_and_tell_with_custom_value_type() {
|
||||
// Ask-and-tell works with non-f64 value types too
|
||||
let study: Study<i32> = Study::new(Direction::Maximize);
|
||||
|
||||
for i in 0..5 {
|
||||
let trial = study.ask();
|
||||
study.tell(trial, Ok::<_, &str>(i * 10));
|
||||
}
|
||||
|
||||
assert_eq!(study.n_trials(), 5);
|
||||
assert_eq!(study.best_value().unwrap(), 40);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests: enqueue trials
|
||||
// =============================================================================
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use optimizer::ParamValue;
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_params_evaluated_first() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
let y = IntParam::new(1, 100);
|
||||
|
||||
// Enqueue a specific configuration
|
||||
study.enqueue(HashMap::from([
|
||||
(x.id(), ParamValue::Float(5.0)),
|
||||
(y.id(), ParamValue::Int(42)),
|
||||
]));
|
||||
|
||||
// The first trial should use the enqueued params
|
||||
let mut trial = study.ask();
|
||||
let x_val = x.suggest(&mut trial).unwrap();
|
||||
let y_val = y.suggest(&mut trial).unwrap();
|
||||
|
||||
assert_eq!(x_val, 5.0);
|
||||
assert_eq!(y_val, 42);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_fifo_order() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Enqueue two configs
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
|
||||
|
||||
// First trial gets first enqueued value
|
||||
let mut trial1 = study.ask();
|
||||
assert_eq!(x.suggest(&mut trial1).unwrap(), 1.0);
|
||||
|
||||
// Second trial gets second enqueued value
|
||||
let mut trial2 = study.ask();
|
||||
assert_eq!(x.suggest(&mut trial2).unwrap(), 2.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_then_normal_sampling_resumes() {
|
||||
let sampler = RandomSampler::with_seed(42);
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Enqueue one config
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(5.0))]));
|
||||
|
||||
// First trial uses enqueued value
|
||||
let mut trial1 = study.ask();
|
||||
assert_eq!(x.suggest(&mut trial1).unwrap(), 5.0);
|
||||
study.tell(trial1, Ok::<_, &str>(25.0));
|
||||
|
||||
// Second trial uses normal sampling (not 5.0)
|
||||
let mut trial2 = study.ask();
|
||||
let x_val = x.suggest(&mut trial2).unwrap();
|
||||
// The sampled value should be in [0, 10] but extremely unlikely to be exactly 5.0
|
||||
assert!((0.0..=10.0).contains(&x_val));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_with_optimize() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
// Enqueue two specific configs
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
|
||||
|
||||
let mut values = Vec::new();
|
||||
|
||||
study
|
||||
.optimize(5, |trial| {
|
||||
let x_val = x.suggest(trial)?;
|
||||
values.push(x_val);
|
||||
Ok::<_, Error>(x_val * x_val)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// First two trials should use enqueued values
|
||||
assert_eq!(values[0], 1.0);
|
||||
assert_eq!(values[1], 2.0);
|
||||
// All 5 trials should have completed
|
||||
assert_eq!(study.n_trials(), 5);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_partial_params_fall_back_to_sampling() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
let y = IntParam::new(1, 100);
|
||||
|
||||
// Enqueue only x, not y
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(3.0))]));
|
||||
|
||||
let mut trial = study.ask();
|
||||
let x_val = x.suggest(&mut trial).unwrap();
|
||||
let y_val = y.suggest(&mut trial).unwrap();
|
||||
|
||||
// x should be the enqueued value
|
||||
assert_eq!(x_val, 3.0);
|
||||
// y should be sampled (within range)
|
||||
assert!((1..=100).contains(&y_val));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_trials_appear_in_completed_trials() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(7.0))]));
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let x_val = x.suggest(trial)?;
|
||||
Ok::<_, Error>(x_val)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let trials = study.trials();
|
||||
assert_eq!(trials.len(), 1);
|
||||
assert_eq!(trials[0].value, 7.0);
|
||||
assert_eq!(
|
||||
*trials[0].params.get(&x.id()).unwrap(),
|
||||
ParamValue::Float(7.0)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_with_ask_and_tell() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(4.0))]));
|
||||
|
||||
let mut trial = study.ask();
|
||||
let x_val = x.suggest(&mut trial).unwrap();
|
||||
assert_eq!(x_val, 4.0);
|
||||
|
||||
study.tell(trial, Ok::<_, &str>(x_val * x_val));
|
||||
assert_eq!(study.n_trials(), 1);
|
||||
assert_eq!(study.best_value().unwrap(), 16.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_n_enqueued() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
assert_eq!(study.n_enqueued(), 0);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
|
||||
assert_eq!(study.n_enqueued(), 1);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
|
||||
assert_eq!(study.n_enqueued(), 2);
|
||||
|
||||
// Creating a trial dequeues one
|
||||
let _ = study.ask();
|
||||
assert_eq!(study.n_enqueued(), 1);
|
||||
|
||||
let _ = study.ask();
|
||||
assert_eq!(study.n_enqueued(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_enqueue_counted_in_n_trials() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 10.0);
|
||||
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(1.0))]));
|
||||
study.enqueue(HashMap::from([(x.id(), ParamValue::Float(2.0))]));
|
||||
|
||||
study
|
||||
.optimize(5, |trial| {
|
||||
let x_val = x.suggest(trial)?;
|
||||
Ok::<_, Error>(x_val)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// All 5 trials count, including the 2 enqueued ones
|
||||
assert_eq!(study.n_trials(), 5);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,228 @@
|
||||
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,
|
||||
HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
// --- 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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &completed));
|
||||
// 2.4 < 2.5 => don't prune
|
||||
let current = vec![(0, 2.4)];
|
||||
assert!(!pruner.should_prune(4, 0, ¤t, &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, ¤t, &completed));
|
||||
// 2.5 < 3.0 => don't prune
|
||||
let current = vec![(0, 2.5)];
|
||||
assert!(!pruner.should_prune(5, 0, ¤t, &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, ¤t, &completed));
|
||||
|
||||
let current = vec![(0, 1.0), (10, 2.0), (100, 3.5)];
|
||||
assert!(!pruner.should_prune(3, 100, ¤t, &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, ¤t, &completed));
|
||||
|
||||
// 6.0 > 5.0 => prune
|
||||
let current = vec![(0, 6.0)];
|
||||
assert!(pruner.should_prune(2, 0, ¤t, &completed));
|
||||
}
|
||||
@@ -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)], &[]));
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
use optimizer::parameter::{FloatParam, Parameter};
|
||||
use optimizer::{AttrValue, Direction, Study};
|
||||
|
||||
#[test]
|
||||
fn set_and_get_float_attr() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("score", 42.5);
|
||||
assert_eq!(trial.user_attr("score"), Some(&AttrValue::Float(42.5)));
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_and_get_int_attr() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("epoch", 42_i64);
|
||||
assert_eq!(trial.user_attr("epoch"), Some(&AttrValue::Int(42)));
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_and_get_string_attr() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("model", "resnet50");
|
||||
assert_eq!(
|
||||
trial.user_attr("model"),
|
||||
Some(&AttrValue::String("resnet50".to_owned()))
|
||||
);
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_and_get_bool_attr() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("converged", true);
|
||||
assert_eq!(trial.user_attr("converged"), Some(&AttrValue::Bool(true)));
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attrs_propagate_to_completed_trial() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("time_secs", 1.5);
|
||||
trial.set_user_attr("tag", "baseline");
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert_eq!(best.user_attr("time_secs"), Some(&AttrValue::Float(1.5)));
|
||||
assert_eq!(
|
||||
best.user_attr("tag"),
|
||||
Some(&AttrValue::String("baseline".to_owned()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overwrite_attr_replaces_value() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("key", "old");
|
||||
trial.set_user_attr("key", "new");
|
||||
assert_eq!(
|
||||
trial.user_attr("key"),
|
||||
Some(&AttrValue::String("new".to_owned()))
|
||||
);
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert_eq!(
|
||||
best.user_attr("key"),
|
||||
Some(&AttrValue::String("new".to_owned()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_attr_returns_none() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
assert_eq!(trial.user_attr("nonexistent"), None);
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert_eq!(best.user_attr("nonexistent"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_attrs_map_returns_all() {
|
||||
let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
let x = FloatParam::new(0.0, 1.0);
|
||||
|
||||
study
|
||||
.optimize(1, |trial| {
|
||||
let _ = x.suggest(trial)?;
|
||||
trial.set_user_attr("a", 1.0);
|
||||
trial.set_user_attr("b", true);
|
||||
assert_eq!(trial.user_attrs().len(), 2);
|
||||
Ok::<_, optimizer::Error>(1.0)
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let best = study.best_trial().unwrap();
|
||||
assert_eq!(best.user_attrs().len(), 2);
|
||||
}
|
||||
Reference in New Issue
Block a user