feat(hyperopt): implement TPE (Tree-Parzen Estimator) optimizer with KDE and LHS
Add a standalone TPE optimizer for Bayesian hyperparameter optimization that replaces PSO by modeling good/bad trial distributions with per-dimension kernel density estimates. Includes Latin Hypercube Sampling for initial exploration, Silverman bandwidth selection, log-sum-exp numerical stability, and JSON-based trial persistence. 15 tests covering splits, bounds, KDE density, convergence on 1D/2D objectives, and history serialization. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -47,6 +47,7 @@ pub mod optimizer;
|
||||
pub mod paths;
|
||||
pub mod sensitivity;
|
||||
pub mod shared_data;
|
||||
pub mod tpe;
|
||||
pub mod traits;
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
733
crates/ml/src/hyperopt/tpe.rs
Normal file
733
crates/ml/src/hyperopt/tpe.rs
Normal file
@@ -0,0 +1,733 @@
|
||||
//! Tree-Parzen Estimator (TPE) for Bayesian Hyperparameter Optimization
|
||||
//!
|
||||
//! TPE builds separate density models for "good" and "bad" trial parameters,
|
||||
//! then suggests new points by maximizing Expected Improvement = l(x)/g(x).
|
||||
//!
|
||||
//! ## Algorithm
|
||||
//!
|
||||
//! 1. Run N initial random trials (Latin Hypercube Sampling)
|
||||
//! 2. After enough data, split trials into "good" (top gamma=25%) and "bad" (bottom 75%)
|
||||
//! 3. Build separate per-dimension kernel density estimates (KDEs) for good and bad groups
|
||||
//! 4. Suggest next point by sampling from the "good" KDE and scoring by EI = l(x)/g(x)
|
||||
//! 5. The optimizer MINIMIZES the objective (lower = better)
|
||||
//!
|
||||
//! ## References
|
||||
//!
|
||||
//! Bergstra, J., Bardenet, R., Bengio, Y., & Kegl, B. (2011).
|
||||
//! "Algorithms for Hyper-Parameter Optimization." NeurIPS.
|
||||
|
||||
use rand::prelude::*;
|
||||
use rand::SeedableRng;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::Path;
|
||||
|
||||
/// A single completed trial recording parameters and objective value.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TrialRecord {
|
||||
/// Parameter values in continuous space.
|
||||
pub params: Vec<f64>,
|
||||
/// Objective value (lower is better).
|
||||
pub objective: f64,
|
||||
}
|
||||
|
||||
/// TPE optimizer configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TpeConfig {
|
||||
/// Number of dimensions in parameter space.
|
||||
pub n_dims: usize,
|
||||
/// Maximum number of trials.
|
||||
pub max_trials: usize,
|
||||
/// Number of initial LHS samples before TPE kicks in.
|
||||
pub n_initial: usize,
|
||||
/// Quantile for good/bad split (0.25 = top 25% are "good").
|
||||
pub gamma: f64,
|
||||
/// Number of candidates to evaluate for EI.
|
||||
pub n_candidates: usize,
|
||||
/// Random seed (optional).
|
||||
pub seed: Option<u64>,
|
||||
}
|
||||
|
||||
impl TpeConfig {
|
||||
/// Create a new TPE configuration with sensible defaults.
|
||||
///
|
||||
/// - `n_initial` is max(5, n_dims) to ensure enough initial exploration.
|
||||
/// - `gamma` defaults to 0.25 (top 25% are "good").
|
||||
/// - `n_candidates` defaults to 100 EI candidates per suggestion.
|
||||
pub fn new(n_dims: usize, max_trials: usize) -> Self {
|
||||
Self {
|
||||
n_dims,
|
||||
max_trials,
|
||||
n_initial: 5.max(n_dims),
|
||||
gamma: 0.25,
|
||||
n_candidates: 100,
|
||||
seed: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tree-Parzen Estimator optimizer.
|
||||
///
|
||||
/// Maintains a history of completed trials and uses kernel density estimation
|
||||
/// to suggest promising new parameter configurations.
|
||||
#[derive(Debug)]
|
||||
pub struct TpeOptimizer {
|
||||
config: TpeConfig,
|
||||
/// Completed trial history.
|
||||
pub trials: Vec<TrialRecord>,
|
||||
rng: StdRng,
|
||||
}
|
||||
|
||||
impl TpeOptimizer {
|
||||
/// Create a new TPE optimizer from configuration.
|
||||
pub fn new(config: TpeConfig) -> Self {
|
||||
let rng = match config.seed {
|
||||
Some(seed) => StdRng::seed_from_u64(seed),
|
||||
None => StdRng::from_entropy(),
|
||||
};
|
||||
Self {
|
||||
config,
|
||||
trials: Vec::new(),
|
||||
rng,
|
||||
}
|
||||
}
|
||||
|
||||
/// Suggest the next parameter configuration to evaluate.
|
||||
///
|
||||
/// During the initial phase (fewer trials than `n_initial`), returns a
|
||||
/// Latin Hypercube sample. After that, uses TPE to suggest the point
|
||||
/// with the highest expected improvement.
|
||||
pub fn suggest(&mut self, bounds: &[(f64, f64)]) -> Vec<f64> {
|
||||
if self.trials.len() < self.config.n_initial {
|
||||
// Initial exploration phase: return single LHS sample
|
||||
let mut samples = self.latin_hypercube_sample(bounds, 1);
|
||||
return samples.pop().unwrap_or_else(|| {
|
||||
// Fallback: uniform random sample
|
||||
bounds.iter().map(|&(lo, hi)| self.rng.gen_range(lo..=hi)).collect()
|
||||
});
|
||||
}
|
||||
|
||||
// TPE phase: split, build KDEs, sample candidates, pick best EI
|
||||
// Clone parameter vectors to break the borrow on self.trials so we can
|
||||
// mutably borrow self.rng via sample_from_kde.
|
||||
let (good, bad) = self.split_trials();
|
||||
let good_params: Vec<Vec<f64>> = good.iter().map(|t| t.params.clone()).collect();
|
||||
let bad_params: Vec<Vec<f64>> = bad.iter().map(|t| t.params.clone()).collect();
|
||||
let good_refs: Vec<&Vec<f64>> = good_params.iter().collect();
|
||||
let bad_refs: Vec<&Vec<f64>> = bad_params.iter().collect();
|
||||
|
||||
// Generate candidates by sampling from the good KDE
|
||||
let mut best_candidate = self.sample_from_kde(&good_refs, bounds);
|
||||
let mut best_ei = Self::expected_improvement_static(&best_candidate, &good_refs, &bad_refs, bounds);
|
||||
|
||||
for _ in 1..self.config.n_candidates {
|
||||
let candidate = self.sample_from_kde(&good_refs, bounds);
|
||||
let ei = Self::expected_improvement_static(&candidate, &good_refs, &bad_refs, bounds);
|
||||
if ei > best_ei {
|
||||
best_ei = ei;
|
||||
best_candidate = candidate;
|
||||
}
|
||||
}
|
||||
|
||||
best_candidate
|
||||
}
|
||||
|
||||
/// Record a completed trial.
|
||||
pub fn add_trial(&mut self, params: Vec<f64>, objective: f64) {
|
||||
self.trials.push(TrialRecord { params, objective });
|
||||
}
|
||||
|
||||
/// Split trials into "good" (top gamma%) and "bad" (bottom 1-gamma%).
|
||||
///
|
||||
/// Trials are sorted by objective ascending (lower = better), so the first
|
||||
/// `ceil(n * gamma)` trials are "good".
|
||||
pub fn split_trials(&self) -> (Vec<&TrialRecord>, Vec<&TrialRecord>) {
|
||||
let mut sorted: Vec<&TrialRecord> = self.trials.iter().collect();
|
||||
sorted.sort_by(|a, b| {
|
||||
a.objective
|
||||
.partial_cmp(&b.objective)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
let n_good = (sorted.len() as f64 * self.config.gamma).ceil() as usize;
|
||||
let n_good = n_good.max(1).min(sorted.len().saturating_sub(1));
|
||||
|
||||
let (good, bad) = sorted.split_at(n_good);
|
||||
(good.to_vec(), bad.to_vec())
|
||||
}
|
||||
|
||||
/// Generate `n` Latin Hypercube samples within the given bounds.
|
||||
///
|
||||
/// For each dimension, divides [0, 1] into `n` equal bins, shuffles them,
|
||||
/// and scales to the parameter bounds.
|
||||
pub fn latin_hypercube_sample(&mut self, bounds: &[(f64, f64)], n: usize) -> Vec<Vec<f64>> {
|
||||
if n == 0 || bounds.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let n_dims = bounds.len();
|
||||
let mut samples = vec![vec![0.0; n_dims]; n];
|
||||
|
||||
for (d, &(lo, hi)) in bounds.iter().enumerate() {
|
||||
// Create permutation of bin indices
|
||||
let mut indices: Vec<usize> = (0..n).collect();
|
||||
indices.shuffle(&mut self.rng);
|
||||
|
||||
for (i, &idx) in indices.iter().enumerate() {
|
||||
// Sample uniformly within the bin
|
||||
let bin_lo = idx as f64 / n as f64;
|
||||
let bin_hi = (idx + 1) as f64 / n as f64;
|
||||
let u = self.rng.gen_range(bin_lo..bin_hi);
|
||||
// Scale to bounds
|
||||
if let Some(s) = samples.get_mut(i) {
|
||||
if let Some(v) = s.get_mut(d) {
|
||||
*v = lo + u * (hi - lo);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
samples
|
||||
}
|
||||
|
||||
/// Compute the log-PDF of a point under an independent per-dimension KDE.
|
||||
///
|
||||
/// Uses Gaussian kernels with Silverman's rule for bandwidth selection.
|
||||
/// Returns the sum of log-PDFs across dimensions (independence assumption).
|
||||
/// Uses log-sum-exp for numerical stability.
|
||||
pub fn kde_log_pdf(samples: &[&Vec<f64>], point: &[f64], bounds: &[(f64, f64)]) -> f64 {
|
||||
if samples.is_empty() {
|
||||
return f64::NEG_INFINITY;
|
||||
}
|
||||
|
||||
let n = samples.len() as f64;
|
||||
let n_dims = point.len();
|
||||
let mut total_log_pdf = 0.0;
|
||||
|
||||
for d in 0..n_dims {
|
||||
// Collect values for this dimension
|
||||
let values: Vec<f64> = samples
|
||||
.iter()
|
||||
.filter_map(|s| s.get(d).copied())
|
||||
.collect();
|
||||
|
||||
if values.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Compute bandwidth via Silverman's rule: h = 1.06 * std * n^(-1/5)
|
||||
let bandwidth = Self::silverman_bandwidth(&values, bounds.get(d).map(|&(lo, hi)| hi - lo));
|
||||
|
||||
// Compute log-PDF at point[d] using log-sum-exp
|
||||
let x = match point.get(d) {
|
||||
Some(&v) => v,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
// log-sum-exp: log(sum(exp(log_k_i))) where log_k_i = -0.5 * ((x - mu_i) / h)^2 - log(h) - 0.5*log(2pi)
|
||||
let log_norm = -0.5_f64 * (2.0 * std::f64::consts::PI).ln() - bandwidth.ln();
|
||||
let log_kernels: Vec<f64> = values
|
||||
.iter()
|
||||
.map(|&mu| {
|
||||
let z = (x - mu) / bandwidth;
|
||||
log_norm - 0.5 * z * z
|
||||
})
|
||||
.collect();
|
||||
|
||||
let log_pdf_d = Self::log_sum_exp(&log_kernels) - n.ln();
|
||||
total_log_pdf += log_pdf_d;
|
||||
}
|
||||
|
||||
total_log_pdf
|
||||
}
|
||||
|
||||
/// Compute the expected improvement score (in log space) for a candidate point.
|
||||
///
|
||||
/// EI = l(x) / g(x), so log(EI) = log_l(x) - log_g(x).
|
||||
/// Higher is better.
|
||||
pub fn expected_improvement(
|
||||
&self,
|
||||
point: &[f64],
|
||||
good_params: &[&Vec<f64>],
|
||||
bad_params: &[&Vec<f64>],
|
||||
bounds: &[(f64, f64)],
|
||||
) -> f64 {
|
||||
Self::expected_improvement_static(point, good_params, bad_params, bounds)
|
||||
}
|
||||
|
||||
/// Static version of expected_improvement (no &self borrow needed).
|
||||
fn expected_improvement_static(
|
||||
point: &[f64],
|
||||
good_params: &[&Vec<f64>],
|
||||
bad_params: &[&Vec<f64>],
|
||||
bounds: &[(f64, f64)],
|
||||
) -> f64 {
|
||||
let log_l = Self::kde_log_pdf(good_params, point, bounds);
|
||||
let log_g = Self::kde_log_pdf(bad_params, point, bounds);
|
||||
log_l - log_g
|
||||
}
|
||||
|
||||
/// Save trial history to a JSON file.
|
||||
pub fn save_history(&self, path: &Path) -> Result<(), std::io::Error> {
|
||||
let json = serde_json::to_string_pretty(&self.trials).map_err(|e| {
|
||||
std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())
|
||||
})?;
|
||||
std::fs::write(path, json)
|
||||
}
|
||||
|
||||
/// Load trial history from a JSON file.
|
||||
///
|
||||
/// Returns the number of trials loaded. Loaded trials are appended to
|
||||
/// any existing trials in the optimizer.
|
||||
pub fn load_history(&mut self, path: &Path) -> Result<usize, std::io::Error> {
|
||||
let json = std::fs::read_to_string(path)?;
|
||||
let records: Vec<TrialRecord> = serde_json::from_str(&json).map_err(|e| {
|
||||
std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string())
|
||||
})?;
|
||||
let count = records.len();
|
||||
self.trials.extend(records);
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Return a reference to the best trial (lowest objective), if any.
|
||||
pub fn best_trial(&self) -> Option<&TrialRecord> {
|
||||
self.trials
|
||||
.iter()
|
||||
.min_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap_or(std::cmp::Ordering::Equal))
|
||||
}
|
||||
|
||||
/// Return the number of completed trials.
|
||||
pub fn n_trials(&self) -> usize {
|
||||
self.trials.len()
|
||||
}
|
||||
|
||||
// --- Private helpers ---
|
||||
|
||||
/// Compute Silverman's bandwidth for a set of 1-D samples.
|
||||
///
|
||||
/// h = 1.06 * std * n^(-1/5), with a floor of range/100 to avoid
|
||||
/// degenerate zero-bandwidth when all samples are identical.
|
||||
fn silverman_bandwidth(values: &[f64], range: Option<f64>) -> f64 {
|
||||
let n = values.len() as f64;
|
||||
if n < 1.0 {
|
||||
return 1.0;
|
||||
}
|
||||
|
||||
let mean = values.iter().sum::<f64>() / n;
|
||||
let variance = values.iter().map(|&v| (v - mean) * (v - mean)).sum::<f64>() / n;
|
||||
let std_dev = variance.sqrt();
|
||||
|
||||
let bandwidth = 1.06 * std_dev * n.powf(-0.2);
|
||||
|
||||
// Floor: avoid zero bandwidth (when all samples are the same)
|
||||
let floor = match range {
|
||||
Some(r) if r > 0.0 => r / 100.0,
|
||||
_ => 0.01,
|
||||
};
|
||||
|
||||
bandwidth.max(floor)
|
||||
}
|
||||
|
||||
/// Numerically stable log-sum-exp: log(sum(exp(values))).
|
||||
fn log_sum_exp(values: &[f64]) -> f64 {
|
||||
if values.is_empty() {
|
||||
return f64::NEG_INFINITY;
|
||||
}
|
||||
|
||||
let max_val = values
|
||||
.iter()
|
||||
.copied()
|
||||
.fold(f64::NEG_INFINITY, f64::max);
|
||||
|
||||
if max_val == f64::NEG_INFINITY {
|
||||
return f64::NEG_INFINITY;
|
||||
}
|
||||
|
||||
let sum_exp: f64 = values.iter().map(|&v| (v - max_val).exp()).sum();
|
||||
max_val + sum_exp.ln()
|
||||
}
|
||||
|
||||
/// Sample a point from a KDE defined by the given samples.
|
||||
///
|
||||
/// For each dimension independently:
|
||||
/// 1. Pick a random sample (kernel center)
|
||||
/// 2. Add Gaussian noise with Silverman bandwidth
|
||||
/// 3. Clamp to bounds
|
||||
fn sample_from_kde(&mut self, samples: &[&Vec<f64>], bounds: &[(f64, f64)]) -> Vec<f64> {
|
||||
let n_dims = bounds.len();
|
||||
let mut point = vec![0.0; n_dims];
|
||||
|
||||
if samples.is_empty() {
|
||||
// Fallback: uniform random
|
||||
for (d, &(lo, hi)) in bounds.iter().enumerate() {
|
||||
if let Some(v) = point.get_mut(d) {
|
||||
*v = self.rng.gen_range(lo..=hi);
|
||||
}
|
||||
}
|
||||
return point;
|
||||
}
|
||||
|
||||
for (d, &(lo, hi)) in bounds.iter().enumerate() {
|
||||
// Collect values for this dimension
|
||||
let values: Vec<f64> = samples
|
||||
.iter()
|
||||
.filter_map(|s| s.get(d).copied())
|
||||
.collect();
|
||||
|
||||
if values.is_empty() {
|
||||
if let Some(v) = point.get_mut(d) {
|
||||
*v = self.rng.gen_range(lo..=hi);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let bandwidth = Self::silverman_bandwidth(&values, Some(hi - lo));
|
||||
|
||||
// Pick a random kernel center
|
||||
let center_idx = self.rng.gen_range(0..values.len());
|
||||
let center = values.get(center_idx).copied().unwrap_or((lo + hi) / 2.0);
|
||||
|
||||
// Sample from Gaussian kernel and clamp to bounds
|
||||
let noise: f64 = self.rng.sample::<f64, _>(rand::distributions::Standard) * bandwidth;
|
||||
let raw = center + noise;
|
||||
if let Some(v) = point.get_mut(d) {
|
||||
*v = raw.clamp(lo, hi);
|
||||
}
|
||||
}
|
||||
|
||||
point
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_tpe_split_good_bad() {
|
||||
let config = TpeConfig::new(2, 10);
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
tpe.add_trial(vec![0.5, 0.5], 10.0);
|
||||
tpe.add_trial(vec![0.3, 0.7], 5.0);
|
||||
tpe.add_trial(vec![0.8, 0.2], 20.0);
|
||||
tpe.add_trial(vec![0.2, 0.8], 1.0);
|
||||
|
||||
let (good, bad) = tpe.split_trials();
|
||||
assert_eq!(good.len(), 1); // top 25% of 4 = 1
|
||||
assert_eq!(bad.len(), 3);
|
||||
assert!((good[0].objective - 1.0).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_suggest_within_bounds() {
|
||||
let config = TpeConfig {
|
||||
n_dims: 2,
|
||||
max_trials: 10,
|
||||
n_initial: 5,
|
||||
gamma: 0.25,
|
||||
n_candidates: 100,
|
||||
seed: Some(123),
|
||||
};
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
let bounds = vec![(0.0, 1.0), (0.0, 1.0)];
|
||||
|
||||
// Add initial trials
|
||||
tpe.add_trial(vec![0.5, 0.5], 5.0);
|
||||
tpe.add_trial(vec![0.3, 0.7], 3.0);
|
||||
tpe.add_trial(vec![0.7, 0.3], 8.0);
|
||||
tpe.add_trial(vec![0.2, 0.8], 1.0);
|
||||
tpe.add_trial(vec![0.1, 0.9], 2.0);
|
||||
tpe.add_trial(vec![0.9, 0.1], 15.0);
|
||||
|
||||
let suggestion = tpe.suggest(&bounds);
|
||||
assert_eq!(suggestion.len(), 2);
|
||||
assert!(
|
||||
suggestion[0] >= 0.0 && suggestion[0] <= 1.0,
|
||||
"dim 0 out of bounds: {}",
|
||||
suggestion[0]
|
||||
);
|
||||
assert!(
|
||||
suggestion[1] >= 0.0 && suggestion[1] <= 1.0,
|
||||
"dim 1 out of bounds: {}",
|
||||
suggestion[1]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_lhs_initial() {
|
||||
let config = TpeConfig {
|
||||
n_dims: 3,
|
||||
max_trials: 10,
|
||||
n_initial: 5,
|
||||
gamma: 0.25,
|
||||
n_candidates: 50,
|
||||
seed: Some(42),
|
||||
};
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
let bounds = vec![(0.0, 10.0), (-1.0, 1.0), (100.0, 200.0)];
|
||||
|
||||
// First suggestion should be LHS (random within bounds)
|
||||
let s = tpe.suggest(&bounds);
|
||||
assert_eq!(s.len(), 3);
|
||||
assert!(s[0] >= 0.0 && s[0] <= 10.0, "dim 0: {}", s[0]);
|
||||
assert!(s[1] >= -1.0 && s[1] <= 1.0, "dim 1: {}", s[1]);
|
||||
assert!(s[2] >= 100.0 && s[2] <= 200.0, "dim 2: {}", s[2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_kde_log_pdf() {
|
||||
// KDE of a single point should peak near that point
|
||||
let samples = vec![vec![0.5, 0.5]];
|
||||
let sample_refs: Vec<&Vec<f64>> = samples.iter().collect();
|
||||
let bounds = vec![(0.0, 1.0), (0.0, 1.0)];
|
||||
|
||||
let pdf_at_sample = TpeOptimizer::kde_log_pdf(&sample_refs, &[0.5, 0.5], &bounds);
|
||||
let pdf_far = TpeOptimizer::kde_log_pdf(&sample_refs, &[0.0, 0.0], &bounds);
|
||||
|
||||
assert!(
|
||||
pdf_at_sample > pdf_far,
|
||||
"PDF should be higher near sample point: at={pdf_at_sample}, far={pdf_far}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_kde_log_pdf_multiple_samples() {
|
||||
// KDE with multiple samples should produce a reasonable density
|
||||
let samples = vec![
|
||||
vec![0.2, 0.8],
|
||||
vec![0.3, 0.7],
|
||||
vec![0.25, 0.75],
|
||||
];
|
||||
let sample_refs: Vec<&Vec<f64>> = samples.iter().collect();
|
||||
let bounds = vec![(0.0, 1.0), (0.0, 1.0)];
|
||||
|
||||
let pdf_near = TpeOptimizer::kde_log_pdf(&sample_refs, &[0.25, 0.75], &bounds);
|
||||
let pdf_far = TpeOptimizer::kde_log_pdf(&sample_refs, &[0.9, 0.1], &bounds);
|
||||
|
||||
assert!(
|
||||
pdf_near > pdf_far,
|
||||
"PDF should be higher near cluster: near={pdf_near}, far={pdf_far}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_convergence_simple() {
|
||||
// TPE should converge toward the minimum of a simple quadratic
|
||||
let config = TpeConfig {
|
||||
n_dims: 1,
|
||||
max_trials: 30,
|
||||
n_initial: 5,
|
||||
gamma: 0.25,
|
||||
n_candidates: 50,
|
||||
seed: Some(42),
|
||||
};
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
let bounds = vec![(0.0, 10.0)];
|
||||
|
||||
// Objective: (x - 3)^2, minimum at x=3
|
||||
for _ in 0..30 {
|
||||
let suggestion = tpe.suggest(&bounds);
|
||||
let x = suggestion[0];
|
||||
let objective = (x - 3.0) * (x - 3.0);
|
||||
tpe.add_trial(suggestion, objective);
|
||||
}
|
||||
|
||||
// Best trial should be near x=3
|
||||
let best = tpe
|
||||
.best_trial()
|
||||
.expect("should have at least one trial");
|
||||
assert!(
|
||||
(best.params[0] - 3.0).abs() < 1.5,
|
||||
"Best should be near 3.0, got {}",
|
||||
best.params[0]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_convergence_2d() {
|
||||
// 2D Rosenbrock-like: (x - 2)^2 + (y - 3)^2
|
||||
let config = TpeConfig {
|
||||
n_dims: 2,
|
||||
max_trials: 50,
|
||||
n_initial: 10,
|
||||
gamma: 0.25,
|
||||
n_candidates: 100,
|
||||
seed: Some(99),
|
||||
};
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
let bounds = vec![(0.0, 5.0), (0.0, 5.0)];
|
||||
|
||||
for _ in 0..50 {
|
||||
let suggestion = tpe.suggest(&bounds);
|
||||
let x = suggestion[0];
|
||||
let y = suggestion[1];
|
||||
let objective = (x - 2.0) * (x - 2.0) + (y - 3.0) * (y - 3.0);
|
||||
tpe.add_trial(suggestion, objective);
|
||||
}
|
||||
|
||||
let best = tpe.best_trial().expect("should have trials");
|
||||
assert!(
|
||||
best.objective < 2.0,
|
||||
"Best objective should be < 2.0, got {}",
|
||||
best.objective
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_history_persistence() {
|
||||
let config = TpeConfig::new(2, 10);
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
tpe.add_trial(vec![0.5, 0.5], 5.0);
|
||||
tpe.add_trial(vec![0.3, 0.7], 3.0);
|
||||
|
||||
let tmp_dir = std::env::temp_dir().join("tpe_test_history");
|
||||
std::fs::create_dir_all(&tmp_dir).unwrap();
|
||||
let path = tmp_dir.join("history.json");
|
||||
|
||||
tpe.save_history(&path).unwrap();
|
||||
|
||||
let config2 = TpeConfig::new(2, 10);
|
||||
let mut tpe2 = TpeOptimizer::new(config2);
|
||||
let loaded = tpe2.load_history(&path).unwrap();
|
||||
assert_eq!(loaded, 2);
|
||||
assert_eq!(tpe2.trials.len(), 2);
|
||||
|
||||
// Verify loaded content
|
||||
assert!((tpe2.trials[0].objective - 5.0).abs() < 1e-10);
|
||||
assert!((tpe2.trials[1].objective - 3.0).abs() < 1e-10);
|
||||
|
||||
std::fs::remove_dir_all(&tmp_dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_lhs_coverage() {
|
||||
// LHS should produce well-distributed samples across bins
|
||||
let config = TpeConfig {
|
||||
n_dims: 1,
|
||||
max_trials: 10,
|
||||
n_initial: 5,
|
||||
gamma: 0.25,
|
||||
n_candidates: 50,
|
||||
seed: Some(7),
|
||||
};
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
let bounds = vec![(0.0, 10.0)];
|
||||
|
||||
let samples = tpe.latin_hypercube_sample(&bounds, 10);
|
||||
assert_eq!(samples.len(), 10);
|
||||
|
||||
// Each sample should be within bounds
|
||||
for s in &samples {
|
||||
assert!(s[0] >= 0.0 && s[0] <= 10.0, "out of bounds: {}", s[0]);
|
||||
}
|
||||
|
||||
// Check that samples cover different bins (not all clustered)
|
||||
let mut bins = [false; 5];
|
||||
for s in &samples {
|
||||
let bin = ((s[0] / 10.0) * 5.0).floor() as usize;
|
||||
let bin = bin.min(4);
|
||||
bins[bin] = true;
|
||||
}
|
||||
let covered = bins.iter().filter(|&&b| b).count();
|
||||
assert!(
|
||||
covered >= 3,
|
||||
"LHS should cover at least 3 of 5 bins, covered {covered}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_empty_bounds() {
|
||||
let config = TpeConfig::new(0, 10);
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
let samples = tpe.latin_hypercube_sample(&[], 5);
|
||||
assert!(samples.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_silverman_bandwidth() {
|
||||
// Identical values should produce floor bandwidth
|
||||
let values = vec![5.0, 5.0, 5.0];
|
||||
let bw = TpeOptimizer::silverman_bandwidth(&values, Some(10.0));
|
||||
assert!(
|
||||
(bw - 0.1).abs() < 1e-10,
|
||||
"Zero-std should use floor bandwidth, got {bw}"
|
||||
);
|
||||
|
||||
// Spread values should produce reasonable bandwidth
|
||||
let values2 = vec![0.0, 5.0, 10.0];
|
||||
let bw2 = TpeOptimizer::silverman_bandwidth(&values2, Some(10.0));
|
||||
assert!(bw2 > 0.1, "Spread values should have larger bandwidth: {bw2}");
|
||||
assert!(bw2 < 10.0, "Bandwidth should be smaller than range: {bw2}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_log_sum_exp() {
|
||||
// log(exp(1) + exp(2)) = log(e + e^2) = log(e^2 * (e^-1 + 1)) = 2 + log(1 + e^-1)
|
||||
let result = TpeOptimizer::log_sum_exp(&[1.0, 2.0]);
|
||||
let expected = (1.0_f64.exp() + 2.0_f64.exp()).ln();
|
||||
assert!(
|
||||
(result - expected).abs() < 1e-10,
|
||||
"log_sum_exp([1, 2]) = {result}, expected {expected}"
|
||||
);
|
||||
|
||||
// Edge case: empty
|
||||
assert_eq!(TpeOptimizer::log_sum_exp(&[]), f64::NEG_INFINITY);
|
||||
|
||||
// Edge case: single value
|
||||
assert!((TpeOptimizer::log_sum_exp(&[5.0]) - 5.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_best_trial() {
|
||||
let config = TpeConfig::new(1, 10);
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
assert!(tpe.best_trial().is_none());
|
||||
|
||||
tpe.add_trial(vec![1.0], 10.0);
|
||||
tpe.add_trial(vec![2.0], 5.0);
|
||||
tpe.add_trial(vec![3.0], 8.0);
|
||||
|
||||
let best = tpe.best_trial().unwrap();
|
||||
assert!((best.objective - 5.0).abs() < 1e-10);
|
||||
assert!((best.params[0] - 2.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_n_trials() {
|
||||
let config = TpeConfig::new(1, 10);
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
assert_eq!(tpe.n_trials(), 0);
|
||||
|
||||
tpe.add_trial(vec![1.0], 1.0);
|
||||
assert_eq!(tpe.n_trials(), 1);
|
||||
|
||||
tpe.add_trial(vec![2.0], 2.0);
|
||||
assert_eq!(tpe.n_trials(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_split_preserves_all_trials() {
|
||||
let config = TpeConfig {
|
||||
n_dims: 1,
|
||||
max_trials: 10,
|
||||
n_initial: 5,
|
||||
gamma: 0.25,
|
||||
n_candidates: 50,
|
||||
seed: None,
|
||||
};
|
||||
let mut tpe = TpeOptimizer::new(config);
|
||||
|
||||
for i in 0..8 {
|
||||
tpe.add_trial(vec![i as f64], i as f64);
|
||||
}
|
||||
|
||||
let (good, bad) = tpe.split_trials();
|
||||
assert_eq!(
|
||||
good.len() + bad.len(),
|
||||
8,
|
||||
"Split should preserve all trials"
|
||||
);
|
||||
// gamma=0.25 of 8 = 2
|
||||
assert_eq!(good.len(), 2);
|
||||
assert_eq!(bad.len(), 6);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user