18 Commits

Author SHA1 Message Date
Manuel Raimann bb79d98027 chore: release v0.6.0 2026-02-11 17:24:10 +01:00
Manuel Raimann f4b0631178 feat: add enqueue_trial() for pre-specified parameter evaluation
Add Study::enqueue() to push specific parameter configurations onto a
FIFO queue. The next call to ask()/create_trial() or the next iteration
in optimize() pops the front entry and injects it into the trial so
that suggest_param() returns the pre-filled value instead of sampling.

Parameters missing from an enqueued map fall back to normal sampling,
and once the queue is drained regular sampler-driven trials resume.
2026-02-11 17:23:25 +01:00
Manuel Raimann 52f3c074dc feat: add trial user attributes for logging and analysis
Add AttrValue enum (Float, Int, String, Bool) with From impls and
set_user_attr/user_attr/user_attrs methods on both Trial and
CompletedTrial. Attributes set during optimization are propagated
through complete_trial and prune_trial. AttrValue is re-exported
at crate root and in the prelude.
2026-02-11 17:18:17 +01:00
Manuel Raimann f67188e3a7 feat: add ask() and tell() methods for ask-and-tell interface 2026-02-11 17:11:38 +01:00
Manuel Raimann 72da883add feat: add Study::top_trials(n) for retrieving best N trials
Returns the top N completed trials sorted by objective value,
respecting the study's optimization direction. Pruned and failed
trials are excluded.
2026-02-11 17:08:10 +01:00
Manuel Raimann bc278d83a3 feat: add timeout-based optimization methods (optimize_until)
Add duration-based variants of all optimization methods that run trials
until a wall-clock deadline rather than for a fixed trial count:

- optimize_until(duration, objective)
- optimize_until_with_callback(duration, objective, callback)
- optimize_until_async(duration, objective)
- optimize_until_parallel(duration, concurrency, objective)
2026-02-11 17:06:08 +01:00
Manuel Raimann c586650df4 feat: add From<RangeInclusive> for FloatParam and IntParam 2026-02-11 17:02:18 +01:00
Manuel Raimann cc78a92ed6 feat: add Study::minimize() and Study::maximize() constructor shortcuts 2026-02-11 17:00:35 +01:00
Manuel Raimann 12c35e7cb4 feat: add WilcoxonPruner for statistics-based trial pruning 2026-02-11 16:59:05 +01:00
Manuel Raimann a885d10436 feat: add HyperbandPruner for multi-bracket trial pruning 2026-02-11 16:52:59 +01:00
Manuel Raimann 41150057d6 feat: add SuccessiveHalvingPruner for SHA-based trial pruning 2026-02-11 16:47:57 +01:00
Manuel Raimann a522fb2f34 feat: add PatientPruner for patience-based trial pruning
Wraps any inner Pruner and requires it to recommend pruning for N
consecutive steps before actually pruning. Useful for noisy
intermediate values where a single bad step shouldn't trigger pruning.
2026-02-11 16:40:00 +01:00
Manuel Raimann ec04757740 feat: add PercentilePruner for configurable percentile-based trial pruning
Generalizes MedianPruner with a configurable percentile threshold.
MedianPruner now reuses the shared compute_percentile function at 50%.
2026-02-11 16:36:57 +01:00
Manuel Raimann 432c74b927 feat: add MedianPruner for statistics-based trial pruning
Prunes trials whose intermediate values are worse than the median
of completed trials at the same step. Supports configurable warmup
steps and minimum trial count before pruning activates.
2026-02-11 16:32:03 +01:00
Manuel Raimann 8199ba6b43 feat: add ThresholdPruner for fixed-bound trial pruning 2026-02-11 16:28:22 +01:00
Manuel Raimann abe15bdd7b feat: add TrialPruned error variant and Pruned trial state
- Add `Pruned` variant to `TrialState`
- Add `Error::TrialPruned` variant and standalone `TrialPruned` struct
  with `From<TrialPruned> for Error` for ergonomic `?` usage
- Add `state` field to `CompletedTrial` (defaults to `Complete`)
- Add `Study::prune_trial()` and `Study::n_pruned_trials()`
- `optimize()` and `optimize_with_callback()` detect `TrialPruned`
  errors via Any downcasting and record pruned trials instead of
  failing them
- `best_trial()` / `best_value()` now filter to only `Complete` trials
- Re-export `TrialPruned` from crate root and prelude
2026-02-11 16:24:57 +01:00
Manuel Raimann 4d8af3242b feat: add intermediate values and pruner integration to Trial
- Add report(step, value) and should_prune() methods to Trial
- Store intermediate_values in Trial, copied to CompletedTrial on completion
- Pass pruner through trial factory so trials can consult it
- Rebuild trial factory when pruner is changed via set_pruner()
2026-02-11 16:08:25 +01:00
Manuel Raimann 2651e61b2c feat: add Pruner trait and NopPruner default implementation
Introduce the pruning extension point for the trial pruning system.
The Pruner trait allows deciding whether to stop a trial early based
on intermediate values. NopPruner (never prunes) is the default.
2026-02-11 15:46:56 +01:00
21 changed files with 4250 additions and 33 deletions
+1 -1
View File
@@ -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"
+34
View File
@@ -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
View File
@@ -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;
}
+43
View File
@@ -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:?}"));
}
}
+559
View File
@@ -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);
}
}
+127
View File
@@ -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);
}
}
+74
View File
@@ -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;
}
+17
View File
@@ -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
}
}
+160
View File
@@ -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)], &[]));
}
}
+294
View File
@@ -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));
}
}
+466
View File
@@ -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);
}
}
+83
View File
@@ -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
}
}
+531
View File
@@ -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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &completed);
let strict_prunes = strict.should_prune(1, 19, &current, &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, &current, &[]));
}
}
+45
View File
@@ -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
View File
@@ -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
View File
@@ -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(&param_id) {
fixed_value
} else {
// Sample using the sampler
self.sample_value(&distribution)
};
let result = param.cast_param_value(&value)?;
// Store distribution, value, and label
+2
View File
@@ -18,4 +18,6 @@ pub enum TrialState {
Complete,
/// The trial failed with an error.
Failed,
/// The trial was pruned (stopped early).
Pruned,
}
+529
View File
@@ -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);
}
+228
View File
@@ -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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &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, &current, &completed));
// 2.4 < 2.5 => don't prune
let current = vec![(0, 2.4)];
assert!(!pruner.should_prune(4, 0, &current, &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, &current, &completed));
// 2.5 < 3.0 => don't prune
let current = vec![(0, 2.5)];
assert!(!pruner.should_prune(5, 0, &current, &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, &current, &completed));
let current = vec![(0, 1.0), (10, 2.0), (100, 3.5)];
assert!(!pruner.should_prune(3, 100, &current, &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, &current, &completed));
// 6.0 > 5.0 => prune
let current = vec![(0, 6.0)];
assert!(pruner.should_prune(2, 0, &current, &completed));
}
+64
View File
@@ -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)], &[]));
}
+148
View File
@@ -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);
}