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
This commit is contained in:
Manuel Raimann
2026-02-11 16:24:57 +01:00
parent 4d8af3242b
commit abe15bdd7b
6 changed files with 144 additions and 23 deletions
+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
}
}
+2 -2
View File
@@ -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,
+5
View File
@@ -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<V = f64> {
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<V> CompletedTrial<V> {
@@ -47,6 +50,7 @@ impl<V> CompletedTrial<V> {
param_labels,
value,
intermediate_values: Vec::new(),
state: TrialState::Complete,
}
}
@@ -66,6 +70,7 @@ impl<V> CompletedTrial<V> {
param_labels,
value,
intermediate_values,
state: TrialState::Complete,
}
}
+96 -21
View File
@@ -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<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();
@@ -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<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();
@@ -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<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)
}
@@ -940,7 +1002,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)
}
@@ -991,3 +1053,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()
}
}
+5
View File
@@ -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
+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,
}