diff --git a/README.md b/README.md index 24e131b..1116318 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,56 @@ let sampler = TpeSampler::builder() let study: Study = Study::with_sampler(Direction::Minimize, sampler); ``` +#### Gamma Strategies + +The gamma parameter controls what fraction of trials are considered "good" when building the TPE model. Instead of a fixed value, you can use adaptive strategies: + +| Strategy | Description | Formula | +|----------|-------------|---------| +| `FixedGamma` | Constant value (default: 0.25) | `γ = constant` | +| `LinearGamma` | Linear interpolation over trials | `γ = γ_min + (γ_max - γ_min) * min(n/n_max, 1)` | +| `SqrtGamma` | Optuna-style inverse sqrt scaling | `γ = min(γ_max, factor/√n / n)` | +| `HyperoptGamma` | Hyperopt-style adaptive | `γ = min(γ_max, (base + 1) / n)` | + +```rust +use optimizer::sampler::tpe::{TpeSampler, SqrtGamma, LinearGamma}; + +// Optuna-style gamma that decreases with more trials +let sampler = TpeSampler::builder() + .gamma_strategy(SqrtGamma::default()) + .build() + .unwrap(); + +// Linear interpolation from 0.1 to 0.3 over 100 trials +let sampler = TpeSampler::builder() + .gamma_strategy(LinearGamma::new(0.1, 0.3, 100).unwrap()) + .build() + .unwrap(); +``` + +You can also implement custom strategies: + +```rust +use optimizer::sampler::tpe::{TpeSampler, GammaStrategy}; + +#[derive(Debug, Clone)] +struct MyGamma { base: f64 } + +impl GammaStrategy for MyGamma { + fn gamma(&self, n_trials: usize) -> f64 { + (self.base + 0.01 * n_trials as f64).min(0.5) + } + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +let sampler = TpeSampler::builder() + .gamma_strategy(MyGamma { base: 0.1 }) + .build() + .unwrap(); +``` + ### Grid Search ```rust diff --git a/src/sampler/tpe.rs b/src/sampler/tpe.rs index b4f5a6e..9fb92e5 100644 --- a/src/sampler/tpe.rs +++ b/src/sampler/tpe.rs @@ -3,6 +3,60 @@ //! 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. +//! +//! # Gamma Strategies +//! +//! The gamma parameter controls what fraction of trials are considered "good". +//! This module provides several built-in strategies via the [`GammaStrategy`] trait: +//! +//! - [`FixedGamma`]: Constant gamma value (default: 0.25) +//! - [`LinearGamma`]: Linear interpolation between min and max based on trial count +//! - [`SqrtGamma`]: Gamma decreases as 1/√n (similar to Optuna) +//! - [`HyperoptGamma`]: Hyperopt-style adaptive gamma +//! +//! You can also implement your own strategy by implementing the [`GammaStrategy`] trait. +//! +//! # Examples +//! +//! Using a built-in gamma strategy: +//! +//! ``` +//! use optimizer::sampler::tpe::{SqrtGamma, TpeSampler}; +//! +//! let sampler = TpeSampler::builder() +//! .gamma_strategy(SqrtGamma::default()) +//! .build() +//! .unwrap(); +//! ``` +//! +//! Implementing a custom gamma strategy: +//! +//! ``` +//! use optimizer::sampler::tpe::{GammaStrategy, TpeSampler}; +//! +//! #[derive(Debug, Clone)] +//! struct MyGamma { +//! base: f64, +//! } +//! +//! impl GammaStrategy for MyGamma { +//! fn gamma(&self, n_trials: usize) -> f64 { +//! (self.base + 0.01 * n_trials as f64).min(0.5) +//! } +//! +//! fn clone_box(&self) -> Box { +//! Box::new(self.clone()) +//! } +//! } +//! +//! let sampler = TpeSampler::builder() +//! .gamma_strategy(MyGamma { base: 0.1 }) +//! .build() +//! .unwrap(); +//! ``` + +use core::fmt::Debug; +use std::sync::Arc; use parking_lot::Mutex; use rand::rngs::StdRng; @@ -14,6 +68,485 @@ use crate::kde::KernelDensityEstimator; use crate::param::ParamValue; use crate::sampler::{CompletedTrial, Sampler}; +// ============================================================================ +// Gamma Strategy Trait and Implementations +// ============================================================================ + +/// A strategy for computing the gamma quantile in TPE. +/// +/// The gamma value determines what fraction of trials are considered "good" +/// when splitting the trial history. Different strategies can adapt this +/// fraction based on the number of completed trials. +/// +/// # Implementation Notes +/// +/// - The returned gamma must be in the range (0.0, 1.0) +/// - Implementations should be deterministic for reproducibility +/// - The `clone_box` method enables trait object cloning +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::GammaStrategy; +/// +/// #[derive(Debug, Clone)] +/// struct ConstantGamma(f64); +/// +/// impl GammaStrategy for ConstantGamma { +/// fn gamma(&self, _n_trials: usize) -> f64 { +/// self.0 +/// } +/// +/// fn clone_box(&self) -> Box { +/// Box::new(self.clone()) +/// } +/// } +/// ``` +pub trait GammaStrategy: Send + Sync + Debug { + /// Computes the gamma quantile based on the number of completed trials. + /// + /// # Arguments + /// + /// * `n_trials` - The number of completed trials in the history. + /// + /// # Returns + /// + /// A gamma value in the range (0.0, 1.0). Values outside this range + /// will be clamped by the sampler. + fn gamma(&self, n_trials: usize) -> f64; + + /// Creates a boxed clone of this strategy. + /// + /// This method enables cloning of trait objects, which is necessary + /// for the builder pattern and sampler configuration. + fn clone_box(&self) -> Box; +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_box() + } +} + +/// A fixed gamma strategy that returns a constant value. +/// +/// This is the simplest strategy and the default behavior of TPE. +/// The gamma value remains constant regardless of the number of trials. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::{FixedGamma, TpeSampler}; +/// +/// // Use 15% of trials as "good" +/// let sampler = TpeSampler::builder() +/// .gamma_strategy(FixedGamma::new(0.15).unwrap()) +/// .build() +/// .unwrap(); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct FixedGamma { + gamma: f64, +} + +impl FixedGamma { + /// Creates a new fixed gamma strategy. + /// + /// # Arguments + /// + /// * `gamma` - The constant gamma value to use. + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0). + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::FixedGamma; + /// + /// let strategy = FixedGamma::new(0.25).unwrap(); + /// assert!((strategy.value() - 0.25).abs() < f64::EPSILON); + /// ``` + pub fn new(gamma: f64) -> Result { + if gamma <= 0.0 || gamma >= 1.0 { + return Err(Error::InvalidGamma(gamma)); + } + Ok(Self { gamma }) + } + + /// Returns the fixed gamma value. + #[must_use] + pub fn value(&self) -> f64 { + self.gamma + } +} + +impl Default for FixedGamma { + /// Creates a fixed gamma strategy with the default value of 0.25. + fn default() -> Self { + Self { gamma: 0.25 } + } +} + +impl GammaStrategy for FixedGamma { + fn gamma(&self, _n_trials: usize) -> f64 { + self.gamma + } + + fn clone_box(&self) -> Box { + Box::new(*self) + } +} + +/// A linear gamma strategy that interpolates between min and max values. +/// +/// The gamma value increases linearly from `gamma_min` to `gamma_max` as the +/// number of trials grows from 0 to `n_trials_max`. Beyond `n_trials_max`, +/// gamma remains at `gamma_max`. +/// +/// This strategy is useful when you want to be more explorative early on +/// (smaller gamma = fewer "good" trials) and more exploitative later +/// (larger gamma = more "good" trials). +/// +/// # Formula +/// +/// ```text +/// gamma = gamma_min + (gamma_max - gamma_min) * min(n_trials / n_trials_max, 1.0) +/// ``` +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::{GammaStrategy, LinearGamma, TpeSampler}; +/// +/// let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap(); +/// +/// // At 0 trials: gamma = 0.1 +/// assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON); +/// +/// // At 50 trials: gamma = 0.25 (midpoint) +/// assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON); +/// +/// // At 100+ trials: gamma = 0.4 +/// assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON); +/// assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct LinearGamma { + gamma_min: f64, + gamma_max: f64, + n_trials_max: usize, +} + +impl LinearGamma { + /// Creates a new linear gamma strategy. + /// + /// # Arguments + /// + /// * `gamma_min` - The minimum gamma value (at 0 trials). + /// * `gamma_max` - The maximum gamma value (at `n_trials_max` trials). + /// * `n_trials_max` - The number of trials at which gamma reaches its maximum. + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if: + /// - `gamma_min` is not in (0.0, 1.0) + /// - `gamma_max` is not in (0.0, 1.0) + /// - `gamma_min > gamma_max` + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::LinearGamma; + /// + /// // Gamma goes from 0.1 to 0.3 over 50 trials + /// let strategy = LinearGamma::new(0.1, 0.3, 50).unwrap(); + /// ``` + pub fn new(gamma_min: f64, gamma_max: f64, n_trials_max: usize) -> Result { + if gamma_min <= 0.0 || gamma_min >= 1.0 { + return Err(Error::InvalidGamma(gamma_min)); + } + if gamma_max <= 0.0 || gamma_max >= 1.0 { + return Err(Error::InvalidGamma(gamma_max)); + } + if gamma_min > gamma_max { + return Err(Error::InvalidGamma(gamma_min)); + } + Ok(Self { + gamma_min, + gamma_max, + n_trials_max, + }) + } + + /// Returns the minimum gamma value. + #[must_use] + pub fn gamma_min(&self) -> f64 { + self.gamma_min + } + + /// Returns the maximum gamma value. + #[must_use] + pub fn gamma_max(&self) -> f64 { + self.gamma_max + } + + /// Returns the number of trials at which gamma reaches its maximum. + #[must_use] + pub fn n_trials_max(&self) -> usize { + self.n_trials_max + } +} + +impl Default for LinearGamma { + /// Creates a linear gamma strategy with default values: + /// - `gamma_min`: 0.10 + /// - `gamma_max`: 0.25 + /// - `n_trials_max`: 100 + fn default() -> Self { + Self { + gamma_min: 0.10, + gamma_max: 0.25, + n_trials_max: 100, + } + } +} + +impl GammaStrategy for LinearGamma { + #[allow(clippy::cast_precision_loss)] + fn gamma(&self, n_trials: usize) -> f64 { + if self.n_trials_max == 0 { + return self.gamma_max; + } + let t = (n_trials as f64 / self.n_trials_max as f64).min(1.0); + self.gamma_min + (self.gamma_max - self.gamma_min) * t + } + + fn clone_box(&self) -> Box { + Box::new(*self) + } +} + +/// A square root gamma strategy inspired by Optuna's default behavior. +/// +/// The gamma value is computed based on the inverse square root of the number +/// of trials, providing a balance between exploration and exploitation that +/// naturally adapts as more data becomes available. +/// +/// # Formula +/// +/// ```text +/// n_good = max(1, floor(gamma_factor / sqrt(n_trials))) +/// gamma = min(gamma_max, n_good / n_trials) +/// ``` +/// +/// When `n_trials` is 0, returns `gamma_max`. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::{GammaStrategy, SqrtGamma, TpeSampler}; +/// +/// let strategy = SqrtGamma::default(); +/// +/// // Gamma decreases as trials increase +/// let g10 = strategy.gamma(10); +/// let g100 = strategy.gamma(100); +/// assert!(g10 > g100, "Gamma should decrease with more trials"); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct SqrtGamma { + gamma_factor: f64, + gamma_max: f64, +} + +impl SqrtGamma { + /// Creates a new square root gamma strategy. + /// + /// # Arguments + /// + /// * `gamma_factor` - The factor controlling how quickly gamma decreases. + /// Higher values mean more "good" trials at any given point. + /// * `gamma_max` - The maximum gamma value (used when `n_trials` is small). + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if: + /// - `gamma_factor` is not positive + /// - `gamma_max` is not in (0.0, 1.0) + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::SqrtGamma; + /// + /// let strategy = SqrtGamma::new(1.0, 0.25).unwrap(); + /// ``` + pub fn new(gamma_factor: f64, gamma_max: f64) -> Result { + if gamma_factor <= 0.0 { + return Err(Error::InvalidGamma(gamma_factor)); + } + if gamma_max <= 0.0 || gamma_max >= 1.0 { + return Err(Error::InvalidGamma(gamma_max)); + } + Ok(Self { + gamma_factor, + gamma_max, + }) + } + + /// Returns the gamma factor. + #[must_use] + pub fn gamma_factor(&self) -> f64 { + self.gamma_factor + } + + /// Returns the maximum gamma value. + #[must_use] + pub fn gamma_max(&self) -> f64 { + self.gamma_max + } +} + +impl Default for SqrtGamma { + /// Creates a square root gamma strategy with default values: + /// - `gamma_factor`: 1.0 + /// - `gamma_max`: 0.25 + fn default() -> Self { + Self { + gamma_factor: 1.0, + gamma_max: 0.25, + } + } +} + +impl GammaStrategy for SqrtGamma { + #[allow(clippy::cast_precision_loss)] + fn gamma(&self, n_trials: usize) -> f64 { + if n_trials == 0 { + return self.gamma_max; + } + let n_good = (self.gamma_factor / (n_trials as f64).sqrt()).max(1.0); + (n_good / n_trials as f64).min(self.gamma_max) + } + + fn clone_box(&self) -> Box { + Box::new(*self) + } +} + +/// A Hyperopt-style gamma strategy. +/// +/// This strategy computes gamma as `min(gamma_max, (gamma_base + 1) / n_trials)`, +/// which is inspired by the original Hyperopt TPE implementation. +/// +/// # Formula +/// +/// ```text +/// gamma = min(gamma_max, (gamma_base + 1) / n_trials) +/// ``` +/// +/// When `n_trials` is 0, returns `gamma_max`. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::{GammaStrategy, HyperoptGamma}; +/// +/// // With gamma_base=24 and gamma_max=0.5: +/// // - At n=25: gamma = min(0.5, 25/25) = 0.5 (capped) +/// // - At n=100: gamma = min(0.5, 25/100) = 0.25 +/// let strategy = HyperoptGamma::new(24.0, 0.5).unwrap(); +/// +/// // Early trials have higher gamma +/// let g50 = strategy.gamma(50); +/// let g200 = strategy.gamma(200); +/// assert!(g50 > g200, "Gamma should decrease with more trials"); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct HyperoptGamma { + gamma_base: f64, + gamma_max: f64, +} + +impl HyperoptGamma { + /// Creates a new Hyperopt-style gamma strategy. + /// + /// # Arguments + /// + /// * `gamma_base` - The base value added to 1 in the numerator. + /// * `gamma_max` - The maximum gamma value. + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if: + /// - `gamma_base` is negative + /// - `gamma_max` is not in (0.0, 1.0) + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::HyperoptGamma; + /// + /// let strategy = HyperoptGamma::new(24.0, 0.25).unwrap(); + /// ``` + pub fn new(gamma_base: f64, gamma_max: f64) -> Result { + if gamma_base < 0.0 { + return Err(Error::InvalidGamma(gamma_base)); + } + if gamma_max <= 0.0 || gamma_max >= 1.0 { + return Err(Error::InvalidGamma(gamma_max)); + } + Ok(Self { + gamma_base, + gamma_max, + }) + } + + /// Returns the gamma base value. + #[must_use] + pub fn gamma_base(&self) -> f64 { + self.gamma_base + } + + /// Returns the maximum gamma value. + #[must_use] + pub fn gamma_max(&self) -> f64 { + self.gamma_max + } +} + +impl Default for HyperoptGamma { + /// Creates a Hyperopt-style gamma strategy with default values: + /// - `gamma_base`: 24.0 + /// - `gamma_max`: 0.25 + fn default() -> Self { + Self { + gamma_base: 24.0, + gamma_max: 0.25, + } + } +} + +impl GammaStrategy for HyperoptGamma { + #[allow(clippy::cast_precision_loss)] + fn gamma(&self, n_trials: usize) -> f64 { + if n_trials == 0 { + return self.gamma_max; + } + ((self.gamma_base + 1.0) / n_trials as f64).min(self.gamma_max) + } + + fn clone_box(&self) -> Box { + Box::new(*self) + } +} + +// ============================================================================ +// TPE Sampler +// ============================================================================ + /// A Tree-Parzen Estimator (TPE) sampler for Bayesian optimization. /// /// TPE works by splitting completed trials into two groups based on their @@ -25,26 +558,47 @@ use crate::sampler::{CompletedTrial, Sampler}; /// During the startup phase (when fewer than `n_startup_trials` are completed), /// TPE falls back to random sampling to gather initial data. /// +/// # Gamma Strategies +/// +/// The gamma quantile can be configured using different strategies via the +/// [`GammaStrategy`] trait. Built-in strategies include: +/// +/// - [`FixedGamma`]: Constant gamma (default: 0.25) +/// - [`LinearGamma`]: Linear interpolation based on trial count +/// - [`SqrtGamma`]: Inverse square root scaling (Optuna-style) +/// - [`HyperoptGamma`]: Hyperopt-style adaptive gamma +/// /// # Examples /// /// ``` /// use optimizer::sampler::tpe::TpeSampler; /// -/// // Create with default settings +/// // Create with default settings (FixedGamma at 0.25) /// let sampler = TpeSampler::new(); /// /// // Create with custom settings using the builder /// let sampler = TpeSampler::builder() -/// .gamma(0.15) +/// .gamma(0.15) // Shorthand for FixedGamma::new(0.15) /// .n_startup_trials(20) /// .n_ei_candidates(32) /// .seed(42) /// .build() /// .unwrap(); /// ``` +/// +/// Using a different gamma strategy: +/// +/// ``` +/// use optimizer::sampler::tpe::{SqrtGamma, TpeSampler}; +/// +/// let sampler = TpeSampler::builder() +/// .gamma_strategy(SqrtGamma::default()) +/// .build() +/// .unwrap(); +/// ``` pub struct TpeSampler { - /// Fraction of trials to consider as "good" (gamma quantile). - gamma: f64, + /// Strategy for computing the gamma quantile. + gamma_strategy: Arc, /// 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. @@ -59,14 +613,14 @@ impl TpeSampler { /// Creates a new TPE sampler with default settings. /// /// Default settings: - /// - gamma: 0.25 (top 25% of trials are considered "good") + /// - gamma strategy: [`FixedGamma`] with gamma = 0.25 /// - `n_startup_trials`: 10 (random sampling for first 10 trials) /// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample) /// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth) #[must_use] pub fn new() -> Self { Self { - gamma: 0.25, + gamma_strategy: Arc::new(FixedGamma::default()), n_startup_trials: 10, n_ei_candidates: 24, kde_bandwidth: None, @@ -96,6 +650,10 @@ impl TpeSampler { /// Creates a new TPE sampler with custom configuration. /// + /// This method uses a fixed gamma value. For more advanced gamma strategies, + /// use [`TpeSampler::with_strategy`] or the builder pattern with + /// [`TpeSamplerBuilder::gamma_strategy`]. + /// /// # Arguments /// /// * `gamma` - Fraction of trials to consider "good" (0.0 to 1.0). @@ -115,9 +673,51 @@ impl TpeSampler { kde_bandwidth: Option, seed: Option, ) -> Result { - if gamma <= 0.0 || gamma >= 1.0 { - return Err(Error::InvalidGamma(gamma)); - } + let gamma_strategy = FixedGamma::new(gamma)?; + Self::with_strategy( + gamma_strategy, + n_startup_trials, + n_ei_candidates, + kde_bandwidth, + seed, + ) + } + + /// Creates a new TPE sampler with a custom gamma strategy. + /// + /// # Arguments + /// + /// * `gamma_strategy` - The strategy for computing the gamma quantile. + /// * `n_startup_trials` - Number of random trials before TPE sampling. + /// * `n_ei_candidates` - Number of candidates to evaluate per sample. + /// * `kde_bandwidth` - Optional fixed bandwidth for KDE. If None, uses Scott's rule. + /// * `seed` - Optional seed for reproducibility. + /// + /// # Errors + /// + /// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::{SqrtGamma, TpeSampler}; + /// + /// let sampler = TpeSampler::with_strategy( + /// SqrtGamma::default(), + /// 10, // n_startup_trials + /// 24, // n_ei_candidates + /// None, // kde_bandwidth + /// Some(42), // seed + /// ) + /// .unwrap(); + /// ``` + pub fn with_strategy( + gamma_strategy: G, + n_startup_trials: usize, + n_ei_candidates: usize, + kde_bandwidth: Option, + seed: Option, + ) -> Result { if let Some(bw) = kde_bandwidth && bw <= 0.0 { @@ -130,7 +730,7 @@ impl TpeSampler { }; Ok(Self { - gamma, + gamma_strategy: Arc::new(gamma_strategy), n_startup_trials, n_ei_candidates, kde_bandwidth, @@ -138,8 +738,16 @@ impl TpeSampler { }) } + /// Returns the gamma strategy used by this sampler. + #[must_use] + pub fn gamma_strategy(&self) -> &dyn GammaStrategy { + self.gamma_strategy.as_ref() + } + /// Splits trials into good and bad groups based on the gamma quantile. /// + /// The gamma value is computed dynamically using the configured [`GammaStrategy`]. + /// /// Returns (`good_trials`, `bad_trials`) where `good_trials` contains trials /// with values below the gamma quantile (for minimization). #[allow( @@ -164,9 +772,15 @@ impl TpeSampler { .unwrap_or(core::cmp::Ordering::Equal) }); + // Compute gamma using the strategy and clamp to valid range + let gamma = self + .gamma_strategy + .gamma(history.len()) + .clamp(f64::EPSILON, 1.0 - f64::EPSILON); + // 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) + let n_good = ((history.len() as f64 * gamma).ceil() as usize) .max(1) .min(history.len() - 1); @@ -421,6 +1035,8 @@ impl Default for TpeSampler { /// /// # Examples /// +/// Using a fixed gamma value: +/// /// ``` /// use optimizer::sampler::tpe::TpeSamplerBuilder; /// @@ -432,9 +1048,23 @@ impl Default for TpeSampler { /// .build() /// .unwrap(); /// ``` +/// +/// Using a custom gamma strategy: +/// +/// ``` +/// use optimizer::sampler::tpe::{SqrtGamma, TpeSamplerBuilder}; +/// +/// let sampler = TpeSamplerBuilder::new() +/// .gamma_strategy(SqrtGamma::default()) +/// .n_startup_trials(20) +/// .build() +/// .unwrap(); +/// ``` #[derive(Debug, Clone)] pub struct TpeSamplerBuilder { - gamma: f64, + gamma_strategy: Box, + /// Raw gamma value for deferred validation (Some if `gamma()` was called) + raw_gamma: Option, n_startup_trials: usize, n_ei_candidates: usize, kde_bandwidth: Option, @@ -445,7 +1075,7 @@ impl TpeSamplerBuilder { /// Creates a new builder with default settings. /// /// Default settings: - /// - gamma: 0.25 (top 25% of trials are considered "good") + /// - gamma strategy: [`FixedGamma`] with gamma = 0.25 /// - `n_startup_trials`: 10 (random sampling for first 10 trials) /// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample) /// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth) @@ -453,7 +1083,8 @@ impl TpeSamplerBuilder { #[must_use] pub fn new() -> Self { Self { - gamma: 0.25, + gamma_strategy: Box::new(FixedGamma::default()), + raw_gamma: None, n_startup_trials: 10, n_ei_candidates: 24, kde_bandwidth: None, @@ -461,7 +1092,10 @@ impl TpeSamplerBuilder { } } - /// Sets the gamma quantile for splitting trials into good/bad groups. + /// Sets a fixed gamma value for splitting trials into good/bad groups. + /// + /// This is a convenience method that creates a [`FixedGamma`] strategy. + /// For more advanced gamma strategies, use [`gamma_strategy`](Self::gamma_strategy). /// /// 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. @@ -487,7 +1121,68 @@ impl TpeSamplerBuilder { /// `build()` will return `Err(Error::InvalidGamma)`. #[must_use] pub fn gamma(mut self, gamma: f64) -> Self { - self.gamma = gamma; + // We defer validation to build() time for consistency with the existing API + // Store the raw value for validation later + self.raw_gamma = Some(gamma); + self + } + + /// Sets a custom gamma strategy for splitting trials into good/bad groups. + /// + /// The gamma strategy determines what fraction of trials are considered + /// "good" based on the number of completed trials. This allows the gamma + /// value to adapt dynamically during optimization. + /// + /// # Arguments + /// + /// * `strategy` - A type implementing [`GammaStrategy`]. + /// + /// # Examples + /// + /// Using built-in strategies: + /// + /// ``` + /// use optimizer::sampler::tpe::{LinearGamma, SqrtGamma, TpeSamplerBuilder}; + /// + /// // Square root strategy (Optuna-style) + /// let sampler = TpeSamplerBuilder::new() + /// .gamma_strategy(SqrtGamma::default()) + /// .build() + /// .unwrap(); + /// + /// // Linear interpolation strategy + /// let sampler = TpeSamplerBuilder::new() + /// .gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap()) + /// .build() + /// .unwrap(); + /// ``` + /// + /// Using a custom strategy: + /// + /// ``` + /// use optimizer::sampler::tpe::{GammaStrategy, TpeSamplerBuilder}; + /// + /// #[derive(Debug, Clone)] + /// struct MyGamma; + /// + /// impl GammaStrategy for MyGamma { + /// fn gamma(&self, n_trials: usize) -> f64 { + /// 0.25 // Always return 0.25 + /// } + /// fn clone_box(&self) -> Box { + /// Box::new(self.clone()) + /// } + /// } + /// + /// let sampler = TpeSamplerBuilder::new() + /// .gamma_strategy(MyGamma) + /// .build() + /// .unwrap(); + /// ``` + #[must_use] + pub fn gamma_strategy(mut self, strategy: G) -> Self { + self.gamma_strategy = Box::new(strategy); + self.raw_gamma = None; // Clear any raw gamma set by gamma() self } @@ -602,7 +1297,7 @@ impl TpeSamplerBuilder { /// /// # Errors /// - /// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0). + /// Returns `Error::InvalidGamma` if a fixed gamma value was set and is not in (0.0, 1.0). /// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive. /// /// # Examples @@ -619,13 +1314,33 @@ impl TpeSamplerBuilder { /// .unwrap(); /// ``` pub fn build(self) -> Result { - TpeSampler::with_config( - self.gamma, - self.n_startup_trials, - self.n_ei_candidates, - self.kde_bandwidth, - self.seed, - ) + // Determine the gamma strategy to use + let gamma_strategy: Arc = if let Some(raw) = self.raw_gamma { + // Validate and create FixedGamma from raw value + Arc::new(FixedGamma::new(raw)?) + } else { + Arc::from(self.gamma_strategy) + }; + + // Validate bandwidth + if let Some(bw) = self.kde_bandwidth + && bw <= 0.0 + { + return Err(Error::InvalidBandwidth(bw)); + } + + let rng = match self.seed { + Some(s) => StdRng::seed_from_u64(s), + None => StdRng::from_os_rng(), + }; + + Ok(TpeSampler { + gamma_strategy, + n_startup_trials: self.n_startup_trials, + n_ei_candidates: self.n_ei_candidates, + kde_bandwidth: self.kde_bandwidth, + rng: Mutex::new(rng), + }) } } @@ -801,7 +1516,8 @@ mod tests { #[test] fn test_tpe_sampler_new() { let sampler = TpeSampler::new(); - assert!((sampler.gamma - 0.25).abs() < f64::EPSILON); + // Default uses FixedGamma with 0.25 + assert!((sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON); assert_eq!(sampler.n_startup_trials, 10); assert_eq!(sampler.n_ei_candidates, 24); } @@ -809,7 +1525,8 @@ mod tests { #[test] fn test_tpe_sampler_with_config() { let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42)).unwrap(); - assert!((sampler.gamma - 0.15).abs() < f64::EPSILON); + // with_config uses FixedGamma + assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON); assert_eq!(sampler.n_startup_trials, 20); assert_eq!(sampler.n_ei_candidates, 32); } @@ -1033,7 +1750,7 @@ mod tests { fn test_tpe_sampler_builder_default() { let builder = TpeSamplerBuilder::new(); let sampler = builder.build().unwrap(); - assert!((sampler.gamma - 0.25).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON); assert_eq!(sampler.n_startup_trials, 10); assert_eq!(sampler.n_ei_candidates, 24); } @@ -1047,7 +1764,7 @@ mod tests { .seed(42) .build() .unwrap(); - assert!((sampler.gamma - 0.15).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON); assert_eq!(sampler.n_startup_trials, 20); assert_eq!(sampler.n_ei_candidates, 32); } @@ -1060,7 +1777,7 @@ mod tests { .n_ei_candidates(48) .build() .unwrap(); - assert!((sampler.gamma - 0.10).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(0) - 0.10).abs() < f64::EPSILON); assert_eq!(sampler.n_startup_trials, 15); assert_eq!(sampler.n_ei_candidates, 48); } @@ -1069,7 +1786,7 @@ mod tests { fn test_tpe_sampler_builder_partial() { // Test setting only some options let sampler = TpeSamplerBuilder::new().gamma(0.20).build().unwrap(); - assert!((sampler.gamma - 0.20).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(0) - 0.20).abs() < f64::EPSILON); assert_eq!(sampler.n_startup_trials, 10); // default assert_eq!(sampler.n_ei_candidates, 24); // default } @@ -1119,4 +1836,284 @@ mod tests { ); } } + + // ======================================================================== + // Gamma Strategy Tests + // ======================================================================== + + #[test] + fn test_fixed_gamma_default() { + let strategy = FixedGamma::default(); + assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); + assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); + assert!((strategy.value() - 0.25).abs() < f64::EPSILON); + } + + #[test] + fn test_fixed_gamma_custom() { + let strategy = FixedGamma::new(0.15).unwrap(); + assert!((strategy.gamma(0) - 0.15).abs() < f64::EPSILON); + assert!((strategy.gamma(50) - 0.15).abs() < f64::EPSILON); + assert!((strategy.gamma(1000) - 0.15).abs() < f64::EPSILON); + } + + #[test] + fn test_fixed_gamma_invalid() { + assert!(FixedGamma::new(0.0).is_err()); + assert!(FixedGamma::new(1.0).is_err()); + assert!(FixedGamma::new(-0.1).is_err()); + assert!(FixedGamma::new(1.5).is_err()); + } + + #[test] + fn test_linear_gamma_default() { + let strategy = LinearGamma::default(); + assert!((strategy.gamma(0) - 0.10).abs() < f64::EPSILON); + assert!((strategy.gamma(50) - 0.175).abs() < f64::EPSILON); // midpoint + assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); + assert!((strategy.gamma(200) - 0.25).abs() < f64::EPSILON); // capped + } + + #[test] + fn test_linear_gamma_custom() { + let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap(); + assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON); + assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON); + assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON); + assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON); + } + + #[test] + fn test_linear_gamma_invalid() { + assert!(LinearGamma::new(0.0, 0.5, 100).is_err()); + assert!(LinearGamma::new(0.1, 1.0, 100).is_err()); + assert!(LinearGamma::new(0.5, 0.2, 100).is_err()); // min > max + } + + #[test] + fn test_sqrt_gamma_default() { + let strategy = SqrtGamma::default(); + // At n=0, returns gamma_max + assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); + + // gamma decreases with more trials + let g10 = strategy.gamma(10); + let g100 = strategy.gamma(100); + assert!(g10 > g100); + } + + #[test] + fn test_sqrt_gamma_custom() { + let strategy = SqrtGamma::new(2.0, 0.5).unwrap(); + assert!((strategy.gamma(0) - 0.5).abs() < f64::EPSILON); + + // At n=4: n_good = max(1, 2/2) = 1, gamma = 1/4 = 0.25 + let g4 = strategy.gamma(4); + assert!((g4 - 0.25).abs() < f64::EPSILON); + } + + #[test] + fn test_sqrt_gamma_invalid() { + assert!(SqrtGamma::new(0.0, 0.25).is_err()); // factor must be positive + assert!(SqrtGamma::new(-1.0, 0.25).is_err()); + assert!(SqrtGamma::new(1.0, 0.0).is_err()); + assert!(SqrtGamma::new(1.0, 1.0).is_err()); + } + + #[test] + fn test_hyperopt_gamma_default() { + let strategy = HyperoptGamma::default(); + // At n=0, returns gamma_max + assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); + + // At n=100: (24+1)/100 = 0.25, so capped to 0.25 + assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); + + // At n=200: (24+1)/200 = 0.125 + assert!((strategy.gamma(200) - 0.125).abs() < f64::EPSILON); + } + + #[test] + fn test_hyperopt_gamma_custom() { + let strategy = HyperoptGamma::new(9.0, 0.5).unwrap(); + // At n=20: (9+1)/20 = 0.5, capped to 0.5 + assert!((strategy.gamma(20) - 0.5).abs() < f64::EPSILON); + + // At n=100: (9+1)/100 = 0.1 + assert!((strategy.gamma(100) - 0.1).abs() < f64::EPSILON); + } + + #[test] + fn test_hyperopt_gamma_invalid() { + assert!(HyperoptGamma::new(-1.0, 0.25).is_err()); + assert!(HyperoptGamma::new(24.0, 0.0).is_err()); + assert!(HyperoptGamma::new(24.0, 1.0).is_err()); + } + + #[test] + fn test_gamma_strategy_clone_box() { + let fixed: Box = Box::new(FixedGamma::new(0.3).unwrap()); + let cloned = fixed.clone(); + assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON); + + let linear: Box = Box::new(LinearGamma::default()); + let cloned = linear.clone(); + assert!((cloned.gamma(0) - 0.10).abs() < f64::EPSILON); + } + + #[test] + fn test_tpe_with_sqrt_gamma_strategy() { + let sampler = TpeSampler::builder() + .gamma_strategy(SqrtGamma::default()) + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + 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, + f64::from(i), + vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())], + ) + }) + .collect(); + + // Should be able to sample with the sqrt gamma strategy + let value = sampler.sample(&dist, 100, &history); + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(&v)); + } else { + panic!("Expected Float value"); + } + } + + #[test] + fn test_tpe_with_linear_gamma_strategy() { + let sampler = TpeSampler::builder() + .gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap()) + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + // Verify the strategy is applied + let g = sampler.gamma_strategy().gamma(25); + assert!((g - 0.2).abs() < f64::EPSILON); // midpoint of 0.1 to 0.3 + } + + #[test] + fn test_tpe_with_hyperopt_gamma_strategy() { + let sampler = TpeSampler::builder() + .gamma_strategy(HyperoptGamma::default()) + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + 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, + f64::from(i), + vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())], + ) + }) + .collect(); + + // Should be able to sample with the hyperopt gamma strategy + let value = sampler.sample(&dist, 100, &history); + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(&v)); + } else { + panic!("Expected Float value"); + } + } + + #[test] + fn test_gamma_overrides_gamma_strategy() { + // When gamma() is called after gamma_strategy(), it should take precedence + let sampler = TpeSampler::builder() + .gamma_strategy(SqrtGamma::default()) + .gamma(0.15) // This should override + .build() + .unwrap(); + + // Should use fixed gamma of 0.15 + assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(100) - 0.15).abs() < f64::EPSILON); + } + + #[test] + fn test_gamma_strategy_overrides_gamma() { + // When gamma_strategy() is called after gamma(), it should take precedence + let sampler = TpeSampler::builder() + .gamma(0.15) + .gamma_strategy(SqrtGamma::default()) // This should override + .build() + .unwrap(); + + // Should use SqrtGamma - gamma decreases with trials + let g10 = sampler.gamma_strategy().gamma(10); + let g100 = sampler.gamma_strategy().gamma(100); + assert!(g10 > g100, "SqrtGamma should decrease with more trials"); + } + + #[test] + fn test_with_strategy_constructor() { + let sampler = TpeSampler::with_strategy( + LinearGamma::new(0.1, 0.4, 100).unwrap(), + 15, + 32, + None, + Some(42), + ) + .unwrap(); + + assert_eq!(sampler.n_startup_trials, 15); + assert_eq!(sampler.n_ei_candidates, 32); + assert!((sampler.gamma_strategy().gamma(0) - 0.1).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(100) - 0.4).abs() < f64::EPSILON); + } + + #[test] + fn test_custom_gamma_strategy() { + #[derive(Debug, Clone)] + struct DoubleGamma; + + impl GammaStrategy for DoubleGamma { + fn gamma(&self, n_trials: usize) -> f64 { + // Double the trial count-based calculation, capped at 0.5 + (0.01 * n_trials as f64).min(0.5) + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + } + + let sampler = TpeSampler::builder() + .gamma_strategy(DoubleGamma) + .build() + .unwrap(); + + assert!((sampler.gamma_strategy().gamma(10) - 0.1).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(50) - 0.5).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(100) - 0.5).abs() < f64::EPSILON); + } }