Manuel Raimann
2120 lines
65 KiB
Rust
2120 lines
65 KiB
Rust
//! 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.
|
|
//!
|
|
//! # 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<dyn GammaStrategy> {
|
|
//! 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;
|
|
use rand::{Rng, SeedableRng};
|
|
|
|
use crate::distribution::Distribution;
|
|
use crate::error::{Error, Result};
|
|
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<dyn GammaStrategy> {
|
|
/// 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<dyn GammaStrategy>;
|
|
}
|
|
|
|
impl Clone for Box<dyn GammaStrategy> {
|
|
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<Self> {
|
|
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<dyn GammaStrategy> {
|
|
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<Self> {
|
|
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<dyn GammaStrategy> {
|
|
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<Self> {
|
|
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<dyn GammaStrategy> {
|
|
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<Self> {
|
|
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<dyn GammaStrategy> {
|
|
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
|
|
/// 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.
|
|
///
|
|
/// # 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 (FixedGamma at 0.25)
|
|
/// let sampler = TpeSampler::new();
|
|
///
|
|
/// // Create with custom settings using the builder
|
|
/// let sampler = TpeSampler::builder()
|
|
/// .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 {
|
|
/// Strategy for computing the gamma quantile.
|
|
gamma_strategy: Arc<dyn GammaStrategy>,
|
|
/// 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,
|
|
/// Optional fixed bandwidth for KDE. If None, uses Scott's rule.
|
|
kde_bandwidth: Option<f64>,
|
|
/// Thread-safe RNG for sampling.
|
|
rng: Mutex<StdRng>,
|
|
}
|
|
|
|
impl TpeSampler {
|
|
/// Creates a new TPE sampler with default settings.
|
|
///
|
|
/// Default settings:
|
|
/// - 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_strategy: Arc::new(FixedGamma::default()),
|
|
n_startup_trials: 10,
|
|
n_ei_candidates: 24,
|
|
kde_bandwidth: None,
|
|
rng: Mutex::new(StdRng::from_os_rng()),
|
|
}
|
|
}
|
|
|
|
/// Creates a builder for configuring a TPE sampler.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::sampler::tpe::TpeSampler;
|
|
///
|
|
/// let sampler = TpeSampler::builder()
|
|
/// .gamma(0.15)
|
|
/// .n_startup_trials(20)
|
|
/// .n_ei_candidates(32)
|
|
/// .seed(42)
|
|
/// .build()
|
|
/// .unwrap();
|
|
/// ```
|
|
#[must_use]
|
|
pub fn builder() -> TpeSamplerBuilder {
|
|
TpeSamplerBuilder::new()
|
|
}
|
|
|
|
/// 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).
|
|
/// * `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::InvalidGamma` if gamma is not in (0.0, 1.0).
|
|
/// Returns `Error::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
|
|
pub fn with_config(
|
|
gamma: f64,
|
|
n_startup_trials: usize,
|
|
n_ei_candidates: usize,
|
|
kde_bandwidth: Option<f64>,
|
|
seed: Option<u64>,
|
|
) -> Result<Self> {
|
|
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<G: GammaStrategy + 'static>(
|
|
gamma_strategy: G,
|
|
n_startup_trials: usize,
|
|
n_ei_candidates: usize,
|
|
kde_bandwidth: Option<f64>,
|
|
seed: Option<u64>,
|
|
) -> Result<Self> {
|
|
if let Some(bw) = kde_bandwidth
|
|
&& bw <= 0.0
|
|
{
|
|
return Err(Error::InvalidBandwidth(bw));
|
|
}
|
|
|
|
let rng = match seed {
|
|
Some(s) => StdRng::seed_from_u64(s),
|
|
None => StdRng::from_os_rng(),
|
|
};
|
|
|
|
Ok(Self {
|
|
gamma_strategy: Arc::new(gamma_strategy),
|
|
n_startup_trials,
|
|
n_ei_candidates,
|
|
kde_bandwidth,
|
|
rng: Mutex::new(rng),
|
|
})
|
|
}
|
|
|
|
/// 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(
|
|
clippy::cast_precision_loss,
|
|
clippy::cast_possible_truncation,
|
|
clippy::cast_sign_loss
|
|
)]
|
|
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<usize> = (0..history.len()).collect();
|
|
sorted_indices.sort_by(|&a, &b| {
|
|
history[a]
|
|
.value
|
|
.partial_cmp(&history[b].value)
|
|
.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 * 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).
|
|
#[allow(
|
|
clippy::cast_possible_truncation,
|
|
clippy::cast_precision_loss,
|
|
clippy::unused_self
|
|
)]
|
|
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<f64>,
|
|
good_values: Vec<f64>,
|
|
bad_values: Vec<f64>,
|
|
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<f64> = good_values.iter().map(|&v| v.ln()).collect();
|
|
let b: Vec<f64> = 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 = match self.kde_bandwidth {
|
|
Some(bw) => KernelDensityEstimator::with_bandwidth(good_internal, bw),
|
|
None => KernelDensityEstimator::new(good_internal),
|
|
};
|
|
let g_kde = match self.kde_bandwidth {
|
|
Some(bw) => KernelDensityEstimator::with_bandwidth(bad_internal, bw),
|
|
None => KernelDensityEstimator::new(bad_internal),
|
|
};
|
|
|
|
// If KDE construction fails, fall back to uniform sampling
|
|
let (Ok(l_kde), Ok(g_kde)) = (l_kde, g_kde) else {
|
|
return rng.random_range(low..=high);
|
|
};
|
|
|
|
// 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,
|
|
clippy::cast_precision_loss,
|
|
clippy::cast_possible_truncation
|
|
)]
|
|
fn sample_tpe_int(
|
|
&self,
|
|
low: i64,
|
|
high: i64,
|
|
log_scale: bool,
|
|
step: Option<i64>,
|
|
good_values: &[i64],
|
|
bad_values: &[i64],
|
|
rng: &mut StdRng,
|
|
) -> i64 {
|
|
// Convert to floats for KDE
|
|
let good_floats: Vec<f64> = good_values.iter().map(|&v| v as f64).collect();
|
|
let bad_floats: Vec<f64> = 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.
|
|
#[allow(clippy::cast_precision_loss, clippy::unused_self)]
|
|
fn sample_tpe_categorical(
|
|
&self,
|
|
n_choices: usize,
|
|
good_indices: &[usize],
|
|
bad_indices: &[usize],
|
|
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::<f64>() * 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
|
|
///
|
|
/// Using a fixed gamma value:
|
|
///
|
|
/// ```
|
|
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
|
///
|
|
/// let sampler = TpeSamplerBuilder::new()
|
|
/// .gamma(0.15)
|
|
/// .n_startup_trials(20)
|
|
/// .n_ei_candidates(32)
|
|
/// .seed(42)
|
|
/// .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_strategy: Box<dyn GammaStrategy>,
|
|
/// Raw gamma value for deferred validation (Some if `gamma()` was called)
|
|
raw_gamma: Option<f64>,
|
|
n_startup_trials: usize,
|
|
n_ei_candidates: usize,
|
|
kde_bandwidth: Option<f64>,
|
|
seed: Option<u64>,
|
|
}
|
|
|
|
impl TpeSamplerBuilder {
|
|
/// Creates a new builder with default settings.
|
|
///
|
|
/// Default settings:
|
|
/// - 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)
|
|
/// - seed: None (use OS-provided entropy)
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self {
|
|
gamma_strategy: Box::new(FixedGamma::default()),
|
|
raw_gamma: None,
|
|
n_startup_trials: 10,
|
|
n_ei_candidates: 24,
|
|
kde_bandwidth: None,
|
|
seed: None,
|
|
}
|
|
}
|
|
|
|
/// 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.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `gamma` - Quantile value, must be in (0.0, 1.0).
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
|
///
|
|
/// let sampler = TpeSamplerBuilder::new()
|
|
/// .gamma(0.10) // Use top 10% as "good" trials
|
|
/// .build()
|
|
/// .unwrap();
|
|
/// ```
|
|
///
|
|
/// # Note
|
|
///
|
|
/// Validation happens at `build()` time. If gamma is not in (0.0, 1.0),
|
|
/// `build()` will return `Err(Error::InvalidGamma)`.
|
|
#[must_use]
|
|
pub fn gamma(mut self, gamma: f64) -> Self {
|
|
// 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<dyn GammaStrategy> {
|
|
/// Box::new(self.clone())
|
|
/// }
|
|
/// }
|
|
///
|
|
/// let sampler = TpeSamplerBuilder::new()
|
|
/// .gamma_strategy(MyGamma)
|
|
/// .build()
|
|
/// .unwrap();
|
|
/// ```
|
|
#[must_use]
|
|
pub fn gamma_strategy<G: GammaStrategy + 'static>(mut self, strategy: G) -> Self {
|
|
self.gamma_strategy = Box::new(strategy);
|
|
self.raw_gamma = None; // Clear any raw gamma set by 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 optimizer::sampler::tpe::TpeSamplerBuilder;
|
|
///
|
|
/// let sampler = TpeSamplerBuilder::new()
|
|
/// .n_startup_trials(20) // Random sample first 20 trials
|
|
/// .build()
|
|
/// .unwrap();
|
|
/// ```
|
|
#[must_use]
|
|
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 optimizer::sampler::tpe::TpeSamplerBuilder;
|
|
///
|
|
/// let sampler = TpeSamplerBuilder::new()
|
|
/// .n_ei_candidates(48) // Evaluate more candidates
|
|
/// .build()
|
|
/// .unwrap();
|
|
/// ```
|
|
#[must_use]
|
|
pub fn n_ei_candidates(mut self, n: usize) -> Self {
|
|
self.n_ei_candidates = n;
|
|
self
|
|
}
|
|
|
|
/// Sets a fixed bandwidth for the kernel density estimator.
|
|
///
|
|
/// By default, TPE uses Scott's rule to automatically select the bandwidth
|
|
/// based on the sample data. Use this method to override with a fixed value.
|
|
///
|
|
/// Smaller bandwidths give more localized, peaky distributions.
|
|
/// Larger bandwidths give smoother, more spread-out distributions.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `bandwidth` - The fixed bandwidth (standard deviation) for Gaussian kernels.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
|
///
|
|
/// let sampler = TpeSamplerBuilder::new()
|
|
/// .kde_bandwidth(0.5) // Fixed bandwidth of 0.5
|
|
/// .build()
|
|
/// .unwrap();
|
|
/// ```
|
|
///
|
|
/// # Note
|
|
///
|
|
/// Validation happens at `build()` time. If bandwidth is not positive,
|
|
/// `build()` will return `Err(Error::InvalidBandwidth)`.
|
|
#[must_use]
|
|
pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self {
|
|
self.kde_bandwidth = Some(bandwidth);
|
|
self
|
|
}
|
|
|
|
/// Sets a seed for reproducible sampling.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `seed` - Seed value for the random number generator.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```
|
|
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
|
///
|
|
/// let sampler = TpeSamplerBuilder::new()
|
|
/// .seed(42) // Reproducible results
|
|
/// .build()
|
|
/// .unwrap();
|
|
/// ```
|
|
#[must_use]
|
|
pub fn seed(mut self, seed: u64) -> Self {
|
|
self.seed = Some(seed);
|
|
self
|
|
}
|
|
|
|
/// Builds the configured [`TpeSampler`].
|
|
///
|
|
/// # Errors
|
|
///
|
|
/// 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
|
|
///
|
|
/// ```
|
|
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
|
///
|
|
/// let sampler = TpeSamplerBuilder::new()
|
|
/// .gamma(0.15)
|
|
/// .n_startup_trials(20)
|
|
/// .n_ei_candidates(32)
|
|
/// .seed(42)
|
|
/// .build()
|
|
/// .unwrap();
|
|
/// ```
|
|
pub fn build(self) -> Result<TpeSampler> {
|
|
// Determine the gamma strategy to use
|
|
let gamma_strategy: Arc<dyn GammaStrategy> = 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),
|
|
})
|
|
}
|
|
}
|
|
|
|
impl Default for TpeSamplerBuilder {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl Sampler for TpeSampler {
|
|
#[allow(clippy::too_many_lines)]
|
|
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<f64> = 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<f64> = 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<i64> = 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<i64> = 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<usize> = 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<usize> = 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)]
|
|
#[allow(
|
|
clippy::similar_names,
|
|
clippy::cast_sign_loss,
|
|
clippy::cast_precision_loss
|
|
)]
|
|
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();
|
|
// 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);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tpe_sampler_with_config() {
|
|
let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42)).unwrap();
|
|
// 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);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tpe_sampler_invalid_gamma_zero() {
|
|
let result = TpeSampler::with_config(0.0, 10, 24, None, None);
|
|
assert!(matches!(result, Err(Error::InvalidGamma(_))));
|
|
}
|
|
|
|
#[test]
|
|
fn test_tpe_sampler_invalid_gamma_one() {
|
|
let result = TpeSampler::with_config(1.0, 10, 24, None, None);
|
|
assert!(matches!(result, Err(Error::InvalidGamma(_))));
|
|
}
|
|
|
|
#[test]
|
|
fn test_tpe_startup_random_sampling() {
|
|
let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)).unwrap();
|
|
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<CompletedTrial> = 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, None, Some(42)).unwrap();
|
|
|
|
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<CompletedTrial> = (0..20)
|
|
.map(|i| {
|
|
create_trial(
|
|
i as u64,
|
|
f64::from(i),
|
|
vec![("x", ParamValue::Float(f64::from(i) / 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, None, Some(42)).unwrap();
|
|
|
|
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<CompletedTrial> = (0..20)
|
|
.map(|i| {
|
|
let x = f64::from(i) / 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::<f64>() / 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, None, Some(42)).unwrap();
|
|
|
|
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 });
|
|
|
|
// Create history where category 1 is consistently good
|
|
let history: Vec<CompletedTrial> = (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, None, Some(42)).unwrap();
|
|
|
|
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<CompletedTrial> = (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<CompletedTrial> = (0..20)
|
|
.map(|i| {
|
|
create_trial(
|
|
i as u64,
|
|
f64::from(i),
|
|
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
let sampler1 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345)).unwrap();
|
|
let sampler2 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345)).unwrap();
|
|
|
|
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().unwrap();
|
|
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);
|
|
}
|
|
|
|
#[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()
|
|
.unwrap();
|
|
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);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tpe_sampler_builder_via_sampler() {
|
|
let sampler = TpeSampler::builder()
|
|
.gamma(0.10)
|
|
.n_startup_trials(15)
|
|
.n_ei_candidates(48)
|
|
.build()
|
|
.unwrap();
|
|
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);
|
|
}
|
|
|
|
#[test]
|
|
fn test_tpe_sampler_builder_partial() {
|
|
// Test setting only some options
|
|
let sampler = TpeSamplerBuilder::new().gamma(0.20).build().unwrap();
|
|
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
|
|
}
|
|
|
|
#[test]
|
|
fn test_tpe_sampler_builder_invalid_gamma() {
|
|
let result = TpeSamplerBuilder::new().gamma(1.5).build();
|
|
assert!(matches!(result, Err(Error::InvalidGamma(_))));
|
|
}
|
|
|
|
#[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<CompletedTrial> = (0..20u32)
|
|
.map(|i| {
|
|
create_trial(
|
|
u64::from(i),
|
|
f64::from(i),
|
|
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
let sampler1 = TpeSampler::builder()
|
|
.seed(99999)
|
|
.n_startup_trials(5)
|
|
.build()
|
|
.unwrap();
|
|
let sampler2 = TpeSampler::builder()
|
|
.seed(99999)
|
|
.n_startup_trials(5)
|
|
.build()
|
|
.unwrap();
|
|
|
|
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"
|
|
);
|
|
}
|
|
}
|
|
|
|
// ========================================================================
|
|
// 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<dyn GammaStrategy> = Box::new(FixedGamma::new(0.3).unwrap());
|
|
let cloned = fixed.clone();
|
|
assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON);
|
|
|
|
let linear: Box<dyn GammaStrategy> = 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<CompletedTrial> = (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<CompletedTrial> = (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<dyn GammaStrategy> {
|
|
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);
|
|
}
|
|
}
|