diff --git a/crates/ml/src/hyperopt/mod.rs b/crates/ml/src/hyperopt/mod.rs index 1b0c86dc2..578bcdf0a 100644 --- a/crates/ml/src/hyperopt/mod.rs +++ b/crates/ml/src/hyperopt/mod.rs @@ -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)] diff --git a/crates/ml/src/hyperopt/tpe.rs b/crates/ml/src/hyperopt/tpe.rs new file mode 100644 index 000000000..46decda2a --- /dev/null +++ b/crates/ml/src/hyperopt/tpe.rs @@ -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, + /// 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, +} + +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, + 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 { + 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> = good.iter().map(|t| t.params.clone()).collect(); + let bad_params: Vec> = bad.iter().map(|t| t.params.clone()).collect(); + let good_refs: Vec<&Vec> = good_params.iter().collect(); + let bad_refs: Vec<&Vec> = 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, 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> { + 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 = (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], 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 = 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 = 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], + bad_params: &[&Vec], + 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], + bad_params: &[&Vec], + 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 { + let json = std::fs::read_to_string(path)?; + let records: Vec = 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 { + let n = values.len() as f64; + if n < 1.0 { + return 1.0; + } + + let mean = values.iter().sum::() / n; + let variance = values.iter().map(|&v| (v - mean) * (v - mean)).sum::() / 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], bounds: &[(f64, f64)]) -> Vec { + 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 = 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::(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> = 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> = 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); + } +}