diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..80c468f --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "optimize" +version = "0.1.0" +edition = "2024" + +[dependencies] +rand = "0.9" +thiserror = "2" +parking_lot = "0.12" +ordered-float = "5" +serde = { version = "1", features = ["derive"], optional = true } +tokio = { version = "1", features = ["sync", "rt-multi-thread"], optional = true } + +[features] +default = [] +serde = ["dep:serde", "ordered-float/serde"] +async = ["dep:tokio"] + +[dev-dependencies] +serde_json = "1" diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..197a4e5 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,10 @@ +# Nightly Features +imports_granularity = "Module" +group_imports = "StdExternalCrate" +#format_strings = true +format_code_in_doc_comments = true + +# Stable Features +merge_derives = true +use_field_init_shorthand = true +use_try_shorthand = true diff --git a/src/distribution.rs b/src/distribution.rs new file mode 100644 index 0000000..ca97532 --- /dev/null +++ b/src/distribution.rs @@ -0,0 +1,52 @@ +//! Parameter distribution types. + +#[cfg(feature = "serde")] +use serde::{Deserialize, Serialize}; + +/// Distribution for floating-point parameters. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct FloatDistribution { + /// Lower bound (inclusive). + pub low: f64, + /// Upper bound (inclusive). + pub high: f64, + /// Whether to sample in log space. + pub log_scale: bool, + /// Optional step size for discretization. + pub step: Option, +} + +/// Distribution for integer parameters. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct IntDistribution { + /// Lower bound (inclusive). + pub low: i64, + /// Upper bound (inclusive). + pub high: i64, + /// Whether to sample in log space. + pub log_scale: bool, + /// Optional step size for discretization. + pub step: Option, +} + +/// Distribution for categorical parameters. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct CategoricalDistribution { + /// Number of choices available. + pub n_choices: usize, +} + +/// Enum wrapping all parameter distribution types. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum Distribution { + /// A floating-point distribution. + Float(FloatDistribution), + /// An integer distribution. + Int(IntDistribution), + /// A categorical distribution. + Categorical(CategoricalDistribution), +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..3503408 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,44 @@ +//! Error types for the optimize library. + +use thiserror::Error; + +/// The error type for TPE operations. +#[derive(Debug, Error)] +pub enum TpeError { + /// Returned when the lower bound is greater than the upper bound. + #[error("invalid bounds: low ({low}) must be less than or equal to high ({high})")] + InvalidBounds { + /// The lower bound value. + low: f64, + /// The upper bound value. + high: f64, + }, + + /// Returned when log scale is used with non-positive bounds. + #[error("invalid log bounds: low must be positive for log scale")] + InvalidLogBounds, + + /// Returned when step size is not positive. + #[error("invalid step: step must be positive")] + InvalidStep, + + /// Returned when categorical choices are empty. + #[error("categorical choices cannot be empty")] + EmptyChoices, + + /// Returned when a parameter is suggested with a different configuration. + #[error("parameter conflict for '{name}': {reason}")] + ParameterConflict { + /// The name of the conflicting parameter. + name: String, + /// The reason for the conflict. + reason: String, + }, + + /// Returned when requesting the best trial but no trials have completed. + #[error("no completed trials available")] + NoCompletedTrials, +} + +/// A specialized Result type for TPE operations. +pub type Result = std::result::Result; diff --git a/src/kde.rs b/src/kde.rs new file mode 100644 index 0000000..d090f56 --- /dev/null +++ b/src/kde.rs @@ -0,0 +1,267 @@ +//! Kernel Density Estimation for continuous parameters. +//! +//! This module provides a Gaussian kernel density estimator used by the TPE +//! sampler to model probability distributions over good and bad trial regions. + +use rand::Rng; + +/// A Gaussian kernel density estimator for continuous distributions. +/// +/// KDE estimates a probability density function from a set of samples by +/// placing Gaussian kernels centered at each sample point. This is used +/// in TPE to model the distributions l(x) (good trials) and g(x) (bad trials). +/// +/// # Examples +/// +/// ```ignore +/// use crate::kde::KernelDensityEstimator; +/// +/// let samples = vec![1.0, 2.0, 3.0, 4.0, 5.0]; +/// let kde = KernelDensityEstimator::new(samples); +/// +/// // Get probability density at a point +/// let density = kde.pdf(2.5); +/// assert!(density > 0.0); +/// +/// // Sample from the estimated distribution +/// let mut rng = rand::rng(); +/// let sample = kde.sample(&mut rng); +/// ``` +#[derive(Clone, Debug)] +pub struct KernelDensityEstimator { + /// The sample points used to construct the KDE. + samples: Vec, + /// The bandwidth (standard deviation) of the Gaussian kernels. + bandwidth: f64, +} + +impl KernelDensityEstimator { + /// Creates a new KDE with automatic bandwidth selection using Scott's rule. + /// + /// Scott's rule sets bandwidth = n^(-1/5) * std_dev, which works well + /// for unimodal distributions close to normal. + /// + /// # Panics + /// + /// Panics if `samples` is empty. + pub fn new(samples: Vec) -> Self { + assert!(!samples.is_empty(), "KDE requires at least one sample"); + + let bandwidth = Self::scotts_rule(&samples); + Self { samples, bandwidth } + } + + /// Creates a new KDE with a specified bandwidth. + /// + /// Use this when you want explicit control over the smoothing parameter. + /// + /// # Panics + /// + /// Panics if `samples` is empty or `bandwidth` is not positive. + pub fn with_bandwidth(samples: Vec, bandwidth: f64) -> Self { + assert!(!samples.is_empty(), "KDE requires at least one sample"); + assert!(bandwidth > 0.0, "Bandwidth must be positive"); + + Self { samples, bandwidth } + } + + /// Computes bandwidth using Scott's rule. + /// + /// Scott's rule: h = n^(-1/5) * sigma + /// where sigma is the sample standard deviation. + fn scotts_rule(samples: &[f64]) -> f64 { + let n = samples.len() as f64; + let std_dev = Self::sample_std_dev(samples); + + // For degenerate case where all samples are identical, + // use a small positive bandwidth + if std_dev < f64::EPSILON { + return 1.0; + } + + n.powf(-0.2) * std_dev + } + + /// Computes the sample standard deviation. + fn sample_std_dev(samples: &[f64]) -> f64 { + let n = samples.len() as f64; + let mean = samples.iter().sum::() / n; + let variance = samples.iter().map(|x| (x - mean).powi(2)).sum::() / n; + variance.sqrt() + } + + /// Returns the probability density at point `x`. + /// + /// The density is computed as the average of Gaussian kernels centered + /// at each sample point: + /// + /// f(x) = (1/n) * sum_i K((x - x_i) / h) + /// + /// where K is the standard Gaussian kernel and h is the bandwidth. + pub fn pdf(&self, x: f64) -> f64 { + let n = self.samples.len() as f64; + let inv_bandwidth = 1.0 / self.bandwidth; + let normalization = inv_bandwidth / (2.0 * std::f64::consts::PI).sqrt(); + + let density: f64 = self + .samples + .iter() + .map(|&xi| { + let z = (x - xi) * inv_bandwidth; + normalization * (-0.5 * z * z).exp() + }) + .sum(); + + density / n + } + + /// Samples a value from the estimated density distribution. + /// + /// Sampling works by: + /// 1. Uniformly selecting one of the kernel centers (samples) + /// 2. Adding Gaussian noise with the bandwidth as standard deviation + pub fn sample(&self, rng: &mut R) -> f64 { + // Select a random sample to center the kernel on + let idx = rng.random_range(0..self.samples.len()); + let center = self.samples[idx]; + + // Add Gaussian noise with bandwidth as standard deviation + // Using Box-Muller transform for Gaussian sampling + let u1: f64 = rng.random(); + let u2: f64 = rng.random(); + + let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos(); + center + z * self.bandwidth + } + + /// Returns the bandwidth of this KDE. + pub fn bandwidth(&self) -> f64 { + self.bandwidth + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_kde_pdf_basic() { + let samples = vec![0.0, 1.0, 2.0]; + let kde = KernelDensityEstimator::new(samples); + + // Density should be positive everywhere + assert!(kde.pdf(0.0) > 0.0); + assert!(kde.pdf(1.0) > 0.0); + assert!(kde.pdf(2.0) > 0.0); + + // Density should be higher near sample points + let mid_density = kde.pdf(1.0); + let far_density = kde.pdf(10.0); + assert!(mid_density > far_density); + } + + #[test] + fn test_kde_pdf_integrates_to_one() { + let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0]; + let kde = KernelDensityEstimator::new(samples); + + // Numerical integration over a wide range + let n_points = 10000; + let low = -10.0; + let high = 15.0; + let dx = (high - low) / n_points as f64; + + let integral: f64 = (0..n_points) + .map(|i| { + let x = low + (i as f64 + 0.5) * dx; + kde.pdf(x) * dx + }) + .sum(); + + // Should be approximately 1.0 (within numerical error) + assert!( + (integral - 1.0).abs() < 0.01, + "Integral = {integral}, expected ~1.0" + ); + } + + #[test] + fn test_kde_with_bandwidth() { + let samples = vec![0.0, 1.0, 2.0]; + let kde = KernelDensityEstimator::with_bandwidth(samples, 0.5); + + assert_eq!(kde.bandwidth(), 0.5); + assert!(kde.pdf(1.0) > 0.0); + } + + #[test] + fn test_kde_sample_in_reasonable_range() { + let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0]; + let kde = KernelDensityEstimator::new(samples); + let mut rng = rand::rng(); + + // Samples should generally be in a reasonable range around the data + for _ in 0..100 { + let s = kde.sample(&mut rng); + // With high probability, samples should be within a few bandwidths + // of the data range. Use a generous range to avoid flaky tests. + assert!(s > -10.0 && s < 15.0, "Sample {s} outside expected range"); + } + } + + #[test] + fn test_kde_single_sample() { + let samples = vec![5.0]; + let kde = KernelDensityEstimator::new(samples); + + // Should have positive density near the sample + assert!(kde.pdf(5.0) > 0.0); + assert!(kde.pdf(4.5) > 0.0); + } + + #[test] + fn test_kde_identical_samples() { + let samples = vec![3.0, 3.0, 3.0, 3.0]; + let kde = KernelDensityEstimator::new(samples); + + // Should handle degenerate case with identical samples + assert!(kde.bandwidth() > 0.0); + assert!(kde.pdf(3.0) > 0.0); + } + + #[test] + fn test_scotts_rule_bandwidth() { + let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]; + let kde = KernelDensityEstimator::new(samples); + + // n = 10, n^(-1/5) ≈ 0.631 + // std_dev ≈ 2.87 + // bandwidth ≈ 0.631 * 2.87 ≈ 1.81 + let bandwidth = kde.bandwidth(); + assert!( + bandwidth > 1.0 && bandwidth < 3.0, + "Bandwidth {bandwidth} outside expected range" + ); + } + + #[test] + #[should_panic(expected = "KDE requires at least one sample")] + fn test_kde_empty_samples() { + let samples: Vec = vec![]; + KernelDensityEstimator::new(samples); + } + + #[test] + #[should_panic(expected = "Bandwidth must be positive")] + fn test_kde_zero_bandwidth() { + let samples = vec![1.0, 2.0, 3.0]; + KernelDensityEstimator::with_bandwidth(samples, 0.0); + } + + #[test] + #[should_panic(expected = "Bandwidth must be positive")] + fn test_kde_negative_bandwidth() { + let samples = vec![1.0, 2.0, 3.0]; + KernelDensityEstimator::with_bandwidth(samples, -1.0); + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..9350a24 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,151 @@ +//! A Tree-Parzen Estimator (TPE) library for black-box optimization. +//! +//! This library provides an Optuna-like API for hyperparameter optimization +//! using the Tree-Parzen Estimator algorithm. It supports: +//! +//! - Float, integer, and categorical parameter types +//! - Log-scale and stepped parameter sampling +//! - Synchronous and async optimization +//! - Parallel trial evaluation with bounded concurrency +//! - Serialization for saving/loading study state +//! +//! # Quick Start +//! +//! ``` +//! use optimize::{Direction, Study, TpeSampler}; +//! +//! // Create a study with TPE sampler +//! let sampler = TpeSampler::builder().seed(42).build(); +//! let study: Study = Study::with_sampler(Direction::Minimize, sampler); +//! +//! // Optimize x^2 for 20 trials +//! study +//! .optimize_with_sampler(20, |trial| { +//! let x = trial.suggest_float("x", -10.0, 10.0)?; +//! Ok::<_, optimize::TpeError>(x * x) +//! }) +//! .unwrap(); +//! +//! // Get the best result +//! let best = study.best_trial().unwrap(); +//! println!("Best value: {} at x={:?}", best.value, best.params); +//! ``` +//! +//! # Creating a Study +//! +//! A [`Study`] manages optimization trials. Create one with an optimization direction: +//! +//! ``` +//! use optimize::{Direction, RandomSampler, Study, TpeSampler}; +//! +//! // Minimize with default random sampler +//! let study: Study = Study::new(Direction::Minimize); +//! +//! // Maximize with TPE sampler +//! let study: Study = Study::with_sampler(Direction::Maximize, TpeSampler::new()); +//! +//! // With seeded sampler for reproducibility +//! let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); +//! ``` +//! +//! # Suggesting Parameters +//! +//! Within the objective function, use [`Trial`] to suggest parameter values: +//! +//! ``` +//! use optimize::{Direction, Study}; +//! +//! let study: Study = Study::new(Direction::Minimize); +//! +//! study +//! .optimize(10, |trial| { +//! // Float parameters +//! let x = trial.suggest_float("x", 0.0, 1.0)?; +//! let lr = trial.suggest_float_log("learning_rate", 1e-5, 1e-1)?; +//! let step = trial.suggest_float_step("step", 0.0, 1.0, 0.1)?; +//! +//! // Integer parameters +//! let n = trial.suggest_int("n_layers", 1, 10)?; +//! let batch = trial.suggest_int_log("batch_size", 16, 256)?; +//! let units = trial.suggest_int_step("units", 32, 512, 32)?; +//! +//! // Categorical parameters +//! let optimizer = trial.suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])?; +//! +//! // Return objective value +//! Ok::<_, optimize::TpeError>(x * n as f64) +//! }) +//! .unwrap(); +//! ``` +//! +//! # Configuring TPE +//! +//! The [`TpeSampler`] can be configured using the builder pattern: +//! +//! ``` +//! use optimize::TpeSampler; +//! +//! let sampler = TpeSampler::builder() +//! .gamma(0.15) // Quantile for good/bad split +//! .n_startup_trials(20) // Random trials before TPE +//! .n_ei_candidates(32) // Candidates to evaluate +//! .seed(42) // Reproducibility +//! .build(); +//! ``` +//! +//! # Async and Parallel Optimization +//! +//! With the `async` feature enabled, you can run trials asynchronously: +//! +//! ```ignore +//! use optimize::{Study, Direction}; +//! +//! // Sequential async +//! study.optimize_async(10, |mut trial| async move { +//! let x = trial.suggest_float("x", 0.0, 1.0)?; +//! Ok((trial, x * x)) +//! }).await?; +//! +//! // Parallel with bounded concurrency +//! study.optimize_parallel(10, 4, |mut trial| async move { +//! let x = trial.suggest_float("x", 0.0, 1.0)?; +//! Ok((trial, x * x)) +//! }).await?; +//! ``` +//! +//! # Serialization +//! +//! With the `serde` feature enabled, studies can be serialized: +//! +//! ```ignore +//! use optimize::{Study, Direction, TpeSampler}; +//! +//! // Save study state +//! let study: Study = Study::new(Direction::Minimize); +//! let json = serde_json::to_string(&study)?; +//! +//! // Load and continue +//! let mut study: Study = serde_json::from_str(&json)?; +//! study.set_sampler(TpeSampler::new()); // Restore sampler +//! study.optimize_with_sampler(10, |trial| { /* ... */ }).unwrap(); +//! ``` +//! +//! # Feature Flags +//! +//! - `serde`: Enable serialization/deserialization of studies and trials +//! - `async`: Enable async optimization methods (requires tokio) + +mod distribution; +mod error; +mod kde; +mod param; +mod sampler; +mod study; +mod trial; +mod types; + +pub use error::{Result, TpeError}; +pub use sampler::{CompletedTrial, RandomSampler, Sampler, TpeSampler, TpeSamplerBuilder}; +pub use study::Study; +pub use trial::Trial; +pub use types::{Direction, TrialState}; diff --git a/src/param.rs b/src/param.rs new file mode 100644 index 0000000..201f6b1 --- /dev/null +++ b/src/param.rs @@ -0,0 +1,20 @@ +//! Parameter value storage types. + +#[cfg(feature = "serde")] +use serde::{Deserialize, Serialize}; + +/// Represents a sampled parameter value. +/// +/// This enum stores different parameter value types uniformly. +/// For categorical parameters, the `Categorical` variant stores +/// the index into the choices array. +#[derive(Clone, Debug, PartialEq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum ParamValue { + /// A floating-point parameter value. + Float(f64), + /// An integer parameter value. + Int(i64), + /// A categorical parameter value, stored as an index into the choices array. + Categorical(usize), +} diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs new file mode 100644 index 0000000..246d8bd --- /dev/null +++ b/src/sampler/mod.rs @@ -0,0 +1,96 @@ +//! Sampler trait and implementations for parameter sampling. + +mod random; +pub mod tpe; + +use std::collections::HashMap; + +pub use random::RandomSampler; +#[cfg(feature = "serde")] +use serde::{Deserialize, Serialize}; +pub use tpe::{TpeSampler, TpeSamplerBuilder}; + +use crate::distribution::Distribution; +use crate::param::ParamValue; + +/// A completed trial with its parameters, distributions, and objective value. +/// +/// This struct stores the results of a completed trial, including all sampled +/// parameter values, their distributions, and the objective value returned +/// by the objective function. +#[derive(Clone, Debug)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct CompletedTrial { + /// The unique identifier for this trial. + pub id: u64, + /// The sampled parameter values, keyed by parameter name. + pub params: HashMap, + /// The parameter distributions used, keyed by parameter name. + pub distributions: HashMap, + /// The objective value returned by the objective function. + pub value: V, +} + +impl CompletedTrial { + /// Creates a new completed trial. + pub fn new( + id: u64, + params: HashMap, + distributions: HashMap, + value: V, + ) -> Self { + Self { + id, + params, + distributions, + value, + } + } +} + +/// Trait for pluggable parameter sampling strategies. +/// +/// Samplers are responsible for generating parameter values based on +/// the distribution and historical trial data. The trait requires +/// `Send + Sync` to support concurrent and async optimization. +/// +/// # Examples +/// +/// Implementing a custom sampler: +/// +/// ```ignore +/// use optimize::{Sampler, ParamValue, Distribution, CompletedTrial}; +/// +/// struct MySampler; +/// +/// impl Sampler for MySampler { +/// fn sample( +/// &self, +/// distribution: &Distribution, +/// trial_id: u64, +/// history: &[CompletedTrial], +/// ) -> ParamValue { +/// // Custom sampling logic here +/// todo!() +/// } +/// } +/// ``` +pub trait Sampler: Send + Sync { + /// Samples a parameter value from the given distribution. + /// + /// # Arguments + /// + /// * `distribution` - The parameter distribution to sample from. + /// * `trial_id` - The unique ID of the trial being sampled for. + /// * `history` - Historical completed trials for informed sampling. + /// + /// # Returns + /// + /// A `ParamValue` sampled from the distribution. + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + history: &[CompletedTrial], + ) -> ParamValue; +} diff --git a/src/sampler/random.rs b/src/sampler/random.rs new file mode 100644 index 0000000..bd2602b --- /dev/null +++ b/src/sampler/random.rs @@ -0,0 +1,275 @@ +//! Random sampler implementation. + +use parking_lot::Mutex; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +use crate::distribution::Distribution; +use crate::param::ParamValue; +use crate::sampler::{CompletedTrial, Sampler}; + +/// A simple random sampler that samples uniformly from distributions. +/// +/// This sampler ignores the trial history and samples uniformly at random, +/// respecting log scale and step size constraints. It serves as a baseline +/// sampler and is used during the startup phase of more sophisticated samplers. +/// +/// # Examples +/// +/// ``` +/// use optimize::RandomSampler; +/// +/// // Create with default RNG +/// let sampler = RandomSampler::new(); +/// +/// // Create with a fixed seed for reproducibility +/// let sampler = RandomSampler::with_seed(42); +/// ``` +pub struct RandomSampler { + rng: Mutex, +} + +impl RandomSampler { + /// Creates a new random sampler with a default random seed. + pub fn new() -> Self { + Self { + rng: Mutex::new(StdRng::from_os_rng()), + } + } + + /// Creates a new random sampler with a fixed seed for reproducibility. + /// + /// Using the same seed will produce the same sequence of sampled values. + pub fn with_seed(seed: u64) -> Self { + Self { + rng: Mutex::new(StdRng::seed_from_u64(seed)), + } + } +} + +impl Default for RandomSampler { + fn default() -> Self { + Self::new() + } +} + +impl Sampler for RandomSampler { + fn sample( + &self, + distribution: &Distribution, + _trial_id: u64, + _history: &[CompletedTrial], + ) -> ParamValue { + let mut rng = self.rng.lock(); + + match distribution { + Distribution::Float(d) => { + let value = if d.log_scale { + // Sample uniformly in log space + let log_low = d.low.ln(); + let log_high = d.high.ln(); + let log_value = rng.random_range(log_low..=log_high); + log_value.exp() + } else if let Some(step) = d.step { + // Sample from step grid + let n_steps = ((d.high - d.low) / step).floor() as i64; + let k = rng.random_range(0..=n_steps); + d.low + (k as f64) * step + } else { + // Uniform sampling + rng.random_range(d.low..=d.high) + }; + ParamValue::Float(value) + } + Distribution::Int(d) => { + let value = if d.log_scale { + // Sample uniformly in log space, then round + let log_low = (d.low as f64).ln(); + let log_high = (d.high as f64).ln(); + let log_value = rng.random_range(log_low..=log_high); + let raw = log_value.exp().round() as i64; + // Clamp to bounds since rounding might push outside + raw.clamp(d.low, d.high) + } else if let Some(step) = d.step { + // Sample from step grid + let n_steps = (d.high - d.low) / step; + let k = rng.random_range(0..=n_steps); + d.low + k * step + } else { + // Uniform sampling + rng.random_range(d.low..=d.high) + }; + ParamValue::Int(value) + } + Distribution::Categorical(d) => { + let index = rng.random_range(0..d.n_choices); + ParamValue::Categorical(index) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + + #[test] + fn test_random_sampler_float() { + let sampler = RandomSampler::with_seed(42); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + for _ in 0..100 { + let value = sampler.sample(&dist, 0, &[]); + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(&v)); + } else { + panic!("Expected Float value"); + } + } + } + + #[test] + fn test_random_sampler_float_log() { + let sampler = RandomSampler::with_seed(42); + let dist = Distribution::Float(FloatDistribution { + low: 1e-5, + high: 1.0, + log_scale: true, + step: None, + }); + + for _ in 0..100 { + let value = sampler.sample(&dist, 0, &[]); + if let ParamValue::Float(v) = value { + assert!((1e-5..=1.0).contains(&v)); + } else { + panic!("Expected Float value"); + } + } + } + + #[test] + fn test_random_sampler_float_step() { + let sampler = RandomSampler::with_seed(42); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: Some(0.25), + }); + + for _ in 0..100 { + let value = sampler.sample(&dist, 0, &[]); + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(&v)); + // Check it's on the step grid + let k = ((v - 0.0) / 0.25).round() as i64; + let expected = 0.0 + k as f64 * 0.25; + assert!((v - expected).abs() < 1e-10); + } else { + panic!("Expected Float value"); + } + } + } + + #[test] + fn test_random_sampler_int() { + let sampler = RandomSampler::with_seed(42); + let dist = Distribution::Int(IntDistribution { + low: 0, + high: 10, + log_scale: false, + step: None, + }); + + for _ in 0..100 { + let value = sampler.sample(&dist, 0, &[]); + if let ParamValue::Int(v) = value { + assert!((0..=10).contains(&v)); + } else { + panic!("Expected Int value"); + } + } + } + + #[test] + fn test_random_sampler_int_log() { + let sampler = RandomSampler::with_seed(42); + let dist = Distribution::Int(IntDistribution { + low: 1, + high: 1000, + log_scale: true, + step: None, + }); + + for _ in 0..100 { + let value = sampler.sample(&dist, 0, &[]); + if let ParamValue::Int(v) = value { + assert!((1..=1000).contains(&v)); + } else { + panic!("Expected Int value"); + } + } + } + + #[test] + fn test_random_sampler_int_step() { + let sampler = RandomSampler::with_seed(42); + let dist = Distribution::Int(IntDistribution { + low: 0, + high: 10, + log_scale: false, + step: Some(2), + }); + + for _ in 0..100 { + let value = sampler.sample(&dist, 0, &[]); + if let ParamValue::Int(v) = value { + assert!((0..=10).contains(&v)); + // Check it's on the step grid: 0, 2, 4, 6, 8, 10 + assert!(v % 2 == 0); + } else { + panic!("Expected Int value"); + } + } + } + + #[test] + fn test_random_sampler_categorical() { + let sampler = RandomSampler::with_seed(42); + let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 }); + + for _ in 0..100 { + let value = sampler.sample(&dist, 0, &[]); + if let ParamValue::Categorical(idx) = value { + assert!(idx < 5); + } else { + panic!("Expected Categorical value"); + } + } + } + + #[test] + fn test_random_sampler_reproducibility() { + let sampler1 = RandomSampler::with_seed(42); + let sampler2 = RandomSampler::with_seed(42); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + for _ in 0..10 { + let v1 = sampler1.sample(&dist, 0, &[]); + let v2 = sampler2.sample(&dist, 0, &[]); + assert_eq!(v1, v2); + } + } +} diff --git a/src/sampler/tpe.rs b/src/sampler/tpe.rs new file mode 100644 index 0000000..2500f55 --- /dev/null +++ b/src/sampler/tpe.rs @@ -0,0 +1,1019 @@ +//! Tree-Parzen Estimator (TPE) sampler implementation. +//! +//! TPE is a Bayesian optimization algorithm that models the objective function +//! using two probability distributions: one for promising (good) parameter values +//! and one for unpromising (bad) parameter values. + +use parking_lot::Mutex; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +use crate::distribution::Distribution; +use crate::kde::KernelDensityEstimator; +use crate::param::ParamValue; +use crate::sampler::{CompletedTrial, Sampler}; + +/// A Tree-Parzen Estimator (TPE) sampler for Bayesian optimization. +/// +/// TPE works by splitting completed trials into two groups based on their +/// objective values: good trials (below the gamma quantile) and bad trials +/// (above the gamma quantile). It then fits kernel density estimators (KDE) +/// to each group and samples new points that maximize the ratio l(x)/g(x), +/// where l(x) is the density of good trials and g(x) is the density of bad trials. +/// +/// During the startup phase (when fewer than `n_startup_trials` are completed), +/// TPE falls back to random sampling to gather initial data. +/// +/// # Examples +/// +/// ``` +/// use optimize::TpeSampler; +/// +/// // Create with default settings +/// let sampler = TpeSampler::new(); +/// +/// // Create with custom settings using the builder +/// let sampler = TpeSampler::builder() +/// .gamma(0.15) +/// .n_startup_trials(20) +/// .n_ei_candidates(32) +/// .seed(42) +/// .build(); +/// ``` +pub struct TpeSampler { + /// Fraction of trials to consider as "good" (gamma quantile). + gamma: f64, + /// Number of trials before TPE kicks in (uses random sampling before this). + n_startup_trials: usize, + /// Number of candidate samples to evaluate when selecting the next point. + n_ei_candidates: usize, + /// Thread-safe RNG for sampling. + rng: Mutex, +} + +impl TpeSampler { + /// Creates a new TPE sampler with default settings. + /// + /// Default settings: + /// - gamma: 0.25 (top 25% of trials are considered "good") + /// - n_startup_trials: 10 (random sampling for first 10 trials) + /// - n_ei_candidates: 24 (evaluate 24 candidates per sample) + pub fn new() -> Self { + Self { + gamma: 0.25, + n_startup_trials: 10, + n_ei_candidates: 24, + rng: Mutex::new(StdRng::from_os_rng()), + } + } + + /// Creates a builder for configuring a TPE sampler. + /// + /// # Examples + /// + /// ``` + /// use optimize::TpeSampler; + /// + /// let sampler = TpeSampler::builder() + /// .gamma(0.15) + /// .n_startup_trials(20) + /// .n_ei_candidates(32) + /// .seed(42) + /// .build(); + /// ``` + pub fn builder() -> TpeSamplerBuilder { + TpeSamplerBuilder::new() + } + + /// Creates a new TPE sampler with custom configuration. + /// + /// # Arguments + /// + /// * `gamma` - Fraction of trials to consider "good" (0.0 to 1.0). + /// * `n_startup_trials` - Number of random trials before TPE sampling. + /// * `n_ei_candidates` - Number of candidates to evaluate per sample. + /// * `seed` - Optional seed for reproducibility. + /// + /// # Panics + /// + /// Panics if gamma is not in (0.0, 1.0). + pub fn with_config( + gamma: f64, + n_startup_trials: usize, + n_ei_candidates: usize, + seed: Option, + ) -> Self { + assert!( + gamma > 0.0 && gamma < 1.0, + "gamma must be in (0.0, 1.0), got {gamma}" + ); + + let rng = match seed { + Some(s) => StdRng::seed_from_u64(s), + None => StdRng::from_os_rng(), + }; + + Self { + gamma, + n_startup_trials, + n_ei_candidates, + rng: Mutex::new(rng), + } + } + + /// Splits trials into good and bad groups based on the gamma quantile. + /// + /// Returns (good_trials, bad_trials) where good_trials contains trials + /// with values below the gamma quantile (for minimization). + fn split_trials<'a>( + &self, + history: &'a [CompletedTrial], + ) -> (Vec<&'a CompletedTrial>, Vec<&'a CompletedTrial>) { + if history.is_empty() { + return (vec![], vec![]); + } + + // Sort trials by value (ascending for minimization) + let mut sorted_indices: Vec = (0..history.len()).collect(); + sorted_indices.sort_by(|&a, &b| { + history[a] + .value + .partial_cmp(&history[b].value) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + // Calculate the split point (gamma quantile) + // Ensure at least 1 trial in each group if possible + let n_good = ((history.len() as f64 * self.gamma).ceil() as usize) + .max(1) + .min(history.len() - 1); + + let good: Vec<_> = sorted_indices[..n_good] + .iter() + .map(|&i| &history[i]) + .collect(); + let bad: Vec<_> = sorted_indices[n_good..] + .iter() + .map(|&i| &history[i]) + .collect(); + + (good, bad) + } + + /// Samples uniformly from a distribution (used during startup phase). + fn sample_uniform(&self, distribution: &Distribution, rng: &mut StdRng) -> ParamValue { + match distribution { + Distribution::Float(d) => { + let value = if d.log_scale { + let log_low = d.low.ln(); + let log_high = d.high.ln(); + rng.random_range(log_low..=log_high).exp() + } else if let Some(step) = d.step { + let n_steps = ((d.high - d.low) / step).floor() as i64; + let k = rng.random_range(0..=n_steps); + d.low + (k as f64) * step + } else { + rng.random_range(d.low..=d.high) + }; + ParamValue::Float(value) + } + Distribution::Int(d) => { + let value = if d.log_scale { + let log_low = (d.low as f64).ln(); + let log_high = (d.high as f64).ln(); + let raw = rng.random_range(log_low..=log_high).exp().round() as i64; + raw.clamp(d.low, d.high) + } else if let Some(step) = d.step { + let n_steps = (d.high - d.low) / step; + let k = rng.random_range(0..=n_steps); + d.low + k * step + } else { + rng.random_range(d.low..=d.high) + }; + ParamValue::Int(value) + } + Distribution::Categorical(d) => { + ParamValue::Categorical(rng.random_range(0..d.n_choices)) + } + } + } + + /// Samples using TPE for float distributions. + #[allow(clippy::too_many_arguments)] + fn sample_tpe_float( + &self, + low: f64, + high: f64, + log_scale: bool, + step: Option, + good_values: Vec, + bad_values: Vec, + rng: &mut StdRng, + ) -> f64 { + // Transform to internal space (log space if needed) + let (internal_low, internal_high, good_internal, bad_internal) = if log_scale { + let i_low = low.ln(); + let i_high = high.ln(); + let g: Vec = good_values.iter().map(|&v| v.ln()).collect(); + let b: Vec = bad_values.iter().map(|&v| v.ln()).collect(); + (i_low, i_high, g, b) + } else { + (low, high, good_values, bad_values) + }; + + // Fit KDEs to good and bad groups + let l_kde = KernelDensityEstimator::new(good_internal); + let g_kde = KernelDensityEstimator::new(bad_internal); + + // Generate candidates from l(x) and select the one with best l(x)/g(x) ratio + let mut best_candidate = internal_low; + let mut best_ratio = f64::NEG_INFINITY; + + for _ in 0..self.n_ei_candidates { + let candidate = l_kde.sample(rng); + + // Clamp to bounds + let candidate = candidate.clamp(internal_low, internal_high); + + let l_density = l_kde.pdf(candidate); + let g_density = g_kde.pdf(candidate); + + // Compute l(x)/g(x) ratio, handling zero density + let ratio = if g_density < f64::EPSILON { + if l_density > f64::EPSILON { + f64::INFINITY + } else { + 0.0 + } + } else { + l_density / g_density + }; + + if ratio > best_ratio { + best_ratio = ratio; + best_candidate = candidate; + } + } + + // Transform back from internal space + let mut value = if log_scale { + best_candidate.exp() + } else { + best_candidate + }; + + // Apply step constraint if present + if let Some(step) = step { + let k = ((value - low) / step).round(); + value = low + k * step; + } + + // Ensure value is within bounds + value.clamp(low, high) + } + + /// Samples using TPE for integer distributions. + #[allow(clippy::too_many_arguments)] + fn sample_tpe_int( + &self, + low: i64, + high: i64, + log_scale: bool, + step: Option, + good_values: Vec, + bad_values: Vec, + rng: &mut StdRng, + ) -> i64 { + // Convert to floats for KDE + let good_floats: Vec = good_values.iter().map(|&v| v as f64).collect(); + let bad_floats: Vec = bad_values.iter().map(|&v| v as f64).collect(); + + // Use float TPE sampling + let float_value = self.sample_tpe_float( + low as f64, + high as f64, + log_scale, + step.map(|s| s as f64), + good_floats, + bad_floats, + rng, + ); + + // Round to nearest integer + let int_value = float_value.round() as i64; + + // Apply step constraint if present + let int_value = if let Some(step) = step { + let k = ((int_value - low) as f64 / step as f64).round() as i64; + low + k * step + } else { + int_value + }; + + // Ensure value is within bounds + int_value.clamp(low, high) + } + + /// Samples using TPE for categorical distributions. + fn sample_tpe_categorical( + &self, + n_choices: usize, + good_indices: Vec, + bad_indices: Vec, + rng: &mut StdRng, + ) -> usize { + // Count occurrences in good and bad groups + let mut good_counts = vec![0usize; n_choices]; + let mut bad_counts = vec![0usize; n_choices]; + + for &idx in &good_indices { + if idx < n_choices { + good_counts[idx] += 1; + } + } + for &idx in &bad_indices { + if idx < n_choices { + bad_counts[idx] += 1; + } + } + + // Add smoothing (Laplace smoothing) to avoid zero probabilities + let good_total = good_indices.len() as f64 + n_choices as f64; + let bad_total = bad_indices.len() as f64 + n_choices as f64; + + // Calculate l(x)/g(x) ratio for each category + let mut weights = vec![0.0f64; n_choices]; + for i in 0..n_choices { + let l_prob = (good_counts[i] as f64 + 1.0) / good_total; + let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total; + weights[i] = l_prob / g_prob; + } + + // Sample proportionally to weights + let total_weight: f64 = weights.iter().sum(); + let threshold = rng.random::() * total_weight; + + let mut cumulative = 0.0; + for (i, &w) in weights.iter().enumerate() { + cumulative += w; + if cumulative >= threshold { + return i; + } + } + + // Fallback to last index (shouldn't happen) + n_choices - 1 + } +} + +impl Default for TpeSampler { + fn default() -> Self { + Self::new() + } +} + +/// Builder for configuring a [`TpeSampler`]. +/// +/// This builder allows fluent configuration of TPE hyperparameters. +/// +/// # Examples +/// +/// ``` +/// use optimize::TpeSamplerBuilder; +/// +/// let sampler = TpeSamplerBuilder::new() +/// .gamma(0.15) +/// .n_startup_trials(20) +/// .n_ei_candidates(32) +/// .seed(42) +/// .build(); +/// ``` +#[derive(Debug, Clone)] +pub struct TpeSamplerBuilder { + gamma: f64, + n_startup_trials: usize, + n_ei_candidates: usize, + seed: Option, +} + +impl TpeSamplerBuilder { + /// Creates a new builder with default settings. + /// + /// Default settings: + /// - gamma: 0.25 (top 25% of trials are considered "good") + /// - n_startup_trials: 10 (random sampling for first 10 trials) + /// - n_ei_candidates: 24 (evaluate 24 candidates per sample) + /// - seed: None (use OS-provided entropy) + pub fn new() -> Self { + Self { + gamma: 0.25, + n_startup_trials: 10, + n_ei_candidates: 24, + seed: None, + } + } + + /// Sets the gamma quantile for splitting trials into good/bad groups. + /// + /// A gamma of 0.25 means the top 25% of trials (by objective value) are + /// considered "good" and used to build the l(x) distribution. + /// + /// # Arguments + /// + /// * `gamma` - Quantile value, must be in (0.0, 1.0). + /// + /// # Panics + /// + /// Panics if gamma is not in (0.0, 1.0). + /// + /// # Examples + /// + /// ``` + /// use optimize::TpeSamplerBuilder; + /// + /// let sampler = TpeSamplerBuilder::new() + /// .gamma(0.10) // Use top 10% as "good" trials + /// .build(); + /// ``` + pub fn gamma(mut self, gamma: f64) -> Self { + assert!( + gamma > 0.0 && gamma < 1.0, + "gamma must be in (0.0, 1.0), got {gamma}" + ); + self.gamma = gamma; + self + } + + /// Sets the number of startup trials before TPE sampling begins. + /// + /// During the startup phase, the sampler uses uniform random sampling + /// to gather initial data. Once `n_startup_trials` have completed, + /// TPE-based sampling begins. + /// + /// # Arguments + /// + /// * `n` - Number of random trials before TPE kicks in. + /// + /// # Examples + /// + /// ``` + /// use optimize::TpeSamplerBuilder; + /// + /// let sampler = TpeSamplerBuilder::new() + /// .n_startup_trials(20) // Random sample first 20 trials + /// .build(); + /// ``` + pub fn n_startup_trials(mut self, n: usize) -> Self { + self.n_startup_trials = n; + self + } + + /// Sets the number of EI (Expected Improvement) candidates to evaluate. + /// + /// When sampling a new point, TPE generates this many candidates from + /// the l(x) distribution and selects the one with the highest l(x)/g(x) + /// ratio. + /// + /// # Arguments + /// + /// * `n` - Number of candidates to evaluate per sample. + /// + /// # Examples + /// + /// ``` + /// use optimize::TpeSamplerBuilder; + /// + /// let sampler = TpeSamplerBuilder::new() + /// .n_ei_candidates(48) // Evaluate more candidates + /// .build(); + /// ``` + pub fn n_ei_candidates(mut self, n: usize) -> Self { + self.n_ei_candidates = n; + self + } + + /// Sets a seed for reproducible sampling. + /// + /// # Arguments + /// + /// * `seed` - Seed value for the random number generator. + /// + /// # Examples + /// + /// ``` + /// use optimize::TpeSamplerBuilder; + /// + /// let sampler = TpeSamplerBuilder::new() + /// .seed(42) // Reproducible results + /// .build(); + /// ``` + pub fn seed(mut self, seed: u64) -> Self { + self.seed = Some(seed); + self + } + + /// Builds the configured [`TpeSampler`]. + /// + /// # Examples + /// + /// ``` + /// use optimize::TpeSamplerBuilder; + /// + /// let sampler = TpeSamplerBuilder::new() + /// .gamma(0.15) + /// .n_startup_trials(20) + /// .n_ei_candidates(32) + /// .seed(42) + /// .build(); + /// ``` + pub fn build(self) -> TpeSampler { + TpeSampler::with_config( + self.gamma, + self.n_startup_trials, + self.n_ei_candidates, + self.seed, + ) + } +} + +impl Default for TpeSamplerBuilder { + fn default() -> Self { + Self::new() + } +} + +impl Sampler for TpeSampler { + fn sample( + &self, + distribution: &Distribution, + _trial_id: u64, + history: &[CompletedTrial], + ) -> ParamValue { + let mut rng = self.rng.lock(); + + // Fall back to random sampling during startup phase + if history.len() < self.n_startup_trials { + return self.sample_uniform(distribution, &mut rng); + } + + // Split trials into good and bad groups + let (good_trials, bad_trials) = self.split_trials(history); + + // Need at least 1 trial in each group for TPE + if good_trials.is_empty() || bad_trials.is_empty() { + return self.sample_uniform(distribution, &mut rng); + } + + // Extract parameter values for this distribution + // Since we don't have the parameter name here, we need to look at all + // trials and find matching distributions + // Note: This is a simplification - in practice, we'd need the param name + // For now, we'll collect values from trials that have this exact distribution type + + match distribution { + Distribution::Float(d) => { + // Collect float values from trials + let good_values: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + // Need values in both groups for TPE + if good_values.is_empty() || bad_values.is_empty() { + return self.sample_uniform(distribution, &mut rng); + } + + let value = self.sample_tpe_float( + d.low, + d.high, + d.log_scale, + d.step, + good_values, + bad_values, + &mut rng, + ); + ParamValue::Float(value) + } + Distribution::Int(d) => { + let good_values: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + if good_values.is_empty() || bad_values.is_empty() { + return self.sample_uniform(distribution, &mut rng); + } + + let value = self.sample_tpe_int( + d.low, + d.high, + d.log_scale, + d.step, + good_values, + bad_values, + &mut rng, + ); + ParamValue::Int(value) + } + Distribution::Categorical(d) => { + let good_indices: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + let bad_indices: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + if good_indices.is_empty() || bad_indices.is_empty() { + return self.sample_uniform(distribution, &mut rng); + } + + let index = + self.sample_tpe_categorical(d.n_choices, good_indices, bad_indices, &mut rng); + ParamValue::Categorical(index) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + + fn create_trial( + id: u64, + value: f64, + params: Vec<(&str, ParamValue, Distribution)>, + ) -> CompletedTrial { + let mut param_map = HashMap::new(); + let mut dist_map = HashMap::new(); + for (name, pv, dist) in params { + param_map.insert(name.to_string(), pv); + dist_map.insert(name.to_string(), dist); + } + CompletedTrial::new(id, param_map, dist_map, value) + } + + #[test] + fn test_tpe_sampler_new() { + let sampler = TpeSampler::new(); + assert_eq!(sampler.gamma, 0.25); + assert_eq!(sampler.n_startup_trials, 10); + assert_eq!(sampler.n_ei_candidates, 24); + } + + #[test] + fn test_tpe_sampler_with_config() { + let sampler = TpeSampler::with_config(0.15, 20, 32, Some(42)); + assert_eq!(sampler.gamma, 0.15); + assert_eq!(sampler.n_startup_trials, 20); + assert_eq!(sampler.n_ei_candidates, 32); + } + + #[test] + #[should_panic(expected = "gamma must be in (0.0, 1.0)")] + fn test_tpe_sampler_invalid_gamma_zero() { + TpeSampler::with_config(0.0, 10, 24, None); + } + + #[test] + #[should_panic(expected = "gamma must be in (0.0, 1.0)")] + fn test_tpe_sampler_invalid_gamma_one() { + TpeSampler::with_config(1.0, 10, 24, None); + } + + #[test] + fn test_tpe_startup_random_sampling() { + let sampler = TpeSampler::with_config(0.25, 10, 24, Some(42)); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // With fewer than n_startup_trials, should use random sampling + let history: Vec = vec![]; + + for _ in 0..100 { + let value = sampler.sample(&dist, 0, &history); + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(&v)); + } else { + panic!("Expected Float value"); + } + } + } + + #[test] + fn test_tpe_split_trials() { + let sampler = TpeSampler::with_config(0.25, 10, 24, Some(42)); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Create 20 trials with values 0..20 + let history: Vec = (0..20) + .map(|i| { + create_trial( + i as u64, + i as f64, + vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())], + ) + }) + .collect(); + + let (good, bad) = sampler.split_trials(&history); + + // With gamma=0.25 and 20 trials, should have 5 good and 15 bad + assert_eq!(good.len(), 5); + assert_eq!(bad.len(), 15); + + // Good trials should have lowest values + for trial in &good { + assert!(trial.value < 5.0); + } + } + + #[test] + fn test_tpe_samples_float_with_history() { + let sampler = TpeSampler::with_config(0.25, 5, 24, Some(42)); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Create history where low values (near 0.2) are "good" + let history: Vec = (0..20) + .map(|i| { + let x = i as f64 / 20.0; + // Objective is (x - 0.2)^2, minimized at x=0.2 + let value = (x - 0.2).powi(2); + create_trial( + i as u64, + value, + vec![("x", ParamValue::Float(x), dist.clone())], + ) + }) + .collect(); + + // TPE should bias toward values near 0.2 + let mut samples = vec![]; + for i in 0..100 { + let value = sampler.sample(&dist, 100 + i, &history); + if let ParamValue::Float(v) = value { + samples.push(v); + } + } + + // Calculate mean of samples - should be closer to 0.2 than 0.5 + let mean: f64 = samples.iter().sum::() / samples.len() as f64; + assert!( + mean < 0.5, + "Mean {mean} should be less than 0.5 (biased toward good region near 0.2)" + ); + } + + #[test] + fn test_tpe_categorical_sampling() { + let sampler = TpeSampler::with_config(0.25, 5, 24, Some(42)); + + let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 }); + + // Create history where category 1 is consistently good + let history: Vec = (0..20) + .map(|i| { + let category = i % 4; + // Category 1 has best (lowest) objective value + let value = if category == 1 { 0.0 } else { 1.0 }; + create_trial( + i as u64, + value, + vec![( + "cat", + ParamValue::Categorical(category as usize), + dist.clone(), + )], + ) + }) + .collect(); + + // TPE should favor category 1 + let mut counts = vec![0usize; 4]; + for i in 0..100 { + let value = sampler.sample(&dist, 100 + i, &history); + if let ParamValue::Categorical(idx) = value { + counts[idx] += 1; + } + } + + // Category 1 should be sampled more often + assert!( + counts[1] > counts[0] && counts[1] > counts[2] && counts[1] > counts[3], + "Category 1 should be most common: {counts:?}" + ); + } + + #[test] + fn test_tpe_int_sampling() { + let sampler = TpeSampler::with_config(0.25, 5, 24, Some(42)); + + let dist = Distribution::Int(IntDistribution { + low: 0, + high: 100, + log_scale: false, + step: None, + }); + + // Create history where values near 30 are good + let history: Vec = (0..20) + .map(|i| { + let x = i * 5; // 0, 5, 10, ..., 95 + let value = ((x as f64) - 30.0).powi(2); + create_trial( + i as u64, + value, + vec![("x", ParamValue::Int(x), dist.clone())], + ) + }) + .collect(); + + // TPE should bias toward values near 30 + for i in 0..50 { + let value = sampler.sample(&dist, 100 + i, &history); + if let ParamValue::Int(v) = value { + assert!((0..=100).contains(&v), "Value {v} out of range"); + } else { + panic!("Expected Int value"); + } + } + } + + #[test] + fn test_tpe_reproducibility() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let history: Vec = (0..20) + .map(|i| { + create_trial( + i as u64, + i as f64, + vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())], + ) + }) + .collect(); + + let sampler1 = TpeSampler::with_config(0.25, 5, 24, Some(12345)); + let sampler2 = TpeSampler::with_config(0.25, 5, 24, Some(12345)); + + for i in 0..10 { + let v1 = sampler1.sample(&dist, i, &history); + let v2 = sampler2.sample(&dist, i, &history); + assert_eq!(v1, v2, "Samples should be identical with same seed"); + } + } + + #[test] + fn test_tpe_sampler_builder_default() { + let builder = TpeSamplerBuilder::new(); + let sampler = builder.build(); + assert_eq!(sampler.gamma, 0.25); + assert_eq!(sampler.n_startup_trials, 10); + assert_eq!(sampler.n_ei_candidates, 24); + } + + #[test] + fn test_tpe_sampler_builder_custom() { + let sampler = TpeSamplerBuilder::new() + .gamma(0.15) + .n_startup_trials(20) + .n_ei_candidates(32) + .seed(42) + .build(); + assert_eq!(sampler.gamma, 0.15); + assert_eq!(sampler.n_startup_trials, 20); + assert_eq!(sampler.n_ei_candidates, 32); + } + + #[test] + fn test_tpe_sampler_builder_via_sampler() { + let sampler = TpeSampler::builder() + .gamma(0.10) + .n_startup_trials(15) + .n_ei_candidates(48) + .build(); + assert_eq!(sampler.gamma, 0.10); + assert_eq!(sampler.n_startup_trials, 15); + assert_eq!(sampler.n_ei_candidates, 48); + } + + #[test] + fn test_tpe_sampler_builder_partial() { + // Test setting only some options + let sampler = TpeSamplerBuilder::new().gamma(0.20).build(); + assert_eq!(sampler.gamma, 0.20); + assert_eq!(sampler.n_startup_trials, 10); // default + assert_eq!(sampler.n_ei_candidates, 24); // default + } + + #[test] + #[should_panic(expected = "gamma must be in (0.0, 1.0)")] + fn test_tpe_sampler_builder_invalid_gamma() { + TpeSamplerBuilder::new().gamma(1.5).build(); + } + + #[test] + fn test_tpe_sampler_builder_reproducibility() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let history: Vec = (0..20) + .map(|i| { + create_trial( + i as u64, + i as f64, + vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())], + ) + }) + .collect(); + + let sampler1 = TpeSampler::builder() + .seed(99999) + .n_startup_trials(5) + .build(); + let sampler2 = TpeSampler::builder() + .seed(99999) + .n_startup_trials(5) + .build(); + + for i in 0..10 { + let v1 = sampler1.sample(&dist, i, &history); + let v2 = sampler2.sample(&dist, i, &history); + assert_eq!( + v1, v2, + "Builder-created samplers with same seed should be identical" + ); + } + } +} diff --git a/src/study.rs b/src/study.rs new file mode 100644 index 0000000..71c5fe5 --- /dev/null +++ b/src/study.rs @@ -0,0 +1,1340 @@ +//! 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 { + 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 optimize::{Direction, Study}; +/// +/// // Create a study to minimize an objective function +/// let study: Study = Study::new(Direction::Minimize); +/// assert_eq!(study.direction(), Direction::Minimize); +/// ``` +pub struct Study +where + V: PartialOrd, +{ + /// The optimization direction. + direction: Direction, + /// The sampler used to generate parameter values. + sampler: Arc, + /// Completed trials (wrapped in Arc for sharing with Trial). + completed_trials: Arc>>>, + /// Counter for generating unique trial IDs. + next_trial_id: AtomicU64, +} + +impl Study +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 optimize::{Direction, Study}; + /// + /// let study: Study = 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 optimize::{Direction, RandomSampler, Study}; + /// + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = 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 optimize::{Direction, Study, TpeSampler}; + /// + /// // After deserializing a study, restore the TPE sampler + /// let mut study: Study = 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`, 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 optimize::{Direction, Study}; + /// + /// let study: Study = 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 optimize::{Direction, Study}; + /// + /// let study: Study = 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 optimize::{Direction, Study}; + /// + /// let study: Study = 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 optimize::{Direction, Study}; + /// + /// let study: Study = 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> + where + V: Clone, + { + self.completed_trials.read().clone() + } + + /// Returns the number of completed trials. + /// + /// Failed trials are not counted. + /// + /// # Examples + /// + /// ``` + /// use optimize::{Direction, Study}; + /// + /// let study: Study = 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 optimize::{Direction, Study}; + /// + /// let study: Study = 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> + 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 optimize::{Direction, Study}; + /// + /// let study: Study = 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 + 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 optimize::{Direction, RandomSampler, Study}; + /// + /// // Minimize x^2 + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); + /// + /// study + /// .optimize(10, |trial| { + /// let x = trial.suggest_float("x", -10.0, 10.0)?; + /// Ok::<_, optimize::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(&self, n_trials: usize, mut objective: F) -> crate::Result<()> + where + F: FnMut(&mut Trial) -> std::result::Result, + 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)`. + /// + /// # Errors + /// + /// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials). + /// + /// # Examples + /// + /// ``` + /// use optimize::{Direction, RandomSampler, Study}; + /// + /// # #[cfg(feature = "async")] + /// # async fn example() -> optimize::Result<()> { + /// // Minimize x^2 with async objective + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = 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::<_, optimize::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( + &self, + n_trials: usize, + objective: F, + ) -> crate::Result<()> + where + F: Fn(Trial) -> Fut, + Fut: Future>, + 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 optimize::{Direction, RandomSampler, Study}; + /// + /// # #[cfg(feature = "async")] + /// # async fn example() -> optimize::Result<()> { + /// // Minimize x^2 with parallel async evaluation + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = 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::<_, optimize::TpeError>((trial, value)) + /// }) + /// .await?; + /// + /// // All trials should have completed + /// assert_eq!(study.n_trials(), 10); + /// # Ok(()) + /// # } + /// ``` + #[cfg(feature = "async")] + pub async fn optimize_parallel( + &self, + n_trials: usize, + concurrency: usize, + objective: F, + ) -> crate::Result<()> + where + F: Fn(Trial) -> Fut + Send + Sync + 'static, + Fut: Future> + 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 optimize::{Direction, RandomSampler, Study}; + /// + /// // Stop early when we find a good enough value + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); + /// + /// study + /// .optimize_with_callback( + /// 100, + /// |trial| { + /// let x = trial.suggest_float("x", -10.0, 10.0)?; + /// Ok::<_, optimize::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( + &self, + n_trials: usize, + mut objective: F, + mut callback: C, + ) -> crate::Result<()> + where + V: Clone, + F: FnMut(&mut Trial) -> std::result::Result, + C: FnMut(&Study, &CompletedTrial) -> 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 that provides full sampler integration. +impl Study { + /// 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`. + /// + /// 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 optimize::{Direction, RandomSampler, Study}; + /// + /// // With a seeded sampler for reproducibility + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = 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 `optimize` 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` + /// 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 optimize::{Direction, RandomSampler, Study}; + /// + /// // Minimize x^2 with sampler integration + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); + /// + /// study + /// .optimize_with_sampler(10, |trial| { + /// let x = trial.suggest_float("x", -10.0, 10.0)?; + /// Ok::<_, optimize::TpeError>(x * x) + /// }) + /// .unwrap(); + /// + /// // At least one trial should have completed + /// assert!(study.n_trials() > 0); + /// ``` + pub fn optimize_with_sampler( + &self, + n_trials: usize, + mut objective: F, + ) -> crate::Result<()> + where + F: FnMut(&mut Trial) -> std::result::Result, + 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 optimize::{Direction, RandomSampler, Study}; + /// + /// // Optimize with sampler integration and early stopping + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = Study::with_sampler(Direction::Minimize, sampler); + /// + /// study + /// .optimize_with_callback_sampler( + /// 100, + /// |trial| { + /// let x = trial.suggest_float("x", -10.0, 10.0)?; + /// Ok::<_, optimize::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( + &self, + n_trials: usize, + mut objective: F, + mut callback: C, + ) -> crate::Result<()> + where + F: FnMut(&mut Trial) -> std::result::Result, + C: FnMut(&Study, &CompletedTrial) -> 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)`. + /// + /// # Errors + /// + /// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials). + /// + /// # Examples + /// + /// ``` + /// use optimize::{Direction, RandomSampler, Study}; + /// + /// # #[cfg(feature = "async")] + /// # async fn example() -> optimize::Result<()> { + /// // Minimize x^2 with async objective and sampler integration + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = 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::<_, optimize::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( + &self, + n_trials: usize, + objective: F, + ) -> crate::Result<()> + where + F: Fn(Trial) -> Fut, + Fut: Future>, + 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 optimize::{Direction, RandomSampler, Study}; + /// + /// # #[cfg(feature = "async")] + /// # async fn example() -> optimize::Result<()> { + /// // Minimize x^2 with parallel async evaluation and sampler integration + /// let sampler = RandomSampler::with_seed(42); + /// let study: Study = 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::<_, optimize::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( + &self, + n_trials: usize, + concurrency: usize, + objective: F, + ) -> crate::Result<()> + where + F: Fn(Trial) -> Fut + Send + Sync + 'static, + Fut: Future> + 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 when serde feature is enabled. +#[cfg(feature = "serde")] +impl Serialize for Study +where + V: PartialOrd + Serialize, +{ + fn serialize(&self, serializer: S) -> Result + 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> + 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 when serde feature is enabled. +#[cfg(feature = "serde")] +impl<'de, V> Deserialize<'de> for Study +where + V: PartialOrd + Deserialize<'de>, +{ + fn deserialize(deserializer: D) -> Result + 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(PhantomData); + + impl<'de, V> Visitor<'de> for StudyVisitor + where + V: PartialOrd + Deserialize<'de>, + { + type Value = Study; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("struct Study") + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut direction = None; + let mut completed_trials: Option>> = 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 = 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 = 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 = 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 = 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 = 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); + } +} diff --git a/src/trial.rs b/src/trial.rs new file mode 100644 index 0000000..ad05f09 --- /dev/null +++ b/src/trial.rs @@ -0,0 +1,793 @@ +//! Trial implementation for tracking sampled parameters and trial state. + +use std::collections::HashMap; +use std::sync::Arc; + +use parking_lot::RwLock; +#[cfg(feature = "serde")] +use serde::{Deserialize, Serialize}; + +use crate::distribution::{ + CategoricalDistribution, Distribution, FloatDistribution, IntDistribution, +}; +use crate::error::{Result, TpeError}; +use crate::param::ParamValue; +use crate::sampler::{CompletedTrial, Sampler}; +use crate::types::TrialState; + +/// A trial represents a single evaluation of the objective function. +/// +/// Each trial has a unique ID and stores the sampled parameters along with +/// their distributions. The trial progresses through states: Running -> Complete/Failed. +/// +/// Trials use a sampler to generate parameter values. When created through +/// `Study::create_trial()`, the trial receives the study's sampler and access +/// to the history of completed trials for informed sampling. +#[derive(Clone)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct Trial { + /// Unique identifier for this trial. + id: u64, + /// Current state of the trial. + state: TrialState, + /// Sampled parameter values, keyed by parameter name. + params: HashMap, + /// Parameter distributions, keyed by parameter name. + distributions: HashMap, + /// The sampler to use for generating parameter values. + #[cfg_attr(feature = "serde", serde(skip))] + sampler: Option>, + /// Access to the history of completed trials (shared with Study). + #[cfg_attr(feature = "serde", serde(skip))] + history: Option>>>>, +} + +impl std::fmt::Debug for Trial { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Trial") + .field("id", &self.id) + .field("state", &self.state) + .field("params", &self.params) + .field("distributions", &self.distributions) + .field("has_sampler", &self.sampler.is_some()) + .field("has_history", &self.history.is_some()) + .finish() + } +} + +impl Trial { + /// Creates a new trial with the given ID. + /// + /// The trial starts in the `Running` state with no parameters sampled. + /// This constructor creates a trial without a sampler, which will use + /// local random sampling for suggest_* methods. + /// + /// For trials that use the study's sampler, use `Trial::with_sampler` instead. + /// + /// # Arguments + /// + /// * `id` - A unique identifier for this trial. + /// + /// # Examples + /// + /// ``` + /// use optimize::Trial; + /// + /// let trial = Trial::new(0); + /// assert_eq!(trial.id(), 0); + /// ``` + pub fn new(id: u64) -> Self { + Self { + id, + state: TrialState::Running, + params: HashMap::new(), + distributions: HashMap::new(), + sampler: None, + history: None, + } + } + + /// Creates a new trial with a sampler and access to trial history. + /// + /// This constructor is used by `Study::create_trial()` to create trials + /// that use the study's sampler for informed parameter suggestions. + /// + /// # Arguments + /// + /// * `id` - A unique identifier for this trial. + /// * `sampler` - The sampler to use for generating parameter values. + /// * `history` - Shared access to the history of completed trials. + pub(crate) fn with_sampler( + id: u64, + sampler: Arc, + history: Arc>>>, + ) -> Self { + Self { + id, + state: TrialState::Running, + params: HashMap::new(), + distributions: HashMap::new(), + sampler: Some(sampler), + history: Some(history), + } + } + + /// Samples a value from the given distribution using the sampler. + /// + /// If the trial has a sampler, it delegates to the sampler's sample method + /// with the history of completed trials. Otherwise, it uses the RandomSampler + /// as a fallback. + fn sample_value(&self, distribution: &Distribution) -> ParamValue { + if let (Some(sampler), Some(history)) = (&self.sampler, &self.history) { + let history_guard = history.read(); + sampler.sample(distribution, self.id, &history_guard) + } else { + // Fallback to RandomSampler when no sampler is configured + use crate::sampler::RandomSampler; + let fallback = RandomSampler::new(); + fallback.sample(distribution, self.id, &[]) + } + } + + /// Returns the unique ID of this trial. + pub fn id(&self) -> u64 { + self.id + } + + /// Returns the current state of this trial. + pub fn state(&self) -> TrialState { + self.state + } + + /// Returns a reference to the sampled parameters. + pub fn params(&self) -> &HashMap { + &self.params + } + + /// Returns a reference to the parameter distributions. + pub fn distributions(&self) -> &HashMap { + &self.distributions + } + + /// Sets the trial state to Complete. + pub(crate) fn set_complete(&mut self) { + self.state = TrialState::Complete; + } + + /// Sets the trial state to Failed. + pub(crate) fn set_failed(&mut self) { + self.state = TrialState::Failed; + } + + /// Suggests a float parameter with the given bounds. + /// + /// If the parameter has already been sampled with the same bounds, the cached value is returned. + /// If the parameter was sampled with different bounds, a `ParameterConflict` error is returned. + /// + /// # Arguments + /// + /// * `name` - The name of the parameter. + /// * `low` - The lower bound (inclusive). + /// * `high` - The upper bound (inclusive). + /// + /// # Errors + /// + /// Returns `InvalidBounds` if `low > high`. + /// Returns `ParameterConflict` if the parameter was previously sampled with different bounds. + /// + /// # Examples + /// + /// ``` + /// use optimize::Trial; + /// + /// let mut trial = Trial::new(0); + /// let x = trial.suggest_float("x", 0.0, 1.0).unwrap(); + /// assert!(x >= 0.0 && x <= 1.0); + /// + /// // Calling again with same bounds returns cached value + /// let x2 = trial.suggest_float("x", 0.0, 1.0).unwrap(); + /// assert_eq!(x, x2); + /// ``` + pub fn suggest_float(&mut self, name: impl Into, low: f64, high: f64) -> Result { + if low > high { + return Err(TpeError::InvalidBounds { low, high }); + } + + let name = name.into(); + let distribution = FloatDistribution { + low, + high, + log_scale: false, + step: None, + }; + + // Check if parameter already exists + if let Some(existing_dist) = self.distributions.get(&name) { + // Verify the distribution matches + if let Distribution::Float(existing) = existing_dist + && existing.low == low + && existing.high == high + && !existing.log_scale + && existing.step.is_none() + { + // Same distribution, return cached value + if let Some(ParamValue::Float(value)) = self.params.get(&name) { + return Ok(*value); + } + } + // Distribution exists but doesn't match + return Err(TpeError::ParameterConflict { + name, + reason: "parameter was previously sampled with different bounds or type" + .to_string(), + }); + } + + // Sample using the sampler + let dist = Distribution::Float(distribution); + let value = match self.sample_value(&dist) { + ParamValue::Float(v) => v, + _ => unreachable!("Float distribution should return Float value"), + }; + + // Store distribution and value + self.distributions.insert(name.clone(), dist); + self.params.insert(name, ParamValue::Float(value)); + + Ok(value) + } + + /// Suggests a float parameter sampled on a logarithmic scale. + /// + /// The value is sampled uniformly in log space, which is useful for parameters + /// that span multiple orders of magnitude (e.g., learning rates). + /// + /// If the parameter has already been sampled with the same bounds and log_scale=true, + /// the cached value is returned. If the parameter was sampled with different configuration, + /// a `ParameterConflict` error is returned. + /// + /// # Arguments + /// + /// * `name` - The name of the parameter. + /// * `low` - The lower bound (inclusive, must be positive). + /// * `high` - The upper bound (inclusive). + /// + /// # Errors + /// + /// Returns `InvalidLogBounds` if `low <= 0`. + /// Returns `InvalidBounds` if `low > high`. + /// Returns `ParameterConflict` if the parameter was previously sampled with different configuration. + /// + /// # Examples + /// + /// ``` + /// use optimize::Trial; + /// + /// let mut trial = Trial::new(0); + /// let lr = trial + /// .suggest_float_log("learning_rate", 1e-5, 1e-1) + /// .unwrap(); + /// assert!(lr >= 1e-5 && lr <= 1e-1); + /// + /// // Calling again with same bounds returns cached value + /// let lr2 = trial + /// .suggest_float_log("learning_rate", 1e-5, 1e-1) + /// .unwrap(); + /// assert_eq!(lr, lr2); + /// ``` + pub fn suggest_float_log( + &mut self, + name: impl Into, + low: f64, + high: f64, + ) -> Result { + if low <= 0.0 { + return Err(TpeError::InvalidLogBounds); + } + + if low > high { + return Err(TpeError::InvalidBounds { low, high }); + } + + let name = name.into(); + let distribution = FloatDistribution { + low, + high, + log_scale: true, + step: None, + }; + + // Check if parameter already exists + if let Some(existing_dist) = self.distributions.get(&name) { + // Verify the distribution matches + if let Distribution::Float(existing) = existing_dist + && existing.low == low + && existing.high == high + && existing.log_scale + && existing.step.is_none() + { + // Same distribution, return cached value + if let Some(ParamValue::Float(value)) = self.params.get(&name) { + return Ok(*value); + } + } + // Distribution exists but doesn't match + return Err(TpeError::ParameterConflict { + name, + reason: "parameter was previously sampled with different bounds or type" + .to_string(), + }); + } + + // Sample using the sampler (sampler handles log-scale transformation) + let dist = Distribution::Float(distribution); + let value = match self.sample_value(&dist) { + ParamValue::Float(v) => v, + _ => unreachable!("Float distribution should return Float value"), + }; + + // Store distribution and value + self.distributions.insert(name.clone(), dist); + self.params.insert(name, ParamValue::Float(value)); + + Ok(value) + } + + /// Suggests a float parameter that snaps to a step grid. + /// + /// The value is sampled from the discrete set {low, low + step, low + 2*step, ...} + /// where each value is <= high. + /// + /// If the parameter has already been sampled with the same configuration, + /// the cached value is returned. If the parameter was sampled with different configuration, + /// a `ParameterConflict` error is returned. + /// + /// # Arguments + /// + /// * `name` - The name of the parameter. + /// * `low` - The lower bound (inclusive). + /// * `high` - The upper bound (inclusive). + /// * `step` - The step size (must be positive). + /// + /// # Errors + /// + /// Returns `InvalidStep` if `step <= 0`. + /// Returns `InvalidBounds` if `low > high`. + /// Returns `ParameterConflict` if the parameter was previously sampled with different configuration. + /// + /// # Examples + /// + /// ``` + /// use optimize::Trial; + /// + /// let mut trial = Trial::new(0); + /// let x = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap(); + /// // x will be one of: 0.0, 0.25, 0.5, 0.75, 1.0 + /// assert!(x >= 0.0 && x <= 1.0); + /// assert!((x / 0.25).fract().abs() < 1e-10 || (x / 0.25).fract().abs() > 1.0 - 1e-10); + /// + /// // Calling again with same bounds returns cached value + /// let x2 = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap(); + /// assert_eq!(x, x2); + /// ``` + pub fn suggest_float_step( + &mut self, + name: impl Into, + low: f64, + high: f64, + step: f64, + ) -> Result { + if step <= 0.0 { + return Err(TpeError::InvalidStep); + } + + if low > high { + return Err(TpeError::InvalidBounds { low, high }); + } + + let name = name.into(); + let distribution = FloatDistribution { + low, + high, + log_scale: false, + step: Some(step), + }; + + // Check if parameter already exists + if let Some(existing_dist) = self.distributions.get(&name) { + // Verify the distribution matches + if let Distribution::Float(existing) = existing_dist + && existing.low == low + && existing.high == high + && !existing.log_scale + && existing.step == Some(step) + { + // Same distribution, return cached value + if let Some(ParamValue::Float(value)) = self.params.get(&name) { + return Ok(*value); + } + } + // Distribution exists but doesn't match + return Err(TpeError::ParameterConflict { + name, + reason: "parameter was previously sampled with different bounds or type" + .to_string(), + }); + } + + // Sample using the sampler (sampler handles step-grid) + let dist = Distribution::Float(distribution); + let value = match self.sample_value(&dist) { + ParamValue::Float(v) => v, + _ => unreachable!("Float distribution should return Float value"), + }; + + // Store distribution and value + self.distributions.insert(name.clone(), dist); + self.params.insert(name, ParamValue::Float(value)); + + Ok(value) + } + + /// Suggests an integer parameter with the given bounds. + /// + /// The value is sampled uniformly from the range [low, high] inclusive. + /// + /// If the parameter has already been sampled with the same bounds, the cached value is returned. + /// If the parameter was sampled with different bounds, a `ParameterConflict` error is returned. + /// + /// # Arguments + /// + /// * `name` - The name of the parameter. + /// * `low` - The lower bound (inclusive). + /// * `high` - The upper bound (inclusive). + /// + /// # Errors + /// + /// Returns `InvalidBounds` if `low > high`. + /// Returns `ParameterConflict` if the parameter was previously sampled with different bounds. + /// + /// # Examples + /// + /// ``` + /// use optimize::Trial; + /// + /// let mut trial = Trial::new(0); + /// let n = trial.suggest_int("n_layers", 1, 10).unwrap(); + /// assert!(n >= 1 && n <= 10); + /// + /// // Calling again with same bounds returns cached value + /// let n2 = trial.suggest_int("n_layers", 1, 10).unwrap(); + /// assert_eq!(n, n2); + /// ``` + pub fn suggest_int(&mut self, name: impl Into, low: i64, high: i64) -> Result { + if low > high { + return Err(TpeError::InvalidBounds { + low: low as f64, + high: high as f64, + }); + } + + let name = name.into(); + let distribution = IntDistribution { + low, + high, + log_scale: false, + step: None, + }; + + // Check if parameter already exists + if let Some(existing_dist) = self.distributions.get(&name) { + // Verify the distribution matches + if let Distribution::Int(existing) = existing_dist + && existing.low == low + && existing.high == high + && !existing.log_scale + && existing.step.is_none() + { + // Same distribution, return cached value + if let Some(ParamValue::Int(value)) = self.params.get(&name) { + return Ok(*value); + } + } + // Distribution exists but doesn't match + return Err(TpeError::ParameterConflict { + name, + reason: "parameter was previously sampled with different bounds or type" + .to_string(), + }); + } + + // Sample using the sampler + let dist = Distribution::Int(distribution); + let value = match self.sample_value(&dist) { + ParamValue::Int(v) => v, + _ => unreachable!("Int distribution should return Int value"), + }; + + // Store distribution and value + self.distributions.insert(name.clone(), dist); + self.params.insert(name, ParamValue::Int(value)); + + Ok(value) + } + + /// Suggests an integer parameter sampled on a logarithmic scale. + /// + /// The value is sampled uniformly in log space, which is useful for parameters + /// that span multiple orders of magnitude (e.g., batch sizes). + /// + /// If the parameter has already been sampled with the same bounds and log_scale=true, + /// the cached value is returned. If the parameter was sampled with different configuration, + /// a `ParameterConflict` error is returned. + /// + /// # Arguments + /// + /// * `name` - The name of the parameter. + /// * `low` - The lower bound (inclusive, must be >= 1). + /// * `high` - The upper bound (inclusive). + /// + /// # Errors + /// + /// Returns `InvalidLogBounds` if `low < 1`. + /// Returns `InvalidBounds` if `low > high`. + /// Returns `ParameterConflict` if the parameter was previously sampled with different configuration. + /// + /// # Examples + /// + /// ``` + /// use optimize::Trial; + /// + /// let mut trial = Trial::new(0); + /// let batch_size = trial.suggest_int_log("batch_size", 1, 1024).unwrap(); + /// assert!(batch_size >= 1 && batch_size <= 1024); + /// + /// // Calling again with same bounds returns cached value + /// let batch_size2 = trial.suggest_int_log("batch_size", 1, 1024).unwrap(); + /// assert_eq!(batch_size, batch_size2); + /// ``` + pub fn suggest_int_log(&mut self, name: impl Into, low: i64, high: i64) -> Result { + if low < 1 { + return Err(TpeError::InvalidLogBounds); + } + + if low > high { + return Err(TpeError::InvalidBounds { + low: low as f64, + high: high as f64, + }); + } + + let name = name.into(); + let distribution = IntDistribution { + low, + high, + log_scale: true, + step: None, + }; + + // Check if parameter already exists + if let Some(existing_dist) = self.distributions.get(&name) { + // Verify the distribution matches + if let Distribution::Int(existing) = existing_dist + && existing.low == low + && existing.high == high + && existing.log_scale + && existing.step.is_none() + { + // Same distribution, return cached value + if let Some(ParamValue::Int(value)) = self.params.get(&name) { + return Ok(*value); + } + } + // Distribution exists but doesn't match + return Err(TpeError::ParameterConflict { + name, + reason: "parameter was previously sampled with different bounds or type" + .to_string(), + }); + } + + // Sample using the sampler (sampler handles log-scale transformation) + let dist = Distribution::Int(distribution); + let value = match self.sample_value(&dist) { + ParamValue::Int(v) => v, + _ => unreachable!("Int distribution should return Int value"), + }; + + // Store distribution and value + self.distributions.insert(name.clone(), dist); + self.params.insert(name, ParamValue::Int(value)); + + Ok(value) + } + + /// Suggests an integer parameter that snaps to a step grid. + /// + /// The value is sampled from the discrete set {low, low + step, low + 2*step, ...} + /// where each value is <= high. + /// + /// If the parameter has already been sampled with the same configuration, + /// the cached value is returned. If the parameter was sampled with different configuration, + /// a `ParameterConflict` error is returned. + /// + /// # Arguments + /// + /// * `name` - The name of the parameter. + /// * `low` - The lower bound (inclusive). + /// * `high` - The upper bound (inclusive). + /// * `step` - The step size (must be positive). + /// + /// # Errors + /// + /// Returns `InvalidStep` if `step <= 0`. + /// Returns `InvalidBounds` if `low > high`. + /// Returns `ParameterConflict` if the parameter was previously sampled with different configuration. + /// + /// # Examples + /// + /// ``` + /// use optimize::Trial; + /// + /// let mut trial = Trial::new(0); + /// let n = trial + /// .suggest_int_step("n_estimators", 100, 500, 50) + /// .unwrap(); + /// // n will be one of: 100, 150, 200, 250, 300, 350, 400, 450, 500 + /// assert!(n >= 100 && n <= 500); + /// assert!((n - 100) % 50 == 0); + /// + /// // Calling again with same bounds returns cached value + /// let n2 = trial + /// .suggest_int_step("n_estimators", 100, 500, 50) + /// .unwrap(); + /// assert_eq!(n, n2); + /// ``` + pub fn suggest_int_step( + &mut self, + name: impl Into, + low: i64, + high: i64, + step: i64, + ) -> Result { + if step <= 0 { + return Err(TpeError::InvalidStep); + } + + if low > high { + return Err(TpeError::InvalidBounds { + low: low as f64, + high: high as f64, + }); + } + + let name = name.into(); + let distribution = IntDistribution { + low, + high, + log_scale: false, + step: Some(step), + }; + + // Check if parameter already exists + if let Some(existing_dist) = self.distributions.get(&name) { + // Verify the distribution matches + if let Distribution::Int(existing) = existing_dist + && existing.low == low + && existing.high == high + && !existing.log_scale + && existing.step == Some(step) + { + // Same distribution, return cached value + if let Some(ParamValue::Int(value)) = self.params.get(&name) { + return Ok(*value); + } + } + // Distribution exists but doesn't match + return Err(TpeError::ParameterConflict { + name, + reason: "parameter was previously sampled with different bounds or type" + .to_string(), + }); + } + + // Sample using the sampler (sampler handles step-grid) + let dist = Distribution::Int(distribution); + let value = match self.sample_value(&dist) { + ParamValue::Int(v) => v, + _ => unreachable!("Int distribution should return Int value"), + }; + + // Store distribution and value + self.distributions.insert(name.clone(), dist); + self.params.insert(name, ParamValue::Int(value)); + + Ok(value) + } + + /// Suggests a categorical parameter from the given choices. + /// + /// The value is selected uniformly at random from the provided choices. + /// Internally, the index of the selected choice is stored. + /// + /// If the parameter has already been sampled with the same number of choices, + /// the cached value (same index) is returned. If the parameter was sampled with + /// a different number of choices, a `ParameterConflict` error is returned. + /// + /// # Arguments + /// + /// * `name` - The name of the parameter. + /// * `choices` - A slice of choices to select from. + /// + /// # Type Parameters + /// + /// * `T` - The type of the choices. Only requires `Clone`. + /// + /// # Errors + /// + /// Returns `EmptyChoices` if `choices` is empty. + /// Returns `ParameterConflict` if the parameter was previously sampled with a different number of choices. + /// + /// # Examples + /// + /// ``` + /// use optimize::Trial; + /// + /// let mut trial = Trial::new(0); + /// let optimizer = trial + /// .suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"]) + /// .unwrap(); + /// assert!(["sgd", "adam", "rmsprop"].contains(&optimizer)); + /// + /// // Calling again with same choices returns cached value + /// let optimizer2 = trial + /// .suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"]) + /// .unwrap(); + /// assert_eq!(optimizer, optimizer2); + /// ``` + pub fn suggest_categorical( + &mut self, + name: impl Into, + choices: &[T], + ) -> Result { + if choices.is_empty() { + return Err(TpeError::EmptyChoices); + } + + let name = name.into(); + let n_choices = choices.len(); + let distribution = CategoricalDistribution { n_choices }; + + // Check if parameter already exists + if let Some(existing_dist) = self.distributions.get(&name) { + // Verify the distribution matches + if let Distribution::Categorical(existing) = existing_dist + && existing.n_choices == n_choices + { + // Same distribution, return cached value + if let Some(ParamValue::Categorical(index)) = self.params.get(&name) { + return Ok(choices[*index].clone()); + } + } + // Distribution exists but doesn't match + return Err(TpeError::ParameterConflict { + name, + reason: "parameter was previously sampled with different number of choices or type" + .to_string(), + }); + } + + // Sample using the sampler + let dist = Distribution::Categorical(distribution); + let index = match self.sample_value(&dist) { + ParamValue::Categorical(idx) => idx, + _ => unreachable!("Categorical distribution should return Categorical value"), + }; + + // Store distribution and value (store the index) + self.distributions.insert(name.clone(), dist); + self.params.insert(name, ParamValue::Categorical(index)); + + Ok(choices[index].clone()) + } +} diff --git a/src/types.rs b/src/types.rs new file mode 100644 index 0000000..ec9a477 --- /dev/null +++ b/src/types.rs @@ -0,0 +1,26 @@ +//! Core types for the optimize library. + +#[cfg(feature = "serde")] +use serde::{Deserialize, Serialize}; + +/// The direction of optimization. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum Direction { + /// Minimize the objective value. + Minimize, + /// Maximize the objective value. + Maximize, +} + +/// The state of a trial in its lifecycle. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum TrialState { + /// The trial is currently running. + Running, + /// The trial completed successfully. + Complete, + /// The trial failed with an error. + Failed, +} diff --git a/tests/integration.rs b/tests/integration.rs new file mode 100644 index 0000000..54f7988 --- /dev/null +++ b/tests/integration.rs @@ -0,0 +1,747 @@ +//! Integration tests for the optimize library. + +use optimize::{Direction, RandomSampler, Study, TpeError, TpeSampler, Trial}; + +// ============================================================================= +// Test: optimize simple quadratic function with TPE, finds near-optimal +// ============================================================================= + +#[test] +fn test_tpe_optimizes_quadratic_function() { + // Minimize f(x) = (x - 3)^2 where x ∈ [-10, 10] + // Optimal: x = 3, f(3) = 0 + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(5) // Quick startup for test + .n_ei_candidates(24) + .build(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(50, |trial| { + let x = trial.suggest_float("x", -10.0, 10.0)?; + Ok::<_, TpeError>((x - 3.0).powi(2)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + // TPE should find a value close to optimal (x ≈ 3) + // We expect the best value to be small (close to 0) + assert!( + best.value < 1.0, + "TPE should find near-optimal: best value {} should be < 1.0", + best.value + ); +} + +#[test] +fn test_tpe_optimizes_multivariate_function() { + // Minimize f(x, y) = x^2 + y^2 where x, y ∈ [-5, 5] + // Optimal: (0, 0), f(0, 0) = 0 + let sampler = TpeSampler::builder().seed(123).n_startup_trials(10).build(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(100, |trial| { + let x = trial.suggest_float("x", -5.0, 5.0)?; + let y = trial.suggest_float("y", -5.0, 5.0)?; + Ok::<_, TpeError>(x * x + y * y) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + // TPE should find a reasonably good solution + assert!( + best.value < 5.0, + "TPE should find near-optimal: best value {} should be < 5.0", + best.value + ); +} + +#[test] +fn test_tpe_maximization() { + // Maximize f(x) = -(x - 2)^2 + 10 where x ∈ [-10, 10] + // Optimal: x = 2, f(2) = 10 + let sampler = TpeSampler::builder().seed(456).n_startup_trials(5).build(); + + let study: Study = Study::with_sampler(Direction::Maximize, sampler); + + study + .optimize_with_sampler(50, |trial| { + let x = trial.suggest_float("x", -10.0, 10.0)?; + Ok::<_, TpeError>(-(x - 2.0).powi(2) + 10.0) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + // For maximization, best value should be better than a random baseline + // The function ranges from -90 (at x=-10 or x=10, when far from x=2) to 10 (at x=2) + // A random approach would average around 0, so finding >5 is a reasonable check + assert!( + best.value > 5.0, + "TPE should find reasonably good solution: best value {} should be > 5.0", + best.value + ); +} + +// ============================================================================= +// Test: RandomSampler samples uniformly across range +// ============================================================================= + +#[test] +fn test_random_sampler_uniform_float_distribution() { + // Test that RandomSampler samples uniformly by running multiple trials + // and checking the distribution of sampled values + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); + + let n_samples = 1000; + let mut samples = Vec::with_capacity(n_samples); + + study + .optimize(n_samples, |trial| { + let x = trial.suggest_float("x", 0.0, 1.0)?; + samples.push(x); + Ok::<_, TpeError>(x) + }) + .unwrap(); + + // All samples should be in range + for &s in &samples { + assert!((0.0..=1.0).contains(&s), "sample {s} out of range [0, 1]"); + } + + // Check distribution is roughly uniform by looking at quartiles + samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); + + let q1 = samples[n_samples / 4]; + let q2 = samples[n_samples / 2]; + let q3 = samples[3 * n_samples / 4]; + + // For uniform distribution, quartiles should be approximately 0.25, 0.5, 0.75 + assert!((q1 - 0.25).abs() < 0.1, "Q1 {q1} should be close to 0.25"); + assert!( + (q2 - 0.5).abs() < 0.1, + "Q2 (median) {q2} should be close to 0.5" + ); + assert!((q3 - 0.75).abs() < 0.1, "Q3 {q3} should be close to 0.75"); +} + +#[test] +fn test_random_sampler_uniform_int_distribution() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(123)); + + let n_samples = 1000; + let mut counts = [0u32; 10]; // counts for values 1-10 + + study + .optimize(n_samples, |trial| { + let n = trial.suggest_int("n", 1, 10)?; + assert!((1..=10).contains(&n), "sample {n} out of range [1, 10]"); + counts[(n - 1) as usize] += 1; + Ok::<_, TpeError>(n as f64) + }) + .unwrap(); + + // Each value should appear roughly n_samples / 10 times + let expected = n_samples as f64 / 10.0; + for (i, &count) in counts.iter().enumerate() { + let diff = (count as f64 - expected).abs() / expected; + assert!( + diff < 0.3, + "value {} appeared {} times, expected ~{}, diff = {:.1}%", + i + 1, + count, + expected, + diff * 100.0 + ); + } +} + +#[test] +fn test_random_sampler_uniform_categorical_distribution() { + let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(456)); + + let n_samples = 1000; + let mut counts = [0u32; 4]; + let choices = ["a", "b", "c", "d"]; + + study + .optimize(n_samples, |trial| { + let choice = trial.suggest_categorical("cat", &choices)?; + let idx = choices.iter().position(|&c| c == choice).unwrap(); + counts[idx] += 1; + Ok::<_, TpeError>(idx as f64) + }) + .unwrap(); + + // Each category should appear roughly n_samples / 4 times + let expected = n_samples as f64 / 4.0; + for (i, &count) in counts.iter().enumerate() { + let diff = (count as f64 - expected).abs() / expected; + assert!( + diff < 0.25, + "category {} appeared {} times, expected ~{}, diff = {:.1}%", + i, + count, + expected, + diff * 100.0 + ); + } +} + +#[test] +fn test_random_sampler_reproducibility() { + // Two studies with the same seed should produce the same sequence + // NOTE: We must use optimize_with_sampler() for Study to get sampler integration + let study1: Study = + Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999)); + let study2: Study = + Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999)); + + let mut values1 = Vec::new(); + let mut values2 = Vec::new(); + + study1 + .optimize_with_sampler(100, |trial| { + let x = trial.suggest_float("x", 0.0, 100.0)?; + values1.push(x); + Ok::<_, TpeError>(x) + }) + .unwrap(); + + study2 + .optimize_with_sampler(100, |trial| { + let x = trial.suggest_float("x", 0.0, 100.0)?; + values2.push(x); + Ok::<_, TpeError>(x) + }) + .unwrap(); + + for (i, (v1, v2)) in values1.iter().zip(values2.iter()).enumerate() { + assert_eq!( + v1, v2, + "values at trial {i} should be identical with same seed: {v1} vs {v2}" + ); + } +} + +// ============================================================================= +// Test: suggest_* methods return cached values on repeated calls +// ============================================================================= + +#[test] +fn test_suggest_float_caching() { + let mut trial = Trial::new(0); + + let x1 = trial.suggest_float("x", 0.0, 10.0).unwrap(); + let x2 = trial.suggest_float("x", 0.0, 10.0).unwrap(); + let x3 = trial.suggest_float("x", 0.0, 10.0).unwrap(); + + assert_eq!(x1, x2, "repeated suggest_float should return cached value"); + assert_eq!(x2, x3, "repeated suggest_float should return cached value"); +} + +#[test] +fn test_suggest_float_log_caching() { + let mut trial = Trial::new(0); + + let x1 = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap(); + let x2 = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap(); + + assert_eq!( + x1, x2, + "repeated suggest_float_log should return cached value" + ); +} + +#[test] +fn test_suggest_float_step_caching() { + let mut trial = Trial::new(0); + + let x1 = trial.suggest_float_step("step", 0.0, 1.0, 0.1).unwrap(); + let x2 = trial.suggest_float_step("step", 0.0, 1.0, 0.1).unwrap(); + + assert_eq!( + x1, x2, + "repeated suggest_float_step should return cached value" + ); +} + +#[test] +fn test_suggest_int_caching() { + let mut trial = Trial::new(0); + + let n1 = trial.suggest_int("n", 1, 100).unwrap(); + let n2 = trial.suggest_int("n", 1, 100).unwrap(); + + assert_eq!(n1, n2, "repeated suggest_int should return cached value"); +} + +#[test] +fn test_suggest_int_log_caching() { + let mut trial = Trial::new(0); + + let n1 = trial.suggest_int_log("batch", 1, 1024).unwrap(); + let n2 = trial.suggest_int_log("batch", 1, 1024).unwrap(); + + assert_eq!( + n1, n2, + "repeated suggest_int_log should return cached value" + ); +} + +#[test] +fn test_suggest_int_step_caching() { + let mut trial = Trial::new(0); + + let n1 = trial.suggest_int_step("units", 32, 512, 32).unwrap(); + let n2 = trial.suggest_int_step("units", 32, 512, 32).unwrap(); + + assert_eq!( + n1, n2, + "repeated suggest_int_step should return cached value" + ); +} + +#[test] +fn test_suggest_categorical_caching() { + let mut trial = Trial::new(0); + + let choices = ["sgd", "adam", "rmsprop"]; + let c1 = trial.suggest_categorical("optimizer", &choices).unwrap(); + let c2 = trial.suggest_categorical("optimizer", &choices).unwrap(); + + assert_eq!( + c1, c2, + "repeated suggest_categorical should return cached value" + ); +} + +#[test] +fn test_multiple_parameters_independent_caching() { + let mut trial = Trial::new(0); + + // Suggest multiple parameters + let x = trial.suggest_float("x", 0.0, 1.0).unwrap(); + let y = trial.suggest_float("y", 0.0, 1.0).unwrap(); + let n = trial.suggest_int("n", 1, 10).unwrap(); + let opt = trial.suggest_categorical("opt", &["a", "b"]).unwrap(); + + // All should be cached independently + assert_eq!(x, trial.suggest_float("x", 0.0, 1.0).unwrap()); + assert_eq!(y, trial.suggest_float("y", 0.0, 1.0).unwrap()); + assert_eq!(n, trial.suggest_int("n", 1, 10).unwrap()); + assert_eq!(opt, trial.suggest_categorical("opt", &["a", "b"]).unwrap()); +} + +// ============================================================================= +// Test: parameter conflict returns error +// ============================================================================= + +#[test] +fn test_parameter_conflict_float_different_bounds() { + let mut trial = Trial::new(0); + + trial.suggest_float("x", 0.0, 1.0).unwrap(); + let result = trial.suggest_float("x", 0.0, 2.0); // Different upper bound + + assert!(matches!(result, Err(TpeError::ParameterConflict { .. }))); +} + +#[test] +fn test_parameter_conflict_float_vs_log() { + let mut trial = Trial::new(0); + + trial.suggest_float("x", 0.1, 1.0).unwrap(); + let result = trial.suggest_float_log("x", 0.1, 1.0); // Same bounds but log scale + + assert!(matches!(result, Err(TpeError::ParameterConflict { .. }))); +} + +#[test] +fn test_parameter_conflict_float_vs_step() { + let mut trial = Trial::new(0); + + trial.suggest_float("x", 0.0, 1.0).unwrap(); + let result = trial.suggest_float_step("x", 0.0, 1.0, 0.1); // Same bounds but with step + + assert!(matches!(result, Err(TpeError::ParameterConflict { .. }))); +} + +#[test] +fn test_parameter_conflict_int_different_bounds() { + let mut trial = Trial::new(0); + + trial.suggest_int("n", 1, 10).unwrap(); + let result = trial.suggest_int("n", 1, 20); // Different upper bound + + assert!(matches!(result, Err(TpeError::ParameterConflict { .. }))); +} + +#[test] +fn test_parameter_conflict_int_vs_log() { + let mut trial = Trial::new(0); + + trial.suggest_int("n", 1, 100).unwrap(); + let result = trial.suggest_int_log("n", 1, 100); // Same bounds but log scale + + assert!(matches!(result, Err(TpeError::ParameterConflict { .. }))); +} + +#[test] +fn test_parameter_conflict_categorical_different_n_choices() { + let mut trial = Trial::new(0); + + trial.suggest_categorical("opt", &["a", "b", "c"]).unwrap(); + let result = trial.suggest_categorical("opt", &["x", "y"]); // Different number of choices + + assert!(matches!(result, Err(TpeError::ParameterConflict { .. }))); +} + +#[test] +fn test_parameter_conflict_float_vs_int() { + let mut trial = Trial::new(0); + + trial.suggest_float("x", 0.0, 10.0).unwrap(); + let result = trial.suggest_int("x", 0, 10); // Different type + + assert!(matches!(result, Err(TpeError::ParameterConflict { .. }))); +} + +#[test] +fn test_parameter_conflict_returns_name() { + let mut trial = Trial::new(0); + + trial.suggest_float("my_param", 0.0, 1.0).unwrap(); + let result = trial.suggest_float("my_param", 0.0, 2.0); + + match result { + Err(TpeError::ParameterConflict { name, .. }) => { + assert_eq!(name, "my_param"); + } + _ => panic!("expected ParameterConflict error"), + } +} + +// ============================================================================= +// Test: empty categorical returns error +// ============================================================================= + +#[test] +fn test_empty_categorical_returns_error() { + let mut trial = Trial::new(0); + let empty: &[&str] = &[]; + + let result = trial.suggest_categorical("opt", empty); + + assert!(matches!(result, Err(TpeError::EmptyChoices))); +} + +#[test] +fn test_empty_categorical_vec_returns_error() { + let mut trial = Trial::new(0); + let empty: Vec = vec![]; + + let result = trial.suggest_categorical("numbers", &empty); + + assert!(matches!(result, Err(TpeError::EmptyChoices))); +} + +// ============================================================================= +// Additional integration tests +// ============================================================================= + +#[test] +fn test_study_basic_workflow() { + let study: Study = Study::new(Direction::Minimize); + + study + .optimize(10, |trial| { + let x = trial.suggest_float("x", -5.0, 5.0)?; + Ok::<_, TpeError>(x * x) + }) + .expect("optimization should succeed"); + + assert_eq!(study.n_trials(), 10); + let best = study.best_trial().expect("should have best trial"); + assert!(best.value >= 0.0, "x^2 should be non-negative"); +} + +#[test] +fn test_study_with_failures() { + let study: Study = Study::new(Direction::Minimize); + + // Every other trial fails + let mut counter = 0; + study + .optimize(10, |trial| { + counter += 1; + if counter % 2 == 0 { + return Err::("intentional failure"); + } + let x = trial + .suggest_float("x", -5.0, 5.0) + .map_err(|_| "param error")?; + Ok(x * x) + }) + .expect("optimization should succeed with some failures"); + + // Only half the trials should have succeeded + assert_eq!(study.n_trials(), 5, "only 5 trials should have completed"); +} + +#[test] +fn test_no_completed_trials_error() { + let study: Study = Study::new(Direction::Minimize); + + let result = study.best_trial(); + assert!(matches!(result, Err(TpeError::NoCompletedTrials))); +} + +#[test] +fn test_invalid_bounds_errors() { + let mut trial = Trial::new(0); + + // low > high for float + let result = trial.suggest_float("x", 10.0, 5.0); + assert!(matches!(result, Err(TpeError::InvalidBounds { .. }))); + + // low > high for int + let result = trial.suggest_int("n", 100, 50); + assert!(matches!(result, Err(TpeError::InvalidBounds { .. }))); +} + +#[test] +fn test_invalid_log_bounds_errors() { + let mut trial = Trial::new(0); + + // low <= 0 for log float + let result = trial.suggest_float_log("x", 0.0, 1.0); + assert!(matches!(result, Err(TpeError::InvalidLogBounds))); + + let result = trial.suggest_float_log("y", -1.0, 1.0); + assert!(matches!(result, Err(TpeError::InvalidLogBounds))); + + // low < 1 for log int + let result = trial.suggest_int_log("n", 0, 100); + assert!(matches!(result, Err(TpeError::InvalidLogBounds))); +} + +#[test] +fn test_invalid_step_errors() { + let mut trial = Trial::new(0); + + // step <= 0 for float + let result = trial.suggest_float_step("x", 0.0, 1.0, 0.0); + assert!(matches!(result, Err(TpeError::InvalidStep))); + + let result = trial.suggest_float_step("y", 0.0, 1.0, -0.1); + assert!(matches!(result, Err(TpeError::InvalidStep))); + + // step <= 0 for int + let result = trial.suggest_int_step("n", 0, 100, 0); + assert!(matches!(result, Err(TpeError::InvalidStep))); +} + +#[test] +fn test_tpe_with_categorical_parameter() { + let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build(); + + let study: Study = Study::with_sampler(Direction::Maximize, sampler); + + // Optimization where the best choice depends on the categorical + study + .optimize_with_sampler(30, |trial| { + let choice = trial.suggest_categorical("model", &["linear", "quadratic", "cubic"])?; + let x = trial.suggest_float("x", 0.0, 2.0)?; + + // cubic model is best at x=1 + let value = match choice { + "linear" => x, + "quadratic" => x * x, + "cubic" => -((x - 1.0).powi(2)) + 10.0, // peak at x=1, max value 10 + _ => unreachable!(), + }; + Ok::<_, TpeError>(value) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have best trial"); + // The optimizer should find that "cubic" with x≈1 is best + assert!( + best.value > 5.0, + "should find good solution, got {}", + best.value + ); +} + +#[test] +fn test_tpe_with_integer_parameters() { + let sampler = TpeSampler::builder().seed(789).n_startup_trials(5).build(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + // Minimize (n - 7)^2 where n ∈ [1, 10] + study + .optimize_with_sampler(30, |trial| { + let n = trial.suggest_int("n", 1, 10)?; + Ok::<_, TpeError>(((n - 7) as f64).powi(2)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have best trial"); + + // Best value should be small (n close to 7) + assert!( + best.value < 5.0, + "should find n close to 7, best value = {}", + best.value + ); +} + +#[test] +fn test_callback_early_stopping() { + use std::cell::Cell; + use std::ops::ControlFlow; + + let study: Study = Study::new(Direction::Minimize); + let trials_run = Cell::new(0); + + study + .optimize_with_callback( + 100, + |trial| { + trials_run.set(trials_run.get() + 1); + let x = trial.suggest_float("x", 0.0, 10.0)?; + Ok::<_, TpeError>(x) + }, + |_study, _trial| { + // Stop after 5 trials + if trials_run.get() >= 5 { + ControlFlow::Break(()) + } else { + ControlFlow::Continue(()) + } + }, + ) + .expect("optimization should succeed"); + + assert_eq!(study.n_trials(), 5, "should have stopped after 5 trials"); +} + +#[test] +fn test_study_trials_iteration() { + let study: Study = Study::new(Direction::Minimize); + + study + .optimize(5, |trial| { + let x = trial.suggest_float("x", 0.0, 1.0)?; + Ok::<_, TpeError>(x) + }) + .unwrap(); + + let trials = study.trials(); + assert_eq!(trials.len(), 5); + + for trial in &trials { + assert!( + !trial.params.is_empty(), + "each trial should have parameters" + ); + assert!( + trial.params.contains_key("x"), + "each trial should have parameter 'x'" + ); + } +} + +#[test] +fn test_study_direction() { + let study_min: Study = Study::new(Direction::Minimize); + assert_eq!(study_min.direction(), Direction::Minimize); + + let study_max: Study = Study::new(Direction::Maximize); + assert_eq!(study_max.direction(), Direction::Maximize); +} + +#[test] +fn test_trial_state() { + use optimize::TrialState; + + let trial = Trial::new(0); + assert_eq!(trial.state(), TrialState::Running); +} + +#[test] +fn test_trial_params_access() { + let mut trial = Trial::new(0); + + trial.suggest_float("x", 0.0, 1.0).unwrap(); + trial.suggest_int("n", 1, 10).unwrap(); + + let params = trial.params(); + assert_eq!(params.len(), 2); + assert!(params.contains_key("x")); + assert!(params.contains_key("n")); +} + +#[test] +fn test_log_scale_float_range() { + let mut trial = Trial::new(0); + + let lr = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap(); + assert!( + (1e-5..=1e-1).contains(&lr), + "log-scale value {lr} out of range" + ); +} + +#[test] +fn test_step_float_snaps_to_grid() { + let mut trial = Trial::new(0); + + let x = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap(); + + // x should be one of: 0.0, 0.25, 0.5, 0.75, 1.0 + let valid_values = [0.0, 0.25, 0.5, 0.75, 1.0]; + let is_valid = valid_values.iter().any(|&v| (x - v).abs() < 1e-10); + assert!(is_valid, "stepped float {x} should snap to grid"); +} + +#[test] +fn test_step_int_snaps_to_grid() { + let mut trial = Trial::new(0); + + let n = trial.suggest_int_step("n", 0, 100, 25).unwrap(); + + // n should be one of: 0, 25, 50, 75, 100 + assert!( + n % 25 == 0 && (0..=100).contains(&n), + "stepped int {n} should snap to grid" + ); +} + +#[test] +fn test_best_value() { + let study: Study = Study::new(Direction::Minimize); + + study + .optimize(10, |trial| { + let x = trial.suggest_float("x", 0.0, 10.0)?; + Ok::<_, TpeError>(x) + }) + .unwrap(); + + let best_value = study.best_value().expect("should have best value"); + let best_trial = study.best_trial().expect("should have best trial"); + + assert_eq!( + best_value, best_trial.value, + "best_value should match best_trial.value" + ); +}