diff --git a/src/study.rs b/src/study.rs index 72dee4d..0042a23 100644 --- a/src/study.rs +++ b/src/study.rs @@ -6,11 +6,14 @@ 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}; @@ -53,6 +56,8 @@ where /// Set automatically for `Study` so that `create_trial()` and all /// optimization methods use the sampler without requiring `_with_sampler` suffixes. trial_factory: Option Trial + Send + Sync>>, + /// Queue of parameter configurations to evaluate next. + enqueued_params: Arc>>>, } impl Study @@ -170,6 +175,7 @@ where completed_trials, next_trial_id: AtomicU64::new(0), trial_factory, + enqueued_params: Arc::new(Mutex::new(VecDeque::new())), } } @@ -261,6 +267,7 @@ where completed_trials, next_trial_id: AtomicU64::new(0), trial_factory, + enqueued_params: Arc::new(Mutex::new(VecDeque::new())), } } @@ -292,6 +299,52 @@ where &*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 = 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) { + 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. pub(crate) fn next_trial_id(&self) -> u64 { self.next_trial_id.fetch_add(1, Ordering::SeqCst) @@ -321,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. diff --git a/src/trial.rs b/src/trial.rs index b0dc05c..3cc33c2 100644 --- a/src/trial.rs +++ b/src/trial.rs @@ -86,6 +86,8 @@ pub struct Trial { pruner: Option>, /// User-defined attributes for logging, debugging, and analysis. user_attrs: HashMap, + /// Pre-filled parameter values from enqueue (used instead of sampling). + fixed_params: HashMap, } impl core::fmt::Debug for Trial { @@ -101,6 +103,7 @@ impl core::fmt::Debug for Trial { .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() } } @@ -139,6 +142,7 @@ impl Trial { intermediate_values: Vec::new(), pruner: None, user_attrs: HashMap::new(), + fixed_params: HashMap::new(), } } @@ -169,9 +173,18 @@ impl Trial { 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) { + 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 @@ -343,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(¶m_id) { + fixed_value + } else { + // Sample using the sampler + self.sample_value(&distribution) + }; + let result = param.cast_param_value(&value)?; // Store distribution, value, and label diff --git a/tests/integration.rs b/tests/integration.rs index 8cfd84d..fa63f24 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -1700,3 +1700,195 @@ fn test_ask_and_tell_with_custom_value_type() { 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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); +}