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:
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -18,4 +18,6 @@ pub enum TrialState {
|
||||
Complete,
|
||||
/// The trial failed with an error.
|
||||
Failed,
|
||||
/// The trial was pruned (stopped early).
|
||||
Pruned,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user