Manuel Raimann
1341 lines
45 KiB
Rust
1341 lines
45 KiB
Rust
//! Study implementation for managing optimization trials.
|
|
|
|
#[cfg(feature = "async")]
|
|
use std::future::Future;
|
|
use std::ops::ControlFlow;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
|
|
use parking_lot::RwLock;
|
|
#[cfg(feature = "serde")]
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
use crate::sampler::{CompletedTrial, RandomSampler, Sampler};
|
|
use crate::trial::Trial;
|
|
use crate::types::Direction;
|
|
|
|
/// Helper function to create default sampler for serde deserialization.
|
|
#[cfg(feature = "serde")]
|
|
fn default_sampler() -> Arc<dyn Sampler> {
|
|
Arc::new(RandomSampler::new())
|
|
}
|
|
|
|
/// A study manages the optimization process, tracking trials and their results.
|
|
///
|
|
/// The study is parameterized by the objective value type `V`, which defaults to `f64`.
|
|
/// The only constraint on `V` is `PartialOrd`, allowing comparison of objective values
|
|
/// to determine which trial is best.
|
|
///
|
|
/// When `V = f64`, the study passes trial history to the sampler for informed
|
|
/// parameter suggestions (e.g., TPE sampler uses history to guide sampling).
|
|
///
|
|
/// # Serialization
|
|
///
|
|
/// When the `serde` feature is enabled, the study can be serialized and deserialized.
|
|
/// The completed trials and trial ID counter are preserved, allowing optimization to
|
|
/// continue after deserialization. The sampler is not serialized; upon deserialization,
|
|
/// a default `RandomSampler` is used. Use `Study::set_sampler()` to restore a custom sampler.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// // Create a study to minimize an objective function
|
|
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
|
/// assert_eq!(study.direction(), Direction::Minimize);
|
|
/// ```
|
|
pub struct Study<V = f64>
|
|
where
|
|
V: PartialOrd,
|
|
{
|
|
/// The optimization direction.
|
|
direction: Direction,
|
|
/// The sampler used to generate parameter values.
|
|
sampler: Arc<dyn Sampler>,
|
|
/// Completed trials (wrapped in Arc for sharing with Trial).
|
|
completed_trials: Arc<RwLock<Vec<CompletedTrial<V>>>>,
|
|
/// Counter for generating unique trial IDs.
|
|
next_trial_id: AtomicU64,
|
|
}
|
|
|
|
impl<V> Study<V>
|
|
where
|
|
V: PartialOrd,
|
|
{
|
|
/// Creates a new study with the given optimization direction.
|
|
///
|
|
/// Uses the default `RandomSampler` for parameter sampling.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `direction` - Whether to minimize or maximize the objective function.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
|
/// assert_eq!(study.direction(), Direction::Minimize);
|
|
/// ```
|
|
pub fn new(direction: Direction) -> Self {
|
|
Self::with_sampler(direction, RandomSampler::new())
|
|
}
|
|
|
|
/// Creates a new study with a custom sampler.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `direction` - Whether to minimize or maximize the objective function.
|
|
/// * `sampler` - The sampler to use for parameter sampling.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
|
|
/// assert_eq!(study.direction(), Direction::Maximize);
|
|
/// ```
|
|
pub fn with_sampler(direction: Direction, sampler: impl Sampler + 'static) -> Self {
|
|
Self {
|
|
direction,
|
|
sampler: Arc::new(sampler),
|
|
completed_trials: Arc::new(RwLock::new(Vec::new())),
|
|
next_trial_id: AtomicU64::new(0),
|
|
}
|
|
}
|
|
|
|
/// Returns the optimization direction.
|
|
pub fn direction(&self) -> Direction {
|
|
self.direction
|
|
}
|
|
|
|
/// Sets a new sampler for the study.
|
|
///
|
|
/// This method is useful after deserializing a study when you want to use
|
|
/// a custom sampler (e.g., TPE) instead of the default `RandomSampler`.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `sampler` - The sampler to use for parameter sampling.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study, TpeSampler};
|
|
///
|
|
/// // After deserializing a study, restore the TPE sampler
|
|
/// let mut study: Study<f64> = Study::new(Direction::Minimize);
|
|
/// study.set_sampler(TpeSampler::new());
|
|
/// ```
|
|
pub fn set_sampler(&mut self, sampler: impl Sampler + 'static) {
|
|
self.sampler = Arc::new(sampler);
|
|
}
|
|
|
|
/// Generates the next unique trial ID.
|
|
pub(crate) fn next_trial_id(&self) -> u64 {
|
|
self.next_trial_id.fetch_add(1, Ordering::SeqCst)
|
|
}
|
|
|
|
/// Creates a new trial with a unique ID.
|
|
///
|
|
/// The trial starts in the `Running` state and can be used to suggest
|
|
/// parameter values. After the objective function is evaluated, call
|
|
/// `complete_trial` or `fail_trial` to record the result.
|
|
///
|
|
/// Note: For `Study<f64>`, this method creates a trial without sampler
|
|
/// integration. Use `create_trial_with_sampler()` to create trials that
|
|
/// use the study's sampler and have access to trial history.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
|
/// let trial = study.create_trial();
|
|
/// assert_eq!(trial.id(), 0);
|
|
///
|
|
/// let trial2 = study.create_trial();
|
|
/// assert_eq!(trial2.id(), 1);
|
|
/// ```
|
|
pub fn create_trial(&self) -> Trial {
|
|
let id = self.next_trial_id();
|
|
Trial::new(id)
|
|
}
|
|
|
|
/// Records a completed trial with its objective value.
|
|
///
|
|
/// This method stores the trial's parameters, distributions, and objective
|
|
/// value in the study's history. The stored data is used by samplers to
|
|
/// inform future parameter suggestions.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `trial` - The trial that was evaluated.
|
|
/// * `value` - The objective value returned by the objective function.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
|
/// let mut trial = study.create_trial();
|
|
/// let x = trial.suggest_float("x", 0.0, 1.0).unwrap();
|
|
/// let objective_value = x * x;
|
|
/// study.complete_trial(trial, objective_value);
|
|
///
|
|
/// assert_eq!(study.n_trials(), 1);
|
|
/// ```
|
|
pub fn complete_trial(&self, mut trial: Trial, value: V) {
|
|
trial.set_complete();
|
|
let completed = CompletedTrial::new(
|
|
trial.id(),
|
|
trial.params().clone(),
|
|
trial.distributions().clone(),
|
|
value,
|
|
);
|
|
self.completed_trials.write().push(completed);
|
|
}
|
|
|
|
/// Records a failed trial with an error message.
|
|
///
|
|
/// Failed trials are not stored in the study's history and do not
|
|
/// contribute to future sampling decisions. This method is useful
|
|
/// when the objective function raises an error that should not stop
|
|
/// the optimization process.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `trial` - The trial that failed.
|
|
/// * `_error` - An error message describing why the trial failed.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
|
/// let trial = study.create_trial();
|
|
/// study.fail_trial(trial, "objective function raised an exception");
|
|
///
|
|
/// // Failed trials are not counted
|
|
/// assert_eq!(study.n_trials(), 0);
|
|
/// ```
|
|
pub fn fail_trial(&self, mut trial: Trial, _error: impl ToString) {
|
|
trial.set_failed();
|
|
// Failed trials are not stored in completed_trials
|
|
// They could be stored in a separate list for debugging if needed
|
|
}
|
|
|
|
/// Returns an iterator over all completed trials.
|
|
///
|
|
/// The iterator yields references to `CompletedTrial` values, which contain
|
|
/// the trial's parameters, distributions, and objective value.
|
|
///
|
|
/// Note: This method acquires a read lock on the completed trials, so the
|
|
/// returned vector is a clone of the internal storage.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
|
/// let mut trial = study.create_trial();
|
|
/// let _ = trial.suggest_float("x", 0.0, 1.0);
|
|
/// study.complete_trial(trial, 0.5);
|
|
///
|
|
/// for completed in study.trials() {
|
|
/// println!("Trial {} has value {:?}", completed.id, completed.value);
|
|
/// }
|
|
/// ```
|
|
pub fn trials(&self) -> Vec<CompletedTrial<V>>
|
|
where
|
|
V: Clone,
|
|
{
|
|
self.completed_trials.read().clone()
|
|
}
|
|
|
|
/// Returns the number of completed trials.
|
|
///
|
|
/// Failed trials are not counted.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
|
/// assert_eq!(study.n_trials(), 0);
|
|
///
|
|
/// let mut trial = study.create_trial();
|
|
/// let _ = trial.suggest_float("x", 0.0, 1.0);
|
|
/// study.complete_trial(trial, 0.5);
|
|
/// assert_eq!(study.n_trials(), 1);
|
|
/// ```
|
|
pub fn n_trials(&self) -> usize {
|
|
self.completed_trials.read().len()
|
|
}
|
|
|
|
/// Returns the trial with the best objective value.
|
|
///
|
|
/// The "best" trial depends on the optimization direction:
|
|
/// - `Direction::Minimize`: Returns the trial with the lowest objective value.
|
|
/// - `Direction::Maximize`: Returns the trial with the highest objective value.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if no trials have been completed.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
|
///
|
|
/// // Error when no trials completed
|
|
/// assert!(study.best_trial().is_err());
|
|
///
|
|
/// let mut trial1 = study.create_trial();
|
|
/// let _ = trial1.suggest_float("x", 0.0, 1.0);
|
|
/// study.complete_trial(trial1, 0.8);
|
|
///
|
|
/// let mut trial2 = study.create_trial();
|
|
/// let _ = trial2.suggest_float("x", 0.0, 1.0);
|
|
/// study.complete_trial(trial2, 0.3);
|
|
///
|
|
/// let best = study.best_trial().unwrap();
|
|
/// assert_eq!(best.value, 0.3); // Minimize: lower is better
|
|
/// ```
|
|
pub fn best_trial(&self) -> crate::Result<CompletedTrial<V>>
|
|
where
|
|
V: Clone,
|
|
{
|
|
let trials = self.completed_trials.read();
|
|
|
|
if trials.is_empty() {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
let best = trials
|
|
.iter()
|
|
.max_by(|a, b| {
|
|
// For Minimize, we want the smallest value to be "max" in ordering
|
|
// For Maximize, we want the largest value to be "max" in ordering
|
|
let ordering = a.value.partial_cmp(&b.value);
|
|
match self.direction {
|
|
Direction::Minimize => {
|
|
// Reverse ordering: smaller values are "greater" for max_by
|
|
ordering
|
|
.map(|o| o.reverse())
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
}
|
|
Direction::Maximize => {
|
|
// Normal ordering: larger values are "greater" for max_by
|
|
ordering.unwrap_or(std::cmp::Ordering::Equal)
|
|
}
|
|
}
|
|
})
|
|
.expect("trials is not empty");
|
|
|
|
Ok(best.clone())
|
|
}
|
|
|
|
/// Returns the best objective value found so far.
|
|
///
|
|
/// The "best" value depends on the optimization direction:
|
|
/// - `Direction::Minimize`: Returns the lowest objective value.
|
|
/// - `Direction::Maximize`: Returns the highest objective value.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if no trials have been completed.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, Study};
|
|
///
|
|
/// let study: Study<f64> = Study::new(Direction::Maximize);
|
|
///
|
|
/// // Error when no trials completed
|
|
/// assert!(study.best_value().is_err());
|
|
///
|
|
/// let mut trial1 = study.create_trial();
|
|
/// let _ = trial1.suggest_float("x", 0.0, 1.0);
|
|
/// study.complete_trial(trial1, 0.3);
|
|
///
|
|
/// let mut trial2 = study.create_trial();
|
|
/// let _ = trial2.suggest_float("x", 0.0, 1.0);
|
|
/// study.complete_trial(trial2, 0.8);
|
|
///
|
|
/// let best = study.best_value().unwrap();
|
|
/// assert_eq!(best, 0.8); // Maximize: higher is better
|
|
/// ```
|
|
pub fn best_value(&self) -> crate::Result<V>
|
|
where
|
|
V: Clone,
|
|
{
|
|
self.best_trial().map(|trial| trial.value)
|
|
}
|
|
|
|
/// Runs optimization with the given objective function.
|
|
///
|
|
/// This method runs `n_trials` evaluations sequentially. For each trial:
|
|
/// 1. A new trial is created
|
|
/// 2. The objective function is called with the trial
|
|
/// 3. If successful, the trial is recorded as completed
|
|
/// 4. If the objective returns an error, the trial is recorded as failed
|
|
///
|
|
/// Failed trials do not stop the optimization; the process continues with
|
|
/// the next trial.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_trials` - The number of trials to run.
|
|
/// * `objective` - A closure that takes a mutable reference to a `Trial` and
|
|
/// returns the objective value or an error.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// // Minimize x^2
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
///
|
|
/// study
|
|
/// .optimize(10, |trial| {
|
|
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
|
/// Ok::<_, optimizer::TpeError>(x * x)
|
|
/// })
|
|
/// .unwrap();
|
|
///
|
|
/// // At least one trial should have completed
|
|
/// assert!(study.n_trials() > 0);
|
|
/// let best = study.best_value().unwrap();
|
|
/// assert!(best >= 0.0);
|
|
/// ```
|
|
pub fn optimize<F, E>(&self, n_trials: usize, mut objective: F) -> crate::Result<()>
|
|
where
|
|
F: FnMut(&mut Trial) -> std::result::Result<V, E>,
|
|
E: ToString,
|
|
{
|
|
for _ in 0..n_trials {
|
|
let mut trial = self.create_trial();
|
|
|
|
match objective(&mut trial) {
|
|
Ok(value) => {
|
|
self.complete_trial(trial, value);
|
|
}
|
|
Err(e) => {
|
|
self.fail_trial(trial, e.to_string());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return error if no trials succeeded
|
|
if self.n_trials() == 0 {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Runs optimization asynchronously with the given objective function.
|
|
///
|
|
/// This method runs `n_trials` evaluations sequentially, but the objective
|
|
/// function can be async (e.g., for I/O-bound operations like network requests
|
|
/// or file operations).
|
|
///
|
|
/// The objective function takes ownership of the `Trial` and must return it
|
|
/// along with the result. This allows async operations to use the trial
|
|
/// across await points.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_trials` - The number of trials to run.
|
|
/// * `objective` - A function that takes a `Trial` and returns a `Future`
|
|
/// that resolves to a tuple of `(Trial, Result<V, E>)`.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// # #[cfg(feature = "async")]
|
|
/// # async fn example() -> optimizer::Result<()> {
|
|
/// // Minimize x^2 with async objective
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
///
|
|
/// study
|
|
/// .optimize_async(10, |mut trial| async move {
|
|
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
|
/// // Simulate async work (e.g., network request)
|
|
/// let value = x * x;
|
|
/// Ok::<_, optimizer::TpeError>((trial, value))
|
|
/// })
|
|
/// .await?;
|
|
///
|
|
/// // At least one trial should have completed
|
|
/// assert!(study.n_trials() > 0);
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
#[cfg(feature = "async")]
|
|
pub async fn optimize_async<F, Fut, E>(
|
|
&self,
|
|
n_trials: usize,
|
|
objective: F,
|
|
) -> crate::Result<()>
|
|
where
|
|
F: Fn(Trial) -> Fut,
|
|
Fut: Future<Output = std::result::Result<(Trial, V), E>>,
|
|
E: ToString,
|
|
{
|
|
for _ in 0..n_trials {
|
|
let trial = self.create_trial();
|
|
|
|
match objective(trial).await {
|
|
Ok((trial, value)) => {
|
|
self.complete_trial(trial, value);
|
|
}
|
|
Err(e) => {
|
|
// For async, we don't have the trial back on error
|
|
// We'll just count this as a failed trial without recording it
|
|
let _ = e.to_string();
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return error if no trials succeeded
|
|
if self.n_trials() == 0 {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Runs optimization with bounded parallelism for concurrent trial evaluation.
|
|
///
|
|
/// This method runs up to `concurrency` trials simultaneously, allowing
|
|
/// efficient use of async I/O-bound objective functions. A semaphore limits
|
|
/// the number of concurrent evaluations.
|
|
///
|
|
/// The objective function takes ownership of the `Trial` and must return it
|
|
/// along with the result. This allows async operations to use the trial
|
|
/// across await points.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_trials` - The total number of trials to run.
|
|
/// * `concurrency` - The maximum number of trials to run simultaneously.
|
|
/// * `objective` - A function that takes a `Trial` and returns a `Future`
|
|
/// that resolves to a tuple of `(Trial, V)` or an error.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// # #[cfg(feature = "async")]
|
|
/// # async fn example() -> optimizer::Result<()> {
|
|
/// // Minimize x^2 with parallel async evaluation
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
///
|
|
/// study
|
|
/// .optimize_parallel(10, 4, |mut trial| async move {
|
|
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
|
/// // Async objective function (e.g., network request)
|
|
/// let value = x * x;
|
|
/// Ok::<_, optimizer::TpeError>((trial, value))
|
|
/// })
|
|
/// .await?;
|
|
///
|
|
/// // All trials should have completed
|
|
/// assert_eq!(study.n_trials(), 10);
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
#[cfg(feature = "async")]
|
|
pub async fn optimize_parallel<F, Fut, E>(
|
|
&self,
|
|
n_trials: usize,
|
|
concurrency: usize,
|
|
objective: F,
|
|
) -> crate::Result<()>
|
|
where
|
|
F: Fn(Trial) -> Fut + Send + Sync + 'static,
|
|
Fut: Future<Output = std::result::Result<(Trial, V), E>> + Send,
|
|
E: ToString + Send + 'static,
|
|
V: Send + 'static,
|
|
{
|
|
use tokio::sync::Semaphore;
|
|
|
|
let semaphore = Arc::new(Semaphore::new(concurrency));
|
|
let objective = Arc::new(objective);
|
|
|
|
let mut handles = Vec::with_capacity(n_trials);
|
|
|
|
for _ in 0..n_trials {
|
|
let permit = semaphore.clone().acquire_owned().await.unwrap();
|
|
let trial = self.create_trial();
|
|
let objective = Arc::clone(&objective);
|
|
|
|
let handle = tokio::spawn(async move {
|
|
let result = objective(trial).await;
|
|
drop(permit); // Release semaphore permit when done
|
|
result
|
|
});
|
|
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Wait for all tasks and record results
|
|
for handle in handles {
|
|
match handle.await.unwrap() {
|
|
Ok((trial, value)) => {
|
|
self.complete_trial(trial, value);
|
|
}
|
|
Err(e) => {
|
|
let _ = e.to_string();
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return error if no trials succeeded
|
|
if self.n_trials() == 0 {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Runs optimization with a callback for monitoring progress.
|
|
///
|
|
/// This method is similar to `optimize`, but calls a callback function after
|
|
/// each completed trial. The callback can inspect the study state and the
|
|
/// completed trial, and can optionally stop optimization early by returning
|
|
/// `ControlFlow::Break(())`.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_trials` - The maximum number of trials to run.
|
|
/// * `objective` - A closure that takes a mutable reference to a `Trial` and
|
|
/// returns the objective value or an error.
|
|
/// * `callback` - A closure called after each successful trial. Returns
|
|
/// `ControlFlow::Continue(())` to proceed or `ControlFlow::Break(())` to stop.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if no trials completed successfully
|
|
/// before optimization stopped (either by completing all trials or early stopping).
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use std::ops::ControlFlow;
|
|
///
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// // Stop early when we find a good enough value
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
///
|
|
/// study
|
|
/// .optimize_with_callback(
|
|
/// 100,
|
|
/// |trial| {
|
|
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
|
/// Ok::<_, optimizer::TpeError>(x * x)
|
|
/// },
|
|
/// |_study, completed_trial| {
|
|
/// // Stop early if we find a value less than 1.0
|
|
/// if completed_trial.value < 1.0 {
|
|
/// ControlFlow::Break(())
|
|
/// } else {
|
|
/// ControlFlow::Continue(())
|
|
/// }
|
|
/// },
|
|
/// )
|
|
/// .unwrap();
|
|
///
|
|
/// // May have stopped early, but should have at least one trial
|
|
/// assert!(study.n_trials() > 0);
|
|
/// ```
|
|
pub fn optimize_with_callback<F, C, E>(
|
|
&self,
|
|
n_trials: usize,
|
|
mut objective: F,
|
|
mut callback: C,
|
|
) -> crate::Result<()>
|
|
where
|
|
V: Clone,
|
|
F: FnMut(&mut Trial) -> std::result::Result<V, E>,
|
|
C: FnMut(&Study<V>, &CompletedTrial<V>) -> ControlFlow<()>,
|
|
E: ToString,
|
|
{
|
|
for _ in 0..n_trials {
|
|
let mut trial = self.create_trial();
|
|
|
|
match objective(&mut trial) {
|
|
Ok(value) => {
|
|
self.complete_trial(trial, value);
|
|
|
|
// Get the just-completed trial for the callback
|
|
let trials = self.completed_trials.read();
|
|
let completed = trials.last().expect("just added a trial");
|
|
|
|
// Call the callback and check if we should stop
|
|
// Note: We need to drop the read lock before calling callback
|
|
// to avoid potential deadlock if callback accesses the study
|
|
let completed_clone = completed.clone();
|
|
drop(trials);
|
|
|
|
if let ControlFlow::Break(()) = callback(self, &completed_clone) {
|
|
break;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
self.fail_trial(trial, e.to_string());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return error if no trials succeeded
|
|
if self.n_trials() == 0 {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// Specialized implementation for Study<f64> that provides full sampler integration.
|
|
impl Study<f64> {
|
|
/// Creates a new trial with sampler integration.
|
|
///
|
|
/// This method creates a trial that uses the study's sampler and has access
|
|
/// to the history of completed trials for informed parameter suggestions.
|
|
/// This is the recommended way to create trials when using `Study<f64>`.
|
|
///
|
|
/// The trial's `suggest_*` methods will delegate to the sampler (e.g., TPE)
|
|
/// which can use historical trial data to make informed sampling decisions.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// // With a seeded sampler for reproducibility
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
/// let mut trial = study.create_trial_with_sampler();
|
|
///
|
|
/// // Parameter suggestions now use the study's sampler and history
|
|
/// let x = trial.suggest_float("x", 0.0, 1.0).unwrap();
|
|
/// ```
|
|
pub fn create_trial_with_sampler(&self) -> Trial {
|
|
let id = self.next_trial_id();
|
|
Trial::with_sampler(
|
|
id,
|
|
Arc::clone(&self.sampler),
|
|
Arc::clone(&self.completed_trials),
|
|
)
|
|
}
|
|
|
|
/// Runs optimization with full sampler integration.
|
|
///
|
|
/// This method is similar to the generic `optimizer` method but creates trials
|
|
/// using `create_trial_with_sampler()`, giving the sampler access to the history
|
|
/// of completed trials for informed parameter suggestions.
|
|
///
|
|
/// This is the recommended way to run optimization when using `Study<f64>`
|
|
/// with advanced samplers like TPE.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_trials` - The number of trials to run.
|
|
/// * `objective` - A closure that takes a mutable reference to a `Trial` and
|
|
/// returns the objective value or an error.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// // Minimize x^2 with sampler integration
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
///
|
|
/// study
|
|
/// .optimize_with_sampler(10, |trial| {
|
|
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
|
/// Ok::<_, optimizer::TpeError>(x * x)
|
|
/// })
|
|
/// .unwrap();
|
|
///
|
|
/// // At least one trial should have completed
|
|
/// assert!(study.n_trials() > 0);
|
|
/// ```
|
|
pub fn optimize_with_sampler<F, E>(
|
|
&self,
|
|
n_trials: usize,
|
|
mut objective: F,
|
|
) -> crate::Result<()>
|
|
where
|
|
F: FnMut(&mut Trial) -> std::result::Result<f64, E>,
|
|
E: ToString,
|
|
{
|
|
for _ in 0..n_trials {
|
|
let mut trial = self.create_trial_with_sampler();
|
|
|
|
match objective(&mut trial) {
|
|
Ok(value) => {
|
|
self.complete_trial(trial, value);
|
|
}
|
|
Err(e) => {
|
|
self.fail_trial(trial, e.to_string());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return error if no trials succeeded
|
|
if self.n_trials() == 0 {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Runs optimization with a callback and full sampler integration.
|
|
///
|
|
/// This method combines the benefits of `optimize_with_sampler` (sampler access
|
|
/// to trial history) with `optimize_with_callback` (progress monitoring and
|
|
/// early stopping).
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_trials` - The maximum number of trials to run.
|
|
/// * `objective` - A closure that takes a mutable reference to a `Trial` and
|
|
/// returns the objective value or an error.
|
|
/// * `callback` - A closure called after each successful trial. Returns
|
|
/// `ControlFlow::Continue(())` to proceed or `ControlFlow::Break(())` to stop.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if no trials completed successfully.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use std::ops::ControlFlow;
|
|
///
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// // Optimize with sampler integration and early stopping
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
///
|
|
/// study
|
|
/// .optimize_with_callback_sampler(
|
|
/// 100,
|
|
/// |trial| {
|
|
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
|
/// Ok::<_, optimizer::TpeError>(x * x)
|
|
/// },
|
|
/// |study, _completed_trial| {
|
|
/// // Stop after finding 5 good trials
|
|
/// if study.n_trials() >= 5 {
|
|
/// ControlFlow::Break(())
|
|
/// } else {
|
|
/// ControlFlow::Continue(())
|
|
/// }
|
|
/// },
|
|
/// )
|
|
/// .unwrap();
|
|
///
|
|
/// assert!(study.n_trials() >= 5);
|
|
/// ```
|
|
pub fn optimize_with_callback_sampler<F, C, E>(
|
|
&self,
|
|
n_trials: usize,
|
|
mut objective: F,
|
|
mut callback: C,
|
|
) -> crate::Result<()>
|
|
where
|
|
F: FnMut(&mut Trial) -> std::result::Result<f64, E>,
|
|
C: FnMut(&Study<f64>, &CompletedTrial<f64>) -> ControlFlow<()>,
|
|
E: ToString,
|
|
{
|
|
for _ in 0..n_trials {
|
|
let mut trial = self.create_trial_with_sampler();
|
|
|
|
match objective(&mut trial) {
|
|
Ok(value) => {
|
|
self.complete_trial(trial, value);
|
|
|
|
// Get the just-completed trial for the callback
|
|
let trials = self.completed_trials.read();
|
|
let completed = trials.last().expect("just added a trial");
|
|
|
|
// Call the callback and check if we should stop
|
|
// Note: We need to drop the read lock before calling callback
|
|
// to avoid potential deadlock if callback accesses the study
|
|
let completed_clone = completed.clone();
|
|
drop(trials);
|
|
|
|
if let ControlFlow::Break(()) = callback(self, &completed_clone) {
|
|
break;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
self.fail_trial(trial, e.to_string());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return error if no trials succeeded
|
|
if self.n_trials() == 0 {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Runs optimization asynchronously with full sampler integration.
|
|
///
|
|
/// This method combines async execution with the TPE sampler's ability to use
|
|
/// historical trial data for informed parameter suggestions.
|
|
///
|
|
/// The objective function takes ownership of the `Trial` and must return it
|
|
/// along with the result. This allows async operations to use the trial
|
|
/// across await points.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_trials` - The number of trials to run.
|
|
/// * `objective` - A function that takes a `Trial` and returns a `Future`
|
|
/// that resolves to a tuple of `(Trial, Result<f64, E>)`.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// # #[cfg(feature = "async")]
|
|
/// # async fn example() -> optimizer::Result<()> {
|
|
/// // Minimize x^2 with async objective and sampler integration
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
///
|
|
/// study
|
|
/// .optimize_async_with_sampler(10, |mut trial| async move {
|
|
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
|
/// // Simulate async work (e.g., network request)
|
|
/// let value = x * x;
|
|
/// Ok::<_, optimizer::TpeError>((trial, value))
|
|
/// })
|
|
/// .await?;
|
|
///
|
|
/// // At least one trial should have completed
|
|
/// assert!(study.n_trials() > 0);
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
#[cfg(feature = "async")]
|
|
pub async fn optimize_async_with_sampler<F, Fut, E>(
|
|
&self,
|
|
n_trials: usize,
|
|
objective: F,
|
|
) -> crate::Result<()>
|
|
where
|
|
F: Fn(Trial) -> Fut,
|
|
Fut: Future<Output = std::result::Result<(Trial, f64), E>>,
|
|
E: ToString,
|
|
{
|
|
for _ in 0..n_trials {
|
|
let trial = self.create_trial_with_sampler();
|
|
|
|
match objective(trial).await {
|
|
Ok((trial, value)) => {
|
|
self.complete_trial(trial, value);
|
|
}
|
|
Err(e) => {
|
|
// For async, we don't have the trial back on error
|
|
// We'll just count this as a failed trial without recording it
|
|
let _ = e.to_string();
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return error if no trials succeeded
|
|
if self.n_trials() == 0 {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Runs optimization with bounded parallelism and full sampler integration.
|
|
///
|
|
/// This method combines parallel async execution with the TPE sampler's ability
|
|
/// to use historical trial data for informed parameter suggestions. Up to
|
|
/// `concurrency` trials run simultaneously.
|
|
///
|
|
/// The objective function takes ownership of the `Trial` and must return it
|
|
/// along with the result. This allows async operations to use the trial
|
|
/// across await points.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `n_trials` - The total number of trials to run.
|
|
/// * `concurrency` - The maximum number of trials to run simultaneously.
|
|
/// * `objective` - A function that takes a `Trial` and returns a `Future`
|
|
/// that resolves to a tuple of `(Trial, f64)` or an error.
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::{Direction, RandomSampler, Study};
|
|
///
|
|
/// # #[cfg(feature = "async")]
|
|
/// # async fn example() -> optimizer::Result<()> {
|
|
/// // Minimize x^2 with parallel async evaluation and sampler integration
|
|
/// let sampler = RandomSampler::with_seed(42);
|
|
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
|
///
|
|
/// study
|
|
/// .optimize_parallel_with_sampler(10, 4, |mut trial| async move {
|
|
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
|
/// // Async objective function (e.g., network request)
|
|
/// let value = x * x;
|
|
/// Ok::<_, optimizer::TpeError>((trial, value))
|
|
/// })
|
|
/// .await?;
|
|
///
|
|
/// // All trials should have completed
|
|
/// assert_eq!(study.n_trials(), 10);
|
|
/// # Ok(())
|
|
/// # }
|
|
/// ```
|
|
#[cfg(feature = "async")]
|
|
pub async fn optimize_parallel_with_sampler<F, Fut, E>(
|
|
&self,
|
|
n_trials: usize,
|
|
concurrency: usize,
|
|
objective: F,
|
|
) -> crate::Result<()>
|
|
where
|
|
F: Fn(Trial) -> Fut + Send + Sync + 'static,
|
|
Fut: Future<Output = std::result::Result<(Trial, f64), E>> + Send,
|
|
E: ToString + Send + 'static,
|
|
{
|
|
use tokio::sync::Semaphore;
|
|
|
|
let semaphore = Arc::new(Semaphore::new(concurrency));
|
|
let objective = Arc::new(objective);
|
|
|
|
let mut handles = Vec::with_capacity(n_trials);
|
|
|
|
for _ in 0..n_trials {
|
|
let permit = semaphore.clone().acquire_owned().await.unwrap();
|
|
let trial = self.create_trial_with_sampler();
|
|
let objective = Arc::clone(&objective);
|
|
|
|
let handle = tokio::spawn(async move {
|
|
let result = objective(trial).await;
|
|
drop(permit); // Release semaphore permit when done
|
|
result
|
|
});
|
|
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Wait for all tasks and record results
|
|
for handle in handles {
|
|
match handle.await.unwrap() {
|
|
Ok((trial, value)) => {
|
|
self.complete_trial(trial, value);
|
|
}
|
|
Err(e) => {
|
|
let _ = e.to_string();
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return error if no trials succeeded
|
|
if self.n_trials() == 0 {
|
|
return Err(crate::TpeError::NoCompletedTrials);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// Manual Serialize implementation for Study<V> when serde feature is enabled.
|
|
#[cfg(feature = "serde")]
|
|
impl<V> Serialize for Study<V>
|
|
where
|
|
V: PartialOrd + Serialize,
|
|
{
|
|
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
|
where
|
|
S: serde::Serializer,
|
|
{
|
|
use serde::ser::SerializeStruct;
|
|
|
|
let mut state = serializer.serialize_struct("Study", 3)?;
|
|
state.serialize_field("direction", &self.direction)?;
|
|
// Serialize the Vec inside the Arc<RwLock<>>
|
|
let trials = self.completed_trials.read();
|
|
state.serialize_field("completed_trials", &*trials)?;
|
|
state.serialize_field("next_trial_id", &self.next_trial_id.load(Ordering::SeqCst))?;
|
|
state.end()
|
|
}
|
|
}
|
|
|
|
// Manual Deserialize implementation for Study<V> when serde feature is enabled.
|
|
#[cfg(feature = "serde")]
|
|
impl<'de, V> Deserialize<'de> for Study<V>
|
|
where
|
|
V: PartialOrd + Deserialize<'de>,
|
|
{
|
|
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
|
where
|
|
D: serde::Deserializer<'de>,
|
|
{
|
|
use std::fmt;
|
|
use std::marker::PhantomData;
|
|
|
|
use serde::de::{self, MapAccess, Visitor};
|
|
|
|
#[derive(serde::Deserialize)]
|
|
#[serde(field_identifier, rename_all = "snake_case")]
|
|
enum Field {
|
|
Direction,
|
|
CompletedTrials,
|
|
NextTrialId,
|
|
}
|
|
|
|
struct StudyVisitor<V>(PhantomData<V>);
|
|
|
|
impl<'de, V> Visitor<'de> for StudyVisitor<V>
|
|
where
|
|
V: PartialOrd + Deserialize<'de>,
|
|
{
|
|
type Value = Study<V>;
|
|
|
|
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
|
|
formatter.write_str("struct Study")
|
|
}
|
|
|
|
fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
|
|
where
|
|
A: MapAccess<'de>,
|
|
{
|
|
let mut direction = None;
|
|
let mut completed_trials: Option<Vec<CompletedTrial<V>>> = None;
|
|
let mut next_trial_id = None;
|
|
|
|
while let Some(key) = map.next_key()? {
|
|
match key {
|
|
Field::Direction => {
|
|
if direction.is_some() {
|
|
return Err(de::Error::duplicate_field("direction"));
|
|
}
|
|
direction = Some(map.next_value()?);
|
|
}
|
|
Field::CompletedTrials => {
|
|
if completed_trials.is_some() {
|
|
return Err(de::Error::duplicate_field("completed_trials"));
|
|
}
|
|
completed_trials = Some(map.next_value()?);
|
|
}
|
|
Field::NextTrialId => {
|
|
if next_trial_id.is_some() {
|
|
return Err(de::Error::duplicate_field("next_trial_id"));
|
|
}
|
|
next_trial_id = Some(map.next_value()?);
|
|
}
|
|
}
|
|
}
|
|
|
|
let direction = direction.ok_or_else(|| de::Error::missing_field("direction"))?;
|
|
let completed_trials =
|
|
completed_trials.ok_or_else(|| de::Error::missing_field("completed_trials"))?;
|
|
let next_trial_id: u64 =
|
|
next_trial_id.ok_or_else(|| de::Error::missing_field("next_trial_id"))?;
|
|
|
|
Ok(Study {
|
|
direction,
|
|
sampler: default_sampler(),
|
|
completed_trials: Arc::new(RwLock::new(completed_trials)),
|
|
next_trial_id: AtomicU64::new(next_trial_id),
|
|
})
|
|
}
|
|
}
|
|
|
|
const FIELDS: &[&str] = &["direction", "completed_trials", "next_trial_id"];
|
|
deserializer.deserialize_struct("Study", FIELDS, StudyVisitor(PhantomData))
|
|
}
|
|
}
|
|
|
|
#[cfg(all(test, feature = "serde"))]
|
|
mod serde_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_study_serde_round_trip() {
|
|
// Create a study and add some trials
|
|
let study: Study<f64> = Study::new(Direction::Minimize);
|
|
|
|
// Run some optimization
|
|
study
|
|
.optimize(5, |trial| {
|
|
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
|
let y = trial.suggest_int("y", 1, 5)?;
|
|
Ok::<_, crate::TpeError>(x + y as f64)
|
|
})
|
|
.unwrap();
|
|
|
|
// Serialize to JSON
|
|
let serialized = serde_json::to_string(&study).unwrap();
|
|
|
|
// Deserialize from JSON
|
|
let deserialized: Study<f64> = serde_json::from_str(&serialized).unwrap();
|
|
|
|
// Verify the data is preserved
|
|
assert_eq!(deserialized.direction(), study.direction());
|
|
assert_eq!(deserialized.n_trials(), study.n_trials());
|
|
|
|
// Verify the best trial is the same
|
|
let original_best = study.best_trial().unwrap();
|
|
let deserialized_best = deserialized.best_trial().unwrap();
|
|
assert_eq!(original_best.id, deserialized_best.id);
|
|
// Use approximate comparison for floats due to JSON serialization precision
|
|
assert!((original_best.value - deserialized_best.value).abs() < 1e-10);
|
|
// Check that all param keys match
|
|
assert_eq!(original_best.params.len(), deserialized_best.params.len());
|
|
for (key, original_val) in &original_best.params {
|
|
let deserialized_val = deserialized_best.params.get(key).unwrap();
|
|
match (original_val, deserialized_val) {
|
|
(crate::param::ParamValue::Float(a), crate::param::ParamValue::Float(b)) => {
|
|
assert!((a - b).abs() < 1e-10, "Float param {key} differs");
|
|
}
|
|
_ => assert_eq!(original_val, deserialized_val),
|
|
}
|
|
}
|
|
|
|
// Verify we can continue optimization on the deserialized study
|
|
let initial_count = deserialized.n_trials();
|
|
deserialized
|
|
.optimize(3, |trial| {
|
|
let x = trial.suggest_float("x", 0.0, 10.0)?;
|
|
let y = trial.suggest_int("y", 1, 5)?;
|
|
Ok::<_, crate::TpeError>(x + y as f64)
|
|
})
|
|
.unwrap();
|
|
|
|
// Verify new trials were added
|
|
assert_eq!(deserialized.n_trials(), initial_count + 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_study_serde_preserves_trial_ids() {
|
|
let study: Study<f64> = Study::new(Direction::Maximize);
|
|
|
|
// Add 5 trials
|
|
study
|
|
.optimize(5, |trial| {
|
|
let x = trial.suggest_float("x", -1.0, 1.0)?;
|
|
Ok::<_, crate::TpeError>(x * x)
|
|
})
|
|
.unwrap();
|
|
|
|
// Serialize and deserialize
|
|
let serialized = serde_json::to_string(&study).unwrap();
|
|
let deserialized: Study<f64> = serde_json::from_str(&serialized).unwrap();
|
|
|
|
// Create a new trial - its ID should continue from where we left off
|
|
let new_trial = deserialized.create_trial();
|
|
assert_eq!(new_trial.id(), 5); // Next trial should be ID 5
|
|
}
|
|
|
|
#[test]
|
|
fn test_completed_trial_serde() {
|
|
use std::collections::HashMap;
|
|
|
|
use crate::distribution::{Distribution, FloatDistribution, IntDistribution};
|
|
use crate::param::ParamValue;
|
|
|
|
let mut params = HashMap::new();
|
|
params.insert("x".to_string(), ParamValue::Float(0.5));
|
|
params.insert("n".to_string(), ParamValue::Int(42));
|
|
|
|
let mut distributions = HashMap::new();
|
|
distributions.insert(
|
|
"x".to_string(),
|
|
Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
}),
|
|
);
|
|
distributions.insert(
|
|
"n".to_string(),
|
|
Distribution::Int(IntDistribution {
|
|
low: 1,
|
|
high: 100,
|
|
log_scale: false,
|
|
step: None,
|
|
}),
|
|
);
|
|
|
|
let completed = CompletedTrial::new(42, params.clone(), distributions.clone(), 0.75);
|
|
|
|
// Serialize and deserialize
|
|
let serialized = serde_json::to_string(&completed).unwrap();
|
|
let deserialized: CompletedTrial<f64> = serde_json::from_str(&serialized).unwrap();
|
|
|
|
assert_eq!(deserialized.id, 42);
|
|
assert_eq!(deserialized.value, 0.75);
|
|
assert_eq!(deserialized.params, params);
|
|
assert_eq!(deserialized.distributions, distributions);
|
|
}
|
|
}
|