From 4d8af3242b80fbdc26e7a21ac7739d5a6d66e094 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Wed, 11 Feb 2026 16:08:25 +0100 Subject: [PATCH] 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() --- src/sampler/mod.rs | 22 ++++++++++++++++++++ src/study.rs | 33 ++++++++++++++++++++++-------- src/trial.rs | 50 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 8 deletions(-) diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 183d7a0..cd04661 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -27,6 +27,8 @@ pub struct CompletedTrial { pub param_labels: HashMap, /// The objective value returned by the objective function. pub value: V, + /// Intermediate objective values reported during the trial. + pub intermediate_values: Vec<(u64, f64)>, } impl CompletedTrial { @@ -44,6 +46,26 @@ impl CompletedTrial { distributions, param_labels, value, + intermediate_values: Vec::new(), + } + } + + /// Creates a new completed trial with intermediate values. + pub fn with_intermediate_values( + id: u64, + params: HashMap, + distributions: HashMap, + param_labels: HashMap, + value: V, + intermediate_values: Vec<(u64, f64)>, + ) -> Self { + Self { + id, + params, + distributions, + param_labels, + value, + intermediate_values, } } diff --git a/src/study.rs b/src/study.rs index bdb85fe..7ceeeb7 100644 --- a/src/study.rs +++ b/src/study.rs @@ -105,14 +105,16 @@ where let sampler: Arc = Arc::new(sampler); let completed_trials = Arc::new(RwLock::new(Vec::new())); + let pruner: Arc = Arc::new(NopPruner); + // For Study, 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: Arc::new(NopPruner), + pruner, completed_trials, next_trial_id: AtomicU64::new(0), trial_factory, @@ -123,6 +125,7 @@ where fn make_trial_factory( sampler: &Arc, completed_trials: &Arc>>>, + pruner: &Arc, ) -> Option Trial + Send + Sync>> where V: 'static, @@ -135,8 +138,14 @@ where f64_trials.map(|trials| { let sampler = Arc::clone(sampler); let trials = Arc::clone(trials); + let pruner = Arc::clone(pruner); let factory: Arc 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 }) @@ -189,13 +198,14 @@ where V: 'static, { let sampler: Arc = Arc::new(sampler); + let pruner: Arc = Arc::new(pruner); let completed_trials = Arc::new(RwLock::new(Vec::new())); - 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: Arc::new(pruner), + pruner, completed_trials, next_trial_id: AtomicU64::new(0), trial_factory, @@ -207,7 +217,8 @@ 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. @@ -215,8 +226,13 @@ where /// # Arguments /// /// * `pruner` - The pruner to use for trial pruning. - pub fn set_pruner(&mut self, pruner: impl Pruner + 'static) { + 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. @@ -288,12 +304,13 @@ where /// ``` pub fn complete_trial(&self, mut trial: Trial, value: V) { trial.set_complete(); - let completed = CompletedTrial::new( + let completed = CompletedTrial::with_intermediate_values( trial.id(), trial.params().clone(), trial.distributions().clone(), trial.param_labels().clone(), value, + trial.intermediate_values().to_vec(), ); self.completed_trials.write().push(completed); } diff --git a/src/trial.rs b/src/trial.rs index 29cbfd2..362a66c 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -9,6 +9,7 @@ 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; @@ -36,6 +37,10 @@ pub struct Trial { sampler: Option>, /// Access to the history of completed trials (shared with Study). history: Option>>>>, + /// 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>, } impl core::fmt::Debug for Trial { @@ -48,6 +53,8 @@ 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()) .finish() } } @@ -83,6 +90,8 @@ impl Trial { param_labels: HashMap::new(), sampler: None, history: None, + intermediate_values: Vec::new(), + pruner: None, } } @@ -100,6 +109,7 @@ impl Trial { id: u64, sampler: Arc, history: Arc>>>, + pruner: Arc, ) -> Self { Self { id, @@ -109,6 +119,8 @@ impl Trial { param_labels: HashMap::new(), sampler: Some(sampler), history: Some(history), + intermediate_values: Vec::new(), + pruner: Some(pruner), } } @@ -159,6 +171,44 @@ 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 the trial state to Complete. pub(crate) fn set_complete(&mut self) { self.state = TrialState::Complete;