diff --git a/src/lib.rs b/src/lib.rs index 93eab99..35fe5c3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -272,9 +272,9 @@ pub use storage::JournalStorage; #[cfg(feature = "sqlite")] pub use storage::SqliteStorage; pub use storage::{MemoryStorage, Storage}; -pub use study::{Study, StudyBuilder}; #[cfg(feature = "serde")] pub use study::StudySnapshot; +pub use study::{Study, StudyBuilder}; pub use trial::{AttrValue, Trial}; pub use types::{Direction, TrialState}; pub use visualization::generate_html_report; @@ -320,9 +320,9 @@ pub mod prelude { #[cfg(feature = "sqlite")] pub use crate::storage::SqliteStorage; pub use crate::storage::{MemoryStorage, Storage}; - pub use crate::study::{Study, StudyBuilder}; #[cfg(feature = "serde")] pub use crate::study::StudySnapshot; + pub use crate::study::{Study, StudyBuilder}; pub use crate::trial::{AttrValue, Trial}; pub use crate::types::Direction; pub use crate::visualization::generate_html_report; diff --git a/src/study.rs b/src/study.rs index 77b6eaf..f2cd352 100644 --- a/src/study.rs +++ b/src/study.rs @@ -4,6 +4,7 @@ use core::any::Any; use core::fmt; #[cfg(feature = "async")] use core::future::Future; +use core::marker::PhantomData; use core::ops::ControlFlow; use core::sync::atomic::{AtomicU64, Ordering}; use core::time::Duration; @@ -89,6 +90,30 @@ where Self::with_sampler(direction, RandomSampler::new()) } + /// Returns a [`StudyBuilder`] for constructing a study with a fluent API. + /// + /// # Examples + /// + /// ``` + /// use optimizer::prelude::*; + /// + /// let study: Study = Study::builder() + /// .minimize() + /// .sampler(TpeSampler::new()) + /// .pruner(NopPruner) + /// .build(); + /// ``` + #[must_use] + pub fn builder() -> StudyBuilder { + StudyBuilder { + direction: Direction::Minimize, + sampler: None, + pruner: None, + storage: None, + _marker: PhantomData, + } + } + /// Creates a study that minimizes the objective value. /// /// This is a shorthand for `Study::with_sampler(Direction::Minimize, sampler)`. @@ -2331,6 +2356,121 @@ impl Study { } } +/// A builder for constructing [`Study`] instances with a fluent API. +/// +/// Created via [`Study::builder()`]. Collects sampler, pruner, direction, +/// and storage options before constructing the study. +/// +/// # Defaults +/// +/// - Direction: [`Minimize`](Direction::Minimize) +/// - Sampler: [`RandomSampler`] +/// - Pruner: [`NopPruner`] +/// - Storage: [`MemoryStorage`](crate::storage::MemoryStorage) +/// +/// # Examples +/// +/// ``` +/// use optimizer::prelude::*; +/// +/// let study: Study = Study::builder() +/// .maximize() +/// .sampler(TpeSampler::new()) +/// .pruner(MedianPruner::new(Direction::Maximize).n_warmup_steps(5)) +/// .build(); +/// +/// assert_eq!(study.direction(), Direction::Maximize); +/// ``` +pub struct StudyBuilder { + direction: Direction, + sampler: Option>, + pruner: Option>, + storage: Option>>, + _marker: PhantomData, +} + +impl StudyBuilder { + /// Sets the optimization direction to minimize. + #[must_use] + pub fn minimize(mut self) -> Self { + self.direction = Direction::Minimize; + self + } + + /// Sets the optimization direction to maximize. + #[must_use] + pub fn maximize(mut self) -> Self { + self.direction = Direction::Maximize; + self + } + + /// Sets the optimization direction. + #[must_use] + pub fn direction(mut self, direction: Direction) -> Self { + self.direction = direction; + self + } + + /// Sets the sampler used for parameter suggestions. + #[must_use] + pub fn sampler(mut self, sampler: impl Sampler + 'static) -> Self { + self.sampler = Some(Box::new(sampler)); + self + } + + /// Sets the pruner used for early stopping of trials. + #[must_use] + pub fn pruner(mut self, pruner: impl Pruner + 'static) -> Self { + self.pruner = Some(Box::new(pruner)); + self + } + + /// Sets a custom storage backend. + #[must_use] + pub fn storage(mut self, storage: impl crate::storage::Storage + 'static) -> Self { + self.storage = Some(Box::new(storage)); + self + } + + /// Builds the [`Study`] with the configured options. + #[must_use] + pub fn build(self) -> Study + where + V: Send + Sync + 'static, + { + let sampler = self + .sampler + .unwrap_or_else(|| Box::new(RandomSampler::new())); + let pruner = self.pruner.unwrap_or_else(|| Box::new(NopPruner)); + let storage = self + .storage + .unwrap_or_else(|| Box::new(crate::storage::MemoryStorage::::new())); + + let sampler: Arc = Arc::from(sampler); + let pruner: Arc = Arc::from(pruner); + let storage: Arc> = Arc::from(storage); + let trial_factory = Study::make_trial_factory(&sampler, &storage, &pruner); + + let next_id = storage + .trials_arc() + .read() + .iter() + .map(|t| t.id) + .max() + .map_or(0, |id| id + 1); + + Study { + direction: self.direction, + sampler, + pruner, + storage, + next_trial_id: AtomicU64::new(next_id), + trial_factory, + enqueued_params: Arc::new(Mutex::new(VecDeque::new())), + } + } +} + #[cfg(feature = "journal")] impl Study where diff --git a/tests/integration.rs b/tests/integration.rs index d323dc8..9e8eb76 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -2245,3 +2245,95 @@ fn test_top_trials_ranks_feasible_above_infeasible() { // Feasible sorted by objective first (5.0, 50.0), then infeasible by violation (0.5, 3.0) assert_eq!(ids, vec![2, 1, 0, 3]); } + +// ============================================================================= +// Test: StudyBuilder +// ============================================================================= + +#[test] +fn test_builder_defaults() { + let study: Study = Study::builder().build(); + assert_eq!(study.direction(), Direction::Minimize); +} + +#[test] +fn test_builder_maximize() { + let study: Study = Study::builder().maximize().build(); + assert_eq!(study.direction(), Direction::Maximize); +} + +#[test] +fn test_builder_minimize() { + let study: Study = Study::builder().minimize().build(); + assert_eq!(study.direction(), Direction::Minimize); +} + +#[test] +fn test_builder_direction() { + let study: Study = Study::builder().direction(Direction::Maximize).build(); + assert_eq!(study.direction(), Direction::Maximize); +} + +#[test] +fn test_builder_with_sampler() { + let x = FloatParam::new(-5.0, 5.0); + let study: Study = Study::builder().sampler(TpeSampler::new()).build(); + + study + .optimize(10, |trial| { + let val = x.suggest(trial)?; + Ok::<_, Error>(val * val) + }) + .unwrap(); + + assert_eq!(study.trials().len(), 10); +} + +#[test] +fn test_builder_with_pruner() { + use optimizer::NopPruner; + + let study: Study = Study::builder().pruner(NopPruner).build(); + + assert_eq!(study.direction(), Direction::Minimize); +} + +#[test] +fn test_builder_chaining() { + let study: Study = Study::builder() + .maximize() + .sampler(RandomSampler::with_seed(42)) + .pruner(optimizer::NopPruner) + .build(); + + assert_eq!(study.direction(), Direction::Maximize); +} + +#[test] +fn test_builder_with_custom_value_type() { + let study: Study = Study::builder().maximize().build(); + assert_eq!(study.direction(), Direction::Maximize); +} + +#[test] +fn test_builder_optimizes_correctly() { + let x = FloatParam::new(-10.0, 10.0); + let study: Study = Study::builder() + .minimize() + .sampler(TpeSampler::builder().seed(42).build().unwrap()) + .build(); + + study + .optimize(100, |trial| { + let val = x.suggest(trial)?; + Ok::<_, Error>((val - 3.0) * (val - 3.0)) + }) + .unwrap(); + + let best = study.best_trial().unwrap(); + assert!( + best.value < 5.0, + "best value should be < 5.0, got {}", + best.value + ); +}