feat: add StudyBuilder for fluent study construction
This commit is contained in:
+2
-2
@@ -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;
|
||||
|
||||
+140
@@ -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<f64> = Study::builder()
|
||||
/// .minimize()
|
||||
/// .sampler(TpeSampler::new())
|
||||
/// .pruner(NopPruner)
|
||||
/// .build();
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn builder() -> StudyBuilder<V> {
|
||||
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<V: PartialOrd + Send + Sync + 'static> Study<V> {
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<f64> = 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<V: PartialOrd = f64> {
|
||||
direction: Direction,
|
||||
sampler: Option<Box<dyn Sampler>>,
|
||||
pruner: Option<Box<dyn Pruner>>,
|
||||
storage: Option<Box<dyn crate::storage::Storage<V>>>,
|
||||
_marker: PhantomData<V>,
|
||||
}
|
||||
|
||||
impl<V: PartialOrd> StudyBuilder<V> {
|
||||
/// 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<V> + 'static) -> Self {
|
||||
self.storage = Some(Box::new(storage));
|
||||
self
|
||||
}
|
||||
|
||||
/// Builds the [`Study`] with the configured options.
|
||||
#[must_use]
|
||||
pub fn build(self) -> Study<V>
|
||||
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::<V>::new()));
|
||||
|
||||
let sampler: Arc<dyn Sampler> = Arc::from(sampler);
|
||||
let pruner: Arc<dyn Pruner> = Arc::from(pruner);
|
||||
let storage: Arc<dyn crate::storage::Storage<V>> = 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<V> Study<V>
|
||||
where
|
||||
|
||||
@@ -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<f64> = Study::builder().build();
|
||||
assert_eq!(study.direction(), Direction::Minimize);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_maximize() {
|
||||
let study: Study<f64> = Study::builder().maximize().build();
|
||||
assert_eq!(study.direction(), Direction::Maximize);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_minimize() {
|
||||
let study: Study<f64> = Study::builder().minimize().build();
|
||||
assert_eq!(study.direction(), Direction::Minimize);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_direction() {
|
||||
let study: Study<f64> = 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<f64> = 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<f64> = Study::builder().pruner(NopPruner).build();
|
||||
|
||||
assert_eq!(study.direction(), Direction::Minimize);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_builder_chaining() {
|
||||
let study: Study<f64> = 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<i32> = 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<f64> = 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
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user