From abe15bdd7bb9890a84d6dd05871cc655c3a7dc02 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 16:24:57 +0100 Subject: [PATCH] feat: add TrialPruned error variant and Pruned trial state - Add `Pruned` variant to `TrialState` - Add `Error::TrialPruned` variant and standalone `TrialPruned` struct with `From 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 --- src/error.rs | 34 +++++++++++++ src/lib.rs | 4 +- src/sampler/mod.rs | 5 ++ src/study.rs | 117 +++++++++++++++++++++++++++++++++++++-------- src/trial.rs | 5 ++ src/types.rs | 2 + 6 files changed, 144 insertions(+), 23 deletions(-) diff --git a/src/error.rs b/src/error.rs index 890c376..44801aa 100644 --- a/src/error.rs +++ b/src/error.rs @@ -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 = core::result::Result; + +/// Convenience type for signalling a pruned trial from an objective function. +/// +/// Implements `Into` so it can be used with `?` in objectives that +/// return `Result`. +/// +/// # Examples +/// +/// ``` +/// use optimizer::{Error, TrialPruned}; +/// +/// fn objective_that_prunes() -> Result { +/// // ... 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 for Error { + fn from(_: TrialPruned) -> Self { + Error::TrialPruned + } +} diff --git a/src/lib.rs b/src/lib.rs index ddbf426..d5c53a1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -194,7 +194,7 @@ 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; @@ -219,7 +219,7 @@ 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, diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index cd04661..65698f2 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -9,6 +9,7 @@ use std::collections::HashMap; use crate::distribution::Distribution; use crate::param::ParamValue; use crate::parameter::{ParamId, Parameter}; +use crate::types::TrialState; /// A completed trial with its parameters, distributions, and objective value. /// @@ -29,6 +30,8 @@ pub struct CompletedTrial { 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, } impl CompletedTrial { @@ -47,6 +50,7 @@ impl CompletedTrial { param_labels, value, intermediate_values: Vec::new(), + state: TrialState::Complete, } } @@ -66,6 +70,7 @@ impl CompletedTrial { param_labels, value, intermediate_values, + state: TrialState::Complete, } } diff --git a/src/study.rs b/src/study.rs index 7ceeeb7..9624cff 100644 --- a/src/study.rs +++ b/src/study.rs @@ -13,7 +13,7 @@ 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. /// @@ -304,7 +304,7 @@ where /// ``` pub fn complete_trial(&self, mut trial: Trial, value: V) { trial.set_complete(); - let completed = CompletedTrial::with_intermediate_values( + let mut completed = CompletedTrial::with_intermediate_values( trial.id(), trial.params().clone(), trial.distributions().clone(), @@ -312,6 +312,7 @@ where value, trial.intermediate_values().to_vec(), ); + completed.state = TrialState::Complete; self.completed_trials.write().push(completed); } @@ -345,6 +346,32 @@ where // They could be stored in a separate list for debugging if needed } + /// 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(), + ); + 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 @@ -399,6 +426,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: @@ -439,12 +475,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 @@ -555,7 +588,8 @@ where pub fn optimize(&self, n_trials: usize, mut objective: F) -> crate::Result<()> where F: FnMut(&mut Trial) -> core::result::Result, - E: ToString, + E: ToString + 'static, + V: Default, { for _ in 0..n_trials { let mut trial = self.create_trial(); @@ -565,13 +599,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); } @@ -656,8 +699,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); } @@ -771,8 +819,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); } @@ -843,10 +896,10 @@ where mut callback: C, ) -> crate::Result<()> where - V: Clone, + V: Clone + Default, F: FnMut(&mut Trial) -> core::result::Result, C: FnMut(&Study, &CompletedTrial) -> ControlFlow<()>, - E: ToString, + E: ToString + 'static, { for _ in 0..n_trials { let mut trial = self.create_trial(); @@ -874,13 +927,22 @@ 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); } @@ -918,7 +980,7 @@ impl Study { pub fn optimize_with_sampler(&self, n_trials: usize, objective: F) -> crate::Result<()> where F: FnMut(&mut Trial) -> core::result::Result, - E: ToString, + E: ToString + 'static, { self.optimize(n_trials, objective) } @@ -940,7 +1002,7 @@ impl Study { where F: FnMut(&mut Trial) -> core::result::Result, C: FnMut(&Study, &CompletedTrial) -> ControlFlow<()>, - E: ToString, + E: ToString + 'static, { self.optimize_with_callback(n_trials, objective, callback) } @@ -991,3 +1053,16 @@ impl Study { .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: &E) -> bool { + let any: &dyn Any = e; + if let Some(err) = any.downcast_ref::() { + matches!(err, crate::Error::TrialPruned) + } else { + any.downcast_ref::().is_some() + } +} diff --git a/src/trial.rs b/src/trial.rs index 362a66c..bc918ca 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -219,6 +219,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 diff --git a/src/types.rs b/src/types.rs index 1d61969..b5e4d96 100644 --- a/src/types.rs +++ b/src/types.rs @@ -18,4 +18,6 @@ pub enum TrialState { Complete, /// The trial failed with an error. Failed, + /// The trial was pruned (stopped early). + Pruned, }