feat: add StudyBuilder for fluent study construction

This commit is contained in:
Manuel Raimann
2026-02-11 23:23:56 +01:00
parent ed80c59795
commit d14f4e0792
3 changed files with 234 additions and 2 deletions
+2 -2
View File
@@ -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
View File
@@ -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
+92
View File
@@ -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
);
}