diff --git a/Cargo.toml b/Cargo.toml index 37fc806..94f64c9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ categories = ["algorithm", "science", "data-structures"] readme = "README.md" [dependencies] +log = "0.4" rand = "0.9" thiserror = "2" parking_lot = "0.12" diff --git a/src/error.rs b/src/error.rs index fe34b47..890c376 100644 --- a/src/error.rs +++ b/src/error.rs @@ -46,6 +46,32 @@ pub enum Error { #[error("KDE requires at least one sample")] EmptySamples, + /// Returned when multivariate KDE samples have zero dimensions. + #[error("multivariate KDE samples must have at least one dimension")] + ZeroDimensions, + + /// Returned when multivariate KDE samples have inconsistent dimensions. + #[error( + "dimension mismatch: expected {expected} dimensions but sample {sample_index} has {got}" + )] + DimensionMismatch { + /// The expected number of dimensions. + expected: usize, + /// The actual number of dimensions in the sample. + got: usize, + /// The index of the sample with mismatched dimensions. + sample_index: usize, + }, + + /// Returned when bandwidth vector length doesn't match the number of dimensions. + #[error("bandwidth dimension mismatch: expected {expected} bandwidths but got {got}")] + BandwidthDimensionMismatch { + /// The expected number of bandwidths. + expected: usize, + /// The actual number of bandwidths provided. + got: usize, + }, + /// Returned when an internal invariant is violated. #[error("internal error: {0}")] Internal(&'static str), diff --git a/src/kde/mod.rs b/src/kde/mod.rs new file mode 100644 index 0000000..276ebdb --- /dev/null +++ b/src/kde/mod.rs @@ -0,0 +1,13 @@ +//! Kernel Density Estimation for parameter distributions. +//! +//! This module provides kernel density estimators used by TPE samplers +//! to model probability distributions over good and bad trial regions. +//! +//! - [`univariate`] - Univariate (single-parameter) KDE +//! - [`multivariate`] - Multivariate (joint-parameter) KDE for capturing parameter dependencies + +mod multivariate; +mod univariate; + +pub(crate) use multivariate::MultivariateKDE; +pub(crate) use univariate::KernelDensityEstimator; diff --git a/src/kde/multivariate.rs b/src/kde/multivariate.rs new file mode 100644 index 0000000..49a5b05 --- /dev/null +++ b/src/kde/multivariate.rs @@ -0,0 +1,921 @@ +//! Multivariate Kernel Density Estimation for joint parameter distributions. +//! +//! This module provides a multivariate kernel density estimator that captures +//! dependencies between parameters. Unlike the univariate KDE which models each +//! parameter independently, the multivariate KDE models the joint distribution +//! to better capture correlations between parameters. + +use rand::Rng; + +use crate::error::{Error, Result}; + +/// A multivariate Gaussian kernel density estimator for joint distributions. +/// +/// `MultivariateKDE` estimates a joint probability density function from a set +/// of multi-dimensional samples. This is used in multivariate TPE to model +/// the joint distributions l(x) and g(x) across multiple parameters simultaneously, +/// capturing their dependencies. +/// +/// # Examples +/// +/// ```ignore +/// use crate::kde::MultivariateKDE; +/// +/// // Create samples: 3 samples with 2 dimensions each +/// let samples = vec![ +/// vec![1.0, 2.0], +/// vec![1.5, 2.5], +/// vec![2.0, 3.0], +/// ]; +/// let kde = MultivariateKDE::new(samples).unwrap(); +/// +/// // Get dimensionality +/// assert_eq!(kde.n_dims(), 2); +/// ``` +#[allow(dead_code)] // Fields and methods will be used in subsequent stories (US-003, US-004) +#[derive(Clone, Debug)] +pub(crate) struct MultivariateKDE { + /// The sample points used to construct the KDE. + /// Each inner Vec is one sample with `n_dims` values. + samples: Vec>, + /// The bandwidth (standard deviation) for each dimension. + /// Uses a diagonal bandwidth matrix (independent bandwidths per dimension). + bandwidths: Vec, + /// The number of dimensions. + n_dims: usize, +} + +#[allow(dead_code)] // Methods will be used in subsequent stories (US-003, US-004) +impl MultivariateKDE { + /// Creates a new multivariate KDE with automatic bandwidth selection using Scott's rule. + /// + /// Scott's rule for multivariate KDE sets bandwidth per dimension as: + /// `h_j = n^(-1/(d+4)) * sigma_j` + /// + /// where n is the number of samples, d is the dimensionality, and `sigma_j` is the + /// standard deviation of the j-th dimension. + /// + /// # Errors + /// + /// Returns `Error::EmptySamples` if `samples` is empty. + /// Returns `Error::DimensionMismatch` if samples have inconsistent dimensions. + /// Returns `Error::ZeroDimensions` if samples have zero dimensions. + #[allow(clippy::cast_precision_loss)] + pub(crate) fn new(samples: Vec>) -> Result { + if samples.is_empty() { + return Err(Error::EmptySamples); + } + + let n_dims = samples[0].len(); + if n_dims == 0 { + return Err(Error::ZeroDimensions); + } + + // Check all samples have the same dimensionality + for (i, sample) in samples.iter().enumerate() { + if sample.len() != n_dims { + return Err(Error::DimensionMismatch { + expected: n_dims, + got: sample.len(), + sample_index: i, + }); + } + } + + let bandwidths = Self::scotts_rule_multivariate(&samples, n_dims); + + Ok(Self { + samples, + bandwidths, + n_dims, + }) + } + + /// Creates a new multivariate KDE with specified bandwidths. + /// + /// Use this when you want explicit control over the smoothing parameters. + /// + /// # Errors + /// + /// Returns `Error::EmptySamples` if `samples` is empty. + /// Returns `Error::DimensionMismatch` if samples have inconsistent dimensions. + /// Returns `Error::ZeroDimensions` if samples have zero dimensions. + /// Returns `Error::BandwidthDimensionMismatch` if bandwidths length doesn't match dimensions. + /// Returns `Error::InvalidBandwidth` if any bandwidth is not positive. + pub(crate) fn with_bandwidths(samples: Vec>, bandwidths: Vec) -> Result { + if samples.is_empty() { + return Err(Error::EmptySamples); + } + + let n_dims = samples[0].len(); + if n_dims == 0 { + return Err(Error::ZeroDimensions); + } + + // Check all samples have the same dimensionality + for (i, sample) in samples.iter().enumerate() { + if sample.len() != n_dims { + return Err(Error::DimensionMismatch { + expected: n_dims, + got: sample.len(), + sample_index: i, + }); + } + } + + // Check bandwidths length matches dimensions + if bandwidths.len() != n_dims { + return Err(Error::BandwidthDimensionMismatch { + expected: n_dims, + got: bandwidths.len(), + }); + } + + // Check all bandwidths are positive + for &bw in &bandwidths { + if bw <= 0.0 { + return Err(Error::InvalidBandwidth(bw)); + } + } + + Ok(Self { + samples, + bandwidths, + n_dims, + }) + } + + /// Computes bandwidths using Scott's rule for multivariate KDE. + /// + /// Scott's rule for d dimensions: `h_j = n^(-1/(d+4)) * sigma_j` + #[allow(clippy::cast_precision_loss)] + fn scotts_rule_multivariate(samples: &[Vec], n_dims: usize) -> Vec { + let n = samples.len() as f64; + let d = n_dims as f64; + + // Scott's rule exponent for multivariate: -1/(d+4) + let exponent = -1.0 / (d + 4.0); + let scale_factor = n.powf(exponent); + + (0..n_dims) + .map(|dim| { + let std_dev = Self::dimension_std_dev(samples, dim); + // For degenerate case where all samples are identical in this dimension, + // use a small positive bandwidth + if std_dev < f64::EPSILON { + 1.0 + } else { + scale_factor * std_dev + } + }) + .collect() + } + + /// Computes the sample standard deviation for a single dimension. + #[allow(clippy::cast_precision_loss)] + fn dimension_std_dev(samples: &[Vec], dim: usize) -> f64 { + let n = samples.len() as f64; + let values: Vec = samples.iter().map(|s| s[dim]).collect(); + let mean = values.iter().sum::() / n; + let variance = values.iter().map(|x| (x - mean).powi(2)).sum::() / n; + variance.sqrt() + } + + /// Returns the number of dimensions. + pub(crate) fn n_dims(&self) -> usize { + self.n_dims + } + + /// Returns the number of samples. + #[allow(dead_code)] // Will be used in subsequent stories + pub(crate) fn n_samples(&self) -> usize { + self.samples.len() + } + + /// Returns the bandwidths for each dimension. + #[cfg(test)] + pub(crate) fn bandwidths(&self) -> &[f64] { + &self.bandwidths + } + + /// Returns a reference to the samples. + #[allow(dead_code)] // Will be used in subsequent stories + pub(crate) fn samples(&self) -> &[Vec] { + &self.samples + } + + /// Returns the log probability density at point `x`. + /// + /// This computes the log-density for numerical stability, using the formula: + /// + /// `log f(x) = log((1/n) * Σ_i Π_j K_hj((x_j - x_ij) / h_j))` + /// + /// The computation uses the log-sum-exp trick for numerical stability: + /// + /// `log(Σ exp(a_i)) = max(a) + log(Σ exp(a_i - max(a)))` + /// + /// # Panics + /// + /// Panics if `x.len() != self.n_dims`. + #[allow(clippy::cast_precision_loss)] + pub(crate) fn log_pdf(&self, x: &[f64]) -> f64 { + assert_eq!( + x.len(), + self.n_dims, + "Point dimension {} doesn't match KDE dimension {}", + x.len(), + self.n_dims + ); + + let n = self.samples.len() as f64; + + // Precompute log normalization constant for each dimension + // For a Gaussian kernel: K_h(z) = (1/(h*sqrt(2*pi))) * exp(-0.5*z^2) + // log(K_h(z)) = -log(h) - 0.5*log(2*pi) - 0.5*z^2 + let log_2pi = (2.0 * core::f64::consts::PI).ln(); + let log_norm_per_dim: Vec = self + .bandwidths + .iter() + .map(|&h| -h.ln() - 0.5 * log_2pi) + .collect(); + + // Compute log of kernel contribution for each sample + // log(prod_j K_hj(z_j)) = sum_j log(K_hj(z_j)) + let log_kernels: Vec = self + .samples + .iter() + .map(|sample| { + let mut log_kernel_sum = 0.0; + for j in 0..self.n_dims { + let z = (x[j] - sample[j]) / self.bandwidths[j]; + log_kernel_sum += log_norm_per_dim[j] - 0.5 * z * z; + } + log_kernel_sum + }) + .collect(); + + // Use log-sum-exp trick for numerical stability + // log(sum(exp(log_kernels))) = max + log(sum(exp(log_kernels - max))) + let max_log_kernel = log_kernels + .iter() + .copied() + .fold(f64::NEG_INFINITY, f64::max); + + // Handle case where all log_kernels are -inf (extremely unlikely point) + if max_log_kernel.is_infinite() && max_log_kernel < 0.0 { + return f64::NEG_INFINITY; + } + + let sum_exp: f64 = log_kernels + .iter() + .map(|&lk| (lk - max_log_kernel).exp()) + .sum(); + + // log((1/n) * sum) = -log(n) + log(sum) + -n.ln() + max_log_kernel + sum_exp.ln() + } + + /// Returns the probability density at point `x`. + /// + /// The density is computed as the average of multivariate Gaussian kernels + /// centered at each sample point: + /// + /// `f(x) = (1/n) * Σ_i Π_j K_hj((x_j - x_ij) / h_j)` + /// + /// where `K_hj` is the univariate Gaussian kernel with bandwidth `h_j`. + /// + /// This method computes in log-space for numerical stability and then + /// exponentiates the result. + /// + /// # Panics + /// + /// Panics if `x.len() != self.n_dims`. + pub(crate) fn pdf(&self, x: &[f64]) -> f64 { + self.log_pdf(x).exp() + } + + /// Samples a point from the estimated joint density distribution. + /// + /// Sampling works by: + /// 1. Uniformly selecting one of the kernel centers (samples) + /// 2. Adding independent Gaussian noise to each dimension with that + /// dimension's bandwidth as the standard deviation + /// + /// This is equivalent to sampling from a mixture of multivariate Gaussians + /// with diagonal covariance matrices, where each mixture component is + /// centered at a sample point. + /// + /// # Returns + /// + /// A `Vec` of length `n_dims` representing a sample from the KDE. + pub(crate) fn sample(&self, rng: &mut R) -> Vec { + // Select a random sample to center the kernel on + let idx = rng.random_range(0..self.samples.len()); + let center = &self.samples[idx]; + + // Add independent Gaussian noise to each dimension + // Using Box-Muller transform for Gaussian sampling + center + .iter() + .zip(self.bandwidths.iter()) + .map(|(¢er_j, &bandwidth_j)| { + let u1: f64 = rng.random(); + let u2: f64 = rng.random(); + + // Box-Muller transform: generates standard normal variate + let z = (-2.0 * u1.ln()).sqrt() * (2.0 * core::f64::consts::PI * u2).cos(); + + // Scale by bandwidth and shift by center + center_j + z * bandwidth_j + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_multivariate_kde_new_basic() { + let samples = vec![vec![1.0, 2.0], vec![1.5, 2.5], vec![2.0, 3.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + assert_eq!(kde.n_dims(), 2); + assert_eq!(kde.n_samples(), 3); + assert_eq!(kde.bandwidths().len(), 2); + } + + #[test] + fn test_multivariate_kde_new_single_sample() { + let samples = vec![vec![1.0, 2.0, 3.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + assert_eq!(kde.n_dims(), 3); + assert_eq!(kde.n_samples(), 1); + // With single sample, bandwidths should default to 1.0 (degenerate case) + for &bw in kde.bandwidths() { + assert!((bw - 1.0).abs() < f64::EPSILON); + } + } + + #[test] + fn test_multivariate_kde_new_single_dimension() { + let samples = vec![vec![1.0], vec![2.0], vec![3.0], vec![4.0], vec![5.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + assert_eq!(kde.n_dims(), 1); + assert_eq!(kde.n_samples(), 5); + assert_eq!(kde.bandwidths().len(), 1); + // Bandwidth should be positive + assert!(kde.bandwidths()[0] > 0.0); + } + + #[test] + fn test_multivariate_kde_scotts_rule() { + // Create samples with known statistics + // 10 samples, 2 dimensions + let samples: Vec> = (0..10) + .map(|i| { + let x = f64::from(i); + vec![x, x * 2.0] // Second dimension has 2x variance + }) + .collect(); + let kde = MultivariateKDE::new(samples).unwrap(); + + // Scott's rule: h = n^(-1/(d+4)) * sigma + // n=10, d=2: exponent = -1/6 ≈ -0.167 + // 10^(-1/6) ≈ 0.681 + // First dim std_dev ≈ 2.87, second ≈ 5.74 + // Expected bandwidths: ~1.95 and ~3.91 + + let bw = kde.bandwidths(); + assert!( + bw[0] > 1.0 && bw[0] < 3.0, + "First bandwidth {} unexpected", + bw[0] + ); + assert!( + bw[1] > 2.0 && bw[1] < 6.0, + "Second bandwidth {} unexpected", + bw[1] + ); + // Second bandwidth should be approximately 2x the first + assert!( + (bw[1] / bw[0] - 2.0).abs() < 0.1, + "Ratio {} not close to 2", + bw[1] / bw[0] + ); + } + + #[test] + fn test_multivariate_kde_with_bandwidths() { + let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; + let bandwidths = vec![0.5, 1.0]; + let kde = MultivariateKDE::with_bandwidths(samples, bandwidths).unwrap(); + + assert_eq!(kde.n_dims(), 2); + assert!((kde.bandwidths()[0] - 0.5).abs() < f64::EPSILON); + assert!((kde.bandwidths()[1] - 1.0).abs() < f64::EPSILON); + } + + #[test] + fn test_multivariate_kde_empty_samples() { + let samples: Vec> = vec![]; + let result = MultivariateKDE::new(samples); + assert!(matches!(result, Err(Error::EmptySamples))); + } + + #[test] + fn test_multivariate_kde_zero_dimensions() { + let samples = vec![vec![], vec![]]; + let result = MultivariateKDE::new(samples); + assert!(matches!(result, Err(Error::ZeroDimensions))); + } + + #[test] + fn test_multivariate_kde_dimension_mismatch() { + let samples = vec![vec![1.0, 2.0], vec![3.0]]; // Second sample has wrong dimensions + let result = MultivariateKDE::new(samples); + assert!(matches!( + result, + Err(Error::DimensionMismatch { + expected: 2, + got: 1, + sample_index: 1 + }) + )); + } + + #[test] + fn test_multivariate_kde_with_bandwidths_wrong_length() { + let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; + let bandwidths = vec![0.5]; // Only 1 bandwidth for 2 dimensions + let result = MultivariateKDE::with_bandwidths(samples, bandwidths); + assert!(matches!( + result, + Err(Error::BandwidthDimensionMismatch { + expected: 2, + got: 1 + }) + )); + } + + #[test] + fn test_multivariate_kde_with_bandwidths_zero() { + let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; + let bandwidths = vec![0.5, 0.0]; // Second bandwidth is zero + let result = MultivariateKDE::with_bandwidths(samples, bandwidths); + assert!(matches!(result, Err(Error::InvalidBandwidth(bw)) if bw == 0.0)); + } + + #[test] + fn test_multivariate_kde_with_bandwidths_negative() { + let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; + let bandwidths = vec![0.5, -1.0]; // Negative bandwidth + let result = MultivariateKDE::with_bandwidths(samples, bandwidths); + assert!( + matches!(result, Err(Error::InvalidBandwidth(bw)) if (bw - (-1.0)).abs() < f64::EPSILON) + ); + } + + #[test] + fn test_multivariate_kde_identical_samples() { + // All samples identical - should handle degenerate case + let samples = vec![vec![5.0, 10.0], vec![5.0, 10.0], vec![5.0, 10.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // Bandwidths should default to 1.0 for degenerate dimensions + for &bw in kde.bandwidths() { + assert!((bw - 1.0).abs() < f64::EPSILON); + } + } + + #[test] + fn test_multivariate_kde_high_dimensional() { + // Test with higher dimensions + let samples: Vec> = (0..20) + .map(|i| { + let x = f64::from(i); + vec![x, x * 0.5, x * 2.0, x * 0.1, x * 10.0] + }) + .collect(); + let kde = MultivariateKDE::new(samples).unwrap(); + + assert_eq!(kde.n_dims(), 5); + assert_eq!(kde.n_samples(), 20); + assert_eq!(kde.bandwidths().len(), 5); + + // All bandwidths should be positive + for &bw in kde.bandwidths() { + assert!(bw > 0.0); + } + } + + #[test] + fn test_multivariate_kde_samples_accessor() { + let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; + let kde = MultivariateKDE::new(samples.clone()).unwrap(); + + assert_eq!(kde.samples(), &samples); + } + + // ==================== PDF Tests ==================== + + #[test] + fn test_multivariate_kde_pdf_basic() { + let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // Density should be positive everywhere + assert!(kde.pdf(&[0.0, 0.0]) > 0.0); + assert!(kde.pdf(&[1.0, 1.0]) > 0.0); + assert!(kde.pdf(&[2.0, 2.0]) > 0.0); + assert!(kde.pdf(&[0.5, 0.5]) > 0.0); + + // Density should be higher near sample points + let near_density = kde.pdf(&[1.0, 1.0]); + let far_density = kde.pdf(&[10.0, 10.0]); + assert!(near_density > far_density); + } + + #[test] + fn test_multivariate_kde_pdf_with_custom_bandwidths() { + let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]]; + let bandwidths = vec![0.5, 0.5]; + let kde = MultivariateKDE::with_bandwidths(samples, bandwidths).unwrap(); + + // Density should be positive + assert!(kde.pdf(&[0.5, 0.5]) > 0.0); + assert!(kde.pdf(&[0.0, 0.0]) > 0.0); + } + + #[test] + fn test_multivariate_kde_pdf_single_sample() { + let samples = vec![vec![5.0, 10.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // Should have positive density near the sample + assert!(kde.pdf(&[5.0, 10.0]) > 0.0); + assert!(kde.pdf(&[4.5, 9.5]) > 0.0); + } + + #[test] + fn test_multivariate_kde_log_pdf_consistency() { + let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // log_pdf and pdf should be consistent: exp(log_pdf(x)) == pdf(x) + let test_points = vec![ + vec![0.0, 0.0], + vec![1.0, 1.0], + vec![0.5, 0.5], + vec![3.0, 3.0], + ]; + + for point in test_points { + let log_p = kde.log_pdf(&point); + let p = kde.pdf(&point); + let p_from_log = log_p.exp(); + assert!( + (p - p_from_log).abs() < 1e-10, + "pdf={p}, exp(log_pdf)={p_from_log}" + ); + } + } + + #[test] + fn test_multivariate_kde_pdf_integrates_to_one_1d() { + // Test with 1D case first (should match univariate KDE behavior) + let samples = vec![vec![0.0], vec![1.0], vec![2.0], vec![3.0], vec![4.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // Numerical integration over a wide range + let n_points = 1000; + let low = -10.0; + let high = 15.0; + let dx = (high - low) / f64::from(n_points); + + let integral: f64 = (0..n_points) + .map(|i| { + let x = low + (f64::from(i) + 0.5) * dx; + kde.pdf(&[x]) * dx + }) + .sum(); + + // Should be approximately 1.0 (within numerical error) + assert!( + (integral - 1.0).abs() < 0.02, + "1D integral = {integral}, expected ~1.0" + ); + } + + #[test] + fn test_multivariate_kde_pdf_integrates_to_one_2d() { + // Test with 2D case + let samples = vec![ + vec![0.0, 0.0], + vec![1.0, 0.0], + vec![0.0, 1.0], + vec![1.0, 1.0], + vec![0.5, 0.5], + ]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // Numerical integration over a 2D grid + let n_points = 100; // 100x100 = 10000 points + let low = -5.0; + let high = 6.0; + let dx = (high - low) / f64::from(n_points); + + let mut integral = 0.0; + for i in 0..n_points { + for j in 0..n_points { + let x = low + (f64::from(i) + 0.5) * dx; + let y = low + (f64::from(j) + 0.5) * dx; + integral += kde.pdf(&[x, y]) * dx * dx; + } + } + + // Should be approximately 1.0 (within numerical error) + // 2D integration has more error, so use larger tolerance + assert!( + (integral - 1.0).abs() < 0.05, + "2D integral = {integral}, expected ~1.0" + ); + } + + #[test] + fn test_multivariate_kde_pdf_symmetry() { + // With symmetric samples, PDF should be symmetric + let samples = vec![ + vec![1.0, 0.0], + vec![-1.0, 0.0], + vec![0.0, 1.0], + vec![0.0, -1.0], + ]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // Density at symmetric points should be equal + let d1 = kde.pdf(&[0.5, 0.0]); + let d2 = kde.pdf(&[-0.5, 0.0]); + assert!( + (d1 - d2).abs() < 1e-10, + "Symmetric points have different densities: {d1} vs {d2}" + ); + + let d3 = kde.pdf(&[0.0, 0.5]); + let d4 = kde.pdf(&[0.0, -0.5]); + assert!( + (d3 - d4).abs() < 1e-10, + "Symmetric points have different densities: {d3} vs {d4}" + ); + } + + #[test] + fn test_multivariate_kde_pdf_high_dimensional() { + // Test with higher dimensions + let samples: Vec> = (0..10) + .map(|i| { + let x = f64::from(i) * 0.1; + vec![x, x, x, x, x] // 5D + }) + .collect(); + let kde = MultivariateKDE::new(samples).unwrap(); + + // Density should be positive + assert!(kde.pdf(&[0.5, 0.5, 0.5, 0.5, 0.5]) > 0.0); + assert!(kde.pdf(&[0.0, 0.0, 0.0, 0.0, 0.0]) > 0.0); + + // Log PDF should be finite + let log_p = kde.log_pdf(&[0.5, 0.5, 0.5, 0.5, 0.5]); + assert!(log_p.is_finite()); + } + + #[test] + fn test_multivariate_kde_pdf_numerical_stability() { + // Test numerical stability with points far from samples + let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // Very far point should have very small but non-negative density + let far_pdf = kde.pdf(&[100.0, 100.0]); + assert!(far_pdf >= 0.0); + assert!(far_pdf.is_finite() || far_pdf == 0.0); + + // Log PDF should be finite (or -inf for zero density) + let far_log_pdf = kde.log_pdf(&[100.0, 100.0]); + assert!(far_log_pdf.is_finite() || far_log_pdf.is_infinite()); + } + + #[test] + #[should_panic(expected = "Point dimension")] + fn test_multivariate_kde_pdf_wrong_dimension() { + let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + + // Should panic with wrong dimension + let _ = kde.pdf(&[0.0]); // Only 1 value for 2D KDE + } + + // ==================== Sampling Tests ==================== + + #[test] + fn test_multivariate_kde_sample_basic() { + let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + let mut rng = rand::rng(); + + // Sample should have correct dimensionality + let sample = kde.sample(&mut rng); + assert_eq!(sample.len(), 2); + } + + #[test] + fn test_multivariate_kde_sample_in_reasonable_range() { + let samples = vec![ + vec![0.0, 0.0], + vec![1.0, 1.0], + vec![2.0, 2.0], + vec![3.0, 3.0], + vec![4.0, 4.0], + ]; + let kde = MultivariateKDE::new(samples).unwrap(); + let mut rng = rand::rng(); + + // Samples should generally be in a reasonable range around the data + for _ in 0..100 { + let s = kde.sample(&mut rng); + // With high probability, samples should be within a few bandwidths + // of the data range. Use a generous range to avoid flaky tests. + assert!( + s[0] > -10.0 && s[0] < 15.0, + "Sample dimension 0: {} outside expected range", + s[0] + ); + assert!( + s[1] > -10.0 && s[1] < 15.0, + "Sample dimension 1: {} outside expected range", + s[1] + ); + } + } + + #[test] + fn test_multivariate_kde_sample_single_sample() { + // When KDE has only one sample, all samples should be centered around it + let samples = vec![vec![5.0, 10.0]]; + let kde = MultivariateKDE::new(samples).unwrap(); + let mut rng = rand::rng(); + + // Generate many samples and check they cluster around (5.0, 10.0) + let n_samples = 100; + let mut sum_x = 0.0; + let mut sum_y = 0.0; + for _ in 0..n_samples { + let s = kde.sample(&mut rng); + sum_x += s[0]; + sum_y += s[1]; + } + + let mean_x = sum_x / f64::from(n_samples); + let mean_y = sum_y / f64::from(n_samples); + + // Mean should be close to the single sample point + // With 100 samples and bandwidth=1.0, we expect mean within ~0.3 of center + assert!( + (mean_x - 5.0).abs() < 1.0, + "Mean x={mean_x}, expected close to 5.0" + ); + assert!( + (mean_y - 10.0).abs() < 1.0, + "Mean y={mean_y}, expected close to 10.0" + ); + } + + #[test] + fn test_multivariate_kde_sample_high_dimensional() { + // Test sampling in higher dimensions + let samples: Vec> = (0..10) + .map(|i| { + let x = f64::from(i) * 0.5; + vec![x, x * 2.0, x * 0.5, x + 1.0, x - 1.0] // 5D + }) + .collect(); + let kde = MultivariateKDE::new(samples).unwrap(); + let mut rng = rand::rng(); + + // Sample should have correct dimensionality + for _ in 0..50 { + let sample = kde.sample(&mut rng); + assert_eq!(sample.len(), 5); + // All values should be finite + for &val in &sample { + assert!(val.is_finite(), "Sample value is not finite: {val}"); + } + } + } + + #[test] + #[allow(clippy::cast_precision_loss)] + fn test_multivariate_kde_sample_respects_bandwidth() { + // Create samples all at origin with large custom bandwidth + let data = vec![vec![0.0, 0.0], vec![0.0, 0.0], vec![0.0, 0.0]]; + let bandwidths = vec![0.1, 10.0]; // Small bandwidth in x, large in y + let kde = MultivariateKDE::with_bandwidths(data, bandwidths).unwrap(); + let mut rng = rand::rng(); + + // Generate samples and check variance in each dimension + let n_samples = 1000; + let mut values_x: Vec = Vec::with_capacity(n_samples); + let mut values_y: Vec = Vec::with_capacity(n_samples); + + for _ in 0..n_samples { + let s = kde.sample(&mut rng); + values_x.push(s[0]); + values_y.push(s[1]); + } + + // Compute sample variances + let n = n_samples as f64; + let mean_x: f64 = values_x.iter().sum::() / n; + let mean_y: f64 = values_y.iter().sum::() / n; + + let var_x: f64 = values_x.iter().map(|x| (x - mean_x).powi(2)).sum::() / n; + let var_y: f64 = values_y.iter().map(|y| (y - mean_y).powi(2)).sum::() / n; + + // Variance should be approximately bandwidth^2 + // x: bandwidth=0.1, expected variance ~0.01 + // y: bandwidth=10.0, expected variance ~100.0 + assert!( + var_x < 0.05, + "X variance {var_x} too large for bandwidth 0.1" + ); + assert!( + var_y > 50.0 && var_y < 200.0, + "Y variance {var_y} unexpected for bandwidth 10.0" + ); + } + + #[test] + fn test_multivariate_kde_sample_distribution_shape() { + // Create samples along a diagonal line + let data = vec![ + vec![0.0, 0.0], + vec![1.0, 1.0], + vec![2.0, 2.0], + vec![3.0, 3.0], + vec![4.0, 4.0], + ]; + let kde = MultivariateKDE::new(data).unwrap(); + let mut rng = rand::rng(); + + // Sample many points and verify the mean is near the center + let n_samples = 500; + let mut sum = [0.0, 0.0]; + + for _ in 0..n_samples { + let s = kde.sample(&mut rng); + sum[0] += s[0]; + sum[1] += s[1]; + } + + let mean_x = sum[0] / f64::from(n_samples); + let mean_y = sum[1] / f64::from(n_samples); + + // Mean should be close to (2.0, 2.0) - the center of the samples + assert!( + (mean_x - 2.0).abs() < 0.5, + "Mean x={mean_x}, expected close to 2.0" + ); + assert!( + (mean_y - 2.0).abs() < 0.5, + "Mean y={mean_y}, expected close to 2.0" + ); + } + + #[test] + fn test_multivariate_kde_sample_deterministic_with_seeded_rng() { + use rand::SeedableRng; + + let data = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; + let kde = MultivariateKDE::new(data).unwrap(); + + // Use a seeded RNG for reproducibility + let mut rng1 = rand::rngs::StdRng::seed_from_u64(42); + let mut rng2 = rand::rngs::StdRng::seed_from_u64(42); + + // Same seed should produce same samples + let result1 = kde.sample(&mut rng1); + let result2 = kde.sample(&mut rng2); + + assert!( + (result1[0] - result2[0]).abs() < f64::EPSILON, + "Samples with same seed differ in dimension 0" + ); + assert!( + (result1[1] - result2[1]).abs() < f64::EPSILON, + "Samples with same seed differ in dimension 1" + ); + } +} diff --git a/src/kde.rs b/src/kde/univariate.rs similarity index 100% rename from src/kde.rs rename to src/kde/univariate.rs diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index b04f05b..0f0b600 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -1,11 +1,16 @@ //! Sampler trait and implementations for parameter sampling. pub mod grid; +pub mod multivariate_tpe; pub mod random; pub mod tpe; use std::collections::HashMap; +pub use multivariate_tpe::{ + ConstantLiarStrategy, MultivariateTpeSampler, MultivariateTpeSamplerBuilder, +}; + use crate::distribution::Distribution; use crate::param::ParamValue; @@ -43,6 +48,57 @@ impl CompletedTrial { } } +/// A pending (running) trial with its parameters and distributions, but no objective value yet. +/// +/// This struct represents a trial that has been started and has sampled parameters, +/// but is still running and hasn't returned an objective value. It is used with the +/// constant liar strategy for parallel optimization. +/// +/// # Examples +/// +/// ```ignore +/// use std::collections::HashMap; +/// use optimizer::sampler::PendingTrial; +/// use optimizer::param::ParamValue; +/// use optimizer::distribution::{Distribution, FloatDistribution}; +/// +/// let mut params = HashMap::new(); +/// params.insert("x".to_string(), ParamValue::Float(0.5)); +/// +/// let mut distributions = HashMap::new(); +/// distributions.insert("x".to_string(), Distribution::Float(FloatDistribution { +/// low: 0.0, high: 1.0, log_scale: false, step: None, +/// })); +/// +/// let pending = PendingTrial::new(1, params, distributions); +/// assert_eq!(pending.id, 1); +/// ``` +#[derive(Clone, Debug)] +pub struct PendingTrial { + /// The unique identifier for this trial. + pub id: u64, + /// The sampled parameter values, keyed by parameter name. + pub params: HashMap, + /// The parameter distributions used, keyed by parameter name. + pub distributions: HashMap, +} + +impl PendingTrial { + /// Creates a new pending trial. + #[must_use] + pub fn new( + id: u64, + params: HashMap, + distributions: HashMap, + ) -> Self { + Self { + id, + params, + distributions, + } + } +} + /// Trait for pluggable parameter sampling strategies. /// /// Samplers are responsible for generating parameter values based on diff --git a/src/sampler/multivariate_tpe.rs b/src/sampler/multivariate_tpe.rs new file mode 100644 index 0000000..8e029df --- /dev/null +++ b/src/sampler/multivariate_tpe.rs @@ -0,0 +1,6522 @@ +//! Multivariate Tree-Parzen Estimator (TPE) sampler implementation. +//! +//! Unlike the standard [`TpeSampler`](super::tpe::TpeSampler) which samples parameters +//! independently, the `MultivariateTpeSampler` models joint distributions over multiple +//! parameters. This allows it to capture correlations between parameters, which can +//! significantly improve optimization performance on problems where parameters interact. +//! +//! # When to Use Multivariate TPE +//! +//! Use `MultivariateTpeSampler` when: +//! - Parameters have strong correlations (e.g., Rosenbrock function) +//! - The optimal value of one parameter depends on another +//! - You have a fixed search space across all trials +//! +//! Use the standard [`TpeSampler`](super::tpe::TpeSampler) when: +//! - Parameters are independent +//! - The search space varies dynamically between trials +//! - You want simpler, faster optimization +//! +//! # Configuration +//! +//! - `gamma_strategy`: Controls what fraction of trials are considered "good" +//! - `n_startup_trials`: Number of random trials before multivariate TPE kicks in +//! - `n_ei_candidates`: Number of candidates to evaluate when selecting the next point +//! - `group`: When true, decomposes search space into independent groups +//! - `warn_independent_sampling`: When true, logs warnings when falling back to independent sampling +//! - `constant_liar`: Strategy for imputing values to pending trials in parallel optimization +//! +//! # Fallback Behavior for Dynamic Search Spaces +//! +//! When the search space varies between trials (e.g., conditional parameters), the sampler +//! gracefully falls back to independent sampling. The fallback hierarchy is: +//! +//! 1. **Multivariate TPE** for parameters that appear in ALL completed trials (the intersection +//! search space). These parameters are modeled jointly using multivariate KDE. +//! +//! 2. **Independent TPE** for parameters that appear in some trials but not the intersection. +//! Each such parameter is modeled with its own univariate KDE. +//! +//! 3. **Uniform random sampling** for new parameters that have never been seen, or during +//! the startup phase before enough trials are collected. +//! +//! When `warn_independent_sampling` is enabled (default), the sampler logs warnings via the +//! `log` crate when fallback occurs, helping you understand why multivariate modeling wasn't +//! possible. +//! +//! # Group Decomposition +//! +//! When `group` is enabled, the sampler analyzes parameter co-occurrence across trials and +//! decomposes the search space into independent groups. Parameters that always appear +//! together form a group, and each group is sampled using multivariate TPE independently. +//! +//! This can improve efficiency when: +//! - Some parameters are truly independent of others +//! - The search space has conditional structure where certain parameters only appear together +//! +//! # Parallel Optimization with Constant Liar +//! +//! For parallel optimization where multiple trials run simultaneously, the constant liar +//! strategy helps avoid redundant exploration. When new samples are requested while trials +//! are pending, the sampler can impute "lie" values for pending trials, allowing them to +//! influence the model. +//! +//! Available strategies via [`ConstantLiarStrategy`]: +//! - `None`: Ignore pending trials (default, suitable for sequential optimization) +//! - `Mean`: Impute the mean of completed trial values +//! - `Best`: Impute the best (minimum) completed value - pessimistic, encourages exploration +//! - `Worst`: Impute the worst (maximum) completed value - optimistic, encourages exploitation +//! - `Custom(f64)`: Impute a specific user-defined value +//! +//! # Examples +//! +//! ## Basic Usage +//! +//! ``` +//! use optimizer::sampler::MultivariateTpeSampler; +//! +//! let sampler = MultivariateTpeSampler::builder() +//! .gamma(0.15) +//! .n_startup_trials(20) +//! .n_ei_candidates(32) +//! .seed(42) +//! .build() +//! .unwrap(); +//! ``` +//! +//! ## With Custom Gamma Strategy and Group Decomposition +//! +//! ``` +//! use optimizer::sampler::MultivariateTpeSampler; +//! use optimizer::sampler::tpe::SqrtGamma; +//! +//! let sampler = MultivariateTpeSampler::builder() +//! .gamma_strategy(SqrtGamma::default()) +//! .group(true) +//! .build() +//! .unwrap(); +//! ``` +//! +//! ## Parallel Optimization with Constant Liar +//! +//! ``` +//! use optimizer::sampler::{ConstantLiarStrategy, MultivariateTpeSampler}; +//! +//! // Use mean imputation for parallel workers +//! let sampler = MultivariateTpeSampler::builder() +//! .constant_liar(ConstantLiarStrategy::Mean) +//! .n_startup_trials(10) +//! .build() +//! .unwrap(); +//! ``` +//! +//! ## Suppressing Fallback Warnings +//! +//! ``` +//! use optimizer::sampler::MultivariateTpeSampler; +//! +//! // Disable fallback warnings for dynamic search spaces +//! let sampler = MultivariateTpeSampler::builder() +//! .warn_independent_sampling(false) +//! .build() +//! .unwrap(); +//! ``` + +use std::collections::HashMap; +use std::sync::Arc; + +use log::warn; +use parking_lot::Mutex; +use rand::SeedableRng; +use rand::rngs::StdRng; + +use super::tpe::{FixedGamma, GammaStrategy}; +use super::{CompletedTrial, PendingTrial, Sampler}; +use crate::distribution::Distribution; +use crate::error::Result; +use crate::param::ParamValue; + +/// Strategy for imputing objective values for pending/running trials during parallel optimization. +/// +/// In parallel optimization, multiple trials may be running simultaneously. The constant liar +/// strategy assigns "lie" values to pending trials so they can be included in the model fitting, +/// which helps avoid redundant exploration of the same region. +/// +/// # Variants +/// +/// - `None`: No imputation; pending trials are ignored (default) +/// - `Mean`: Impute with the mean of completed trial values +/// - `Best`: Impute with the best (minimum for minimization) completed value +/// - `Worst`: Impute with the worst (maximum for minimization) completed value +/// - `Custom(f64)`: Impute with a specific user-provided value +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::ConstantLiarStrategy; +/// +/// // Use mean imputation for parallel optimization +/// let strategy = ConstantLiarStrategy::Mean; +/// +/// // Use a custom lie value +/// let strategy = ConstantLiarStrategy::Custom(0.5); +/// ``` +#[derive(Debug, Clone, Copy, PartialEq, Default)] +pub enum ConstantLiarStrategy { + /// No imputation; pending trials are ignored. + #[default] + None, + /// Impute with the mean of completed trial values. + Mean, + /// Impute with the best (minimum for minimization) completed value. + Best, + /// Impute with the worst (maximum for minimization) completed value. + Worst, + /// Impute with a specific user-provided value. + Custom(f64), +} + +/// A Multivariate Tree-Parzen Estimator (TPE) sampler for Bayesian optimization. +/// +/// This sampler extends the standard TPE approach by modeling joint distributions +/// over all parameters, allowing it to capture parameter correlations. +/// +/// # Fields +/// +/// - `gamma_strategy`: Strategy for computing the gamma quantile +/// - `n_startup_trials`: Number of random trials before TPE sampling begins +/// - `n_ei_candidates`: Number of candidates to evaluate per joint sample +/// - `group`: Whether to decompose search space into independent groups +/// - `warn_independent_sampling`: Whether to log warnings when falling back to independent sampling +pub struct MultivariateTpeSampler { + /// Strategy for computing the gamma quantile. + gamma_strategy: Arc, + /// Number of trials before multivariate 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, + /// Whether to decompose search space into independent groups based on parameter co-occurrence. + group: bool, + /// Whether to log warnings when falling back to independent sampling. + warn_independent_sampling: bool, + /// Strategy for imputing objective values for pending trials in parallel optimization. + constant_liar: ConstantLiarStrategy, + /// Thread-safe RNG for sampling. + rng: Mutex, + /// Cache for joint samples to maintain consistency across parameters within the same trial. + /// The tuple contains (`trial_id`, cached joint sample). + joint_sample_cache: Mutex)>>, +} + +impl MultivariateTpeSampler { + /// Creates a new Multivariate 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) + /// - `group`: false (no group decomposition) + /// - `warn_independent_sampling`: true (log warnings on fallback) + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::MultivariateTpeSampler; + /// + /// let sampler = MultivariateTpeSampler::new(); + /// ``` + #[must_use] + pub fn new() -> Self { + Self { + gamma_strategy: Arc::new(FixedGamma::default()), + n_startup_trials: 10, + n_ei_candidates: 24, + group: false, + warn_independent_sampling: true, + constant_liar: ConstantLiarStrategy::None, + rng: Mutex::new(StdRng::from_os_rng()), + joint_sample_cache: Mutex::new(None), + } + } + + /// Creates a builder for configuring a Multivariate TPE sampler. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::MultivariateTpeSampler; + /// + /// let sampler = MultivariateTpeSampler::builder() + /// .gamma(0.15) + /// .n_startup_trials(20) + /// .n_ei_candidates(32) + /// .seed(42) + /// .build() + /// .unwrap(); + /// ``` + #[must_use] + pub fn builder() -> MultivariateTpeSamplerBuilder { + MultivariateTpeSamplerBuilder::new() + } + + /// Returns the gamma strategy used by this sampler. + #[must_use] + pub fn gamma_strategy(&self) -> &dyn GammaStrategy { + self.gamma_strategy.as_ref() + } + + /// Returns the number of startup trials. + #[must_use] + pub fn n_startup_trials(&self) -> usize { + self.n_startup_trials + } + + /// Returns the number of EI candidates. + #[must_use] + pub fn n_ei_candidates(&self) -> usize { + self.n_ei_candidates + } + + /// Returns whether group decomposition is enabled. + #[must_use] + pub fn group(&self) -> bool { + self.group + } + + /// Returns whether independent sampling warnings are enabled. + #[must_use] + pub fn warn_independent_sampling(&self) -> bool { + self.warn_independent_sampling + } + + /// Returns the constant liar strategy for parallel optimization. + #[must_use] + pub fn constant_liar(&self) -> &ConstantLiarStrategy { + &self.constant_liar + } + + /// Imputes objective values for pending trials based on the constant liar strategy. + /// + /// In parallel optimization, multiple trials may be running simultaneously. This method + /// assigns "lie" values to pending trials so they can be included in the model fitting, + /// which helps avoid redundant exploration of the same region. + /// + /// # Arguments + /// + /// * `pending_trials` - Trials that are currently running and have no objective value yet. + /// * `completed_trials` - Trials that have completed and have objective values. + /// + /// # Returns + /// + /// A vector of `CompletedTrial` objects containing both the original completed trials + /// and the pending trials with imputed values. If the strategy is `None`, returns + /// only the completed trials (pending trials are ignored). + /// + /// # Imputation Strategies + /// + /// - `None`: Pending trials are ignored (returns only completed trials) + /// - `Mean`: Pending trials get the mean of completed objective values + /// - `Best`: Pending trials get the minimum completed objective value (for minimization) + /// - `Worst`: Pending trials get the maximum completed objective value (for minimization) + /// - `Custom(v)`: Pending trials get the specified value `v` + /// + /// # Examples + /// + /// ```ignore + /// use std::collections::HashMap; + /// use optimizer::sampler::{ + /// ConstantLiarStrategy, MultivariateTpeSampler, CompletedTrial, PendingTrial, + /// }; + /// use optimizer::param::ParamValue; + /// use optimizer::distribution::{Distribution, FloatDistribution}; + /// + /// // Create a sampler with mean imputation + /// let sampler = MultivariateTpeSampler::builder() + /// .constant_liar(ConstantLiarStrategy::Mean) + /// .build() + /// .unwrap(); + /// + /// // Create some completed trials + /// let dist = Distribution::Float(FloatDistribution { + /// low: 0.0, high: 1.0, log_scale: false, step: None, + /// }); + /// let completed = vec![ + /// CompletedTrial::new( + /// 0, + /// [("x".to_string(), ParamValue::Float(0.2))].into_iter().collect(), + /// [("x".to_string(), dist.clone())].into_iter().collect(), + /// 1.0, + /// ), + /// CompletedTrial::new( + /// 1, + /// [("x".to_string(), ParamValue::Float(0.8))].into_iter().collect(), + /// [("x".to_string(), dist.clone())].into_iter().collect(), + /// 3.0, + /// ), + /// ]; + /// + /// // Create a pending trial + /// let pending = vec![ + /// PendingTrial::new( + /// 2, + /// [("x".to_string(), ParamValue::Float(0.5))].into_iter().collect(), + /// [("x".to_string(), dist.clone())].into_iter().collect(), + /// ), + /// ]; + /// + /// // Impute values + /// let augmented = sampler.impute_pending_trials(&pending, &completed); + /// + /// // Should have 3 trials total + /// assert_eq!(augmented.len(), 3); + /// + /// // The pending trial should have the mean value (1.0 + 3.0) / 2 = 2.0 + /// let imputed = augmented.iter().find(|t| t.id == 2).unwrap(); + /// assert!((imputed.value - 2.0).abs() < f64::EPSILON); + /// ``` + #[must_use] + pub fn impute_pending_trials( + &self, + pending_trials: &[PendingTrial], + completed_trials: &[CompletedTrial], + ) -> Vec { + // Start with a copy of completed trials + let mut result: Vec = completed_trials.to_vec(); + + // If strategy is None or no pending trials, just return completed trials + if matches!(self.constant_liar, ConstantLiarStrategy::None) || pending_trials.is_empty() { + return result; + } + + // Compute the imputation value based on strategy + let imputed_value = self.compute_imputation_value(completed_trials); + + // Convert pending trials to completed trials with imputed values + for pending in pending_trials { + result.push(CompletedTrial::new( + pending.id, + pending.params.clone(), + pending.distributions.clone(), + imputed_value, + )); + } + + result + } + + /// Computes the imputation value based on the constant liar strategy. + /// + /// This is a helper method used by [`impute_pending_trials`](Self::impute_pending_trials). + /// + /// # Arguments + /// + /// * `completed_trials` - The completed trials to compute the imputation value from. + /// + /// # Returns + /// + /// The imputed value based on the strategy. Returns 0.0 if there are no completed + /// trials (except for `Custom` strategy which returns its specified value). + #[allow(clippy::cast_precision_loss)] + fn compute_imputation_value(&self, completed_trials: &[CompletedTrial]) -> f64 { + match self.constant_liar { + ConstantLiarStrategy::None => 0.0, // This case is handled before calling this method + ConstantLiarStrategy::Mean => { + if completed_trials.is_empty() { + 0.0 + } else { + let sum: f64 = completed_trials.iter().map(|t| t.value).sum(); + sum / completed_trials.len() as f64 + } + } + ConstantLiarStrategy::Best => { + // Best means minimum for minimization problems + completed_trials + .iter() + .map(|t| t.value) + .fold(f64::INFINITY, f64::min) + } + ConstantLiarStrategy::Worst => { + // Worst means maximum for minimization problems + completed_trials + .iter() + .map(|t| t.value) + .fold(f64::NEG_INFINITY, f64::max) + } + ConstantLiarStrategy::Custom(v) => v, + } + } + + /// Filters trials to those containing all parameters in the search space. + /// + /// This method is used to identify trials that can be used for multivariate + /// KDE fitting. Only trials that contain ALL parameters in the search space + /// are included, ensuring we can model the joint distribution over all + /// parameters. + /// + /// # Arguments + /// + /// * `history` - All completed trials from the optimization history. + /// * `search_space` - The intersection search space containing parameters that + /// appear in all trials. + /// + /// # Returns + /// + /// A vector of references to trials that contain all parameters in the search space. + /// + /// # Examples + /// + /// ```ignore + /// use std::collections::HashMap; + /// use optimizer::sampler::MultivariateTpeSampler; + /// use optimizer::sampler::tpe::IntersectionSearchSpace; + /// + /// let sampler = MultivariateTpeSampler::new(); + /// let trials = vec![/* ... completed trials ... */]; + /// let search_space = IntersectionSearchSpace::calculate(&trials); + /// let filtered = sampler.filter_trials(&trials, &search_space); + /// ``` + #[must_use] + pub fn filter_trials<'a>( + &self, + history: &'a [CompletedTrial], + search_space: &HashMap, + ) -> Vec<&'a CompletedTrial> { + history + .iter() + .filter(|trial| { + // Include trial only if it has ALL parameters in the search space + search_space + .keys() + .all(|param| trial.params.contains_key(param)) + }) + .collect() + } + + /// Splits filtered trials into good and bad groups based on the gamma quantile. + /// + /// This method is similar to [`TpeSampler::split_trials`](super::tpe::TpeSampler) + /// but operates on pre-filtered trials that have been selected for multivariate + /// sampling. + /// + /// The gamma value is computed dynamically using the configured [`GammaStrategy`]. + /// Trials are sorted by objective value (ascending for minimization), and the + /// gamma quantile determines the split point. + /// + /// # Arguments + /// + /// * `trials` - Filtered trials to split (typically from [`filter_trials`](Self::filter_trials)). + /// + /// # Returns + /// + /// A tuple `(good_trials, bad_trials)` where: + /// - `good_trials` contains trials with values below the gamma quantile + /// - `bad_trials` contains trials with values at or above the gamma quantile + /// + /// Both vectors are guaranteed to be non-empty when the input has at least 2 trials. + /// If the input has fewer than 2 trials, both vectors may be empty or one may + /// contain the single trial. + /// + /// # Examples + /// + /// ```ignore + /// use std::collections::HashMap; + /// use optimizer::sampler::MultivariateTpeSampler; + /// use optimizer::sampler::tpe::IntersectionSearchSpace; + /// + /// let sampler = MultivariateTpeSampler::new(); + /// let trials = vec![/* ... completed trials ... */]; + /// let search_space = IntersectionSearchSpace::calculate(&trials); + /// let filtered = sampler.filter_trials(&trials, &search_space); + /// let (good, bad) = sampler.split_trials(&filtered); + /// + /// // good contains trials with lowest objective values + /// // bad contains trials with higher objective values + /// ``` + #[allow( + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss + )] + #[must_use] + pub fn split_trials<'a>( + &self, + trials: &[&'a CompletedTrial], + ) -> (Vec<&'a CompletedTrial>, Vec<&'a CompletedTrial>) { + if trials.is_empty() { + return (vec![], vec![]); + } + + // Sort trials by objective value (ascending for minimization) + let mut sorted_indices: Vec = (0..trials.len()).collect(); + sorted_indices.sort_by(|&a, &b| { + trials[a] + .value + .partial_cmp(&trials[b].value) + .unwrap_or(core::cmp::Ordering::Equal) + }); + + // Compute gamma using the strategy and clamp to valid range + let gamma = self + .gamma_strategy + .gamma(trials.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 = ((trials.len() as f64 * gamma).ceil() as usize) + .max(1) + .min(trials.len().saturating_sub(1)); + + // Handle edge case: if we have only 1 trial, put it in good + if trials.len() == 1 { + return (vec![trials[0]], vec![]); + } + + let good: Vec<_> = sorted_indices[..n_good] + .iter() + .map(|&i| trials[i]) + .collect(); + let bad: Vec<_> = sorted_indices[n_good..] + .iter() + .map(|&i| trials[i]) + .collect(); + + (good, bad) + } + + /// Extracts parameter values from trials as a numeric observation matrix. + /// + /// This method converts trial parameter values into a matrix format suitable + /// for multivariate KDE fitting. Each row in the output represents one trial's + /// parameter values in the specified order. + /// + /// # Arguments + /// + /// * `trials` - Trials to extract observations from. + /// * `param_order` - The order of parameters in the output vectors. This ensures + /// consistent column ordering across the observation matrix. + /// + /// # Returns + /// + /// A `Vec>` where: + /// - Outer vec has one entry per trial + /// - Inner vec has one entry per parameter (in `param_order` order) + /// - Float values are used directly + /// - Int values are converted to f64 + /// - Categorical parameters are skipped (not included in output) + /// + /// # Panics + /// + /// This method does not panic. If a parameter is missing from a trial or has + /// an unsupported type (Categorical), it is simply skipped. + /// + /// # Examples + /// + /// ```ignore + /// use std::collections::HashMap; + /// use optimizer::sampler::MultivariateTpeSampler; + /// + /// let sampler = MultivariateTpeSampler::new(); + /// let trials = vec![/* ... completed trials ... */]; + /// let filtered = sampler.filter_trials(&trials, &search_space); + /// + /// // Extract observations for x and y in that order + /// let param_order = vec!["x".to_string(), "y".to_string()]; + /// let observations = sampler.extract_observations(&filtered, ¶m_order); + /// + /// // observations[i][0] is the x value for trial i + /// // observations[i][1] is the y value for trial i + /// ``` + #[must_use] + #[allow(clippy::cast_precision_loss)] + pub fn extract_observations( + &self, + trials: &[&CompletedTrial], + param_order: &[String], + ) -> Vec> { + trials + .iter() + .map(|trial| { + param_order + .iter() + .filter_map(|param_name| { + trial.params.get(param_name).and_then(|value| match value { + crate::param::ParamValue::Float(f) => Some(*f), + crate::param::ParamValue::Int(i) => Some(*i as f64), + crate::param::ParamValue::Categorical(_) => None, // Skip categorical + }) + }) + .collect() + }) + .collect() + } + + /// Selects the best candidate from a set of samples using the joint acquisition function. + /// + /// This method implements the core TPE selection criterion: it generates candidates + /// from the "good" KDE (l(x)) and selects the one that maximizes the ratio l(x)/g(x), + /// which is equivalent to maximizing `log(l(x)) - log(g(x))`. + /// + /// # Arguments + /// + /// * `good_kde` - The KDE fitted on the "good" trials (low objective values). + /// * `bad_kde` - The KDE fitted on the "bad" trials (high objective values). + /// + /// # Returns + /// + /// A `Vec` representing the selected candidate point in the parameter space. + /// The point is chosen to maximize the expected improvement proxy l(x)/g(x). + /// + /// # Algorithm + /// + /// 1. Generate `n_ei_candidates` samples from `good_kde` + /// 2. For each candidate, compute `log(l(x)) - log(g(x))` + /// 3. Select the candidate with the highest ratio + /// + /// # Edge Cases + /// + /// - If `g(x)` is very small (near zero), the log-space computation handles this + /// gracefully without division issues. + /// - If all candidates have `-inf` log ratios, the first candidate is returned. + /// + /// # Examples + /// + /// ```ignore + /// use optimizer::sampler::MultivariateTpeSampler; + /// use optimizer::kde::MultivariateKDE; + /// + /// let sampler = MultivariateTpeSampler::builder() + /// .n_ei_candidates(24) + /// .seed(42) + /// .build() + /// .unwrap(); + /// + /// let good_obs = vec![vec![0.1, 0.2], vec![0.2, 0.3], vec![0.15, 0.25]]; + /// let bad_obs = vec![vec![0.8, 0.9], vec![0.7, 0.85], vec![0.9, 0.95]]; + /// + /// let good_kde = MultivariateKDE::new(good_obs).unwrap(); + /// let bad_kde = MultivariateKDE::new(bad_obs).unwrap(); + /// + /// let selected = sampler.select_candidate(&good_kde, &bad_kde); + /// // selected should be a point likely in the "good" region + /// ``` + #[must_use] + #[allow(dead_code)] // Used by tests + pub(crate) fn select_candidate( + &self, + good_kde: &crate::kde::MultivariateKDE, + bad_kde: &crate::kde::MultivariateKDE, + ) -> Vec { + let mut rng = self.rng.lock(); + + // Generate candidates from the good distribution + let candidates: Vec> = (0..self.n_ei_candidates) + .map(|_| good_kde.sample(&mut *rng)) + .collect(); + + // Compute log(l(x)) - log(g(x)) for each candidate + // This is equivalent to log(l(x)/g(x)) which we want to maximize + let log_ratios: Vec = candidates + .iter() + .map(|candidate| { + let log_l = good_kde.log_pdf(candidate); + let log_g = bad_kde.log_pdf(candidate); + log_l - log_g + }) + .collect(); + + // Find the candidate with the maximum log ratio + let mut best_idx = 0; + let mut best_ratio = f64::NEG_INFINITY; + + for (idx, &ratio) in log_ratios.iter().enumerate() { + // Handle NaN by treating it as worse than any finite value + if ratio > best_ratio || (best_ratio.is_nan() && !ratio.is_nan()) { + best_ratio = ratio; + best_idx = idx; + } + } + + candidates.into_iter().nth(best_idx).unwrap_or_default() + } + + /// Selects the best candidate using an external RNG. + /// + /// This variant accepts an external RNG, used when the caller already holds the lock. + fn select_candidate_with_rng( + &self, + good_kde: &crate::kde::MultivariateKDE, + bad_kde: &crate::kde::MultivariateKDE, + rng: &mut StdRng, + ) -> Vec { + // Generate candidates from the good distribution + let candidates: Vec> = (0..self.n_ei_candidates) + .map(|_| good_kde.sample(rng)) + .collect(); + + // Compute log(l(x)) - log(g(x)) for each candidate + let log_ratios: Vec = candidates + .iter() + .map(|candidate| { + let log_l = good_kde.log_pdf(candidate); + let log_g = bad_kde.log_pdf(candidate); + log_l - log_g + }) + .collect(); + + // Find the candidate with the maximum log ratio + let mut best_idx = 0; + let mut best_ratio = f64::NEG_INFINITY; + + for (idx, &ratio) in log_ratios.iter().enumerate() { + if ratio > best_ratio || (best_ratio.is_nan() && !ratio.is_nan()) { + best_ratio = ratio; + best_idx = idx; + } + } + + candidates.into_iter().nth(best_idx).unwrap_or_default() + } + + /// Samples all parameters uniformly at random. + /// + /// This is a fallback method used when multivariate TPE cannot be applied. + #[allow(clippy::unused_self)] + fn sample_all_uniform( + &self, + search_space: &HashMap, + rng: &mut rand::rngs::StdRng, + ) -> HashMap { + search_space + .iter() + .map(|(name, dist)| (name.clone(), Self::sample_uniform_single(dist, rng))) + .collect() + } + + /// Samples a single parameter uniformly at random from its distribution. + fn sample_uniform_single( + distribution: &Distribution, + rng: &mut rand::rngs::StdRng, + ) -> ParamValue { + use rand::Rng; + + 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 { + #[allow(clippy::cast_possible_truncation)] + let n_steps = ((d.high - d.low) / step).floor() as i64; + let k = rng.random_range(0..=n_steps); + #[allow(clippy::cast_precision_loss)] + let result = d.low + (k as f64) * step; + result + } else { + rng.random_range(d.low..=d.high) + }; + ParamValue::Float(value) + } + Distribution::Int(d) => { + #[allow(clippy::cast_precision_loss)] + let value = if d.log_scale { + let log_low = (d.low as f64).ln(); + let log_high = (d.high as f64).ln(); + #[allow(clippy::cast_possible_truncation)] + 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 { + #[allow(clippy::cast_possible_truncation)] + 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) => { + let index = rng.random_range(0..d.n_choices); + ParamValue::Categorical(index) + } + } + } + + /// Samples parameters jointly using multivariate TPE. + /// + /// This method samples all parameters in the search space jointly, capturing + /// correlations between them. During the startup phase (before `n_startup_trials` + /// have completed), it falls back to uniform random sampling. + /// + /// # Arguments + /// + /// * `search_space` - A map of parameter names to their distributions. + /// * `history` - Historical completed trials for informed sampling. + /// + /// # Returns + /// + /// A `HashMap` containing sampled values for all parameters + /// in the search space. + /// + /// # Algorithm + /// + /// 1. During startup phase: sample each parameter uniformly at random + /// 2. After startup: + /// - Compute intersection search space (parameters common to all trials) + /// - Filter trials to those containing all intersection parameters + /// - Split into good/bad trials based on gamma quantile + /// - Fit multivariate KDEs on continuous parameters (Float and Int) + /// - Select best candidate using l(x)/g(x) acquisition function + /// - Sample categorical parameters independently (for now) + /// + /// # Notes + /// + /// - Categorical parameters are currently sampled independently, not jointly. + /// - If there are not enough continuous parameters for multivariate modeling, + /// falls back to independent sampling. + /// + /// # Examples + /// + /// ```ignore + /// use std::collections::HashMap; + /// use optimizer::sampler::MultivariateTpeSampler; + /// use optimizer::distribution::{Distribution, FloatDistribution}; + /// + /// let sampler = MultivariateTpeSampler::builder() + /// .n_startup_trials(10) + /// .seed(42) + /// .build() + /// .unwrap(); + /// + /// let mut search_space = HashMap::new(); + /// search_space.insert("x".to_string(), Distribution::Float(FloatDistribution { + /// low: 0.0, high: 1.0, log_scale: false, step: None, + /// })); + /// search_space.insert("y".to_string(), Distribution::Float(FloatDistribution { + /// low: 0.0, high: 1.0, log_scale: false, step: None, + /// })); + /// + /// let history = vec![]; // No history yet + /// let params = sampler.sample_joint(&search_space, &history); + /// ``` + #[must_use] + #[allow(clippy::too_many_lines)] + pub fn sample_joint( + &self, + search_space: &HashMap, + history: &[CompletedTrial], + ) -> HashMap { + let mut rng = self.rng.lock(); + + // Early returns for cases requiring random sampling + if history.len() < self.n_startup_trials { + return self.sample_all_uniform(search_space, &mut rng); + } + + // If group mode is enabled, decompose search space into independent groups + if self.group { + drop(rng); + return self.sample_with_groups(search_space, history); + } + + // Non-grouped mode: use the original single-group logic + self.sample_single_group(search_space, history, &mut rng) + } + + /// Samples parameters by decomposing the search space into independent groups. + /// + /// When `group=true`, this method analyzes the trial history to identify groups of + /// parameters that always appear together, then samples each group independently + /// using multivariate TPE. This is more efficient when parameters naturally partition + /// into independent subsets (e.g., due to conditional search spaces). + /// + /// # Arguments + /// + /// * `search_space` - The full search space containing all parameters to sample. + /// * `history` - Completed trials from the optimization history. + /// + /// # Returns + /// + /// A `HashMap` mapping parameter names to their sampled values. + fn sample_with_groups( + &self, + search_space: &HashMap, + history: &[CompletedTrial], + ) -> HashMap { + use std::collections::HashSet; + + use crate::sampler::tpe::GroupDecomposedSearchSpace; + + // Decompose the search space into independent parameter groups + let groups = GroupDecomposedSearchSpace::calculate(history); + + let mut result: HashMap = HashMap::new(); + + // Sample each group independently + for group in &groups { + // Build a sub-search space for this group + let group_search_space: HashMap = search_space + .iter() + .filter(|(name, _)| group.contains(*name)) + .map(|(name, dist)| (name.clone(), dist.clone())) + .collect(); + + if group_search_space.is_empty() { + continue; + } + + // Filter history to trials that have at least one parameter in this group + let group_history: Vec<&CompletedTrial> = history + .iter() + .filter(|trial| { + trial + .distributions + .keys() + .any(|param| group.contains(param)) + }) + .collect(); + + // Build completed trials from references for the group + // We need to create a temporary slice for sample_group_internal + let group_history_owned: Vec = + group_history.iter().map(|t| (*t).clone()).collect(); + + // Sample this group using multivariate TPE + let mut rng = self.rng.lock(); + let group_result = + self.sample_single_group(&group_search_space, &group_history_owned, &mut rng); + drop(rng); + + // Merge group results into the main result + for (name, value) in group_result { + result.insert(name, value); + } + } + + // Handle parameters not in any group (sample independently) + let grouped_params: HashSet = groups.iter().flatten().cloned().collect(); + let ungrouped_params: HashMap = search_space + .iter() + .filter(|(name, _)| !grouped_params.contains(*name) && !result.contains_key(*name)) + .map(|(name, dist)| (name.clone(), dist.clone())) + .collect(); + + if !ungrouped_params.is_empty() { + if self.warn_independent_sampling { + let param_names: Vec<&str> = ungrouped_params.keys().map(String::as_str).collect(); + warn!( + "MultivariateTpeSampler: Parameters {param_names:?} are not in any group. \ + Sampling independently." + ); + } + + // Sample ungrouped parameters uniformly (no history for them) + let mut rng = self.rng.lock(); + for (name, dist) in &ungrouped_params { + let value = Self::sample_uniform_single(dist, &mut rng); + result.insert(name.clone(), value); + } + } + + result + } + + /// Samples parameters as a single group using multivariate TPE. + /// + /// This is the core multivariate TPE sampling logic, used both in non-grouped mode + /// and for sampling individual groups in grouped mode. + /// + /// # Arguments + /// + /// * `search_space` - The search space for this group of parameters. + /// * `history` - Completed trials to use for model fitting. + /// * `rng` - Random number generator (caller must hold lock). + /// + /// # Returns + /// + /// A `HashMap` mapping parameter names to their sampled values. + #[allow(clippy::too_many_lines)] + fn sample_single_group( + &self, + search_space: &HashMap, + history: &[CompletedTrial], + rng: &mut StdRng, + ) -> HashMap { + use crate::kde::MultivariateKDE; + use crate::sampler::tpe::IntersectionSearchSpace; + + // Early returns for cases requiring random sampling + if history.len() < self.n_startup_trials { + return self.sample_all_uniform(search_space, rng); + } + + let intersection = IntersectionSearchSpace::calculate(history); + if intersection.is_empty() { + if self.warn_independent_sampling { + warn!( + "MultivariateTpeSampler: No common parameters found across trials. \ + Falling back to fully independent sampling." + ); + } + return self.sample_all_independent_with_rng(search_space, history, rng); + } + + let filtered = self.filter_trials(history, &intersection); + if filtered.len() < 2 { + // Not enough trials in intersection - use independent TPE on full history + if self.warn_independent_sampling { + warn!( + "MultivariateTpeSampler: Only {} trial(s) with all intersection parameters. \ + Falling back to independent sampling.", + filtered.len() + ); + } + return self.sample_all_independent_with_rng(search_space, history, rng); + } + + let (good, bad) = self.split_trials(&filtered); + if good.is_empty() || bad.is_empty() { + // Can't split trials properly - use independent TPE on full history + if self.warn_independent_sampling { + warn!( + "MultivariateTpeSampler: Unable to split trials into good/bad groups. \ + Falling back to independent sampling." + ); + } + return self.sample_all_independent_with_rng(search_space, history, rng); + } + + // Sample categorical parameters using TPE with l(x)/g(x) ratio + let mut result: HashMap = HashMap::new(); + for (name, dist) in &intersection { + if let Distribution::Categorical(d) = dist { + let good_indices = Self::extract_categorical_indices(&good, name); + let bad_indices = Self::extract_categorical_indices(&bad, name); + let idx = + Self::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); + result.insert(name.clone(), ParamValue::Categorical(idx)); + } + } + + // Collect continuous parameters + let mut param_order: Vec = intersection + .iter() + .filter(|(_, dist)| !matches!(dist, Distribution::Categorical(_))) + .map(|(name, _)| name.clone()) + .collect(); + + if param_order.is_empty() { + // Only categorical parameters in intersection - fill remaining with independent TPE + self.fill_remaining_independent_with_rng( + search_space, + &intersection, + history, + &mut result, + rng, + ); + return result; + } + + param_order.sort(); + + // Extract observations and validate + let good_obs = self.extract_observations(&good, ¶m_order); + let bad_obs = self.extract_observations(&bad, ¶m_order); + let expected_dims = param_order.len(); + + let valid = !good_obs.is_empty() + && !bad_obs.is_empty() + && good_obs.iter().all(|obs| obs.len() == expected_dims) + && bad_obs.iter().all(|obs| obs.len() == expected_dims); + + if !valid { + // Observations invalid - fill remaining with independent TPE + self.fill_remaining_independent_with_rng( + search_space, + &intersection, + history, + &mut result, + rng, + ); + return result; + } + + // Fit KDEs using let...else pattern + let Ok(good_kde) = MultivariateKDE::new(good_obs) else { + // KDE construction failed - fill remaining with independent TPE + self.fill_remaining_independent_with_rng( + search_space, + &intersection, + history, + &mut result, + rng, + ); + return result; + }; + + let Ok(bad_kde) = MultivariateKDE::new(bad_obs) else { + // KDE construction failed - fill remaining with independent TPE + self.fill_remaining_independent_with_rng( + search_space, + &intersection, + history, + &mut result, + rng, + ); + return result; + }; + + let selected = self.select_candidate_with_rng(&good_kde, &bad_kde, rng); + + // Map selected values to parameter names + for (idx, param_name) in param_order.iter().enumerate() { + if let Some(dist) = intersection.get(param_name) { + let value = selected[idx]; + let param_value = self.convert_to_param_value(value, dist); + if let Some(pv) = param_value { + result.insert(param_name.clone(), pv); + } + } + } + + // Fill remaining parameters using independent TPE sampling + self.fill_remaining_independent_with_rng( + search_space, + &intersection, + history, + &mut result, + rng, + ); + result + } + + /// Converts a raw f64 value to a `ParamValue` based on the distribution. + #[allow(clippy::unused_self)] + fn convert_to_param_value(&self, value: f64, dist: &Distribution) -> Option { + match dist { + Distribution::Float(d) => { + let clamped = value.clamp(d.low, d.high); + let stepped = if let Some(step) = d.step { + let steps = ((clamped - d.low) / step).round(); + (d.low + steps * step).clamp(d.low, d.high) + } else { + clamped + }; + Some(ParamValue::Float(stepped)) + } + Distribution::Int(d) => { + #[allow(clippy::cast_possible_truncation)] + let int_value = value.round() as i64; + let clamped = int_value.clamp(d.low, d.high); + let stepped = if let Some(step) = d.step { + let steps = (clamped - d.low) / step; + (d.low + steps * step).clamp(d.low, d.high) + } else { + clamped + }; + Some(ParamValue::Int(stepped)) + } + Distribution::Categorical(_) => None, + } + } + + /// Fills remaining parameters in result using independent TPE sampling. + /// + /// This method is used to sample parameters that are not in the intersection + /// search space. It uses independent univariate TPE sampling for each parameter, + /// similar to the standard [`TpeSampler`](super::tpe::TpeSampler). + /// + /// When there isn't enough history for a parameter, falls back to uniform sampling. + #[allow(dead_code)] + fn fill_remaining_independent( + &self, + search_space: &HashMap, + intersection: &HashMap, + history: &[CompletedTrial], + result: &mut HashMap, + ) { + // Identify parameters not in result (and not in intersection) + let missing_params: Vec<(&String, &Distribution)> = search_space + .iter() + .filter(|(name, _)| !result.contains_key(*name)) + .collect(); + + if missing_params.is_empty() { + return; + } + + // Log warning if we have missing params and warnings are enabled + if self.warn_independent_sampling && !intersection.is_empty() { + let param_names: Vec<&str> = missing_params.iter().map(|(n, _)| n.as_str()).collect(); + warn!( + "MultivariateTpeSampler: Parameters {param_names:?} are not in the intersection \ + search space. Sampling independently." + ); + } + + // Split trials for independent sampling + let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); + + let mut rng = self.rng.lock(); + + for (name, dist) in missing_params { + let value = + self.sample_independent_tpe(name, dist, &good_trials, &bad_trials, &mut rng); + result.insert(name.clone(), value); + } + } + + /// Fills remaining parameters using independent TPE sampling with an external RNG. + /// + /// This variant accepts an external RNG, used when the caller already holds the lock. + fn fill_remaining_independent_with_rng( + &self, + search_space: &HashMap, + intersection: &HashMap, + history: &[CompletedTrial], + result: &mut HashMap, + rng: &mut StdRng, + ) { + // Identify parameters not in result (and not in intersection) + let missing_params: Vec<(&String, &Distribution)> = search_space + .iter() + .filter(|(name, _)| !result.contains_key(*name)) + .collect(); + + if missing_params.is_empty() { + return; + } + + // Log warning if we have missing params and warnings are enabled + if self.warn_independent_sampling && !intersection.is_empty() { + let param_names: Vec<&str> = missing_params.iter().map(|(n, _)| n.as_str()).collect(); + warn!( + "MultivariateTpeSampler: Parameters {param_names:?} are not in the intersection \ + search space. Sampling independently." + ); + } + + // Split trials for independent sampling + let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); + + for (name, dist) in missing_params { + let value = self.sample_independent_tpe(name, dist, &good_trials, &bad_trials, rng); + result.insert(name.clone(), value); + } + } + + /// Samples a single parameter using independent TPE. + /// + /// This method extracts values for the given parameter from good and bad trials, + /// fits univariate KDEs, and samples using the TPE acquisition function. + #[allow(clippy::too_many_lines)] + fn sample_independent_tpe( + &self, + param_name: &str, + distribution: &Distribution, + good_trials: &[&CompletedTrial], + bad_trials: &[&CompletedTrial], + rng: &mut StdRng, + ) -> ParamValue { + match distribution { + Distribution::Float(d) => { + let good_values: Vec = good_trials + .iter() + .filter_map(|t| t.params.get(param_name)) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .filter_map(|t| t.params.get(param_name)) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + if good_values.is_empty() || bad_values.is_empty() { + return Self::sample_uniform_single(distribution, rng); + } + + let value = self.sample_tpe_float( + d.low, + d.high, + d.log_scale, + d.step, + good_values, + bad_values, + rng, + ); + ParamValue::Float(value) + } + Distribution::Int(d) => { + let good_values: Vec = good_trials + .iter() + .filter_map(|t| t.params.get(param_name)) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .filter_map(|t| t.params.get(param_name)) + .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_single(distribution, rng); + } + + let value = self.sample_tpe_int( + d.low, + d.high, + d.log_scale, + d.step, + &good_values, + &bad_values, + rng, + ); + ParamValue::Int(value) + } + Distribution::Categorical(d) => { + let good_indices: Vec = good_trials + .iter() + .filter_map(|t| t.params.get(param_name)) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + let bad_indices: Vec = bad_trials + .iter() + .filter_map(|t| t.params.get(param_name)) + .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_single(distribution, rng); + } + + let idx = + Self::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); + ParamValue::Categorical(idx) + } + } + } + + /// Samples using TPE for float distributions. + #[allow(clippy::too_many_arguments)] + fn sample_tpe_float( + &self, + low: f64, + high: f64, + log_scale: bool, + step: Option, + good_values: Vec, + bad_values: Vec, + rng: &mut StdRng, + ) -> f64 { + use rand::Rng; + + use crate::kde::KernelDensityEstimator; + + // Transform to internal space (log space if needed) + let (internal_low, internal_high, good_internal, bad_internal) = if log_scale { + let i_low = low.ln(); + let i_high = high.ln(); + let g: Vec = good_values.iter().map(|&v| v.ln()).collect(); + let b: Vec = bad_values.iter().map(|&v| v.ln()).collect(); + (i_low, i_high, g, b) + } else { + (low, high, good_values, bad_values) + }; + + // Fit KDEs to good and bad groups + let l_kde = KernelDensityEstimator::new(good_internal); + let g_kde = KernelDensityEstimator::new(bad_internal); + + // 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, + good_values: &[i64], + bad_values: &[i64], + rng: &mut StdRng, + ) -> i64 { + // Convert to floats for KDE + let good_floats: Vec = good_values.iter().map(|&v| v as f64).collect(); + let bad_floats: Vec = bad_values.iter().map(|&v| v as f64).collect(); + + // Use float TPE sampling + let float_value = self.sample_tpe_float( + low as f64, + high as f64, + log_scale, + step.map(|s| s as f64), + good_floats, + bad_floats, + rng, + ); + + // Round to nearest integer + let int_value = float_value.round() as i64; + + // Apply step constraint if present + let int_value = if let Some(step) = step { + let k = ((int_value - low) as f64 / step as f64).round() as i64; + low + k * step + } else { + int_value + }; + + // Ensure value is within bounds + int_value.clamp(low, high) + } + + /// Samples all parameters using independent TPE sampling. + /// + /// This is used as a complete fallback when no intersection search space exists. + #[allow(dead_code)] // Used by tests + fn sample_all_independent( + &self, + search_space: &HashMap, + history: &[CompletedTrial], + ) -> HashMap { + // Split trials for independent sampling + let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); + + let mut rng = self.rng.lock(); + let mut result = HashMap::new(); + + for (name, dist) in search_space { + let value = + self.sample_independent_tpe(name, dist, &good_trials, &bad_trials, &mut rng); + result.insert(name.clone(), value); + } + + result + } + + /// Samples all parameters using independent TPE sampling with an external RNG. + /// + /// This variant accepts an external RNG, used when the caller already holds the lock. + fn sample_all_independent_with_rng( + &self, + search_space: &HashMap, + history: &[CompletedTrial], + rng: &mut StdRng, + ) -> HashMap { + // Split trials for independent sampling + let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); + + let mut result = HashMap::new(); + + for (name, dist) in search_space { + let value = self.sample_independent_tpe(name, dist, &good_trials, &bad_trials, rng); + result.insert(name.clone(), value); + } + + result + } + + /// Samples using TPE for categorical distributions. + /// + /// This method computes kernel-weighted category probabilities based on + /// the good and bad trial groups, then samples proportionally to the + /// `l(x)/g(x)` ratio for each category. + /// + /// # Arguments + /// + /// * `n_choices` - The number of categories in the categorical distribution. + /// * `good_indices` - Category indices from the "good" trials. + /// * `bad_indices` - Category indices from the "bad" trials. + /// * `rng` - Random number generator for sampling. + /// + /// # Returns + /// + /// The selected category index. + /// + /// # Algorithm + /// + /// 1. Count occurrences of each category in good and bad groups + /// 2. Apply Laplace smoothing (add 1 to each count) to avoid zero probabilities + /// 3. Compute `l(x)/g(x)` ratio for each category + /// 4. Sample proportionally to the computed weights + #[allow(clippy::cast_precision_loss)] + fn sample_tpe_categorical( + n_choices: usize, + good_indices: &[usize], + bad_indices: &[usize], + rng: &mut rand::rngs::StdRng, + ) -> usize { + use rand::Rng; + + // 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 Laplace smoothing to avoid zero probabilities + let good_total = good_indices.len() as f64 + n_choices as f64; + let bad_total = bad_indices.len() as f64 + n_choices as f64; + + // Calculate l(x)/g(x) ratio for each category + let mut weights = vec![0.0f64; n_choices]; + for i in 0..n_choices { + let l_prob = (good_counts[i] as f64 + 1.0) / good_total; + let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total; + weights[i] = l_prob / g_prob; + } + + // Sample proportionally to weights + let total_weight: f64 = weights.iter().sum(); + let threshold = rng.random::() * total_weight; + + let mut cumulative = 0.0; + for (i, &w) in weights.iter().enumerate() { + cumulative += w; + if cumulative >= threshold { + return i; + } + } + + // Fallback to last index (shouldn't happen) + n_choices - 1 + } + + /// Extracts categorical indices from trials for a specific parameter. + /// + /// # Arguments + /// + /// * `trials` - Trials to extract from. + /// * `param_name` - The name of the categorical parameter. + /// + /// # Returns + /// + /// A vector of category indices from the trials. + fn extract_categorical_indices(trials: &[&CompletedTrial], param_name: &str) -> Vec { + trials + .iter() + .filter_map(|trial| { + trial.params.get(param_name).and_then(|value| { + if let ParamValue::Categorical(idx) = value { + Some(*idx) + } else { + None + } + }) + }) + .collect() + } +} + +impl Default for MultivariateTpeSampler { + fn default() -> Self { + Self::new() + } +} + +impl Sampler for MultivariateTpeSampler { + /// Samples a parameter value from the given distribution using multivariate TPE. + /// + /// This method integrates with the multivariate sampling strategy by: + /// 1. On first call for a trial: generating a joint sample for all parameters + /// 2. On subsequent calls for the same trial: returning cached values + /// + /// This ensures consistency across parameters within the same trial while + /// still conforming to the single-parameter [`Sampler`] interface. + /// + /// # Arguments + /// + /// * `distribution` - The parameter distribution to sample from. + /// * `trial_id` - The unique ID of the trial being sampled for. + /// * `history` - Historical completed trials for informed sampling. + /// + /// # Returns + /// + /// A [`ParamValue`] sampled from the distribution. + /// + /// # Note + /// + /// Since the [`Sampler`] trait doesn't provide the parameter name, this method + /// must infer the parameter from the distribution. For best results with + /// multivariate sampling, use [`sample_joint`](Self::sample_joint) directly + /// when you have access to the full search space. + fn sample( + &self, + distribution: &Distribution, + trial_id: u64, + history: &[CompletedTrial], + ) -> ParamValue { + // Check if we have a cached joint sample for this trial + { + let cache = self.joint_sample_cache.lock(); + if let Some((cached_trial_id, ref cached_sample)) = *cache + && cached_trial_id == trial_id + { + // Try to find a matching parameter from the cached sample + if let Some(value) = Self::find_matching_param(distribution, cached_sample) { + return value; + } + } + } + + // Build the search space from history to get parameter names + let search_space = Self::build_search_space_from_history(distribution, history); + + // Generate a joint sample + let joint_sample = self.sample_joint(&search_space, history); + + // Cache the joint sample for this trial + { + let mut cache = self.joint_sample_cache.lock(); + *cache = Some((trial_id, joint_sample.clone())); + } + + // Find and return the value for the requested distribution + Self::find_matching_param(distribution, &joint_sample).unwrap_or_else(|| { + // Fallback to uniform sampling if no match found + let mut rng = self.rng.lock(); + Self::sample_uniform_single(distribution, &mut rng) + }) + } +} + +impl MultivariateTpeSampler { + /// Finds a matching parameter value from the cached sample based on distribution. + /// + /// This is an associated function that matches parameters by comparing + /// distribution bounds and types. + fn find_matching_param( + distribution: &Distribution, + cached_sample: &HashMap, + ) -> Option { + // Match by distribution type and value compatibility + for value in cached_sample.values() { + match (distribution, value) { + (Distribution::Float(d), ParamValue::Float(v)) => { + if *v >= d.low && *v <= d.high { + return Some(value.clone()); + } + } + (Distribution::Int(d), ParamValue::Int(v)) => { + if *v >= d.low && *v <= d.high { + return Some(value.clone()); + } + } + (Distribution::Categorical(d), ParamValue::Categorical(v)) => { + if *v < d.n_choices { + return Some(value.clone()); + } + } + _ => {} + } + } + None + } + + /// Builds a search space from history and the current distribution. + /// + /// This is an associated function that extracts parameter distributions + /// from historical trials and includes the current distribution to form + /// a complete search space. + fn build_search_space_from_history( + current_distribution: &Distribution, + history: &[CompletedTrial], + ) -> HashMap { + let mut search_space = HashMap::new(); + + // Collect distributions from history + for trial in history { + for (name, dist) in &trial.distributions { + search_space + .entry(name.clone()) + .or_insert_with(|| dist.clone()); + } + } + + // If the search space is empty, create a placeholder for the current distribution + if search_space.is_empty() { + search_space.insert("_current".to_string(), current_distribution.clone()); + } + + search_space + } +} + +/// Builder for configuring a [`MultivariateTpeSampler`]. +/// +/// This builder allows fluent configuration of multivariate TPE hyperparameters. +/// +/// # Examples +/// +/// Using a fixed gamma value: +/// +/// ``` +/// use optimizer::sampler::MultivariateTpeSamplerBuilder; +/// +/// let sampler = MultivariateTpeSamplerBuilder::new() +/// .gamma(0.15) +/// .n_startup_trials(20) +/// .n_ei_candidates(32) +/// .seed(42) +/// .build() +/// .unwrap(); +/// ``` +/// +/// Using a custom gamma strategy: +/// +/// ``` +/// use optimizer::sampler::MultivariateTpeSamplerBuilder; +/// use optimizer::sampler::tpe::SqrtGamma; +/// +/// let sampler = MultivariateTpeSamplerBuilder::new() +/// .gamma_strategy(SqrtGamma::default()) +/// .group(true) +/// .build() +/// .unwrap(); +/// ``` +#[derive(Debug, Clone)] +pub struct MultivariateTpeSamplerBuilder { + gamma_strategy: Box, + /// Raw gamma value for deferred validation (Some if `gamma()` was called). + raw_gamma: Option, + n_startup_trials: usize, + n_ei_candidates: usize, + group: bool, + warn_independent_sampling: bool, + constant_liar: ConstantLiarStrategy, + seed: Option, +} + +impl MultivariateTpeSamplerBuilder { + /// 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) + /// - `group`: false (no group decomposition) + /// - `warn_independent_sampling`: true (log warnings on fallback) + /// - `constant_liar`: None (no imputation for pending trials) + /// - 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, + group: false, + warn_independent_sampling: true, + constant_liar: ConstantLiarStrategy::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::MultivariateTpeSamplerBuilder; + /// + /// let sampler = MultivariateTpeSamplerBuilder::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 { + 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 + /// + /// ``` + /// use optimizer::sampler::MultivariateTpeSamplerBuilder; + /// use optimizer::sampler::tpe::{LinearGamma, SqrtGamma}; + /// + /// // Square root strategy (Optuna-style) + /// let sampler = MultivariateTpeSamplerBuilder::new() + /// .gamma_strategy(SqrtGamma::default()) + /// .build() + /// .unwrap(); + /// + /// // Linear interpolation strategy + /// let sampler = MultivariateTpeSamplerBuilder::new() + /// .gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap()) + /// .build() + /// .unwrap(); + /// ``` + #[must_use] + pub fn gamma_strategy(mut self, strategy: G) -> Self { + self.gamma_strategy = Box::new(strategy); + self.raw_gamma = None; // Clear any raw gamma set by gamma() + self + } + + /// 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, + /// multivariate TPE-based sampling begins. + /// + /// # Arguments + /// + /// * `n` - Number of random trials before TPE kicks in. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::MultivariateTpeSamplerBuilder; + /// + /// let sampler = MultivariateTpeSamplerBuilder::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, multivariate 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::MultivariateTpeSamplerBuilder; + /// + /// let sampler = MultivariateTpeSamplerBuilder::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 + } + + /// Enables or disables group decomposition. + /// + /// When enabled, the sampler analyzes parameter co-occurrence in trial history + /// and decomposes the search space into independent groups. Each group is then + /// sampled using multivariate TPE independently, which can improve efficiency + /// when some parameters are truly independent. + /// + /// # Arguments + /// + /// * `group` - Whether to enable group decomposition. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::MultivariateTpeSamplerBuilder; + /// + /// let sampler = MultivariateTpeSamplerBuilder::new() + /// .group(true) // Enable group decomposition + /// .build() + /// .unwrap(); + /// ``` + #[must_use] + pub fn group(mut self, group: bool) -> Self { + self.group = group; + self + } + + /// Enables or disables warnings when falling back to independent sampling. + /// + /// When multivariate TPE cannot model all parameters jointly (e.g., due to + /// dynamic search spaces), it falls back to independent sampling for some + /// parameters. This setting controls whether a warning is logged when this + /// occurs. + /// + /// # Arguments + /// + /// * `warn` - Whether to log warnings on fallback. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::MultivariateTpeSamplerBuilder; + /// + /// let sampler = MultivariateTpeSamplerBuilder::new() + /// .warn_independent_sampling(false) // Suppress fallback warnings + /// .build() + /// .unwrap(); + /// ``` + #[must_use] + pub fn warn_independent_sampling(mut self, warn: bool) -> Self { + self.warn_independent_sampling = warn; + self + } + + /// Sets the constant liar strategy for parallel optimization. + /// + /// The constant liar strategy determines how to impute objective values for + /// pending (not-yet-completed) trials. This is useful in parallel optimization + /// where multiple trials may be running simultaneously. + /// + /// # Arguments + /// + /// * `strategy` - The [`ConstantLiarStrategy`] to use. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::{ConstantLiarStrategy, MultivariateTpeSamplerBuilder}; + /// + /// // Use mean imputation for pending trials + /// let sampler = MultivariateTpeSamplerBuilder::new() + /// .constant_liar(ConstantLiarStrategy::Mean) + /// .build() + /// .unwrap(); + /// + /// // Use worst-case imputation (pessimistic) + /// let sampler = MultivariateTpeSamplerBuilder::new() + /// .constant_liar(ConstantLiarStrategy::Worst) + /// .build() + /// .unwrap(); + /// + /// // Use a custom imputation value + /// let sampler = MultivariateTpeSamplerBuilder::new() + /// .constant_liar(ConstantLiarStrategy::Custom(0.5)) + /// .build() + /// .unwrap(); + /// ``` + #[must_use] + pub fn constant_liar(mut self, strategy: ConstantLiarStrategy) -> Self { + self.constant_liar = strategy; + self + } + + /// Sets a seed for reproducible sampling. + /// + /// # Arguments + /// + /// * `seed` - Seed value for the random number generator. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::MultivariateTpeSamplerBuilder; + /// + /// let sampler = MultivariateTpeSamplerBuilder::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 [`MultivariateTpeSampler`]. + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if a fixed gamma value was set and is not in (0.0, 1.0). + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::MultivariateTpeSamplerBuilder; + /// + /// let sampler = MultivariateTpeSamplerBuilder::new() + /// .gamma(0.15) + /// .n_startup_trials(20) + /// .n_ei_candidates(32) + /// .seed(42) + /// .build() + /// .unwrap(); + /// ``` + pub fn build(self) -> Result { + // Determine the gamma strategy to use + let gamma_strategy: Arc = if let Some(raw) = self.raw_gamma { + // Validate and create FixedGamma from raw value + Arc::new(FixedGamma::new(raw)?) + } else { + Arc::from(self.gamma_strategy) + }; + + let rng = match self.seed { + Some(s) => StdRng::seed_from_u64(s), + None => StdRng::from_os_rng(), + }; + + Ok(MultivariateTpeSampler { + gamma_strategy, + n_startup_trials: self.n_startup_trials, + n_ei_candidates: self.n_ei_candidates, + group: self.group, + warn_independent_sampling: self.warn_independent_sampling, + constant_liar: self.constant_liar, + rng: Mutex::new(rng), + joint_sample_cache: Mutex::new(None), + }) + } +} + +impl Default for MultivariateTpeSamplerBuilder { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::error::Error; + use crate::sampler::tpe::{HyperoptGamma, LinearGamma, SqrtGamma}; + + #[test] + fn test_multivariate_tpe_sampler_new() { + let sampler = MultivariateTpeSampler::new(); + + // Check default values + assert!( + (sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON, + "Default gamma should be 0.25" + ); + assert_eq!(sampler.n_startup_trials(), 10); + assert_eq!(sampler.n_ei_candidates(), 24); + assert!(!sampler.group()); + assert!(sampler.warn_independent_sampling()); + } + + #[test] + fn test_multivariate_tpe_sampler_default() { + let sampler = MultivariateTpeSampler::default(); + + // Default should match new() + assert!( + (sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON, + "Default gamma should be 0.25" + ); + assert_eq!(sampler.n_startup_trials(), 10); + assert_eq!(sampler.n_ei_candidates(), 24); + assert!(!sampler.group()); + assert!(sampler.warn_independent_sampling()); + } + + // ======================================================================== + // Builder Tests + // ======================================================================== + + #[test] + fn test_builder_default() { + let builder = MultivariateTpeSamplerBuilder::new(); + let sampler = builder.build().unwrap(); + + assert!( + (sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON, + "Default gamma should be 0.25" + ); + assert_eq!(sampler.n_startup_trials(), 10); + assert_eq!(sampler.n_ei_candidates(), 24); + assert!(!sampler.group()); + assert!(sampler.warn_independent_sampling()); + } + + #[test] + fn test_builder_default_impl() { + let builder = MultivariateTpeSamplerBuilder::default(); + let sampler = builder.build().unwrap(); + + assert!( + (sampler.gamma_strategy().gamma(0) - 0.25).abs() < f64::EPSILON, + "Default gamma should be 0.25" + ); + assert_eq!(sampler.n_startup_trials(), 10); + } + + #[test] + fn test_builder_gamma() { + let sampler = MultivariateTpeSamplerBuilder::new() + .gamma(0.15) + .build() + .unwrap(); + + assert!( + (sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON, + "Gamma should be 0.15" + ); + // Gamma should be constant across trial counts (FixedGamma) + assert!( + (sampler.gamma_strategy().gamma(100) - 0.15).abs() < f64::EPSILON, + "Gamma should be constant" + ); + } + + #[test] + fn test_builder_gamma_invalid_zero() { + let result = MultivariateTpeSamplerBuilder::new().gamma(0.0).build(); + assert!(matches!(result, Err(Error::InvalidGamma(_)))); + } + + #[test] + fn test_builder_gamma_invalid_one() { + let result = MultivariateTpeSamplerBuilder::new().gamma(1.0).build(); + assert!(matches!(result, Err(Error::InvalidGamma(_)))); + } + + #[test] + fn test_builder_gamma_invalid_negative() { + let result = MultivariateTpeSamplerBuilder::new().gamma(-0.1).build(); + assert!(matches!(result, Err(Error::InvalidGamma(_)))); + } + + #[test] + fn test_builder_gamma_invalid_greater_than_one() { + let result = MultivariateTpeSamplerBuilder::new().gamma(1.5).build(); + assert!(matches!(result, Err(Error::InvalidGamma(_)))); + } + + #[test] + fn test_builder_gamma_strategy_sqrt() { + let sampler = MultivariateTpeSamplerBuilder::new() + .gamma_strategy(SqrtGamma::default()) + .build() + .unwrap(); + + // SqrtGamma decreases with more 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_builder_gamma_strategy_linear() { + let sampler = MultivariateTpeSamplerBuilder::new() + .gamma_strategy(LinearGamma::new(0.1, 0.4, 100).unwrap()) + .build() + .unwrap(); + + assert!((sampler.gamma_strategy().gamma(0) - 0.1).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(50) - 0.25).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(100) - 0.4).abs() < f64::EPSILON); + } + + #[test] + fn test_builder_gamma_strategy_hyperopt() { + let sampler = MultivariateTpeSamplerBuilder::new() + .gamma_strategy(HyperoptGamma::default()) + .build() + .unwrap(); + + // HyperoptGamma decreases with more trials + let g50 = sampler.gamma_strategy().gamma(50); + let g200 = sampler.gamma_strategy().gamma(200); + assert!(g50 > g200, "HyperoptGamma should decrease with more trials"); + } + + #[test] + fn test_builder_n_startup_trials() { + let sampler = MultivariateTpeSamplerBuilder::new() + .n_startup_trials(20) + .build() + .unwrap(); + + assert_eq!(sampler.n_startup_trials(), 20); + } + + #[test] + fn test_builder_n_ei_candidates() { + let sampler = MultivariateTpeSamplerBuilder::new() + .n_ei_candidates(48) + .build() + .unwrap(); + + assert_eq!(sampler.n_ei_candidates(), 48); + } + + #[test] + fn test_builder_group() { + let sampler = MultivariateTpeSamplerBuilder::new() + .group(true) + .build() + .unwrap(); + + assert!(sampler.group()); + } + + #[test] + fn test_builder_warn_independent_sampling() { + let sampler = MultivariateTpeSamplerBuilder::new() + .warn_independent_sampling(false) + .build() + .unwrap(); + + assert!(!sampler.warn_independent_sampling()); + } + + #[test] + fn test_builder_seed() { + // Two samplers with the same seed should produce the same sequence + // We can't directly test RNG output, but we verify build succeeds + let sampler1 = MultivariateTpeSamplerBuilder::new() + .seed(42) + .build() + .unwrap(); + + let sampler2 = MultivariateTpeSamplerBuilder::new() + .seed(42) + .build() + .unwrap(); + + // Verify both built successfully with same config + assert_eq!(sampler1.n_startup_trials(), sampler2.n_startup_trials()); + } + + #[test] + fn test_builder_all_options() { + let sampler = MultivariateTpeSamplerBuilder::new() + .gamma(0.20) + .n_startup_trials(15) + .n_ei_candidates(32) + .group(true) + .warn_independent_sampling(false) + .seed(12345) + .build() + .unwrap(); + + assert!((sampler.gamma_strategy().gamma(0) - 0.20).abs() < f64::EPSILON); + assert_eq!(sampler.n_startup_trials(), 15); + assert_eq!(sampler.n_ei_candidates(), 32); + assert!(sampler.group()); + assert!(!sampler.warn_independent_sampling()); + } + + #[test] + fn test_builder_gamma_overrides_gamma_strategy() { + // When gamma() is called after gamma_strategy(), it should take precedence + let sampler = MultivariateTpeSamplerBuilder::new() + .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_builder_gamma_strategy_overrides_gamma() { + // When gamma_strategy() is called after gamma(), it should take precedence + let sampler = MultivariateTpeSamplerBuilder::new() + .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_builder_via_sampler() { + let sampler = MultivariateTpeSampler::builder() + .gamma(0.10) + .n_startup_trials(25) + .n_ei_candidates(64) + .group(true) + .build() + .unwrap(); + + assert!((sampler.gamma_strategy().gamma(0) - 0.10).abs() < f64::EPSILON); + assert_eq!(sampler.n_startup_trials(), 25); + assert_eq!(sampler.n_ei_candidates(), 64); + assert!(sampler.group()); + } + + #[test] + fn test_builder_custom_gamma_strategy() { + #[derive(Debug, Clone)] + struct ConstantGamma(f64); + + impl GammaStrategy for ConstantGamma { + fn gamma(&self, _n_trials: usize) -> f64 { + self.0 + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + } + + let sampler = MultivariateTpeSamplerBuilder::new() + .gamma_strategy(ConstantGamma(0.33)) + .build() + .unwrap(); + + assert!((sampler.gamma_strategy().gamma(0) - 0.33).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(100) - 0.33).abs() < f64::EPSILON); + } + + #[test] + fn test_builder_clone() { + let builder = MultivariateTpeSamplerBuilder::new() + .gamma(0.20) + .n_startup_trials(15); + + let builder2 = builder.clone(); + + let sampler1 = builder.build().unwrap(); + let sampler2 = builder2.build().unwrap(); + + assert!((sampler1.gamma_strategy().gamma(0) - 0.20).abs() < f64::EPSILON); + assert!((sampler2.gamma_strategy().gamma(0) - 0.20).abs() < f64::EPSILON); + assert_eq!(sampler1.n_startup_trials(), 15); + assert_eq!(sampler2.n_startup_trials(), 15); + } + + #[test] + fn test_builder_constant_liar_default() { + let sampler = MultivariateTpeSamplerBuilder::new().build().unwrap(); + + assert_eq!(*sampler.constant_liar(), ConstantLiarStrategy::None); + } + + #[test] + fn test_builder_constant_liar_mean() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Mean) + .build() + .unwrap(); + + assert_eq!(*sampler.constant_liar(), ConstantLiarStrategy::Mean); + } + + #[test] + fn test_builder_constant_liar_best() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Best) + .build() + .unwrap(); + + assert_eq!(*sampler.constant_liar(), ConstantLiarStrategy::Best); + } + + #[test] + fn test_builder_constant_liar_worst() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Worst) + .build() + .unwrap(); + + assert_eq!(*sampler.constant_liar(), ConstantLiarStrategy::Worst); + } + + #[test] + fn test_builder_constant_liar_custom() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Custom(0.5)) + .build() + .unwrap(); + + match sampler.constant_liar() { + ConstantLiarStrategy::Custom(v) => { + assert!((v - 0.5).abs() < f64::EPSILON); + } + _ => panic!("Expected Custom variant"), + } + } + + #[test] + fn test_builder_constant_liar_with_other_options() { + let sampler = MultivariateTpeSamplerBuilder::new() + .gamma(0.20) + .n_startup_trials(15) + .n_ei_candidates(32) + .constant_liar(ConstantLiarStrategy::Worst) + .seed(42) + .build() + .unwrap(); + + assert!((sampler.gamma_strategy().gamma(0) - 0.20).abs() < f64::EPSILON); + assert_eq!(sampler.n_startup_trials(), 15); + assert_eq!(sampler.n_ei_candidates(), 32); + assert_eq!(*sampler.constant_liar(), ConstantLiarStrategy::Worst); + } + + // ======================================================================== + // impute_pending_trials Tests + // ======================================================================== + + mod impute_pending_trials_tests { + use std::collections::HashMap; + + use super::*; + use crate::distribution::FloatDistribution; + use crate::param::ParamValue; + use crate::sampler::{CompletedTrial, PendingTrial}; + + fn float_dist() -> Distribution { + Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }) + } + + fn create_completed_trial(id: u64, x_value: f64, objective: f64) -> CompletedTrial { + let mut params = HashMap::new(); + params.insert("x".to_string(), ParamValue::Float(x_value)); + let mut distributions = HashMap::new(); + distributions.insert("x".to_string(), float_dist()); + CompletedTrial::new(id, params, distributions, objective) + } + + fn create_pending_trial(id: u64, x_value: f64) -> PendingTrial { + let mut params = HashMap::new(); + params.insert("x".to_string(), ParamValue::Float(x_value)); + let mut distributions = HashMap::new(); + distributions.insert("x".to_string(), float_dist()); + PendingTrial::new(id, params, distributions) + } + + #[test] + fn test_impute_none_strategy_ignores_pending() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::None) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, 1.0), + create_completed_trial(1, 0.8, 3.0), + ]; + let pending = vec![create_pending_trial(2, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Should only have the completed trials + assert_eq!(result.len(), 2); + assert!(result.iter().all(|t| t.id == 0 || t.id == 1)); + } + + #[test] + fn test_impute_mean_strategy() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Mean) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, 1.0), + create_completed_trial(1, 0.8, 3.0), + ]; + let pending = vec![create_pending_trial(2, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Should have 3 trials + assert_eq!(result.len(), 3); + + // The pending trial should have mean value (1.0 + 3.0) / 2 = 2.0 + let imputed = result.iter().find(|t| t.id == 2).unwrap(); + assert!((imputed.value - 2.0).abs() < f64::EPSILON); + } + + #[test] + fn test_impute_best_strategy() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Best) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, 1.0), + create_completed_trial(1, 0.8, 3.0), + create_completed_trial(2, 0.5, 2.0), + ]; + let pending = vec![create_pending_trial(3, 0.6)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + assert_eq!(result.len(), 4); + + // Best (minimum) is 1.0 + let imputed = result.iter().find(|t| t.id == 3).unwrap(); + assert!((imputed.value - 1.0).abs() < f64::EPSILON); + } + + #[test] + fn test_impute_worst_strategy() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Worst) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, 1.0), + create_completed_trial(1, 0.8, 3.0), + create_completed_trial(2, 0.5, 2.0), + ]; + let pending = vec![create_pending_trial(3, 0.6)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + assert_eq!(result.len(), 4); + + // Worst (maximum) is 3.0 + let imputed = result.iter().find(|t| t.id == 3).unwrap(); + assert!((imputed.value - 3.0).abs() < f64::EPSILON); + } + + #[test] + fn test_impute_custom_strategy() { + let custom_value = 42.0; + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Custom(custom_value)) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, 1.0), + create_completed_trial(1, 0.8, 3.0), + ]; + let pending = vec![create_pending_trial(2, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + assert_eq!(result.len(), 3); + + // Should use custom value regardless of completed values + let imputed = result.iter().find(|t| t.id == 2).unwrap(); + assert!((imputed.value - custom_value).abs() < f64::EPSILON); + } + + #[test] + fn test_impute_multiple_pending_trials() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Mean) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, 1.0), + create_completed_trial(1, 0.8, 5.0), + ]; + let pending = vec![ + create_pending_trial(2, 0.3), + create_pending_trial(3, 0.7), + create_pending_trial(4, 0.5), + ]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Should have 5 trials total + assert_eq!(result.len(), 5); + + // All pending trials should have mean value (1.0 + 5.0) / 2 = 3.0 + let mean_value = 3.0; + for id in [2, 3, 4] { + let imputed = result.iter().find(|t| t.id == id).unwrap(); + assert!( + (imputed.value - mean_value).abs() < f64::EPSILON, + "Trial {id} should have imputed value {mean_value}" + ); + } + } + + #[test] + fn test_impute_empty_pending_trials() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Mean) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, 1.0), + create_completed_trial(1, 0.8, 3.0), + ]; + let pending: Vec = vec![]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Should just return the completed trials unchanged + assert_eq!(result.len(), 2); + } + + #[test] + fn test_impute_empty_completed_trials_mean() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Mean) + .build() + .unwrap(); + + let completed: Vec = vec![]; + let pending = vec![create_pending_trial(0, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Should have 1 trial with imputed value 0.0 (mean of empty is 0) + assert_eq!(result.len(), 1); + assert!((result[0].value - 0.0).abs() < f64::EPSILON); + } + + #[test] + fn test_impute_empty_completed_trials_best() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Best) + .build() + .unwrap(); + + let completed: Vec = vec![]; + let pending = vec![create_pending_trial(0, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Best of empty is INFINITY + assert_eq!(result.len(), 1); + assert!(result[0].value.is_infinite() && result[0].value > 0.0); + } + + #[test] + fn test_impute_empty_completed_trials_worst() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Worst) + .build() + .unwrap(); + + let completed: Vec = vec![]; + let pending = vec![create_pending_trial(0, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Worst of empty is NEG_INFINITY + assert_eq!(result.len(), 1); + assert!(result[0].value.is_infinite() && result[0].value < 0.0); + } + + #[test] + fn test_impute_empty_completed_trials_custom() { + let custom_value = 100.0; + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Custom(custom_value)) + .build() + .unwrap(); + + let completed: Vec = vec![]; + let pending = vec![create_pending_trial(0, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Custom value is used regardless of completed trials + assert_eq!(result.len(), 1); + assert!((result[0].value - custom_value).abs() < f64::EPSILON); + } + + #[test] + fn test_impute_preserves_pending_trial_params() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Mean) + .build() + .unwrap(); + + let completed = vec![create_completed_trial(0, 0.2, 1.0)]; + + // Create a pending trial with specific parameter value + let mut params = HashMap::new(); + params.insert("x".to_string(), ParamValue::Float(0.777)); + let mut distributions = HashMap::new(); + distributions.insert("x".to_string(), float_dist()); + let pending = vec![PendingTrial::new(1, params, distributions)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + assert_eq!(result.len(), 2); + + let imputed = result.iter().find(|t| t.id == 1).unwrap(); + + // Parameter value should be preserved + if let Some(ParamValue::Float(v)) = imputed.params.get("x") { + assert!((*v - 0.777).abs() < f64::EPSILON); + } else { + panic!("Expected Float parameter 'x'"); + } + + // Distribution should be preserved + assert!(imputed.distributions.contains_key("x")); + } + + #[test] + fn test_impute_with_negative_values() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Mean) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, -5.0), + create_completed_trial(1, 0.8, 3.0), + ]; + let pending = vec![create_pending_trial(2, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Mean is (-5.0 + 3.0) / 2 = -1.0 + let imputed = result.iter().find(|t| t.id == 2).unwrap(); + assert!((imputed.value - (-1.0)).abs() < f64::EPSILON); + } + + #[test] + fn test_impute_best_with_negative_values() { + let sampler = MultivariateTpeSamplerBuilder::new() + .constant_liar(ConstantLiarStrategy::Best) + .build() + .unwrap(); + + let completed = vec![ + create_completed_trial(0, 0.2, -5.0), + create_completed_trial(1, 0.8, 3.0), + ]; + let pending = vec![create_pending_trial(2, 0.5)]; + + let result = sampler.impute_pending_trials(&pending, &completed); + + // Best (minimum) is -5.0 + let imputed = result.iter().find(|t| t.id == 2).unwrap(); + assert!((imputed.value - (-5.0)).abs() < f64::EPSILON); + } + } + + // ======================================================================== + // filter_trials Tests + // ======================================================================== + + mod filter_trials_tests { + use std::collections::HashMap; + + use super::*; + use crate::distribution::{FloatDistribution, IntDistribution}; + use crate::param::ParamValue; + use crate::sampler::CompletedTrial; + + fn create_trial( + id: u64, + params: Vec<(&str, ParamValue, Distribution)>, + value: f64, + ) -> 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) + } + + fn float_dist() -> Distribution { + Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }) + } + + fn int_dist() -> Distribution { + Distribution::Int(IntDistribution { + low: 1, + high: 10, + log_scale: false, + step: None, + }) + } + + #[test] + fn test_filter_trials_empty_history() { + let sampler = MultivariateTpeSampler::new(); + let history: Vec = vec![]; + let search_space: HashMap = HashMap::new(); + + let filtered = sampler.filter_trials(&history, &search_space); + assert!(filtered.is_empty()); + } + + #[test] + fn test_filter_trials_empty_search_space() { + let sampler = MultivariateTpeSampler::new(); + let history = vec![ + create_trial(0, vec![("x", ParamValue::Float(0.5), float_dist())], 1.0), + create_trial(1, vec![("y", ParamValue::Float(0.3), float_dist())], 0.5), + ]; + let search_space: HashMap = HashMap::new(); + + // With empty search space, all trials should pass (vacuously true) + let filtered = sampler.filter_trials(&history, &search_space); + assert_eq!(filtered.len(), 2); + } + + #[test] + fn test_filter_trials_all_match() { + let sampler = MultivariateTpeSampler::new(); + let history = vec![ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist()), + ("y", ParamValue::Float(0.3), float_dist()), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("x", ParamValue::Float(0.7), float_dist()), + ("y", ParamValue::Float(0.2), float_dist()), + ], + 0.5, + ), + ]; + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist()); + search_space.insert("y".to_string(), float_dist()); + + let filtered = sampler.filter_trials(&history, &search_space); + assert_eq!(filtered.len(), 2); + assert_eq!(filtered[0].id, 0); + assert_eq!(filtered[1].id, 1); + } + + #[test] + fn test_filter_trials_partial_match() { + let sampler = MultivariateTpeSampler::new(); + + // Trial 0: has x and y + // Trial 1: has only x + // Trial 2: has x and y + let history = vec![ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist()), + ("y", ParamValue::Float(0.3), float_dist()), + ], + 1.0, + ), + create_trial(1, vec![("x", ParamValue::Float(0.7), float_dist())], 0.5), + create_trial( + 2, + vec![ + ("x", ParamValue::Float(0.6), float_dist()), + ("y", ParamValue::Float(0.4), float_dist()), + ], + 0.8, + ), + ]; + + // Search space requires both x and y + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist()); + search_space.insert("y".to_string(), float_dist()); + + let filtered = sampler.filter_trials(&history, &search_space); + + // Only trials 0 and 2 should match (trial 1 is missing y) + assert_eq!(filtered.len(), 2); + assert_eq!(filtered[0].id, 0); + assert_eq!(filtered[1].id, 2); + } + + #[test] + fn test_filter_trials_none_match() { + let sampler = MultivariateTpeSampler::new(); + + // All trials have x, but search space requires both x and z + let history = vec![ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist()), + ("y", ParamValue::Float(0.3), float_dist()), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("x", ParamValue::Float(0.7), float_dist()), + ("y", ParamValue::Float(0.2), float_dist()), + ], + 0.5, + ), + ]; + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist()); + search_space.insert("z".to_string(), float_dist()); // No trial has z + + let filtered = sampler.filter_trials(&history, &search_space); + + // No trials should match since none have z + assert!(filtered.is_empty()); + } + + #[test] + fn test_filter_trials_mixed_param_types() { + let sampler = MultivariateTpeSampler::new(); + + // Trials with mixed parameter types + let history = vec![ + create_trial( + 0, + vec![ + ("learning_rate", ParamValue::Float(0.01), float_dist()), + ("n_layers", ParamValue::Int(3), int_dist()), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("learning_rate", ParamValue::Float(0.001), float_dist()), + ("n_layers", ParamValue::Int(5), int_dist()), + ("dropout", ParamValue::Float(0.2), float_dist()), // Extra param + ], + 0.8, + ), + create_trial( + 2, + vec![ + ("learning_rate", ParamValue::Float(0.005), float_dist()), + // Missing n_layers + ("dropout", ParamValue::Float(0.1), float_dist()), + ], + 0.9, + ), + ]; + + // Search space requires learning_rate and n_layers + let mut search_space = HashMap::new(); + search_space.insert("learning_rate".to_string(), float_dist()); + search_space.insert("n_layers".to_string(), int_dist()); + + let filtered = sampler.filter_trials(&history, &search_space); + + // Trials 0 and 1 have both params, trial 2 is missing n_layers + assert_eq!(filtered.len(), 2); + assert_eq!(filtered[0].id, 0); + assert_eq!(filtered[1].id, 1); + } + + #[test] + fn test_filter_trials_superset_params_accepted() { + let sampler = MultivariateTpeSampler::new(); + + // All trials have more params than the search space requires + let history = vec![ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist()), + ("y", ParamValue::Float(0.3), float_dist()), + ("z", ParamValue::Float(0.1), float_dist()), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("x", ParamValue::Float(0.7), float_dist()), + ("y", ParamValue::Float(0.2), float_dist()), + ("w", ParamValue::Float(0.9), float_dist()), // Different extra param + ], + 0.5, + ), + ]; + + // Search space only requires x + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist()); + + let filtered = sampler.filter_trials(&history, &search_space); + + // Both trials should be accepted (they have x, even though they have extras) + assert_eq!(filtered.len(), 2); + } + + #[test] + fn test_filter_trials_preserves_order() { + let sampler = MultivariateTpeSampler::new(); + + let history: Vec = (0..10) + .map(|i| create_trial(i, vec![("x", ParamValue::Float(0.5), float_dist())], 1.0)) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist()); + + let filtered = sampler.filter_trials(&history, &search_space); + + // Order should be preserved + for (i, trial) in filtered.iter().enumerate() { + #[allow(clippy::cast_possible_truncation)] + let expected_id = i as u64; + assert_eq!(trial.id, expected_id); + } + } + + #[test] + fn test_filter_trials_single_param_search_space() { + let sampler = MultivariateTpeSampler::new(); + + // Some trials have the required param, some don't + let history = vec![ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist()), + ("y", ParamValue::Float(0.3), float_dist()), + ], + 1.0, + ), + create_trial(1, vec![("y", ParamValue::Float(0.7), float_dist())], 0.5), + create_trial(2, vec![("x", ParamValue::Float(0.6), float_dist())], 0.8), + ]; + + // Search space only requires x + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist()); + + let filtered = sampler.filter_trials(&history, &search_space); + + // Trials 0 and 2 have x + assert_eq!(filtered.len(), 2); + assert_eq!(filtered[0].id, 0); + assert_eq!(filtered[1].id, 2); + } + } + + // ======================================================================== + // split_trials Tests + // ======================================================================== + + #[allow(clippy::cast_precision_loss)] + mod split_trials_tests { + use std::collections::HashMap; + + use super::*; + use crate::distribution::FloatDistribution; + use crate::param::ParamValue; + use crate::sampler::CompletedTrial; + + fn create_trial(id: u64, value: f64) -> CompletedTrial { + let float_dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let mut params = HashMap::new(); + params.insert("x".to_string(), ParamValue::Float(0.5)); + let mut distributions = HashMap::new(); + distributions.insert("x".to_string(), float_dist); + CompletedTrial::new(id, params, distributions, value) + } + + #[test] + fn test_split_trials_empty() { + let sampler = MultivariateTpeSampler::new(); + let trials: Vec<&CompletedTrial> = vec![]; + + let (good, bad) = sampler.split_trials(&trials); + + assert!(good.is_empty()); + assert!(bad.is_empty()); + } + + #[test] + fn test_split_trials_single_trial() { + let sampler = MultivariateTpeSampler::new(); + let trial = create_trial(0, 1.0); + let trials: Vec<&CompletedTrial> = vec![&trial]; + + let (good, bad) = sampler.split_trials(&trials); + + // Single trial should go to good + assert_eq!(good.len(), 1); + assert!(bad.is_empty()); + assert_eq!(good[0].id, 0); + } + + #[test] + fn test_split_trials_two_trials() { + let sampler = MultivariateTpeSampler::new(); + let good_trial = create_trial(0, 0.5); // Better (lower) + let bad_trial = create_trial(1, 1.0); // Worse (higher) + let trial_refs: Vec<&CompletedTrial> = vec![&good_trial, &bad_trial]; + + let (good, bad) = sampler.split_trials(&trial_refs); + + // With 2 trials and gamma=0.25, ceil(2*0.25)=1, so 1 good, 1 bad + assert_eq!(good.len(), 1); + assert_eq!(bad.len(), 1); + assert_eq!(good[0].id, 0); // Lower value in good + assert_eq!(bad[0].id, 1); // Higher value in bad + } + + #[test] + fn test_split_trials_many_trials_default_gamma() { + // Default gamma is 0.25, so with 20 trials, ceil(20*0.25)=5 good + let sampler = MultivariateTpeSampler::new(); + + // Create 20 trials with values 0..20 + let trial_data: Vec = + (0..20).map(|i| create_trial(i, i as f64)).collect(); + let trial_refs: Vec<&CompletedTrial> = trial_data.iter().collect(); + + let (good, bad) = sampler.split_trials(&trial_refs); + + // With gamma=0.25 and 20 trials: ceil(20 * 0.25) = 5 good + assert_eq!(good.len(), 5); + assert_eq!(bad.len(), 15); + + // Good trials should have lowest values (0, 1, 2, 3, 4) + for trial in &good { + assert!( + trial.value < 5.0, + "Good trial has value {}, expected < 5.0", + trial.value + ); + } + + // Bad trials should have higher values (5..20) + for trial in &bad { + assert!( + trial.value >= 5.0, + "Bad trial has value {}, expected >= 5.0", + trial.value + ); + } + } + + #[test] + fn test_split_trials_custom_gamma() { + // Create sampler with gamma = 0.10 + let sampler = MultivariateTpeSampler::builder() + .gamma(0.10) + .build() + .unwrap(); + + // Create 20 trials with values 0..20 + let trial_data: Vec = + (0..20).map(|i| create_trial(i, i as f64)).collect(); + let trial_refs: Vec<&CompletedTrial> = trial_data.iter().collect(); + + let (good, bad) = sampler.split_trials(&trial_refs); + + // With gamma=0.10 and 20 trials: ceil(20 * 0.10) = 2 good + assert_eq!(good.len(), 2); + assert_eq!(bad.len(), 18); + + // Good trials should have lowest values (0, 1) + for trial in &good { + assert!( + trial.value < 2.0, + "Good trial has value {}, expected < 2.0", + trial.value + ); + } + } + + #[test] + fn test_split_trials_unsorted_input() { + let sampler = MultivariateTpeSampler::new(); + + // Create trials in non-sorted order + let trial_data = [ + create_trial(0, 5.0), + create_trial(1, 1.0), + create_trial(2, 8.0), + create_trial(3, 0.5), + create_trial(4, 3.0), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + + let (good, bad) = sampler.split_trials(&trials); + + // With 5 trials and gamma=0.25: ceil(5*0.25)=2 good, 3 bad + assert_eq!(good.len(), 2); + assert_eq!(bad.len(), 3); + + // Good should contain trials 3 (0.5) and 1 (1.0) - lowest values + let good_ids: Vec = good.iter().map(|t| t.id).collect(); + assert!( + good_ids.contains(&3), + "Trial 3 (value=0.5) should be in good" + ); + assert!( + good_ids.contains(&1), + "Trial 1 (value=1.0) should be in good" + ); + + // Bad should contain trials 0 (5.0), 2 (8.0), 4 (3.0) + let bad_ids: Vec = bad.iter().map(|t| t.id).collect(); + assert!(bad_ids.contains(&0), "Trial 0 (value=5.0) should be in bad"); + assert!(bad_ids.contains(&2), "Trial 2 (value=8.0) should be in bad"); + assert!(bad_ids.contains(&4), "Trial 4 (value=3.0) should be in bad"); + } + + #[test] + fn test_split_trials_with_ties() { + let sampler = MultivariateTpeSampler::new(); + + // Create trials with tied values + let trial_data = [ + create_trial(0, 1.0), + create_trial(1, 1.0), + create_trial(2, 2.0), + create_trial(3, 2.0), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + + let (good, bad) = sampler.split_trials(&trials); + + // With 4 trials and gamma=0.25: ceil(4*0.25)=1 good, 3 bad + assert_eq!(good.len(), 1); + assert_eq!(bad.len(), 3); + + // Good should contain one of the trials with value 1.0 + assert!( + (good[0].value - 1.0).abs() < f64::EPSILON, + "Good trial should have value 1.0" + ); + } + + #[test] + fn test_split_trials_ensures_both_groups_nonempty() { + // Even with extreme gamma values, we should have at least 1 in each group + // when there are at least 2 trials + + // Very small gamma (would give 0 good without clamping) + let sampler = MultivariateTpeSampler::builder() + .gamma(0.01) + .build() + .unwrap(); + + let trial_data = [create_trial(0, 0.5), create_trial(1, 1.0)]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + + let (good, bad) = sampler.split_trials(&trials); + + // Should still have 1 in each group + assert_eq!(good.len(), 1); + assert_eq!(bad.len(), 1); + } + + #[test] + fn test_split_trials_large_gamma() { + // Large gamma (would give all good without clamping) + let sampler = MultivariateTpeSampler::builder() + .gamma(0.99) + .build() + .unwrap(); + + let trial_data: Vec = + (0..10).map(|i| create_trial(i, i as f64)).collect(); + let trial_refs: Vec<&CompletedTrial> = trial_data.iter().collect(); + + let (good, bad) = sampler.split_trials(&trial_refs); + + // With gamma=0.99: ceil(10*0.99)=10, but min ensures at least 1 in bad + // Actually: n_good = min(10, 10-1) = 9 + assert_eq!(good.len(), 9); + assert_eq!(bad.len(), 1); + + // Bad should contain the trial with highest value + assert!( + (bad[0].value - 9.0).abs() < f64::EPSILON, + "Bad trial should have highest value" + ); + } + + #[test] + fn test_split_trials_nan_handling() { + let sampler = MultivariateTpeSampler::new(); + + // Create trials with NaN values - they should be sorted consistently + let trial_data = [ + create_trial(0, 1.0), + create_trial(1, f64::NAN), + create_trial(2, 0.5), + create_trial(3, 2.0), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + + let (good, bad) = sampler.split_trials(&trials); + + // Should complete without panic + // Total should equal input + assert_eq!(good.len() + bad.len(), 4); + } + + #[test] + fn test_split_trials_preserves_trial_references() { + let sampler = MultivariateTpeSampler::new(); + + let trial_data: Vec = + (0..5).map(|i| create_trial(i, i as f64)).collect(); + let trial_refs: Vec<&CompletedTrial> = trial_data.iter().collect(); + + let (good, bad) = sampler.split_trials(&trial_refs); + + // Verify that references point to the original data + for trial in good.iter().chain(bad.iter()) { + // Find the original trial by ID + let original = trial_data.iter().find(|t| t.id == trial.id).unwrap(); + assert!( + core::ptr::eq(*trial, original), + "Reference should point to original trial" + ); + } + } + + #[test] + fn test_split_trials_integration_with_filter() { + // Test that split_trials works correctly with filter_trials output + let sampler = MultivariateTpeSampler::new(); + + let float_dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Create trials with varying parameters + let mut trials_vec = vec![]; + for i in 0..10 { + let mut params = HashMap::new(); + params.insert("x".to_string(), ParamValue::Float(i as f64 / 10.0)); + params.insert("y".to_string(), ParamValue::Float(i as f64 / 10.0)); + let mut distributions = HashMap::new(); + distributions.insert("x".to_string(), float_dist.clone()); + distributions.insert("y".to_string(), float_dist.clone()); + trials_vec.push(CompletedTrial::new(i, params, distributions, i as f64)); + } + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist.clone()); + search_space.insert("y".to_string(), float_dist); + + // Filter trials + let filtered = sampler.filter_trials(&trials_vec, &search_space); + assert_eq!(filtered.len(), 10); + + // Split filtered trials + let (good, bad) = sampler.split_trials(&filtered); + + // With 10 trials and gamma=0.25: ceil(10*0.25)=3 good + assert_eq!(good.len(), 3); + assert_eq!(bad.len(), 7); + + // Verify good trials have lowest values + for trial in &good { + assert!(trial.value < 3.0); + } + } + } + + // ======================================================================== + // extract_observations Tests + // ======================================================================== + + #[allow(clippy::cast_precision_loss, clippy::cast_possible_wrap)] + mod extract_observations_tests { + use std::collections::HashMap; + + use super::*; + use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + use crate::param::ParamValue; + use crate::sampler::CompletedTrial; + + fn float_dist() -> Distribution { + Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }) + } + + fn int_dist() -> Distribution { + Distribution::Int(IntDistribution { + low: 1, + high: 100, + log_scale: false, + step: None, + }) + } + + fn categorical_dist(n: usize) -> Distribution { + Distribution::Categorical(CategoricalDistribution { n_choices: n }) + } + + fn create_trial( + id: u64, + params: Vec<(&str, ParamValue, Distribution)>, + value: f64, + ) -> 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_extract_observations_empty_trials() { + let sampler = MultivariateTpeSampler::new(); + let trials: Vec<&CompletedTrial> = vec![]; + let param_order = vec!["x".to_string(), "y".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + assert!(observations.is_empty()); + } + + #[test] + fn test_extract_observations_empty_param_order() { + let sampler = MultivariateTpeSampler::new(); + let trial = create_trial(0, vec![("x", ParamValue::Float(0.5), float_dist())], 1.0); + let trials: Vec<&CompletedTrial> = vec![&trial]; + let param_order: Vec = vec![]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + // Should have one row (for the trial) with zero columns + assert_eq!(observations.len(), 1); + assert!(observations[0].is_empty()); + } + + #[test] + fn test_extract_observations_single_float_param() { + let sampler = MultivariateTpeSampler::new(); + let trial_data = [ + create_trial(0, vec![("x", ParamValue::Float(0.1), float_dist())], 1.0), + create_trial(1, vec![("x", ParamValue::Float(0.5), float_dist())], 0.5), + create_trial(2, vec![("x", ParamValue::Float(0.9), float_dist())], 0.8), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + let param_order = vec!["x".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + assert_eq!(observations.len(), 3); + assert_eq!(observations[0].len(), 1); + assert!((observations[0][0] - 0.1).abs() < f64::EPSILON); + assert!((observations[1][0] - 0.5).abs() < f64::EPSILON); + assert!((observations[2][0] - 0.9).abs() < f64::EPSILON); + } + + #[test] + fn test_extract_observations_multiple_float_params() { + let sampler = MultivariateTpeSampler::new(); + let trial_data = [ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.1), float_dist()), + ("y", ParamValue::Float(0.2), float_dist()), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("x", ParamValue::Float(0.3), float_dist()), + ("y", ParamValue::Float(0.4), float_dist()), + ], + 0.5, + ), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + let param_order = vec!["x".to_string(), "y".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + assert_eq!(observations.len(), 2); + assert_eq!(observations[0].len(), 2); + assert!((observations[0][0] - 0.1).abs() < f64::EPSILON); + assert!((observations[0][1] - 0.2).abs() < f64::EPSILON); + assert!((observations[1][0] - 0.3).abs() < f64::EPSILON); + assert!((observations[1][1] - 0.4).abs() < f64::EPSILON); + } + + #[test] + fn test_extract_observations_respects_param_order() { + let sampler = MultivariateTpeSampler::new(); + let trial = create_trial( + 0, + vec![ + ("a", ParamValue::Float(1.0), float_dist()), + ("b", ParamValue::Float(2.0), float_dist()), + ("c", ParamValue::Float(3.0), float_dist()), + ], + 1.0, + ); + let trials: Vec<&CompletedTrial> = vec![&trial]; + + // Different orderings + let order_abc = vec!["a".to_string(), "b".to_string(), "c".to_string()]; + let order_cba = vec!["c".to_string(), "b".to_string(), "a".to_string()]; + let order_bac = vec!["b".to_string(), "a".to_string(), "c".to_string()]; + + let obs_abc = sampler.extract_observations(&trials, &order_abc); + let obs_cba = sampler.extract_observations(&trials, &order_cba); + let obs_bac = sampler.extract_observations(&trials, &order_bac); + + assert!((obs_abc[0][0] - 1.0).abs() < f64::EPSILON); + assert!((obs_abc[0][1] - 2.0).abs() < f64::EPSILON); + assert!((obs_abc[0][2] - 3.0).abs() < f64::EPSILON); + + assert!((obs_cba[0][0] - 3.0).abs() < f64::EPSILON); + assert!((obs_cba[0][1] - 2.0).abs() < f64::EPSILON); + assert!((obs_cba[0][2] - 1.0).abs() < f64::EPSILON); + + assert!((obs_bac[0][0] - 2.0).abs() < f64::EPSILON); + assert!((obs_bac[0][1] - 1.0).abs() < f64::EPSILON); + assert!((obs_bac[0][2] - 3.0).abs() < f64::EPSILON); + } + + #[test] + fn test_extract_observations_int_conversion() { + let sampler = MultivariateTpeSampler::new(); + let trial_data = [ + create_trial(0, vec![("n_layers", ParamValue::Int(3), int_dist())], 1.0), + create_trial(1, vec![("n_layers", ParamValue::Int(5), int_dist())], 0.5), + create_trial(2, vec![("n_layers", ParamValue::Int(10), int_dist())], 0.8), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + let param_order = vec!["n_layers".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + assert_eq!(observations.len(), 3); + assert!((observations[0][0] - 3.0).abs() < f64::EPSILON); + assert!((observations[1][0] - 5.0).abs() < f64::EPSILON); + assert!((observations[2][0] - 10.0).abs() < f64::EPSILON); + } + + #[test] + fn test_extract_observations_mixed_float_and_int() { + let sampler = MultivariateTpeSampler::new(); + let trial_data = [ + create_trial( + 0, + vec![ + ("learning_rate", ParamValue::Float(0.01), float_dist()), + ("n_layers", ParamValue::Int(3), int_dist()), + ("batch_size", ParamValue::Int(32), int_dist()), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("learning_rate", ParamValue::Float(0.001), float_dist()), + ("n_layers", ParamValue::Int(5), int_dist()), + ("batch_size", ParamValue::Int(64), int_dist()), + ], + 0.5, + ), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + let param_order = vec![ + "learning_rate".to_string(), + "n_layers".to_string(), + "batch_size".to_string(), + ]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + assert_eq!(observations.len(), 2); + assert_eq!(observations[0].len(), 3); + + assert!((observations[0][0] - 0.01).abs() < f64::EPSILON); + assert!((observations[0][1] - 3.0).abs() < f64::EPSILON); + assert!((observations[0][2] - 32.0).abs() < f64::EPSILON); + + assert!((observations[1][0] - 0.001).abs() < f64::EPSILON); + assert!((observations[1][1] - 5.0).abs() < f64::EPSILON); + assert!((observations[1][2] - 64.0).abs() < f64::EPSILON); + } + + #[test] + fn test_extract_observations_skips_categorical() { + let sampler = MultivariateTpeSampler::new(); + let trial_data = [ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist()), + ("optimizer", ParamValue::Categorical(1), categorical_dist(3)), + ("y", ParamValue::Float(0.3), float_dist()), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("x", ParamValue::Float(0.7), float_dist()), + ("optimizer", ParamValue::Categorical(0), categorical_dist(3)), + ("y", ParamValue::Float(0.2), float_dist()), + ], + 0.5, + ), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + + // Include categorical in param order - it should be skipped + let param_order = vec!["x".to_string(), "optimizer".to_string(), "y".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + // Each observation should have only 2 values (x and y), not 3 + assert_eq!(observations.len(), 2); + assert_eq!(observations[0].len(), 2); + assert_eq!(observations[1].len(), 2); + + // First row: x=0.5, y=0.3 + assert!((observations[0][0] - 0.5).abs() < f64::EPSILON); + assert!((observations[0][1] - 0.3).abs() < f64::EPSILON); + + // Second row: x=0.7, y=0.2 + assert!((observations[1][0] - 0.7).abs() < f64::EPSILON); + assert!((observations[1][1] - 0.2).abs() < f64::EPSILON); + } + + #[test] + fn test_extract_observations_only_categorical_in_order() { + let sampler = MultivariateTpeSampler::new(); + let trial = create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist()), + ("optimizer", ParamValue::Categorical(1), categorical_dist(3)), + ], + 1.0, + ); + let trials: Vec<&CompletedTrial> = vec![&trial]; + + // Only request categorical param + let param_order = vec!["optimizer".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + // Should have one row with zero columns (categorical skipped) + assert_eq!(observations.len(), 1); + assert!(observations[0].is_empty()); + } + + #[test] + fn test_extract_observations_missing_param() { + let sampler = MultivariateTpeSampler::new(); + let trial_data = [ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist()), + ("y", ParamValue::Float(0.3), float_dist()), + ], + 1.0, + ), + // This trial is missing "y" + create_trial(1, vec![("x", ParamValue::Float(0.7), float_dist())], 0.5), + ]; + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + let param_order = vec!["x".to_string(), "y".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + // First trial should have both params + assert_eq!(observations[0].len(), 2); + // Second trial is missing y, so it only has x + assert_eq!(observations[1].len(), 1); + } + + #[test] + fn test_extract_observations_large_int_values() { + let sampler = MultivariateTpeSampler::new(); + // Test with large integer values to verify precision + let trial = create_trial( + 0, + vec![ + ("small_int", ParamValue::Int(1), int_dist()), + ("medium_int", ParamValue::Int(1_000_000), int_dist()), + ("negative_int", ParamValue::Int(-42), int_dist()), + ], + 1.0, + ); + let trials: Vec<&CompletedTrial> = vec![&trial]; + let param_order = vec![ + "small_int".to_string(), + "medium_int".to_string(), + "negative_int".to_string(), + ]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + assert_eq!(observations.len(), 1); + assert_eq!(observations[0].len(), 3); + assert!((observations[0][0] - 1.0).abs() < f64::EPSILON); + assert!((observations[0][1] - 1_000_000.0).abs() < f64::EPSILON); + assert!((observations[0][2] - (-42.0)).abs() < f64::EPSILON); + } + + #[test] + fn test_extract_observations_many_trials() { + let sampler = MultivariateTpeSampler::new(); + + // Create 100 trials with predictable values + let trial_data: Vec = (0..100) + .map(|i| { + create_trial( + i, + vec![ + ("x", ParamValue::Float(i as f64 / 100.0), float_dist()), + ("y", ParamValue::Int(i as i64), int_dist()), + ], + i as f64, + ) + }) + .collect(); + let trials: Vec<&CompletedTrial> = trial_data.iter().collect(); + let param_order = vec!["x".to_string(), "y".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + assert_eq!(observations.len(), 100); + for (i, obs) in observations.iter().enumerate() { + assert_eq!(obs.len(), 2); + assert!((obs[0] - i as f64 / 100.0).abs() < f64::EPSILON); + assert!((obs[1] - i as f64).abs() < f64::EPSILON); + } + } + + #[test] + fn test_extract_observations_subset_of_params() { + let sampler = MultivariateTpeSampler::new(); + let trial = create_trial( + 0, + vec![ + ("a", ParamValue::Float(1.0), float_dist()), + ("b", ParamValue::Float(2.0), float_dist()), + ("c", ParamValue::Float(3.0), float_dist()), + ("d", ParamValue::Float(4.0), float_dist()), + ], + 1.0, + ); + let trials: Vec<&CompletedTrial> = vec![&trial]; + + // Only extract a subset of params + let param_order = vec!["b".to_string(), "d".to_string()]; + + let observations = sampler.extract_observations(&trials, ¶m_order); + + assert_eq!(observations.len(), 1); + assert_eq!(observations[0].len(), 2); + assert!((observations[0][0] - 2.0).abs() < f64::EPSILON); // b + assert!((observations[0][1] - 4.0).abs() < f64::EPSILON); // d + } + + #[test] + fn test_extract_observations_integration_with_pipeline() { + // Test full pipeline: filter -> split -> extract + let sampler = MultivariateTpeSampler::new(); + + let float_dist_val = float_dist(); + let int_dist_val = int_dist(); + + let trial_data: Vec = (0..20) + .map(|i| { + let mut params = HashMap::new(); + params.insert("x".to_string(), ParamValue::Float(i as f64 / 20.0)); + params.insert("n".to_string(), ParamValue::Int(i as i64)); + let mut distributions = HashMap::new(); + distributions.insert("x".to_string(), float_dist_val.clone()); + distributions.insert("n".to_string(), int_dist_val.clone()); + CompletedTrial::new(i, params, distributions, i as f64) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist_val); + search_space.insert("n".to_string(), int_dist_val); + + // Filter + let filtered = sampler.filter_trials(&trial_data, &search_space); + assert_eq!(filtered.len(), 20); + + // Split + let (good, bad) = sampler.split_trials(&filtered); + assert_eq!(good.len(), 5); // gamma=0.25, ceil(20*0.25)=5 + assert_eq!(bad.len(), 15); + + // Extract from good trials + let param_order = vec!["x".to_string(), "n".to_string()]; + let good_obs = sampler.extract_observations(&good, ¶m_order); + let bad_obs = sampler.extract_observations(&bad, ¶m_order); + + assert_eq!(good_obs.len(), 5); + assert_eq!(bad_obs.len(), 15); + + // Good observations should all have low values (trials 0-4) + for obs in &good_obs { + assert_eq!(obs.len(), 2); + assert!(obs[0] < 0.25); // x < 5/20 = 0.25 + assert!(obs[1] < 5.0); // n < 5 + } + + // Bad observations should have higher values + for obs in &bad_obs { + assert_eq!(obs.len(), 2); + assert!(obs[0] >= 0.25 || obs[1] >= 5.0); + } + } + } + + // ======================================================================== + // select_candidate Tests + // ======================================================================== + + mod select_candidate_tests { + use super::*; + use crate::kde::MultivariateKDE; + + #[test] + fn test_select_candidate_basic() { + // Create a sampler with a fixed seed for reproducibility + let sampler = MultivariateTpeSampler::builder() + .n_ei_candidates(24) + .seed(42) + .build() + .unwrap(); + + // Good samples clustered around (0.1, 0.1) + let good_samples = vec![ + vec![0.1, 0.1], + vec![0.15, 0.12], + vec![0.08, 0.11], + vec![0.12, 0.09], + vec![0.11, 0.13], + ]; + + // Bad samples clustered around (0.9, 0.9) + let bad_samples = vec![ + vec![0.9, 0.9], + vec![0.85, 0.88], + vec![0.92, 0.91], + vec![0.88, 0.93], + vec![0.91, 0.87], + ]; + + let good_kde = MultivariateKDE::new(good_samples).unwrap(); + let bad_kde = MultivariateKDE::new(bad_samples).unwrap(); + + let selected = sampler.select_candidate(&good_kde, &bad_kde); + + // The selected candidate should have 2 dimensions + assert_eq!(selected.len(), 2); + + // The selected point should be closer to the good region than the bad region + // (though not always perfectly so due to stochasticity) + let dist_to_good = ((selected[0] - 0.1).powi(2) + (selected[1] - 0.1).powi(2)).sqrt(); + let dist_to_bad = ((selected[0] - 0.9).powi(2) + (selected[1] - 0.9).powi(2)).sqrt(); + + assert!( + dist_to_good < dist_to_bad, + "Selected point ({}, {}) is closer to bad region than good region", + selected[0], + selected[1] + ); + } + + #[test] + fn test_select_candidate_returns_correct_dimension() { + let sampler = MultivariateTpeSampler::builder() + .n_ei_candidates(10) + .seed(123) + .build() + .unwrap(); + + // 3D case + let good_samples = vec![ + vec![0.1, 0.2, 0.3], + vec![0.15, 0.25, 0.35], + vec![0.12, 0.22, 0.32], + ]; + let bad_samples = vec![ + vec![0.8, 0.7, 0.6], + vec![0.85, 0.75, 0.65], + vec![0.82, 0.72, 0.62], + ]; + + let good_kde = MultivariateKDE::new(good_samples).unwrap(); + let bad_kde = MultivariateKDE::new(bad_samples).unwrap(); + + let selected = sampler.select_candidate(&good_kde, &bad_kde); + assert_eq!(selected.len(), 3); + } + + #[test] + fn test_select_candidate_one_dimension() { + let sampler = MultivariateTpeSampler::builder() + .n_ei_candidates(20) + .seed(456) + .build() + .unwrap(); + + // 1D case - good samples near 0, bad samples near 10 + let good_samples = vec![vec![0.0], vec![0.5], vec![1.0], vec![0.3], vec![0.7]]; + let bad_samples = vec![vec![8.0], vec![9.0], vec![10.0], vec![8.5], vec![9.5]]; + + let good_kde = MultivariateKDE::new(good_samples).unwrap(); + let bad_kde = MultivariateKDE::new(bad_samples).unwrap(); + + let selected = sampler.select_candidate(&good_kde, &bad_kde); + assert_eq!(selected.len(), 1); + + // Selected value should be closer to the good region + assert!( + selected[0] < 5.0, + "Selected value {} should be closer to good region (< 5.0)", + selected[0] + ); + } + + #[test] + fn test_select_candidate_reproducibility() { + // Two samplers with the same seed should produce the same results + let sampler1 = MultivariateTpeSampler::builder() + .n_ei_candidates(24) + .seed(999) + .build() + .unwrap(); + + let sampler2 = MultivariateTpeSampler::builder() + .n_ei_candidates(24) + .seed(999) + .build() + .unwrap(); + + let good_samples = vec![vec![1.0, 2.0], vec![1.5, 2.5], vec![1.2, 2.2]]; + let bad_samples = vec![vec![8.0, 9.0], vec![8.5, 9.5], vec![8.2, 9.2]]; + + let good_kde = MultivariateKDE::new(good_samples.clone()).unwrap(); + let bad_kde = MultivariateKDE::new(bad_samples.clone()).unwrap(); + + let selected1 = sampler1.select_candidate(&good_kde, &bad_kde); + + // Need to recreate KDEs for second sampler since we consumed them + let good_kde2 = MultivariateKDE::new(good_samples).unwrap(); + let bad_kde2 = MultivariateKDE::new(bad_samples).unwrap(); + + let selected2 = sampler2.select_candidate(&good_kde2, &bad_kde2); + + // With same seed, should get same result + assert!( + (selected1[0] - selected2[0]).abs() < f64::EPSILON, + "Dimension 0: {} vs {}", + selected1[0], + selected2[0] + ); + assert!( + (selected1[1] - selected2[1]).abs() < f64::EPSILON, + "Dimension 1: {} vs {}", + selected1[1], + selected2[1] + ); + } + + #[test] + fn test_select_candidate_with_single_candidate() { + // With n_ei_candidates=1, should still work + let sampler = MultivariateTpeSampler::builder() + .n_ei_candidates(1) + .seed(789) + .build() + .unwrap(); + + let good_samples = vec![vec![0.0, 0.0], vec![0.1, 0.1]]; + let bad_samples = vec![vec![5.0, 5.0], vec![5.1, 5.1]]; + + let good_kde = MultivariateKDE::new(good_samples).unwrap(); + let bad_kde = MultivariateKDE::new(bad_samples).unwrap(); + + let selected = sampler.select_candidate(&good_kde, &bad_kde); + assert_eq!(selected.len(), 2); + } + + #[test] + fn test_select_candidate_many_candidates() { + // With many candidates, should find a good point + let sampler = MultivariateTpeSampler::builder() + .n_ei_candidates(100) + .seed(111) + .build() + .unwrap(); + + let good_samples = vec![ + vec![0.0, 0.0], + vec![0.1, 0.1], + vec![0.2, 0.2], + vec![0.05, 0.15], + vec![0.15, 0.05], + ]; + let bad_samples = vec![ + vec![10.0, 10.0], + vec![10.1, 10.1], + vec![10.2, 10.2], + vec![10.05, 10.15], + vec![10.15, 10.05], + ]; + + let good_kde = MultivariateKDE::new(good_samples).unwrap(); + let bad_kde = MultivariateKDE::new(bad_samples).unwrap(); + + let selected = sampler.select_candidate(&good_kde, &bad_kde); + + // With more candidates, should definitely find a point in the good region + assert!( + selected[0] < 5.0 && selected[1] < 5.0, + "Selected point ({}, {}) should be in good region", + selected[0], + selected[1] + ); + } + + #[test] + fn test_select_candidate_overlapping_distributions() { + // When distributions overlap, selection should still work + let sampler = MultivariateTpeSampler::builder() + .n_ei_candidates(24) + .seed(222) + .build() + .unwrap(); + + // Overlapping distributions - good centered at 0, bad centered at 1 + let good_samples = vec![ + vec![0.0, 0.0], + vec![0.5, 0.5], + vec![-0.5, -0.5], + vec![0.3, -0.3], + vec![-0.3, 0.3], + ]; + let bad_samples = vec![ + vec![1.0, 1.0], + vec![1.5, 1.5], + vec![0.5, 0.5], // This overlaps with good + vec![1.3, 0.7], + vec![0.7, 1.3], + ]; + + let good_kde = MultivariateKDE::new(good_samples).unwrap(); + let bad_kde = MultivariateKDE::new(bad_samples).unwrap(); + + let selected = sampler.select_candidate(&good_kde, &bad_kde); + + // Should still return a valid point + assert_eq!(selected.len(), 2); + assert!(selected[0].is_finite()); + assert!(selected[1].is_finite()); + } + + #[test] + fn test_select_candidate_high_dimensional() { + // Test with higher dimensions (5D) + let sampler = MultivariateTpeSampler::builder() + .n_ei_candidates(50) + .seed(333) + .build() + .unwrap(); + + // Good samples near origin + let good_samples: Vec> = (0..10) + .map(|i| { + let offset = f64::from(i) * 0.01; + vec![offset, offset, offset, offset, offset] + }) + .collect(); + + // Bad samples far from origin + let bad_samples: Vec> = (0..10) + .map(|i| { + let offset = 10.0 + f64::from(i) * 0.01; + vec![offset, offset, offset, offset, offset] + }) + .collect(); + + let good_kde = MultivariateKDE::new(good_samples).unwrap(); + let bad_kde = MultivariateKDE::new(bad_samples).unwrap(); + + let selected = sampler.select_candidate(&good_kde, &bad_kde); + + assert_eq!(selected.len(), 5); + + // All dimensions should be closer to 0 than to 10 + for (dim, &val) in selected.iter().enumerate() { + assert!( + val.abs() < 5.0, + "Dimension {dim} value {val} should be closer to origin" + ); + } + } + + #[test] + fn test_select_candidate_integration_with_pipeline() { + // Full integration test: extract observations -> fit KDEs -> select candidate + use crate::distribution::FloatDistribution; + use crate::param::ParamValue; + + let sampler = MultivariateTpeSampler::builder() + .gamma(0.25) + .n_ei_candidates(24) + .seed(444) + .build() + .unwrap(); + + let float_dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 10.0, + log_scale: false, + step: None, + }); + + // Create trials with objective values equal to x + y + // So good trials have low x and y values + let trial_data: Vec = (0..20) + .map(|i| { + #[allow(clippy::cast_precision_loss)] + let x = (i as f64) / 2.0; + #[allow(clippy::cast_precision_loss)] + let y = (i as f64) / 2.0; + let mut params = HashMap::new(); + params.insert("x".to_string(), ParamValue::Float(x)); + params.insert("y".to_string(), ParamValue::Float(y)); + let mut distributions = HashMap::new(); + distributions.insert("x".to_string(), float_dist.clone()); + distributions.insert("y".to_string(), float_dist.clone()); + CompletedTrial::new(i, params, distributions, x + y) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist.clone()); + search_space.insert("y".to_string(), float_dist); + + // Filter and split + let filtered = sampler.filter_trials(&trial_data, &search_space); + let (good, bad) = sampler.split_trials(&filtered); + + // Extract observations + let param_order = vec!["x".to_string(), "y".to_string()]; + let good_obs = sampler.extract_observations(&good, ¶m_order); + let bad_obs = sampler.extract_observations(&bad, ¶m_order); + + // Fit KDEs + let good_kde = MultivariateKDE::new(good_obs).unwrap(); + let bad_kde = MultivariateKDE::new(bad_obs).unwrap(); + + // Select candidate + let selected = sampler.select_candidate(&good_kde, &bad_kde); + + assert_eq!(selected.len(), 2); + + // Selected point should be in the "good" region (low x, low y) + // Good trials have x, y in [0, 2.5), bad trials have x, y in [2.5, 10) + assert!( + selected[0] < 5.0 && selected[1] < 5.0, + "Selected point ({}, {}) should be in good region", + selected[0], + selected[1] + ); + } + } + + // ======================================================================== + // sample_joint Tests + // ======================================================================== + + #[allow( + clippy::cast_precision_loss, + clippy::cast_possible_wrap, + clippy::cast_possible_truncation, + clippy::cast_lossless, + clippy::collapsible_if + )] + mod sample_joint_tests { + use std::collections::HashMap; + + use super::*; + use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + use crate::param::ParamValue; + use crate::sampler::CompletedTrial; + + fn float_dist(low: f64, high: f64) -> Distribution { + Distribution::Float(FloatDistribution { + low, + high, + log_scale: false, + step: None, + }) + } + + fn int_dist(low: i64, high: i64) -> Distribution { + Distribution::Int(IntDistribution { + low, + high, + log_scale: false, + step: None, + }) + } + + fn categorical_dist(n: usize) -> Distribution { + Distribution::Categorical(CategoricalDistribution { n_choices: n }) + } + + fn create_trial( + id: u64, + params: Vec<(&str, ParamValue, Distribution)>, + value: f64, + ) -> 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_sample_joint_empty_history_returns_all_params() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(10) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + + let history: Vec = vec![]; + + let result = sampler.sample_joint(&search_space, &history); + + assert_eq!(result.len(), 2); + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + } + + #[test] + fn test_sample_joint_startup_phase_uses_random() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(10) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + + // Create 5 trials (less than n_startup_trials=10) + let history: Vec = (0..5) + .map(|i| { + create_trial( + i, + vec![( + "x", + ParamValue::Float(i as f64 / 10.0), + float_dist(0.0, 1.0), + )], + i as f64, + ) + }) + .collect(); + + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("x")); + if let Some(ParamValue::Float(v)) = result.get("x") { + assert!(*v >= 0.0 && *v <= 1.0); + } else { + panic!("Expected Float value for x"); + } + } + + #[test] + fn test_sample_joint_after_startup_uses_tpe() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .n_ei_candidates(24) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 10.0)); + search_space.insert("y".to_string(), float_dist(0.0, 10.0)); + + // Create 20 trials with good values near origin + let history: Vec = (0..20) + .map(|i| { + let x = i as f64 / 2.0; + let y = i as f64 / 2.0; + create_trial( + i, + vec![ + ("x", ParamValue::Float(x), float_dist(0.0, 10.0)), + ("y", ParamValue::Float(y), float_dist(0.0, 10.0)), + ], + x + y, // Objective: lower is better + ) + }) + .collect(); + + let result = sampler.sample_joint(&search_space, &history); + + assert_eq!(result.len(), 2); + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + + // Values should be in the distribution bounds + if let Some(ParamValue::Float(x)) = result.get("x") { + assert!(*x >= 0.0 && *x <= 10.0, "x={x} should be within [0, 10]"); + } + if let Some(ParamValue::Float(y)) = result.get("y") { + assert!(*y >= 0.0 && *y <= 10.0, "y={y} should be within [0, 10]"); + } + } + + #[test] + fn test_sample_joint_biases_toward_good_region() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .n_ei_candidates(50) + .seed(123) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 10.0)); + + // Create trials: good ones near 0, bad ones near 10 + let mut history: Vec = vec![]; + // Good trials (low value = good) + for i in 0..5 { + history.push(create_trial( + i, + vec![( + "x", + ParamValue::Float(i as f64 * 0.5), + float_dist(0.0, 10.0), + )], + i as f64 * 0.5, // Low objective + )); + } + // Bad trials + for i in 5..15 { + history.push(create_trial( + i, + vec![( + "x", + ParamValue::Float(5.0 + (i as f64 - 5.0) * 0.5), + float_dist(0.0, 10.0), + )], + 5.0 + (i as f64 - 5.0) * 0.5, // High objective + )); + } + + // Sample multiple times and check bias + let mut low_count = 0; + for _ in 0..20 { + let result = sampler.sample_joint(&search_space, &history); + if let Some(ParamValue::Float(x)) = result.get("x") { + if *x < 5.0 { + low_count += 1; + } + } + } + + // Should be biased toward lower values + assert!( + low_count > 10, + "Expected more samples in good region, got {low_count}/20" + ); + } + + #[test] + fn test_sample_joint_with_int_params() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("n_layers".to_string(), int_dist(1, 10)); + + // Create history + let history: Vec = (0..10) + .map(|i| { + create_trial( + i, + vec![("n_layers", ParamValue::Int(i as i64 + 1), int_dist(1, 10))], + i as f64, + ) + }) + .collect(); + + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("n_layers")); + if let Some(ParamValue::Int(v)) = result.get("n_layers") { + assert!(*v >= 1 && *v <= 10, "n_layers={v} should be within [1, 10]"); + } else { + panic!("Expected Int value for n_layers"); + } + } + + #[test] + fn test_sample_joint_with_mixed_params() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("lr".to_string(), float_dist(0.0, 1.0)); + search_space.insert("n_layers".to_string(), int_dist(1, 5)); + + // Create history with both params + let history: Vec = (0..15) + .map(|i| { + create_trial( + i, + vec![ + ( + "lr", + ParamValue::Float(i as f64 / 15.0), + float_dist(0.0, 1.0), + ), + ( + "n_layers", + ParamValue::Int((i % 5) as i64 + 1), + int_dist(1, 5), + ), + ], + i as f64, + ) + }) + .collect(); + + let result = sampler.sample_joint(&search_space, &history); + + assert_eq!(result.len(), 2); + assert!(result.contains_key("lr")); + assert!(result.contains_key("n_layers")); + } + + #[test] + fn test_sample_joint_with_categorical_params() { + // Categorical params are currently sampled independently + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert("optimizer".to_string(), categorical_dist(3)); + + let history: Vec = (0..15) + .map(|i| { + create_trial( + i, + vec![ + ( + "x", + ParamValue::Float(i as f64 / 15.0), + float_dist(0.0, 1.0), + ), + ( + "optimizer", + ParamValue::Categorical(i as usize % 3), + categorical_dist(3), + ), + ], + i as f64, + ) + }) + .collect(); + + let result = sampler.sample_joint(&search_space, &history); + + assert_eq!(result.len(), 2); + assert!(result.contains_key("x")); + assert!(result.contains_key("optimizer")); + + if let Some(ParamValue::Categorical(v)) = result.get("optimizer") { + assert!(*v < 3, "optimizer={v} should be in [0, 3)"); + } else { + panic!("Expected Categorical value for optimizer"); + } + } + + #[test] + fn test_sample_joint_reproducibility() { + let search_space = { + let mut s = HashMap::new(); + s.insert("x".to_string(), float_dist(0.0, 1.0)); + s.insert("y".to_string(), float_dist(0.0, 1.0)); + s + }; + + let history: Vec = (0..15) + .map(|i| { + create_trial( + i, + vec![ + ( + "x", + ParamValue::Float(i as f64 / 15.0), + float_dist(0.0, 1.0), + ), + ( + "y", + ParamValue::Float(i as f64 / 15.0), + float_dist(0.0, 1.0), + ), + ], + i as f64, + ) + }) + .collect(); + + let sampler1 = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(999) + .build() + .unwrap(); + + let sampler2 = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(999) + .build() + .unwrap(); + + let result1 = sampler1.sample_joint(&search_space, &history); + let result2 = sampler2.sample_joint(&search_space, &history); + + // With same seed, should get same results + assert_eq!(result1, result2); + } + + #[test] + fn test_sample_joint_clamps_to_bounds() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(2) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + // Narrow distribution + search_space.insert("x".to_string(), float_dist(0.0, 0.1)); + + // Create trials at the edge + let history: Vec = (0..10) + .map(|i| { + create_trial( + i, + vec![( + "x", + ParamValue::Float(i as f64 / 100.0), + float_dist(0.0, 0.1), + )], + i as f64, + ) + }) + .collect(); + + // Sample multiple times + for _ in 0..10 { + let result = sampler.sample_joint(&search_space, &history); + if let Some(ParamValue::Float(x)) = result.get("x") { + assert!( + *x >= 0.0 && *x <= 0.1, + "x={x} should be clamped to [0.0, 0.1]" + ); + } + } + } + + #[test] + fn test_sample_joint_handles_empty_intersection() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(2) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + + // Create trials with non-overlapping parameters + let history = vec![ + create_trial( + 0, + vec![("x", ParamValue::Float(0.5), float_dist(0.0, 1.0))], + 1.0, + ), + create_trial( + 1, + vec![("y", ParamValue::Float(0.5), float_dist(0.0, 1.0))], + 1.0, + ), + ]; + + // Should fall back to random sampling since no common params + let result = sampler.sample_joint(&search_space, &history); + + assert_eq!(result.len(), 2); + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + } + + #[test] + fn test_sample_joint_integration_many_dimensions() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(10) + .n_ei_candidates(24) + .seed(42) + .build() + .unwrap(); + + // 5D search space + let mut search_space = HashMap::new(); + for i in 0..5 { + search_space.insert(format!("x{i}"), float_dist(0.0, 10.0)); + } + + // Create history + let history: Vec = (0..30) + .map(|trial_id| { + let mut param_map = HashMap::new(); + let mut dist_map = HashMap::new(); + for dim in 0..5 { + let name = format!("x{dim}"); + let value = (trial_id as f64 + dim as f64) / 3.0; + param_map.insert(name.clone(), ParamValue::Float(value)); + dist_map.insert(name, float_dist(0.0, 10.0)); + } + CompletedTrial::new(trial_id, param_map, dist_map, trial_id as f64) + }) + .collect(); + + let result = sampler.sample_joint(&search_space, &history); + + assert_eq!(result.len(), 5); + for i in 0..5 { + let name = format!("x{i}"); + assert!(result.contains_key(&name), "Missing parameter {name}"); + if let Some(ParamValue::Float(v)) = result.get(&name) { + assert!( + *v >= 0.0 && *v <= 10.0, + "{name}={v} should be within [0, 10]" + ); + } + } + } + + // ==================================================================== + // Categorical Parameter Tests (US-014) + // ==================================================================== + + #[test] + fn test_sample_tpe_categorical_basic() { + use rand::SeedableRng; + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + + // Category 0 is good (appears more in good trials) + let good_indices = vec![0, 0, 0, 1]; + let bad_indices = vec![1, 1, 2, 2]; + + // Sample many times and check bias toward good category + let mut counts = [0usize; 3]; + for _ in 0..1000 { + let idx = MultivariateTpeSampler::sample_tpe_categorical( + 3, + &good_indices, + &bad_indices, + &mut rng, + ); + counts[idx] += 1; + } + + // Category 0 should be sampled most frequently + assert!( + counts[0] > counts[1], + "Category 0 (good) should be sampled more than category 1: {} vs {}", + counts[0], + counts[1] + ); + assert!( + counts[0] > counts[2], + "Category 0 (good) should be sampled more than category 2: {} vs {}", + counts[0], + counts[2] + ); + } + + #[test] + fn test_sample_tpe_categorical_laplace_smoothing() { + use rand::SeedableRng; + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + + // Category 2 never appears, but should still be sampled due to Laplace smoothing + let good_indices = vec![0, 0, 1]; + let bad_indices = vec![0, 1, 1]; + + let mut sampled_two = false; + for _ in 0..1000 { + let idx = MultivariateTpeSampler::sample_tpe_categorical( + 3, + &good_indices, + &bad_indices, + &mut rng, + ); + if idx == 2 { + sampled_two = true; + break; + } + } + + assert!( + sampled_two, + "Category 2 should be sampled occasionally due to Laplace smoothing" + ); + } + + #[test] + fn test_sample_tpe_categorical_empty_good() { + use rand::SeedableRng; + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + + // Empty good group - all categories should have equal probability + let good_indices: Vec = vec![]; + let bad_indices = vec![0, 1, 2]; + + let mut counts = [0usize; 3]; + for _ in 0..1000 { + let idx = MultivariateTpeSampler::sample_tpe_categorical( + 3, + &good_indices, + &bad_indices, + &mut rng, + ); + counts[idx] += 1; + } + + // All categories should be sampled (roughly uniformly with Laplace) + assert!(counts[0] > 0, "Category 0 should be sampled"); + assert!(counts[1] > 0, "Category 1 should be sampled"); + assert!(counts[2] > 0, "Category 2 should be sampled"); + } + + #[test] + fn test_sample_tpe_categorical_all_indices_valid() { + use rand::SeedableRng; + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + + let n_choices = 4; + let good_indices = vec![0, 1, 2, 3]; + let bad_indices = vec![0, 1, 2, 3]; + + // All samples should be valid indices + for _ in 0..100 { + let idx = MultivariateTpeSampler::sample_tpe_categorical( + n_choices, + &good_indices, + &bad_indices, + &mut rng, + ); + assert!(idx < n_choices, "Index {idx} should be < {n_choices}"); + } + } + + #[test] + fn test_extract_categorical_indices_basic() { + let trials = [ + create_trial( + 0, + vec![("cat", ParamValue::Categorical(1), categorical_dist(3))], + 0.5, + ), + create_trial( + 1, + vec![("cat", ParamValue::Categorical(0), categorical_dist(3))], + 1.0, + ), + create_trial( + 2, + vec![("cat", ParamValue::Categorical(2), categorical_dist(3))], + 1.5, + ), + ]; + + let trial_refs: Vec<&CompletedTrial> = trials.iter().collect(); + let indices = MultivariateTpeSampler::extract_categorical_indices(&trial_refs, "cat"); + + assert_eq!(indices, vec![1, 0, 2]); + } + + #[test] + fn test_extract_categorical_indices_missing_param() { + let trials = [ + create_trial( + 0, + vec![("cat", ParamValue::Categorical(1), categorical_dist(3))], + 0.5, + ), + create_trial( + 1, + vec![("other", ParamValue::Float(1.0), float_dist(0.0, 2.0))], + 1.0, + ), + ]; + + let trial_refs: Vec<&CompletedTrial> = trials.iter().collect(); + let indices = MultivariateTpeSampler::extract_categorical_indices(&trial_refs, "cat"); + + // Only the trial with "cat" should be included + assert_eq!(indices, vec![1]); + } + + #[test] + fn test_extract_categorical_indices_wrong_type() { + let trials = [create_trial( + 0, + vec![("param", ParamValue::Float(1.0), float_dist(0.0, 2.0))], + 0.5, + )]; + + let trial_refs: Vec<&CompletedTrial> = trials.iter().collect(); + let indices = MultivariateTpeSampler::extract_categorical_indices(&trial_refs, "param"); + + // Should be empty since param is Float, not Categorical + assert!(indices.is_empty()); + } + + #[test] + fn test_sample_joint_categorical_tpe_sampling() { + let mut search_space = HashMap::new(); + search_space.insert("cat".to_string(), categorical_dist(3)); + + // Create history where category 0 is consistently good + let mut history = Vec::new(); + // Good trials (low objective): category 0 + for i in 0..5 { + history.push(create_trial( + i, + vec![("cat", ParamValue::Categorical(0), categorical_dist(3))], + 0.1, + )); + } + // Bad trials (high objective): categories 1 and 2 + for i in 5u64..15 { + let cat = if i.is_multiple_of(2) { 1 } else { 2 }; + history.push(create_trial( + i, + vec![("cat", ParamValue::Categorical(cat), categorical_dist(3))], + 10.0, + )); + } + + // Sample many times and count + let mut counts = [0usize; 3]; + for seed in 0..100 { + let sampler_test = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(seed) + .build() + .unwrap(); + let result = sampler_test.sample_joint(&search_space, &history); + if let Some(ParamValue::Categorical(idx)) = result.get("cat") { + counts[*idx] += 1; + } + } + + // Category 0 should be sampled most frequently (it's in the good group) + assert!( + counts[0] > counts[1] && counts[0] > counts[2], + "Category 0 (good) should be sampled most: {counts:?}" + ); + } + + #[test] + fn test_sample_joint_mixed_continuous_and_categorical() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert("cat".to_string(), categorical_dist(3)); + + // Create history with both types + let mut history = Vec::new(); + // Good trials: low x values, category 0 + for i in 0..5 { + history.push(create_trial( + i, + vec![ + ( + "x", + ParamValue::Float(f64::from(i as u32) * 0.05), + float_dist(0.0, 1.0), + ), + ("cat", ParamValue::Categorical(0), categorical_dist(3)), + ], + f64::from(i as u32) * 0.1, + )); + } + // Bad trials: high x values, categories 1 and 2 + for i in 5u64..15 { + let cat = if i.is_multiple_of(2) { 1 } else { 2 }; + history.push(create_trial( + i, + vec![ + ( + "x", + ParamValue::Float(0.5 + f64::from(i as u32 - 5) * 0.05), + float_dist(0.0, 1.0), + ), + ("cat", ParamValue::Categorical(cat), categorical_dist(3)), + ], + 5.0 + f64::from(i as u32), + )); + } + + let result = sampler.sample_joint(&search_space, &history); + + // Both parameters should be present + assert!(result.contains_key("x")); + assert!(result.contains_key("cat")); + + // x should be Float + assert!( + matches!(result.get("x"), Some(ParamValue::Float(_))), + "x should be Float" + ); + + // cat should be Categorical + assert!( + matches!(result.get("cat"), Some(ParamValue::Categorical(_))), + "cat should be Categorical" + ); + } + + #[test] + fn test_sample_joint_only_categorical_params() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let mut search_space = HashMap::new(); + search_space.insert("cat1".to_string(), categorical_dist(2)); + search_space.insert("cat2".to_string(), categorical_dist(3)); + + // Create history with only categorical parameters + let mut history = Vec::new(); + for i in 0u64..15 { + let cat1 = usize::from(i >= 5); + let cat2 = (i % 3) as usize; + let value = if i < 5 { 0.1 } else { 10.0 }; + history.push(create_trial( + i, + vec![ + ("cat1", ParamValue::Categorical(cat1), categorical_dist(2)), + ("cat2", ParamValue::Categorical(cat2), categorical_dist(3)), + ], + value, + )); + } + + let result = sampler.sample_joint(&search_space, &history); + + assert_eq!(result.len(), 2); + assert!(matches!( + result.get("cat1"), + Some(ParamValue::Categorical(_)) + )); + assert!(matches!( + result.get("cat2"), + Some(ParamValue::Categorical(_)) + )); + } + } + + // ======================================================================== + // Sampler Trait Implementation Tests + // ======================================================================== + #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)] + mod sampler_trait_tests { + use super::*; + + fn float_dist(low: f64, high: f64) -> Distribution { + Distribution::Float(crate::distribution::FloatDistribution { + low, + high, + log_scale: false, + step: None, + }) + } + + fn int_dist(low: i64, high: i64) -> Distribution { + Distribution::Int(crate::distribution::IntDistribution { + low, + high, + log_scale: false, + step: None, + }) + } + + fn categorical_dist(n_choices: usize) -> Distribution { + Distribution::Categorical(crate::distribution::CategoricalDistribution { n_choices }) + } + + fn create_trial( + id: u64, + params: Vec<(&str, ParamValue, Distribution)>, + value: f64, + ) -> CompletedTrial { + let mut param_map = HashMap::new(); + let mut dist_map = HashMap::new(); + for (name, param, dist) in params { + param_map.insert(name.to_string(), param); + dist_map.insert(name.to_string(), dist); + } + CompletedTrial::new(id, param_map, dist_map, value) + } + + #[test] + fn test_sampler_trait_basic() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let dist = float_dist(0.0, 1.0); + let history: Vec = vec![]; + + // During startup phase, should sample uniformly + let value = sampler.sample(&dist, 0, &history); + match value { + ParamValue::Float(v) => { + assert!((0.0..=1.0).contains(&v), "Value {v} should be in [0, 1]"); + } + _ => panic!("Expected Float value"), + } + } + + #[test] + fn test_sampler_trait_with_history() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .n_ei_candidates(24) + .seed(42) + .build() + .unwrap(); + + let dist = float_dist(0.0, 1.0); + + // Create enough history to trigger TPE sampling + let history: Vec = (0..10) + .map(|i| { + create_trial( + i, + vec![( + "x", + ParamValue::Float(f64::from(i as u32) / 10.0), + float_dist(0.0, 1.0), + )], + f64::from(i as u32), + ) + }) + .collect(); + + let value = sampler.sample(&dist, 10, &history); + match value { + ParamValue::Float(v) => { + assert!((0.0..=1.0).contains(&v), "Value {v} should be in [0, 1]"); + } + _ => panic!("Expected Float value"), + } + } + + #[test] + fn test_sampler_trait_cache_consistency() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let dist = float_dist(0.0, 1.0); + let history: Vec = vec![]; + + // Sample twice with the same trial_id + let value1 = sampler.sample(&dist, 0, &history); + let value2 = sampler.sample(&dist, 0, &history); + + // Should return cached value for same trial_id + match (&value1, &value2) { + (ParamValue::Float(v1), ParamValue::Float(v2)) => { + assert!( + (*v1 - *v2).abs() < f64::EPSILON, + "Same trial_id should return same value: {v1} vs {v2}" + ); + } + _ => panic!("Expected Float values"), + } + } + + #[test] + fn test_sampler_trait_different_trial_ids() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let dist = float_dist(0.0, 1.0); + let history: Vec = vec![]; + + // Sample with different trial_ids + let value1 = sampler.sample(&dist, 0, &history); + let value2 = sampler.sample(&dist, 1, &history); + + // Values may differ (new joint sample for each trial_id) + // Just verify both are valid + match (&value1, &value2) { + (ParamValue::Float(v1), ParamValue::Float(v2)) => { + assert!((0.0..=1.0).contains(v1), "Value {v1} should be in [0, 1]"); + assert!((0.0..=1.0).contains(v2), "Value {v2} should be in [0, 1]"); + } + _ => panic!("Expected Float values"), + } + } + + #[test] + fn test_sampler_trait_int_distribution() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let dist = int_dist(0, 10); + let history: Vec = vec![]; + + let value = sampler.sample(&dist, 0, &history); + match value { + ParamValue::Int(v) => { + assert!((0..=10).contains(&v), "Value {v} should be in [0, 10]"); + } + _ => panic!("Expected Int value"), + } + } + + #[test] + fn test_sampler_trait_categorical_distribution() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .build() + .unwrap(); + + let dist = categorical_dist(3); + let history: Vec = vec![]; + + let value = sampler.sample(&dist, 0, &history); + match value { + ParamValue::Categorical(v) => { + assert!(v < 3, "Value {v} should be < 3"); + } + _ => panic!("Expected Categorical value"), + } + } + + #[test] + fn test_sampler_trait_with_multivariate_history() { + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .n_ei_candidates(24) + .seed(42) + .build() + .unwrap(); + + let dist_x = float_dist(0.0, 1.0); + let dist_y = float_dist(0.0, 1.0); + + // Create history with two parameters + let history: Vec = (0..10) + .map(|i| { + create_trial( + i, + vec![ + ( + "x", + ParamValue::Float(f64::from(i as u32) / 10.0), + float_dist(0.0, 1.0), + ), + ( + "y", + ParamValue::Float(f64::from((10 - i) as u32) / 10.0), + float_dist(0.0, 1.0), + ), + ], + f64::from(i as u32), + ) + }) + .collect(); + + // Sample x and y for the same trial + let value_x = sampler.sample(&dist_x, 10, &history); + let value_y = sampler.sample(&dist_y, 10, &history); + + // Both should be valid Float values + match (&value_x, &value_y) { + (ParamValue::Float(vx), ParamValue::Float(vy)) => { + assert!((0.0..=1.0).contains(vx), "Value {vx} should be in [0, 1]"); + assert!((0.0..=1.0).contains(vy), "Value {vy} should be in [0, 1]"); + } + _ => panic!("Expected Float values"), + } + } + + #[test] + fn test_find_matching_param_float() { + let mut cached = HashMap::new(); + cached.insert("x".to_string(), ParamValue::Float(0.5)); + cached.insert("y".to_string(), ParamValue::Float(0.8)); + + let dist = float_dist(0.0, 1.0); + let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); + + assert!(result.is_some()); + if let Some(ParamValue::Float(v)) = result { + assert!((0.0..=1.0).contains(&v)); + } + } + + #[test] + fn test_find_matching_param_int() { + let mut cached = HashMap::new(); + cached.insert("n".to_string(), ParamValue::Int(5)); + + let dist = int_dist(0, 10); + let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); + + assert!(result.is_some()); + if let Some(ParamValue::Int(v)) = result { + assert!((0..=10).contains(&v)); + } + } + + #[test] + fn test_find_matching_param_categorical() { + let mut cached = HashMap::new(); + cached.insert("choice".to_string(), ParamValue::Categorical(1)); + + let dist = categorical_dist(3); + let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); + + assert!(result.is_some()); + if let Some(ParamValue::Categorical(v)) = result { + assert!(v < 3); + } + } + + #[test] + fn test_find_matching_param_no_match() { + let mut cached = HashMap::new(); + cached.insert("x".to_string(), ParamValue::Float(0.5)); + + // Looking for Int, but only Float in cache + let dist = int_dist(0, 10); + let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); + + assert!(result.is_none()); + } + + #[test] + fn test_find_matching_param_out_of_bounds() { + let mut cached = HashMap::new(); + cached.insert("x".to_string(), ParamValue::Float(5.0)); // Out of bounds + + let dist = float_dist(0.0, 1.0); + let result = MultivariateTpeSampler::find_matching_param(&dist, &cached); + + assert!(result.is_none()); + } + + #[test] + fn test_build_search_space_from_history() { + let history = vec![ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), float_dist(0.0, 1.0)), + ("y", ParamValue::Int(5), int_dist(0, 10)), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("x", ParamValue::Float(0.3), float_dist(0.0, 1.0)), + ("y", ParamValue::Int(3), int_dist(0, 10)), + ], + 2.0, + ), + ]; + + let current_dist = float_dist(0.0, 1.0); + let search_space = + MultivariateTpeSampler::build_search_space_from_history(¤t_dist, &history); + + assert!(search_space.contains_key("x")); + assert!(search_space.contains_key("y")); + } + + #[test] + fn test_build_search_space_empty_history() { + let history: Vec = vec![]; + + let current_dist = float_dist(0.0, 1.0); + let search_space = + MultivariateTpeSampler::build_search_space_from_history(¤t_dist, &history); + + // Should have a placeholder for current distribution + assert!(!search_space.is_empty()); + assert!(search_space.contains_key("_current")); + } + } + + // ======================================================================== + // Independent Fallback Sampling Tests + // ======================================================================== + + #[allow( + clippy::cast_precision_loss, + clippy::cast_possible_wrap, + clippy::cast_possible_truncation, + clippy::cast_lossless + )] + mod independent_fallback_tests { + use super::*; + use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + + fn create_trial( + id: u64, + params: Vec<(&str, ParamValue, Distribution)>, + value: f64, + ) -> 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) + } + + fn float_dist(low: f64, high: f64) -> Distribution { + Distribution::Float(FloatDistribution { + low, + high, + log_scale: false, + step: None, + }) + } + + fn int_dist(low: i64, high: i64) -> Distribution { + Distribution::Int(IntDistribution { + low, + high, + log_scale: false, + step: None, + }) + } + + fn categorical_dist(n_choices: usize) -> Distribution { + Distribution::Categorical(CategoricalDistribution { n_choices }) + } + + #[test] + fn test_fallback_on_empty_intersection() { + // Create trials with completely different parameters + let history = vec![ + create_trial( + 0, + vec![("x", ParamValue::Float(0.1), float_dist(0.0, 1.0))], + 1.0, + ), + create_trial( + 1, + vec![("y", ParamValue::Float(0.2), float_dist(0.0, 1.0))], + 2.0, + ), + create_trial( + 2, + vec![("z", ParamValue::Float(0.3), float_dist(0.0, 1.0))], + 3.0, + ), + ]; + + // Request parameters that none of the trials have in common + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(1) + .seed(42) + .warn_independent_sampling(false) + .build() + .unwrap(); + + // This should fall back to independent sampling since no common params + let result = sampler.sample_joint(&search_space, &history); + + // Should have all requested parameters + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + + // Values should be within bounds + if let ParamValue::Float(x) = result.get("x").unwrap() { + assert!((0.0..=1.0).contains(x)); + } else { + panic!("Expected Float for x"); + } + if let ParamValue::Float(y) = result.get("y").unwrap() { + assert!((0.0..=1.0).contains(y)); + } else { + panic!("Expected Float for y"); + } + } + + #[test] + fn test_fallback_fills_missing_params() { + // Trials have x and y, but we request x, y, and z + let history = vec![ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.2), float_dist(0.0, 1.0)), + ("y", ParamValue::Float(0.3), float_dist(0.0, 1.0)), + ], + 1.0, + ), + create_trial( + 1, + vec![ + ("x", ParamValue::Float(0.4), float_dist(0.0, 1.0)), + ("y", ParamValue::Float(0.5), float_dist(0.0, 1.0)), + ], + 2.0, + ), + create_trial( + 2, + vec![ + ("x", ParamValue::Float(0.6), float_dist(0.0, 1.0)), + ("y", ParamValue::Float(0.7), float_dist(0.0, 1.0)), + ], + 3.0, + ), + ]; + + // Request x, y (which are in intersection) and z (which is not) + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + search_space.insert("z".to_string(), float_dist(0.0, 1.0)); + + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(1) + .seed(42) + .warn_independent_sampling(false) + .build() + .unwrap(); + + let result = sampler.sample_joint(&search_space, &history); + + // Should have all three parameters + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + assert!(result.contains_key("z")); + + // All values should be within bounds + for (name, value) in &result { + if let ParamValue::Float(v) = value { + assert!( + (0.0..=1.0).contains(v), + "Parameter {name} has value {v} out of bounds" + ); + } + } + } + + #[test] + fn test_independent_tpe_sampling_with_int() { + // Create trials with int parameters + let history: Vec = (0..20) + .map(|i| { + create_trial( + i, + vec![("n", ParamValue::Int((i % 10) as i64), int_dist(0, 10))], + (i as f64) * 0.1, + ) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("n".to_string(), int_dist(0, 10)); + search_space.insert("m".to_string(), int_dist(0, 5)); // Not in history + + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .warn_independent_sampling(false) + .build() + .unwrap(); + + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("n")); + assert!(result.contains_key("m")); + + if let ParamValue::Int(n) = result.get("n").unwrap() { + assert!((0..=10).contains(n)); + } else { + panic!("Expected Int for n"); + } + if let ParamValue::Int(m) = result.get("m").unwrap() { + assert!((0..=5).contains(m)); + } else { + panic!("Expected Int for m"); + } + } + + #[test] + fn test_independent_tpe_sampling_with_categorical() { + // Create trials with categorical parameters + let history: Vec = (0..20) + .map(|i| { + create_trial( + i, + vec![( + "cat", + ParamValue::Categorical(i as usize % 3), + categorical_dist(3), + )], + (i as f64) * 0.1, + ) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("cat".to_string(), categorical_dist(3)); + search_space.insert("other_cat".to_string(), categorical_dist(4)); // Not in history + + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .warn_independent_sampling(false) + .build() + .unwrap(); + + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("cat")); + assert!(result.contains_key("other_cat")); + + if let ParamValue::Categorical(cat) = result.get("cat").unwrap() { + assert!(*cat < 3); + } else { + panic!("Expected Categorical for cat"); + } + if let ParamValue::Categorical(other) = result.get("other_cat").unwrap() { + assert!(*other < 4); + } else { + panic!("Expected Categorical for other_cat"); + } + } + + #[test] + fn test_sample_all_independent_uses_tpe() { + // Create trials with a clear pattern - good values clustered low + let history: Vec = (0..30) + .map(|i| { + let x = if i < 10 { + 0.1 + (i as f64) * 0.02 // Good trials: 0.1-0.28 + } else { + 0.6 + ((i - 10) as f64) * 0.02 // Bad trials: 0.6-0.98 + }; + let value = if i < 10 { 0.0 } else { 1.0 }; + create_trial( + 0, + vec![("x", ParamValue::Float(x), float_dist(0.0, 1.0))], + value, + ) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .n_ei_candidates(48) + .seed(123) + .warn_independent_sampling(false) + .build() + .unwrap(); + + // Sample multiple times and check that we tend to sample from the good region + let mut low_count = 0; + let n_samples = 100; + for _ in 0..n_samples { + let result = sampler.sample_all_independent(&search_space, &history); + if let Some(ParamValue::Float(x)) = result.get("x") + && *x < 0.5 + { + low_count += 1; + } + } + + // TPE should bias towards lower values where good trials are clustered + // With random sampling we'd expect ~50%, with TPE we should see more bias + assert!( + low_count > 60, + "Expected TPE to bias towards good region, but only got {low_count} out of {n_samples} in low region" + ); + } + + #[test] + fn test_fallback_with_few_filtered_trials() { + // Create trials where only 1 has all the required parameters + let history = vec![ + create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.1), float_dist(0.0, 1.0)), + ("y", ParamValue::Float(0.2), float_dist(0.0, 1.0)), + ], + 1.0, + ), + create_trial( + 1, + vec![("x", ParamValue::Float(0.3), float_dist(0.0, 1.0))], + 2.0, + ), + create_trial( + 2, + vec![("x", ParamValue::Float(0.5), float_dist(0.0, 1.0))], + 3.0, + ), + ]; + + // Request both x and y - only trial 0 has both + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert("y".to_string(), float_dist(0.0, 1.0)); + + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(1) + .seed(42) + .warn_independent_sampling(false) + .build() + .unwrap(); + + // Should fall back since only 1 trial has both params (need at least 2) + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + } + + #[test] + fn test_warn_independent_sampling_flag() { + // This test verifies the flag exists and can be set + let sampler_with_warn = MultivariateTpeSampler::builder() + .warn_independent_sampling(true) + .build() + .unwrap(); + assert!(sampler_with_warn.warn_independent_sampling()); + + let sampler_without_warn = MultivariateTpeSampler::builder() + .warn_independent_sampling(false) + .build() + .unwrap(); + assert!(!sampler_without_warn.warn_independent_sampling()); + } + + #[test] + fn test_fill_remaining_uniform_fallback() { + // During startup phase, should use uniform sampling + let history: Vec = vec![]; // No history + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(10) // High startup means we'll use random + .seed(42) + .build() + .unwrap(); + + // With no history and high startup threshold, should sample uniformly + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("x")); + if let ParamValue::Float(x) = result.get("x").unwrap() { + assert!((0.0..=1.0).contains(x)); + } + } + + #[test] + fn test_mixed_params_intersection_and_not() { + // Create history with x,y - both will be in intersection + let history: Vec = (0..15) + .map(|i| { + create_trial( + i, + vec![ + ( + "x", + ParamValue::Float((i as f64) * 0.05), + float_dist(0.0, 1.0), + ), + ("y", ParamValue::Int(i as i64 % 5), int_dist(0, 10)), + ], + (i as f64) * 0.1, + ) + }) + .collect(); + + // Request x, y (in intersection) and z, cat (not in intersection) + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), float_dist(0.0, 1.0)); + search_space.insert("y".to_string(), int_dist(0, 10)); + search_space.insert("z".to_string(), float_dist(-1.0, 1.0)); + search_space.insert("cat".to_string(), categorical_dist(5)); + + let sampler = MultivariateTpeSampler::builder() + .n_startup_trials(5) + .seed(42) + .warn_independent_sampling(false) + .build() + .unwrap(); + + let result = sampler.sample_joint(&search_space, &history); + + // All parameters should be present + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + assert!(result.contains_key("z")); + assert!(result.contains_key("cat")); + + // Validate bounds + if let ParamValue::Float(x) = result.get("x").unwrap() { + assert!((0.0..=1.0).contains(x)); + } + if let ParamValue::Int(y) = result.get("y").unwrap() { + assert!((0..=10).contains(y)); + } + if let ParamValue::Float(z) = result.get("z").unwrap() { + assert!((-1.0..=1.0).contains(z)); + } + if let ParamValue::Categorical(cat) = result.get("cat").unwrap() { + assert!(*cat < 5); + } + } + } + + /// Tests for group decomposition integration (US-020). + mod group_sampling_tests { + use super::*; + use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + + fn create_trial( + id: u64, + params: Vec<(&str, ParamValue, Distribution)>, + value: f64, + ) -> 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_group_mode_disabled_samples_all_together() { + // When group=false, should sample all params jointly + let sampler = MultivariateTpeSamplerBuilder::new() + .group(false) + .n_startup_trials(3) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Create history with all params together + let history: Vec = (0..10) + .map(|i| { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.1; + create_trial( + i, + vec![ + ("x", ParamValue::Float(v), dist.clone()), + ("y", ParamValue::Float(v + 0.05), dist.clone()), + ], + v * v, + ) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist); + + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + assert_eq!(result.len(), 2); + } + + #[test] + fn test_group_mode_enabled_samples_groups_independently() { + // When group=true, should decompose into groups based on co-occurrence + let sampler = MultivariateTpeSamplerBuilder::new() + .group(true) + .n_startup_trials(3) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Create history with two independent groups: + // Group 1: x, y appear together + // Group 2: a, b appear together + // Groups never co-occur + let mut history = Vec::new(); + for i in 0..5 { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.1; + history.push(create_trial( + i, + vec![ + ("x", ParamValue::Float(v), dist.clone()), + ("y", ParamValue::Float(v + 0.05), dist.clone()), + ], + v * v, + )); + } + for i in 5..10 { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.05; + history.push(create_trial( + i, + vec![ + ("a", ParamValue::Float(v), dist.clone()), + ("b", ParamValue::Float(v + 0.1), dist.clone()), + ], + v + 0.5, + )); + } + + // Search space contains params from both groups + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist.clone()); + search_space.insert("a".to_string(), dist.clone()); + search_space.insert("b".to_string(), dist); + + let result = sampler.sample_joint(&search_space, &history); + + // All params should be sampled + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + assert!(result.contains_key("a")); + assert!(result.contains_key("b")); + assert_eq!(result.len(), 4); + + // Values should be within bounds + for value in result.values() { + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(v)); + } + } + } + + #[test] + fn test_group_mode_with_single_group() { + // When all params co-occur, group mode should behave like non-group mode + let sampler_grouped = MultivariateTpeSamplerBuilder::new() + .group(true) + .n_startup_trials(3) + .seed(123) + .build() + .unwrap(); + + let sampler_ungrouped = MultivariateTpeSamplerBuilder::new() + .group(false) + .n_startup_trials(3) + .seed(123) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // All params appear together in all trials + let history: Vec = (0..10) + .map(|i| { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.1; + create_trial( + i, + vec![ + ("x", ParamValue::Float(v), dist.clone()), + ("y", ParamValue::Float(v + 0.05), dist.clone()), + ("z", ParamValue::Float(v + 0.1), dist.clone()), + ], + v * v, + ) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist.clone()); + search_space.insert("z".to_string(), dist); + + let result_grouped = sampler_grouped.sample_joint(&search_space, &history); + let result_ungrouped = sampler_ungrouped.sample_joint(&search_space, &history); + + // Both should produce valid samples with all params + assert_eq!(result_grouped.len(), 3); + assert_eq!(result_ungrouped.len(), 3); + + // Both should have valid bounds + for result in [&result_grouped, &result_ungrouped] { + for value in result.values() { + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(v)); + } + } + } + } + + #[test] + fn test_group_mode_handles_isolated_params() { + // Test handling of parameters that only appear alone + let sampler = MultivariateTpeSamplerBuilder::new() + .group(true) + .n_startup_trials(3) + .warn_independent_sampling(false) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Each param appears alone (forms its own isolated group) + let history = vec![ + create_trial(0, vec![("x", ParamValue::Float(0.3), dist.clone())], 1.0), + create_trial(1, vec![("y", ParamValue::Float(0.5), dist.clone())], 0.5), + create_trial(2, vec![("z", ParamValue::Float(0.7), dist.clone())], 0.8), + create_trial(3, vec![("x", ParamValue::Float(0.2), dist.clone())], 1.2), + create_trial(4, vec![("y", ParamValue::Float(0.6), dist.clone())], 0.4), + create_trial(5, vec![("z", ParamValue::Float(0.8), dist.clone())], 0.7), + ]; + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist.clone()); + search_space.insert("z".to_string(), dist); + + let result = sampler.sample_joint(&search_space, &history); + + // All params should be sampled + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + assert!(result.contains_key("z")); + + // Values should be within bounds + for value in result.values() { + if let ParamValue::Float(v) = value { + assert!((0.0..=1.0).contains(v)); + } + } + } + + #[test] + fn test_group_mode_during_startup() { + // During startup phase, should sample uniformly regardless of group mode + let sampler = MultivariateTpeSamplerBuilder::new() + .group(true) + .n_startup_trials(10) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Only 5 trials (less than n_startup_trials=10) + let history: Vec = (0..5) + .map(|i| { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.1; + create_trial(i, vec![("x", ParamValue::Float(v), dist.clone())], v * v) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist); + + let result = sampler.sample_joint(&search_space, &history); + + // Should sample uniformly for all params + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + } + + #[test] + fn test_group_mode_with_mixed_distribution_types() { + // Test group mode with Float, Int, and Categorical distributions + let sampler = MultivariateTpeSamplerBuilder::new() + .group(true) + .n_startup_trials(3) + .seed(42) + .build() + .unwrap(); + + let dist_float = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_int = Distribution::Int(IntDistribution { + low: 0, + high: 10, + log_scale: false, + step: None, + }); + let dist_cat = Distribution::Categorical(CategoricalDistribution { n_choices: 3 }); + + // Group 1: float, int co-occur + // Group 2: categorical alone + let mut history = Vec::new(); + for i in 0..5 { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.1; + #[allow(clippy::cast_possible_wrap)] + let int_v = (i % 10) as i64; + history.push(create_trial( + i, + vec![ + ("lr", ParamValue::Float(v), dist_float.clone()), + ("layers", ParamValue::Int(int_v), dist_int.clone()), + ], + v * v, + )); + } + for i in 5..10 { + history.push(create_trial( + i, + vec![( + "opt", + ParamValue::Categorical((i % 3) as usize), + dist_cat.clone(), + )], + 1.0, + )); + } + + let mut search_space = HashMap::new(); + search_space.insert("lr".to_string(), dist_float); + search_space.insert("layers".to_string(), dist_int); + search_space.insert("opt".to_string(), dist_cat); + + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("lr")); + assert!(result.contains_key("layers")); + assert!(result.contains_key("opt")); + + // Check types + assert!(matches!(result.get("lr"), Some(ParamValue::Float(_)))); + assert!(matches!(result.get("layers"), Some(ParamValue::Int(_)))); + assert!(matches!( + result.get("opt"), + Some(ParamValue::Categorical(_)) + )); + } + + #[test] + fn test_group_mode_empty_history() { + // With empty history (startup phase), should sample uniformly + let sampler = MultivariateTpeSamplerBuilder::new() + .group(true) + .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 = vec![]; + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist); + + let result = sampler.sample_joint(&search_space, &history); + + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + } + + #[test] + fn test_non_grouped_sampling_with_different_groups() { + // When group=false but history has separate groups, + // intersection will be empty and fall back to independent + let sampler = MultivariateTpeSamplerBuilder::new() + .group(false) + .n_startup_trials(3) + .warn_independent_sampling(false) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Two non-overlapping groups + let mut history = Vec::new(); + for i in 0..5 { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.1; + history.push(create_trial( + i, + vec![("x", ParamValue::Float(v), dist.clone())], + v, + )); + } + for i in 5..10 { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.05; + history.push(create_trial( + i, + vec![("y", ParamValue::Float(v), dist.clone())], + v, + )); + } + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist); + + let result = sampler.sample_joint(&search_space, &history); + + // Should still produce valid results via fallback + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + } + + #[test] + fn test_group_mode_deterministic_with_seed() { + // Same seed should produce same results + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let history: Vec = (0..10) + .map(|i| { + #[allow(clippy::cast_precision_loss)] + let v = (i as f64) * 0.1; + create_trial( + i, + vec![ + ("x", ParamValue::Float(v), dist.clone()), + ("y", ParamValue::Float(v + 0.05), dist.clone()), + ], + v * v, + ) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist); + + // Same seed should produce same results + let sampler1 = MultivariateTpeSamplerBuilder::new() + .group(true) + .seed(999) + .build() + .unwrap(); + + let sampler2 = MultivariateTpeSamplerBuilder::new() + .group(true) + .seed(999) + .build() + .unwrap(); + + let result1 = sampler1.sample_joint(&search_space, &history); + let result2 = sampler2.sample_joint(&search_space, &history); + + assert_eq!(result1, result2); + } + + #[test] + #[allow(clippy::cast_precision_loss)] + fn test_group_mode_handles_ungrouped_params() { + // Test that params not in any group get sampled independently + let sampler = MultivariateTpeSamplerBuilder::new() + .group(true) + .n_startup_trials(3) + .warn_independent_sampling(false) + .seed(42) + .build() + .unwrap(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // History only has x, but search space has x and y + let history: Vec = (0..10) + .map(|i| { + let v = (i as f64) * 0.1; + create_trial(i, vec![("x", ParamValue::Float(v), dist.clone())], v * v) + }) + .collect(); + + let mut search_space = HashMap::new(); + search_space.insert("x".to_string(), dist.clone()); + search_space.insert("y".to_string(), dist); // Not in history + + let result = sampler.sample_joint(&search_space, &history); + + // Both should be sampled + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + + // y should be sampled uniformly (not in any group) + if let ParamValue::Float(v) = result.get("y").unwrap() { + assert!((0.0..=1.0).contains(v)); + } + } + } +} diff --git a/src/sampler/tpe/gamma.rs b/src/sampler/tpe/gamma.rs new file mode 100644 index 0000000..900c92a --- /dev/null +++ b/src/sampler/tpe/gamma.rs @@ -0,0 +1,670 @@ +use core::fmt::Debug; + +use crate::Error; + +/// A strategy for computing the gamma quantile in TPE. +/// +/// The gamma value determines what fraction of trials are considered "good" +/// when splitting the trial history. Different strategies can adapt this +/// fraction based on the number of completed trials. +/// +/// # Implementation Notes +/// +/// - The returned gamma must be in the range (0.0, 1.0) +/// - Implementations should be deterministic for reproducibility +/// - The `clone_box` method enables trait object cloning +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::GammaStrategy; +/// +/// #[derive(Debug, Clone)] +/// struct ConstantGamma(f64); +/// +/// impl GammaStrategy for ConstantGamma { +/// fn gamma(&self, _n_trials: usize) -> f64 { +/// self.0 +/// } +/// +/// fn clone_box(&self) -> Box { +/// Box::new(self.clone()) +/// } +/// } +/// ``` +pub trait GammaStrategy: Send + Sync + Debug { + /// Computes the gamma quantile based on the number of completed trials. + /// + /// # Arguments + /// + /// * `n_trials` - The number of completed trials in the history. + /// + /// # Returns + /// + /// A gamma value in the range (0.0, 1.0). Values outside this range + /// will be clamped by the sampler. + fn gamma(&self, n_trials: usize) -> f64; + + /// Creates a boxed clone of this strategy. + /// + /// This method enables cloning of trait objects, which is necessary + /// for the builder pattern and sampler configuration. + fn clone_box(&self) -> Box; +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_box() + } +} + +/// A fixed gamma strategy that returns a constant value. +/// +/// This is the simplest strategy and the default behavior of TPE. +/// The gamma value remains constant regardless of the number of trials. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::{FixedGamma, TpeSampler}; +/// +/// // Use 15% of trials as "good" +/// let sampler = TpeSampler::builder() +/// .gamma_strategy(FixedGamma::new(0.15).unwrap()) +/// .build() +/// .unwrap(); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct FixedGamma { + gamma: f64, +} + +impl FixedGamma { + /// Creates a new fixed gamma strategy. + /// + /// # Arguments + /// + /// * `gamma` - The constant gamma value to use. + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0). + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::FixedGamma; + /// + /// let strategy = FixedGamma::new(0.25).unwrap(); + /// assert!((strategy.value() - 0.25).abs() < f64::EPSILON); + /// ``` + pub fn new(gamma: f64) -> crate::Result { + if gamma <= 0.0 || gamma >= 1.0 { + return Err(Error::InvalidGamma(gamma)); + } + Ok(Self { gamma }) + } + + /// Returns the fixed gamma value. + #[must_use] + pub fn value(&self) -> f64 { + self.gamma + } +} + +impl Default for FixedGamma { + /// Creates a fixed gamma strategy with the default value of 0.25. + fn default() -> Self { + Self { gamma: 0.25 } + } +} + +impl GammaStrategy for FixedGamma { + fn gamma(&self, _n_trials: usize) -> f64 { + self.gamma + } + + fn clone_box(&self) -> Box { + Box::new(*self) + } +} + +/// A linear gamma strategy that interpolates between min and max values. +/// +/// The gamma value increases linearly from `gamma_min` to `gamma_max` as the +/// number of trials grows from 0 to `n_trials_max`. Beyond `n_trials_max`, +/// gamma remains at `gamma_max`. +/// +/// This strategy is useful when you want to be more explorative early on +/// (smaller gamma = fewer "good" trials) and more exploitative later +/// (larger gamma = more "good" trials). +/// +/// # Formula +/// +/// ```text +/// gamma = gamma_min + (gamma_max - gamma_min) * min(n_trials / n_trials_max, 1.0) +/// ``` +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::{GammaStrategy, LinearGamma, TpeSampler}; +/// +/// let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap(); +/// +/// // At 0 trials: gamma = 0.1 +/// assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON); +/// +/// // At 50 trials: gamma = 0.25 (midpoint) +/// assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON); +/// +/// // At 100+ trials: gamma = 0.4 +/// assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON); +/// assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct LinearGamma { + gamma_min: f64, + gamma_max: f64, + n_trials_max: usize, +} + +impl LinearGamma { + /// Creates a new linear gamma strategy. + /// + /// # Arguments + /// + /// * `gamma_min` - The minimum gamma value (at 0 trials). + /// * `gamma_max` - The maximum gamma value (at `n_trials_max` trials). + /// * `n_trials_max` - The number of trials at which gamma reaches its maximum. + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if: + /// - `gamma_min` is not in (0.0, 1.0) + /// - `gamma_max` is not in (0.0, 1.0) + /// - `gamma_min > gamma_max` + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::LinearGamma; + /// + /// // Gamma goes from 0.1 to 0.3 over 50 trials + /// let strategy = LinearGamma::new(0.1, 0.3, 50).unwrap(); + /// ``` + pub fn new(gamma_min: f64, gamma_max: f64, n_trials_max: usize) -> crate::Result { + if gamma_min <= 0.0 || gamma_min >= 1.0 { + return Err(Error::InvalidGamma(gamma_min)); + } + if gamma_max <= 0.0 || gamma_max >= 1.0 { + return Err(Error::InvalidGamma(gamma_max)); + } + if gamma_min > gamma_max { + return Err(Error::InvalidGamma(gamma_min)); + } + Ok(Self { + gamma_min, + gamma_max, + n_trials_max, + }) + } + + /// Returns the minimum gamma value. + #[must_use] + pub fn gamma_min(&self) -> f64 { + self.gamma_min + } + + /// Returns the maximum gamma value. + #[must_use] + pub fn gamma_max(&self) -> f64 { + self.gamma_max + } + + /// Returns the number of trials at which gamma reaches its maximum. + #[must_use] + pub fn n_trials_max(&self) -> usize { + self.n_trials_max + } +} + +impl Default for LinearGamma { + /// Creates a linear gamma strategy with default values: + /// - `gamma_min`: 0.10 + /// - `gamma_max`: 0.25 + /// - `n_trials_max`: 100 + fn default() -> Self { + Self { + gamma_min: 0.10, + gamma_max: 0.25, + n_trials_max: 100, + } + } +} + +impl GammaStrategy for LinearGamma { + #[allow(clippy::cast_precision_loss)] + fn gamma(&self, n_trials: usize) -> f64 { + if self.n_trials_max == 0 { + return self.gamma_max; + } + let t = (n_trials as f64 / self.n_trials_max as f64).min(1.0); + self.gamma_min + (self.gamma_max - self.gamma_min) * t + } + + fn clone_box(&self) -> Box { + Box::new(*self) + } +} + +/// A square root gamma strategy inspired by Optuna's default behavior. +/// +/// The gamma value is computed based on the inverse square root of the number +/// of trials, providing a balance between exploration and exploitation that +/// naturally adapts as more data becomes available. +/// +/// # Formula +/// +/// ```text +/// n_good = max(1, floor(gamma_factor / sqrt(n_trials))) +/// gamma = min(gamma_max, n_good / n_trials) +/// ``` +/// +/// When `n_trials` is 0, returns `gamma_max`. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::{GammaStrategy, SqrtGamma, TpeSampler}; +/// +/// let strategy = SqrtGamma::default(); +/// +/// // Gamma decreases as trials increase +/// let g10 = strategy.gamma(10); +/// let g100 = strategy.gamma(100); +/// assert!(g10 > g100, "Gamma should decrease with more trials"); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct SqrtGamma { + gamma_factor: f64, + gamma_max: f64, +} + +impl SqrtGamma { + /// Creates a new square root gamma strategy. + /// + /// # Arguments + /// + /// * `gamma_factor` - The factor controlling how quickly gamma decreases. + /// Higher values mean more "good" trials at any given point. + /// * `gamma_max` - The maximum gamma value (used when `n_trials` is small). + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if: + /// - `gamma_factor` is not positive + /// - `gamma_max` is not in (0.0, 1.0) + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::SqrtGamma; + /// + /// let strategy = SqrtGamma::new(1.0, 0.25).unwrap(); + /// ``` + pub fn new(gamma_factor: f64, gamma_max: f64) -> crate::Result { + if gamma_factor <= 0.0 { + return Err(Error::InvalidGamma(gamma_factor)); + } + if gamma_max <= 0.0 || gamma_max >= 1.0 { + return Err(Error::InvalidGamma(gamma_max)); + } + Ok(Self { + gamma_factor, + gamma_max, + }) + } + + /// Returns the gamma factor. + #[must_use] + pub fn gamma_factor(&self) -> f64 { + self.gamma_factor + } + + /// Returns the maximum gamma value. + #[must_use] + pub fn gamma_max(&self) -> f64 { + self.gamma_max + } +} + +impl Default for SqrtGamma { + /// Creates a square root gamma strategy with default values: + /// - `gamma_factor`: 1.0 + /// - `gamma_max`: 0.25 + fn default() -> Self { + Self { + gamma_factor: 1.0, + gamma_max: 0.25, + } + } +} + +impl GammaStrategy for SqrtGamma { + #[allow(clippy::cast_precision_loss)] + fn gamma(&self, n_trials: usize) -> f64 { + if n_trials == 0 { + return self.gamma_max; + } + let n_good = (self.gamma_factor / (n_trials as f64).sqrt()).max(1.0); + (n_good / n_trials as f64).min(self.gamma_max) + } + + fn clone_box(&self) -> Box { + Box::new(*self) + } +} + +/// A Hyperopt-style gamma strategy. +/// +/// This strategy computes gamma as `min(gamma_max, (gamma_base + 1) / n_trials)`, +/// which is inspired by the original Hyperopt TPE implementation. +/// +/// # Formula +/// +/// ```text +/// gamma = min(gamma_max, (gamma_base + 1) / n_trials) +/// ``` +/// +/// When `n_trials` is 0, returns `gamma_max`. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::tpe::{GammaStrategy, HyperoptGamma}; +/// +/// // With gamma_base=24 and gamma_max=0.5: +/// // - At n=25: gamma = min(0.5, 25/25) = 0.5 (capped) +/// // - At n=100: gamma = min(0.5, 25/100) = 0.25 +/// let strategy = HyperoptGamma::new(24.0, 0.5).unwrap(); +/// +/// // Early trials have higher gamma +/// let g50 = strategy.gamma(50); +/// let g200 = strategy.gamma(200); +/// assert!(g50 > g200, "Gamma should decrease with more trials"); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct HyperoptGamma { + gamma_base: f64, + gamma_max: f64, +} + +impl HyperoptGamma { + /// Creates a new Hyperopt-style gamma strategy. + /// + /// # Arguments + /// + /// * `gamma_base` - The base value added to 1 in the numerator. + /// * `gamma_max` - The maximum gamma value. + /// + /// # Errors + /// + /// Returns `Error::InvalidGamma` if: + /// - `gamma_base` is negative + /// - `gamma_max` is not in (0.0, 1.0) + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::tpe::HyperoptGamma; + /// + /// let strategy = HyperoptGamma::new(24.0, 0.25).unwrap(); + /// ``` + pub fn new(gamma_base: f64, gamma_max: f64) -> crate::Result { + if gamma_base < 0.0 { + return Err(Error::InvalidGamma(gamma_base)); + } + if gamma_max <= 0.0 || gamma_max >= 1.0 { + return Err(Error::InvalidGamma(gamma_max)); + } + Ok(Self { + gamma_base, + gamma_max, + }) + } + + /// Returns the gamma base value. + #[must_use] + pub fn gamma_base(&self) -> f64 { + self.gamma_base + } + + /// Returns the maximum gamma value. + #[must_use] + pub fn gamma_max(&self) -> f64 { + self.gamma_max + } +} + +impl Default for HyperoptGamma { + /// Creates a Hyperopt-style gamma strategy with default values: + /// - `gamma_base`: 24.0 + /// - `gamma_max`: 0.25 + fn default() -> Self { + Self { + gamma_base: 24.0, + gamma_max: 0.25, + } + } +} + +impl GammaStrategy for HyperoptGamma { + #[allow(clippy::cast_precision_loss)] + fn gamma(&self, n_trials: usize) -> f64 { + if n_trials == 0 { + return self.gamma_max; + } + ((self.gamma_base + 1.0) / n_trials as f64).min(self.gamma_max) + } + + fn clone_box(&self) -> Box { + Box::new(*self) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sampler::tpe::TpeSampler; + + #[test] + fn test_fixed_gamma_default() { + let strategy = FixedGamma::default(); + assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); + assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); + assert!((strategy.value() - 0.25).abs() < f64::EPSILON); + } + + #[test] + fn test_fixed_gamma_custom() { + let strategy = FixedGamma::new(0.15).unwrap(); + assert!((strategy.gamma(0) - 0.15).abs() < f64::EPSILON); + assert!((strategy.gamma(50) - 0.15).abs() < f64::EPSILON); + assert!((strategy.gamma(1000) - 0.15).abs() < f64::EPSILON); + } + + #[test] + fn test_fixed_gamma_invalid() { + assert!(FixedGamma::new(0.0).is_err()); + assert!(FixedGamma::new(1.0).is_err()); + assert!(FixedGamma::new(-0.1).is_err()); + assert!(FixedGamma::new(1.5).is_err()); + } + + #[test] + fn test_linear_gamma_default() { + let strategy = LinearGamma::default(); + assert!((strategy.gamma(0) - 0.10).abs() < f64::EPSILON); + assert!((strategy.gamma(50) - 0.175).abs() < f64::EPSILON); // midpoint + assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); + assert!((strategy.gamma(200) - 0.25).abs() < f64::EPSILON); // capped + } + + #[test] + fn test_linear_gamma_custom() { + let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap(); + assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON); + assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON); + assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON); + assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON); + } + + #[test] + fn test_linear_gamma_invalid() { + assert!(LinearGamma::new(0.0, 0.5, 100).is_err()); + assert!(LinearGamma::new(0.1, 1.0, 100).is_err()); + assert!(LinearGamma::new(0.5, 0.2, 100).is_err()); // min > max + } + + #[test] + fn test_sqrt_gamma_default() { + let strategy = SqrtGamma::default(); + // At n=0, returns gamma_max + assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); + + // gamma decreases with more trials + let g10 = strategy.gamma(10); + let g100 = strategy.gamma(100); + assert!(g10 > g100); + } + + #[test] + fn test_sqrt_gamma_custom() { + let strategy = SqrtGamma::new(2.0, 0.5).unwrap(); + assert!((strategy.gamma(0) - 0.5).abs() < f64::EPSILON); + + // At n=4: n_good = max(1, 2/2) = 1, gamma = 1/4 = 0.25 + let g4 = strategy.gamma(4); + assert!((g4 - 0.25).abs() < f64::EPSILON); + } + + #[test] + fn test_sqrt_gamma_invalid() { + assert!(SqrtGamma::new(0.0, 0.25).is_err()); // factor must be positive + assert!(SqrtGamma::new(-1.0, 0.25).is_err()); + assert!(SqrtGamma::new(1.0, 0.0).is_err()); + assert!(SqrtGamma::new(1.0, 1.0).is_err()); + } + + #[test] + fn test_hyperopt_gamma_default() { + let strategy = HyperoptGamma::default(); + // At n=0, returns gamma_max + assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); + + // At n=100: (24+1)/100 = 0.25, so capped to 0.25 + assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); + + // At n=200: (24+1)/200 = 0.125 + assert!((strategy.gamma(200) - 0.125).abs() < f64::EPSILON); + } + + #[test] + fn test_hyperopt_gamma_custom() { + let strategy = HyperoptGamma::new(9.0, 0.5).unwrap(); + // At n=20: (9+1)/20 = 0.5, capped to 0.5 + assert!((strategy.gamma(20) - 0.5).abs() < f64::EPSILON); + + // At n=100: (9+1)/100 = 0.1 + assert!((strategy.gamma(100) - 0.1).abs() < f64::EPSILON); + } + + #[test] + fn test_hyperopt_gamma_invalid() { + assert!(HyperoptGamma::new(-1.0, 0.25).is_err()); + assert!(HyperoptGamma::new(24.0, 0.0).is_err()); + assert!(HyperoptGamma::new(24.0, 1.0).is_err()); + } + + #[test] + fn test_gamma_strategy_clone_box() { + let fixed: Box = Box::new(FixedGamma::new(0.3).unwrap()); + let cloned = fixed.clone(); + assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON); + + let linear: Box = Box::new(LinearGamma::default()); + let cloned = linear.clone(); + assert!((cloned.gamma(0) - 0.10).abs() < f64::EPSILON); + } + + #[test] + fn test_tpe_with_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_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_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 + #[allow(clippy::cast_precision_loss)] + (0.01 * n_trials as f64).min(0.5) + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + } + + let sampler = TpeSampler::builder() + .gamma_strategy(DoubleGamma) + .build() + .unwrap(); + + assert!((sampler.gamma_strategy().gamma(10) - 0.1).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(50) - 0.5).abs() < f64::EPSILON); + assert!((sampler.gamma_strategy().gamma(100) - 0.5).abs() < f64::EPSILON); + } +} diff --git a/src/sampler/tpe/mod.rs b/src/sampler/tpe/mod.rs new file mode 100644 index 0000000..407fa9b --- /dev/null +++ b/src/sampler/tpe/mod.rs @@ -0,0 +1,12 @@ +//! Tree-Parzen Estimator (TPE) sampler implementation and utilities. +//! +//! This module provides TPE-based sampling for Bayesian optimization, +//! including support for intersection search space calculation. + +mod gamma; +mod sampler; +pub mod search_space; + +pub use gamma::{FixedGamma, GammaStrategy, HyperoptGamma, LinearGamma, SqrtGamma}; +pub use sampler::{TpeSampler, TpeSamplerBuilder}; +pub use search_space::{GroupDecomposedSearchSpace, IntersectionSearchSpace}; diff --git a/src/sampler/tpe.rs b/src/sampler/tpe/sampler.rs similarity index 66% rename from src/sampler/tpe.rs rename to src/sampler/tpe/sampler.rs index 9fb92e5..610d40b 100644 --- a/src/sampler/tpe.rs +++ b/src/sampler/tpe/sampler.rs @@ -66,483 +66,13 @@ use crate::distribution::Distribution; use crate::error::{Error, Result}; use crate::kde::KernelDensityEstimator; use crate::param::ParamValue; +use crate::sampler::tpe::gamma::{FixedGamma, GammaStrategy}; use crate::sampler::{CompletedTrial, Sampler}; // ============================================================================ // Gamma Strategy Trait and Implementations // ============================================================================ -/// A strategy for computing the gamma quantile in TPE. -/// -/// The gamma value determines what fraction of trials are considered "good" -/// when splitting the trial history. Different strategies can adapt this -/// fraction based on the number of completed trials. -/// -/// # Implementation Notes -/// -/// - The returned gamma must be in the range (0.0, 1.0) -/// - Implementations should be deterministic for reproducibility -/// - The `clone_box` method enables trait object cloning -/// -/// # Examples -/// -/// ``` -/// use optimizer::sampler::tpe::GammaStrategy; -/// -/// #[derive(Debug, Clone)] -/// struct ConstantGamma(f64); -/// -/// impl GammaStrategy for ConstantGamma { -/// fn gamma(&self, _n_trials: usize) -> f64 { -/// self.0 -/// } -/// -/// fn clone_box(&self) -> Box { -/// Box::new(self.clone()) -/// } -/// } -/// ``` -pub trait GammaStrategy: Send + Sync + Debug { - /// Computes the gamma quantile based on the number of completed trials. - /// - /// # Arguments - /// - /// * `n_trials` - The number of completed trials in the history. - /// - /// # Returns - /// - /// A gamma value in the range (0.0, 1.0). Values outside this range - /// will be clamped by the sampler. - fn gamma(&self, n_trials: usize) -> f64; - - /// Creates a boxed clone of this strategy. - /// - /// This method enables cloning of trait objects, which is necessary - /// for the builder pattern and sampler configuration. - fn clone_box(&self) -> Box; -} - -impl Clone for Box { - fn clone(&self) -> Self { - self.clone_box() - } -} - -/// A fixed gamma strategy that returns a constant value. -/// -/// This is the simplest strategy and the default behavior of TPE. -/// The gamma value remains constant regardless of the number of trials. -/// -/// # Examples -/// -/// ``` -/// use optimizer::sampler::tpe::{FixedGamma, TpeSampler}; -/// -/// // Use 15% of trials as "good" -/// let sampler = TpeSampler::builder() -/// .gamma_strategy(FixedGamma::new(0.15).unwrap()) -/// .build() -/// .unwrap(); -/// ``` -#[derive(Debug, Clone, Copy)] -pub struct FixedGamma { - gamma: f64, -} - -impl FixedGamma { - /// Creates a new fixed gamma strategy. - /// - /// # Arguments - /// - /// * `gamma` - The constant gamma value to use. - /// - /// # Errors - /// - /// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0). - /// - /// # Examples - /// - /// ``` - /// use optimizer::sampler::tpe::FixedGamma; - /// - /// let strategy = FixedGamma::new(0.25).unwrap(); - /// assert!((strategy.value() - 0.25).abs() < f64::EPSILON); - /// ``` - pub fn new(gamma: f64) -> Result { - if gamma <= 0.0 || gamma >= 1.0 { - return Err(Error::InvalidGamma(gamma)); - } - Ok(Self { gamma }) - } - - /// Returns the fixed gamma value. - #[must_use] - pub fn value(&self) -> f64 { - self.gamma - } -} - -impl Default for FixedGamma { - /// Creates a fixed gamma strategy with the default value of 0.25. - fn default() -> Self { - Self { gamma: 0.25 } - } -} - -impl GammaStrategy for FixedGamma { - fn gamma(&self, _n_trials: usize) -> f64 { - self.gamma - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -/// A linear gamma strategy that interpolates between min and max values. -/// -/// The gamma value increases linearly from `gamma_min` to `gamma_max` as the -/// number of trials grows from 0 to `n_trials_max`. Beyond `n_trials_max`, -/// gamma remains at `gamma_max`. -/// -/// This strategy is useful when you want to be more explorative early on -/// (smaller gamma = fewer "good" trials) and more exploitative later -/// (larger gamma = more "good" trials). -/// -/// # Formula -/// -/// ```text -/// gamma = gamma_min + (gamma_max - gamma_min) * min(n_trials / n_trials_max, 1.0) -/// ``` -/// -/// # Examples -/// -/// ``` -/// use optimizer::sampler::tpe::{GammaStrategy, LinearGamma, TpeSampler}; -/// -/// let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap(); -/// -/// // At 0 trials: gamma = 0.1 -/// assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON); -/// -/// // At 50 trials: gamma = 0.25 (midpoint) -/// assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON); -/// -/// // At 100+ trials: gamma = 0.4 -/// assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON); -/// assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON); -/// ``` -#[derive(Debug, Clone, Copy)] -pub struct LinearGamma { - gamma_min: f64, - gamma_max: f64, - n_trials_max: usize, -} - -impl LinearGamma { - /// Creates a new linear gamma strategy. - /// - /// # Arguments - /// - /// * `gamma_min` - The minimum gamma value (at 0 trials). - /// * `gamma_max` - The maximum gamma value (at `n_trials_max` trials). - /// * `n_trials_max` - The number of trials at which gamma reaches its maximum. - /// - /// # Errors - /// - /// Returns `Error::InvalidGamma` if: - /// - `gamma_min` is not in (0.0, 1.0) - /// - `gamma_max` is not in (0.0, 1.0) - /// - `gamma_min > gamma_max` - /// - /// # Examples - /// - /// ``` - /// use optimizer::sampler::tpe::LinearGamma; - /// - /// // Gamma goes from 0.1 to 0.3 over 50 trials - /// let strategy = LinearGamma::new(0.1, 0.3, 50).unwrap(); - /// ``` - pub fn new(gamma_min: f64, gamma_max: f64, n_trials_max: usize) -> Result { - if gamma_min <= 0.0 || gamma_min >= 1.0 { - return Err(Error::InvalidGamma(gamma_min)); - } - if gamma_max <= 0.0 || gamma_max >= 1.0 { - return Err(Error::InvalidGamma(gamma_max)); - } - if gamma_min > gamma_max { - return Err(Error::InvalidGamma(gamma_min)); - } - Ok(Self { - gamma_min, - gamma_max, - n_trials_max, - }) - } - - /// Returns the minimum gamma value. - #[must_use] - pub fn gamma_min(&self) -> f64 { - self.gamma_min - } - - /// Returns the maximum gamma value. - #[must_use] - pub fn gamma_max(&self) -> f64 { - self.gamma_max - } - - /// Returns the number of trials at which gamma reaches its maximum. - #[must_use] - pub fn n_trials_max(&self) -> usize { - self.n_trials_max - } -} - -impl Default for LinearGamma { - /// Creates a linear gamma strategy with default values: - /// - `gamma_min`: 0.10 - /// - `gamma_max`: 0.25 - /// - `n_trials_max`: 100 - fn default() -> Self { - Self { - gamma_min: 0.10, - gamma_max: 0.25, - n_trials_max: 100, - } - } -} - -impl GammaStrategy for LinearGamma { - #[allow(clippy::cast_precision_loss)] - fn gamma(&self, n_trials: usize) -> f64 { - if self.n_trials_max == 0 { - return self.gamma_max; - } - let t = (n_trials as f64 / self.n_trials_max as f64).min(1.0); - self.gamma_min + (self.gamma_max - self.gamma_min) * t - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -/// A square root gamma strategy inspired by Optuna's default behavior. -/// -/// The gamma value is computed based on the inverse square root of the number -/// of trials, providing a balance between exploration and exploitation that -/// naturally adapts as more data becomes available. -/// -/// # Formula -/// -/// ```text -/// n_good = max(1, floor(gamma_factor / sqrt(n_trials))) -/// gamma = min(gamma_max, n_good / n_trials) -/// ``` -/// -/// When `n_trials` is 0, returns `gamma_max`. -/// -/// # Examples -/// -/// ``` -/// use optimizer::sampler::tpe::{GammaStrategy, SqrtGamma, TpeSampler}; -/// -/// let strategy = SqrtGamma::default(); -/// -/// // Gamma decreases as trials increase -/// let g10 = strategy.gamma(10); -/// let g100 = strategy.gamma(100); -/// assert!(g10 > g100, "Gamma should decrease with more trials"); -/// ``` -#[derive(Debug, Clone, Copy)] -pub struct SqrtGamma { - gamma_factor: f64, - gamma_max: f64, -} - -impl SqrtGamma { - /// Creates a new square root gamma strategy. - /// - /// # Arguments - /// - /// * `gamma_factor` - The factor controlling how quickly gamma decreases. - /// Higher values mean more "good" trials at any given point. - /// * `gamma_max` - The maximum gamma value (used when `n_trials` is small). - /// - /// # Errors - /// - /// Returns `Error::InvalidGamma` if: - /// - `gamma_factor` is not positive - /// - `gamma_max` is not in (0.0, 1.0) - /// - /// # Examples - /// - /// ``` - /// use optimizer::sampler::tpe::SqrtGamma; - /// - /// let strategy = SqrtGamma::new(1.0, 0.25).unwrap(); - /// ``` - pub fn new(gamma_factor: f64, gamma_max: f64) -> Result { - if gamma_factor <= 0.0 { - return Err(Error::InvalidGamma(gamma_factor)); - } - if gamma_max <= 0.0 || gamma_max >= 1.0 { - return Err(Error::InvalidGamma(gamma_max)); - } - Ok(Self { - gamma_factor, - gamma_max, - }) - } - - /// Returns the gamma factor. - #[must_use] - pub fn gamma_factor(&self) -> f64 { - self.gamma_factor - } - - /// Returns the maximum gamma value. - #[must_use] - pub fn gamma_max(&self) -> f64 { - self.gamma_max - } -} - -impl Default for SqrtGamma { - /// Creates a square root gamma strategy with default values: - /// - `gamma_factor`: 1.0 - /// - `gamma_max`: 0.25 - fn default() -> Self { - Self { - gamma_factor: 1.0, - gamma_max: 0.25, - } - } -} - -impl GammaStrategy for SqrtGamma { - #[allow(clippy::cast_precision_loss)] - fn gamma(&self, n_trials: usize) -> f64 { - if n_trials == 0 { - return self.gamma_max; - } - let n_good = (self.gamma_factor / (n_trials as f64).sqrt()).max(1.0); - (n_good / n_trials as f64).min(self.gamma_max) - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - -/// A Hyperopt-style gamma strategy. -/// -/// This strategy computes gamma as `min(gamma_max, (gamma_base + 1) / n_trials)`, -/// which is inspired by the original Hyperopt TPE implementation. -/// -/// # Formula -/// -/// ```text -/// gamma = min(gamma_max, (gamma_base + 1) / n_trials) -/// ``` -/// -/// When `n_trials` is 0, returns `gamma_max`. -/// -/// # Examples -/// -/// ``` -/// use optimizer::sampler::tpe::{GammaStrategy, HyperoptGamma}; -/// -/// // With gamma_base=24 and gamma_max=0.5: -/// // - At n=25: gamma = min(0.5, 25/25) = 0.5 (capped) -/// // - At n=100: gamma = min(0.5, 25/100) = 0.25 -/// let strategy = HyperoptGamma::new(24.0, 0.5).unwrap(); -/// -/// // Early trials have higher gamma -/// let g50 = strategy.gamma(50); -/// let g200 = strategy.gamma(200); -/// assert!(g50 > g200, "Gamma should decrease with more trials"); -/// ``` -#[derive(Debug, Clone, Copy)] -pub struct HyperoptGamma { - gamma_base: f64, - gamma_max: f64, -} - -impl HyperoptGamma { - /// Creates a new Hyperopt-style gamma strategy. - /// - /// # Arguments - /// - /// * `gamma_base` - The base value added to 1 in the numerator. - /// * `gamma_max` - The maximum gamma value. - /// - /// # Errors - /// - /// Returns `Error::InvalidGamma` if: - /// - `gamma_base` is negative - /// - `gamma_max` is not in (0.0, 1.0) - /// - /// # Examples - /// - /// ``` - /// use optimizer::sampler::tpe::HyperoptGamma; - /// - /// let strategy = HyperoptGamma::new(24.0, 0.25).unwrap(); - /// ``` - pub fn new(gamma_base: f64, gamma_max: f64) -> Result { - if gamma_base < 0.0 { - return Err(Error::InvalidGamma(gamma_base)); - } - if gamma_max <= 0.0 || gamma_max >= 1.0 { - return Err(Error::InvalidGamma(gamma_max)); - } - Ok(Self { - gamma_base, - gamma_max, - }) - } - - /// Returns the gamma base value. - #[must_use] - pub fn gamma_base(&self) -> f64 { - self.gamma_base - } - - /// Returns the maximum gamma value. - #[must_use] - pub fn gamma_max(&self) -> f64 { - self.gamma_max - } -} - -impl Default for HyperoptGamma { - /// Creates a Hyperopt-style gamma strategy with default values: - /// - `gamma_base`: 24.0 - /// - `gamma_max`: 0.25 - fn default() -> Self { - Self { - gamma_base: 24.0, - gamma_max: 0.25, - } - } -} - -impl GammaStrategy for HyperoptGamma { - #[allow(clippy::cast_precision_loss)] - fn gamma(&self, n_trials: usize) -> f64 { - if n_trials == 0 { - return self.gamma_max; - } - ((self.gamma_base + 1.0) / n_trials as f64).min(self.gamma_max) - } - - fn clone_box(&self) -> Box { - Box::new(*self) - } -} - // ============================================================================ // TPE Sampler // ============================================================================ @@ -578,7 +108,7 @@ impl GammaStrategy for HyperoptGamma { /// /// // Create with custom settings using the builder /// let sampler = TpeSampler::builder() -/// .gamma(0.15) // Shorthand for FixedGamma::new(0.15) +/// .gamma(0.15) // Shorthand for Fixednew(0.15) /// .n_startup_trials(20) /// .n_ei_candidates(32) /// .seed(42) @@ -1836,284 +1366,4 @@ mod tests { ); } } - - // ======================================================================== - // Gamma Strategy Tests - // ======================================================================== - - #[test] - fn test_fixed_gamma_default() { - let strategy = FixedGamma::default(); - assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); - assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); - assert!((strategy.value() - 0.25).abs() < f64::EPSILON); - } - - #[test] - fn test_fixed_gamma_custom() { - let strategy = FixedGamma::new(0.15).unwrap(); - assert!((strategy.gamma(0) - 0.15).abs() < f64::EPSILON); - assert!((strategy.gamma(50) - 0.15).abs() < f64::EPSILON); - assert!((strategy.gamma(1000) - 0.15).abs() < f64::EPSILON); - } - - #[test] - fn test_fixed_gamma_invalid() { - assert!(FixedGamma::new(0.0).is_err()); - assert!(FixedGamma::new(1.0).is_err()); - assert!(FixedGamma::new(-0.1).is_err()); - assert!(FixedGamma::new(1.5).is_err()); - } - - #[test] - fn test_linear_gamma_default() { - let strategy = LinearGamma::default(); - assert!((strategy.gamma(0) - 0.10).abs() < f64::EPSILON); - assert!((strategy.gamma(50) - 0.175).abs() < f64::EPSILON); // midpoint - assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); - assert!((strategy.gamma(200) - 0.25).abs() < f64::EPSILON); // capped - } - - #[test] - fn test_linear_gamma_custom() { - let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap(); - assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON); - assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON); - assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON); - assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON); - } - - #[test] - fn test_linear_gamma_invalid() { - assert!(LinearGamma::new(0.0, 0.5, 100).is_err()); - assert!(LinearGamma::new(0.1, 1.0, 100).is_err()); - assert!(LinearGamma::new(0.5, 0.2, 100).is_err()); // min > max - } - - #[test] - fn test_sqrt_gamma_default() { - let strategy = SqrtGamma::default(); - // At n=0, returns gamma_max - assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); - - // gamma decreases with more trials - let g10 = strategy.gamma(10); - let g100 = strategy.gamma(100); - assert!(g10 > g100); - } - - #[test] - fn test_sqrt_gamma_custom() { - let strategy = SqrtGamma::new(2.0, 0.5).unwrap(); - assert!((strategy.gamma(0) - 0.5).abs() < f64::EPSILON); - - // At n=4: n_good = max(1, 2/2) = 1, gamma = 1/4 = 0.25 - let g4 = strategy.gamma(4); - assert!((g4 - 0.25).abs() < f64::EPSILON); - } - - #[test] - fn test_sqrt_gamma_invalid() { - assert!(SqrtGamma::new(0.0, 0.25).is_err()); // factor must be positive - assert!(SqrtGamma::new(-1.0, 0.25).is_err()); - assert!(SqrtGamma::new(1.0, 0.0).is_err()); - assert!(SqrtGamma::new(1.0, 1.0).is_err()); - } - - #[test] - fn test_hyperopt_gamma_default() { - let strategy = HyperoptGamma::default(); - // At n=0, returns gamma_max - assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON); - - // At n=100: (24+1)/100 = 0.25, so capped to 0.25 - assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON); - - // At n=200: (24+1)/200 = 0.125 - assert!((strategy.gamma(200) - 0.125).abs() < f64::EPSILON); - } - - #[test] - fn test_hyperopt_gamma_custom() { - let strategy = HyperoptGamma::new(9.0, 0.5).unwrap(); - // At n=20: (9+1)/20 = 0.5, capped to 0.5 - assert!((strategy.gamma(20) - 0.5).abs() < f64::EPSILON); - - // At n=100: (9+1)/100 = 0.1 - assert!((strategy.gamma(100) - 0.1).abs() < f64::EPSILON); - } - - #[test] - fn test_hyperopt_gamma_invalid() { - assert!(HyperoptGamma::new(-1.0, 0.25).is_err()); - assert!(HyperoptGamma::new(24.0, 0.0).is_err()); - assert!(HyperoptGamma::new(24.0, 1.0).is_err()); - } - - #[test] - fn test_gamma_strategy_clone_box() { - let fixed: Box = Box::new(FixedGamma::new(0.3).unwrap()); - let cloned = fixed.clone(); - assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON); - - let linear: Box = Box::new(LinearGamma::default()); - let cloned = linear.clone(); - assert!((cloned.gamma(0) - 0.10).abs() < f64::EPSILON); - } - - #[test] - fn test_tpe_with_sqrt_gamma_strategy() { - let sampler = TpeSampler::builder() - .gamma_strategy(SqrtGamma::default()) - .n_startup_trials(5) - .seed(42) - .build() - .unwrap(); - - let dist = Distribution::Float(FloatDistribution { - low: 0.0, - high: 1.0, - log_scale: false, - step: None, - }); - - let history: Vec = (0..20) - .map(|i| { - create_trial( - i as u64, - f64::from(i), - vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())], - ) - }) - .collect(); - - // Should be able to sample with the sqrt gamma strategy - let value = sampler.sample(&dist, 100, &history); - if let ParamValue::Float(v) = value { - assert!((0.0..=1.0).contains(&v)); - } else { - panic!("Expected Float value"); - } - } - - #[test] - fn test_tpe_with_linear_gamma_strategy() { - let sampler = TpeSampler::builder() - .gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap()) - .n_startup_trials(5) - .seed(42) - .build() - .unwrap(); - - // Verify the strategy is applied - let g = sampler.gamma_strategy().gamma(25); - assert!((g - 0.2).abs() < f64::EPSILON); // midpoint of 0.1 to 0.3 - } - - #[test] - fn test_tpe_with_hyperopt_gamma_strategy() { - let sampler = TpeSampler::builder() - .gamma_strategy(HyperoptGamma::default()) - .n_startup_trials(5) - .seed(42) - .build() - .unwrap(); - - let dist = Distribution::Float(FloatDistribution { - low: 0.0, - high: 1.0, - log_scale: false, - step: None, - }); - - let history: Vec = (0..20) - .map(|i| { - create_trial( - i as u64, - f64::from(i), - vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())], - ) - }) - .collect(); - - // Should be able to sample with the hyperopt gamma strategy - let value = sampler.sample(&dist, 100, &history); - if let ParamValue::Float(v) = value { - assert!((0.0..=1.0).contains(&v)); - } else { - panic!("Expected Float value"); - } - } - - #[test] - fn test_gamma_overrides_gamma_strategy() { - // When gamma() is called after gamma_strategy(), it should take precedence - let sampler = TpeSampler::builder() - .gamma_strategy(SqrtGamma::default()) - .gamma(0.15) // This should override - .build() - .unwrap(); - - // Should use fixed gamma of 0.15 - assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON); - assert!((sampler.gamma_strategy().gamma(100) - 0.15).abs() < f64::EPSILON); - } - - #[test] - fn test_gamma_strategy_overrides_gamma() { - // When gamma_strategy() is called after gamma(), it should take precedence - let sampler = TpeSampler::builder() - .gamma(0.15) - .gamma_strategy(SqrtGamma::default()) // This should override - .build() - .unwrap(); - - // Should use SqrtGamma - gamma decreases with trials - let g10 = sampler.gamma_strategy().gamma(10); - let g100 = sampler.gamma_strategy().gamma(100); - assert!(g10 > g100, "SqrtGamma should decrease with more trials"); - } - - #[test] - fn test_with_strategy_constructor() { - let sampler = TpeSampler::with_strategy( - LinearGamma::new(0.1, 0.4, 100).unwrap(), - 15, - 32, - None, - Some(42), - ) - .unwrap(); - - assert_eq!(sampler.n_startup_trials, 15); - assert_eq!(sampler.n_ei_candidates, 32); - assert!((sampler.gamma_strategy().gamma(0) - 0.1).abs() < f64::EPSILON); - assert!((sampler.gamma_strategy().gamma(100) - 0.4).abs() < f64::EPSILON); - } - - #[test] - fn test_custom_gamma_strategy() { - #[derive(Debug, Clone)] - struct DoubleGamma; - - impl GammaStrategy for DoubleGamma { - fn gamma(&self, n_trials: usize) -> f64 { - // Double the trial count-based calculation, capped at 0.5 - (0.01 * n_trials as f64).min(0.5) - } - - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } - } - - let sampler = TpeSampler::builder() - .gamma_strategy(DoubleGamma) - .build() - .unwrap(); - - assert!((sampler.gamma_strategy().gamma(10) - 0.1).abs() < f64::EPSILON); - assert!((sampler.gamma_strategy().gamma(50) - 0.5).abs() < f64::EPSILON); - assert!((sampler.gamma_strategy().gamma(100) - 0.5).abs() < f64::EPSILON); - } } diff --git a/src/sampler/tpe/search_space.rs b/src/sampler/tpe/search_space.rs new file mode 100644 index 0000000..197eeea --- /dev/null +++ b/src/sampler/tpe/search_space.rs @@ -0,0 +1,1086 @@ +//! Search space intersection utilities for multivariate TPE sampling. +//! +//! This module provides utilities for computing the intersection of search spaces +//! across completed trials, which is necessary for multivariate TPE sampling where +//! we need to identify parameters that appear in all trials. +//! +//! It also provides utilities for decomposing the search space into independent +//! groups based on parameter co-occurrence patterns, which enables more efficient +//! multivariate modeling when parameters naturally partition into independent subsets. + +use std::collections::{HashMap, HashSet, VecDeque}; + +use crate::distribution::Distribution; +use crate::sampler::CompletedTrial; + +/// Computes the intersection of search spaces across completed trials. +/// +/// In multivariate TPE sampling, we need to model the joint distribution of parameters. +/// However, trials may have different sets of parameters due to dynamic search spaces +/// (e.g., conditional parameters). The intersection search space contains only those +/// parameters that appear in ALL completed trials, allowing us to fit a joint KDE. +/// +/// # Example +/// +/// ```ignore +/// use std::collections::HashMap; +/// use optimizer::distribution::{Distribution, FloatDistribution}; +/// use optimizer::param::ParamValue; +/// use optimizer::sampler::CompletedTrial; +/// use optimizer::sampler::tpe::IntersectionSearchSpace; +/// +/// // Create two trials with overlapping parameters +/// let dist_x = Distribution::Float(FloatDistribution { +/// low: 0.0, high: 1.0, log_scale: false, step: None, +/// }); +/// let dist_y = Distribution::Float(FloatDistribution { +/// low: 0.0, high: 1.0, log_scale: false, step: None, +/// }); +/// +/// let mut params1 = HashMap::new(); +/// params1.insert("x".to_string(), ParamValue::Float(0.5)); +/// params1.insert("y".to_string(), ParamValue::Float(0.3)); +/// let mut dists1 = HashMap::new(); +/// dists1.insert("x".to_string(), dist_x.clone()); +/// dists1.insert("y".to_string(), dist_y.clone()); +/// let trial1 = CompletedTrial::new(0, params1, dists1, 1.0); +/// +/// let mut params2 = HashMap::new(); +/// params2.insert("x".to_string(), ParamValue::Float(0.7)); +/// // Note: trial2 doesn't have "y" +/// let mut dists2 = HashMap::new(); +/// dists2.insert("x".to_string(), dist_x.clone()); +/// let trial2 = CompletedTrial::new(1, params2, dists2, 0.5); +/// +/// let trials = vec![trial1, trial2]; +/// let intersection = IntersectionSearchSpace::calculate(&trials); +/// +/// // Only "x" appears in both trials +/// assert!(intersection.contains_key("x")); +/// assert!(!intersection.contains_key("y")); +/// ``` +#[derive(Debug, Clone, Default)] +pub struct IntersectionSearchSpace; + +impl IntersectionSearchSpace { + /// Calculates the intersection of search spaces across all completed trials. + /// + /// This method identifies parameters that appear in ALL completed trials and returns + /// a mapping from parameter names to their distributions. If any parameter has + /// different distributions across trials (e.g., different bounds), the first + /// encountered distribution is used. + /// + /// # Arguments + /// + /// * `trials` - A slice of completed trials to compute the intersection from. + /// + /// # Returns + /// + /// A `HashMap` mapping parameter names to their distributions, containing only + /// parameters that appear in every trial. Returns an empty map if `trials` is empty. + /// + /// # Notes + /// + /// - Parameters must appear in ALL trials to be included in the intersection. + /// - If a parameter has different distributions in different trials, the distribution + /// from the first trial containing that parameter is used. + /// - For dynamic search spaces where not all parameters are sampled in every trial, + /// this helps identify the "stable" set of parameters that can be modeled jointly. + #[must_use] + pub fn calculate(trials: &[CompletedTrial]) -> HashMap { + if trials.is_empty() { + return HashMap::new(); + } + + // Get parameter names from the first trial as the initial candidate set + let first_trial = &trials[0]; + let mut candidate_params: HashSet<&str> = first_trial + .distributions + .keys() + .map(String::as_str) + .collect(); + + // Intersect with parameter sets from all other trials + for trial in trials.iter().skip(1) { + let trial_params: HashSet<&str> = + trial.distributions.keys().map(String::as_str).collect(); + candidate_params.retain(|param| trial_params.contains(param)); + } + + // Build the result map using distributions from the first trial + // that contains each parameter + let mut result = HashMap::new(); + for param_name in candidate_params { + // Find the first trial that has this parameter and use its distribution + for trial in trials { + if let Some(dist) = trial.distributions.get(param_name) { + result.insert(param_name.to_string(), dist.clone()); + break; + } + } + } + + result + } +} + +/// Decomposes the search space into independent parameter groups based on co-occurrence. +/// +/// In multivariate TPE, we model the joint distribution of parameters. However, when +/// certain parameters never appear together (e.g., they are in different branches of +/// a conditional), modeling them jointly is wasteful and can hurt performance. +/// +/// `GroupDecomposedSearchSpace` analyzes trial history to identify groups of parameters +/// that always appear together. By building a co-occurrence graph (where edges connect +/// parameters that appear in the same trial) and finding its connected components, we +/// can partition parameters into independent groups that can be sampled separately. +/// +/// # Example +/// +/// ```ignore +/// use std::collections::HashMap; +/// use optimizer::distribution::{Distribution, FloatDistribution}; +/// use optimizer::param::ParamValue; +/// use optimizer::sampler::CompletedTrial; +/// use optimizer::sampler::tpe::GroupDecomposedSearchSpace; +/// +/// // Trials with two independent parameter groups: +/// // Group 1: "x", "y" always appear together +/// // Group 2: "a", "b" always appear together +/// // But "x"/"y" never appear with "a"/"b" +/// +/// let dist = Distribution::Float(FloatDistribution { +/// low: 0.0, high: 1.0, log_scale: false, step: None, +/// }); +/// +/// let mut params1 = HashMap::new(); +/// params1.insert("x".to_string(), ParamValue::Float(0.1)); +/// params1.insert("y".to_string(), ParamValue::Float(0.2)); +/// let mut dists1 = HashMap::new(); +/// dists1.insert("x".to_string(), dist.clone()); +/// dists1.insert("y".to_string(), dist.clone()); +/// let trial1 = CompletedTrial::new(0, params1, dists1, 1.0); +/// +/// let mut params2 = HashMap::new(); +/// params2.insert("a".to_string(), ParamValue::Float(0.3)); +/// params2.insert("b".to_string(), ParamValue::Float(0.4)); +/// let mut dists2 = HashMap::new(); +/// dists2.insert("a".to_string(), dist.clone()); +/// dists2.insert("b".to_string(), dist.clone()); +/// let trial2 = CompletedTrial::new(1, params2, dists2, 0.5); +/// +/// let trials = vec![trial1, trial2]; +/// let groups = GroupDecomposedSearchSpace::calculate(&trials); +/// +/// // Should produce two groups: {"x", "y"} and {"a", "b"} +/// assert_eq!(groups.len(), 2); +/// ``` +#[derive(Debug, Clone, Default)] +pub struct GroupDecomposedSearchSpace; + +impl GroupDecomposedSearchSpace { + /// Calculates parameter groups based on co-occurrence in completed trials. + /// + /// This method builds a co-occurrence graph where: + /// - Each parameter is a node + /// - An edge exists between two parameters if they appear in the same trial + /// + /// Then finds connected components to partition parameters into independent groups. + /// + /// # Arguments + /// + /// * `trials` - A slice of completed trials to analyze. + /// + /// # Returns + /// + /// A `Vec>` where each `HashSet` represents an independent group + /// of parameters that co-occur. Parameters within the same group have appeared + /// together in at least one trial (directly or transitively). Parameters in + /// different groups have never appeared in the same trial. + /// + /// Returns an empty vector if `trials` is empty. + /// + /// # Panics + /// + /// This function does not panic under normal circumstances. The internal + /// `.expect()` calls are guarded by the algorithm's invariants. + /// + /// # Notes + /// + /// - Single-parameter groups are included in the output. + /// - The order of groups in the output vector is not guaranteed. + /// - This is useful for `MultivariateTpeSampler` when `group=true` to sample + /// independent parameter groups separately. + #[must_use] + pub fn calculate(trials: &[CompletedTrial]) -> Vec> { + if trials.is_empty() { + return Vec::new(); + } + + // Collect all unique parameter names + let mut all_params: HashSet = HashSet::new(); + for trial in trials { + for param_name in trial.distributions.keys() { + all_params.insert(param_name.clone()); + } + } + + if all_params.is_empty() { + return Vec::new(); + } + + // Build adjacency list for co-occurrence graph + // Two parameters are adjacent if they appear in the same trial + let mut adjacency: HashMap> = HashMap::new(); + for param in &all_params { + adjacency.insert(param.clone(), HashSet::new()); + } + + for trial in trials { + let trial_params: Vec<&String> = trial.distributions.keys().collect(); + // Connect all pairs of parameters in this trial + for (i, param1) in trial_params.iter().enumerate() { + for param2 in trial_params.iter().skip(i + 1) { + adjacency + .get_mut(*param1) + .expect("param should exist in adjacency map") + .insert((*param2).clone()); + adjacency + .get_mut(*param2) + .expect("param should exist in adjacency map") + .insert((*param1).clone()); + } + } + } + + // Find connected components using BFS + let mut visited: HashSet = HashSet::new(); + let mut groups: Vec> = Vec::new(); + + for param in &all_params { + if visited.contains(param) { + continue; + } + + // BFS to find all parameters in this component + let mut component: HashSet = HashSet::new(); + let mut queue: VecDeque = VecDeque::new(); + queue.push_back(param.clone()); + visited.insert(param.clone()); + + while let Some(current) = queue.pop_front() { + component.insert(current.clone()); + + if let Some(neighbors) = adjacency.get(¤t) { + for neighbor in neighbors { + if !visited.contains(neighbor) { + visited.insert(neighbor.clone()); + queue.push_back(neighbor.clone()); + } + } + } + } + + groups.push(component); + } + + groups + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; + use crate::param::ParamValue; + + fn create_trial( + id: u64, + params: Vec<(&str, ParamValue, Distribution)>, + value: f64, + ) -> 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_empty_trials() { + let trials: Vec = vec![]; + let result = IntersectionSearchSpace::calculate(&trials); + assert!(result.is_empty()); + } + + #[test] + fn test_single_trial() { + let dist_x = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_y = Distribution::Int(IntDistribution { + low: 1, + high: 10, + log_scale: false, + step: None, + }); + + let trial = create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), dist_x.clone()), + ("y", ParamValue::Int(5), dist_y.clone()), + ], + 1.0, + ); + + let result = IntersectionSearchSpace::calculate(&[trial]); + assert_eq!(result.len(), 2); + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + assert_eq!(result.get("x"), Some(&dist_x)); + assert_eq!(result.get("y"), Some(&dist_y)); + } + + #[test] + fn test_all_trials_same_params() { + let dist_x = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_y = Distribution::Float(FloatDistribution { + low: -1.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let trials: Vec = (0..5) + .map(|i| { + #[allow(clippy::cast_precision_loss)] + let val = i as f64 * 0.1; + create_trial( + i, + vec![ + ("x", ParamValue::Float(val), dist_x.clone()), + ("y", ParamValue::Float(val - 0.5), dist_y.clone()), + ], + val * val, + ) + }) + .collect(); + + let result = IntersectionSearchSpace::calculate(&trials); + assert_eq!(result.len(), 2); + assert!(result.contains_key("x")); + assert!(result.contains_key("y")); + } + + #[test] + fn test_partial_overlap() { + let dist_x = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_y = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_z = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Trial 1: x, y + let trial1 = create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), dist_x.clone()), + ("y", ParamValue::Float(0.3), dist_y.clone()), + ], + 1.0, + ); + + // Trial 2: x, z (no y) + let trial2 = create_trial( + 1, + vec![ + ("x", ParamValue::Float(0.7), dist_x.clone()), + ("z", ParamValue::Float(0.2), dist_z.clone()), + ], + 0.5, + ); + + // Trial 3: x, y, z (has all) + let trial3 = create_trial( + 2, + vec![ + ("x", ParamValue::Float(0.6), dist_x.clone()), + ("y", ParamValue::Float(0.4), dist_y.clone()), + ("z", ParamValue::Float(0.1), dist_z.clone()), + ], + 0.8, + ); + + let result = IntersectionSearchSpace::calculate(&[trial1, trial2, trial3]); + + // Only "x" appears in all three trials + assert_eq!(result.len(), 1); + assert!(result.contains_key("x")); + assert!(!result.contains_key("y")); + assert!(!result.contains_key("z")); + } + + #[test] + fn test_no_common_params() { + let dist_x = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_y = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Trial 1: only x + let trial1 = create_trial(0, vec![("x", ParamValue::Float(0.5), dist_x.clone())], 1.0); + + // Trial 2: only y + let trial2 = create_trial(1, vec![("y", ParamValue::Float(0.3), dist_y.clone())], 0.5); + + let result = IntersectionSearchSpace::calculate(&[trial1, trial2]); + + // No common parameters + assert!(result.is_empty()); + } + + #[test] + fn test_mixed_distribution_types() { + let dist_float = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_int = Distribution::Int(IntDistribution { + low: 1, + high: 100, + log_scale: false, + step: None, + }); + let dist_cat = Distribution::Categorical(CategoricalDistribution { n_choices: 3 }); + + let trial1 = create_trial( + 0, + vec![ + ("learning_rate", ParamValue::Float(0.01), dist_float.clone()), + ("n_layers", ParamValue::Int(3), dist_int.clone()), + ("optimizer", ParamValue::Categorical(0), dist_cat.clone()), + ], + 1.0, + ); + + let trial2 = create_trial( + 1, + vec![ + ( + "learning_rate", + ParamValue::Float(0.001), + dist_float.clone(), + ), + ("n_layers", ParamValue::Int(5), dist_int.clone()), + ("optimizer", ParamValue::Categorical(1), dist_cat.clone()), + ], + 0.8, + ); + + let result = IntersectionSearchSpace::calculate(&[trial1, trial2]); + + assert_eq!(result.len(), 3); + assert!(matches!( + result.get("learning_rate"), + Some(Distribution::Float(_)) + )); + assert!(matches!(result.get("n_layers"), Some(Distribution::Int(_)))); + assert!(matches!( + result.get("optimizer"), + Some(Distribution::Categorical(_)) + )); + } + + #[test] + fn test_distribution_from_first_trial() { + // Test that when distributions differ, the first trial's distribution is used + let dist_x_v1 = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_x_v2 = Distribution::Float(FloatDistribution { + low: 0.0, + high: 10.0, // Different upper bound + log_scale: false, + step: None, + }); + + let trial1 = create_trial( + 0, + vec![("x", ParamValue::Float(0.5), dist_x_v1.clone())], + 1.0, + ); + + let trial2 = create_trial( + 1, + vec![("x", ParamValue::Float(5.0), dist_x_v2.clone())], + 0.5, + ); + + let result = IntersectionSearchSpace::calculate(&[trial1, trial2]); + + assert_eq!(result.len(), 1); + // Should use the distribution from the first trial + assert_eq!(result.get("x"), Some(&dist_x_v1)); + } + + #[test] + fn test_many_trials_with_conditional_params() { + // Simulate a scenario with conditional parameters + // e.g., "use_dropout" is a bool, and "dropout_rate" only exists when use_dropout=true + let dist_lr = Distribution::Float(FloatDistribution { + low: 1e-5, + high: 1e-1, + log_scale: true, + step: None, + }); + let dist_dropout = Distribution::Categorical(CategoricalDistribution { n_choices: 2 }); + let dist_dropout_rate = Distribution::Float(FloatDistribution { + low: 0.0, + high: 0.5, + log_scale: false, + step: None, + }); + + // Trial 1: use_dropout=true, has dropout_rate + let trial1 = create_trial( + 0, + vec![ + ("lr", ParamValue::Float(0.01), dist_lr.clone()), + ( + "use_dropout", + ParamValue::Categorical(1), + dist_dropout.clone(), + ), + ( + "dropout_rate", + ParamValue::Float(0.2), + dist_dropout_rate.clone(), + ), + ], + 1.0, + ); + + // Trial 2: use_dropout=false, no dropout_rate + let trial2 = create_trial( + 1, + vec![ + ("lr", ParamValue::Float(0.001), dist_lr.clone()), + ( + "use_dropout", + ParamValue::Categorical(0), + dist_dropout.clone(), + ), + ], + 0.8, + ); + + // Trial 3: use_dropout=true, has dropout_rate + let trial3 = create_trial( + 2, + vec![ + ("lr", ParamValue::Float(0.005), dist_lr.clone()), + ( + "use_dropout", + ParamValue::Categorical(1), + dist_dropout.clone(), + ), + ( + "dropout_rate", + ParamValue::Float(0.3), + dist_dropout_rate.clone(), + ), + ], + 0.9, + ); + + let result = IntersectionSearchSpace::calculate(&[trial1, trial2, trial3]); + + // Only "lr" and "use_dropout" appear in all trials + assert_eq!(result.len(), 2); + assert!(result.contains_key("lr")); + assert!(result.contains_key("use_dropout")); + assert!(!result.contains_key("dropout_rate")); // Not in trial2 + } + + // ==================== GroupDecomposedSearchSpace Tests ==================== + + #[test] + fn test_group_empty_trials() { + let trials: Vec = vec![]; + let groups = GroupDecomposedSearchSpace::calculate(&trials); + assert!(groups.is_empty()); + } + + #[test] + fn test_group_single_trial_single_param() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let trial = create_trial(0, vec![("x", ParamValue::Float(0.5), dist)], 1.0); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial]); + + assert_eq!(groups.len(), 1); + assert!(groups[0].contains("x")); + assert_eq!(groups[0].len(), 1); + } + + #[test] + fn test_group_single_trial_multiple_params() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + let trial = create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.5), dist.clone()), + ("y", ParamValue::Float(0.3), dist.clone()), + ("z", ParamValue::Float(0.7), dist), + ], + 1.0, + ); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial]); + + // All params appear together, so one group + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].len(), 3); + assert!(groups[0].contains("x")); + assert!(groups[0].contains("y")); + assert!(groups[0].contains("z")); + } + + #[test] + fn test_group_two_independent_groups() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Trial 1: x, y (group 1) + let trial1 = create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.1), dist.clone()), + ("y", ParamValue::Float(0.2), dist.clone()), + ], + 1.0, + ); + + // Trial 2: a, b (group 2 - never appears with x or y) + let trial2 = create_trial( + 1, + vec![ + ("a", ParamValue::Float(0.3), dist.clone()), + ("b", ParamValue::Float(0.4), dist), + ], + 0.5, + ); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2]); + + // Should have 2 independent groups + assert_eq!(groups.len(), 2); + + // Find which group has x/y and which has a/b + let group_xy = groups.iter().find(|g| g.contains("x")); + let group_ab = groups.iter().find(|g| g.contains("a")); + + assert!(group_xy.is_some()); + assert!(group_ab.is_some()); + + let group_xy = group_xy.expect("group with x should exist"); + let group_ab = group_ab.expect("group with a should exist"); + + assert_eq!(group_xy.len(), 2); + assert!(group_xy.contains("x")); + assert!(group_xy.contains("y")); + + assert_eq!(group_ab.len(), 2); + assert!(group_ab.contains("a")); + assert!(group_ab.contains("b")); + } + + #[test] + fn test_group_transitive_connection() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Trial 1: x, y + let trial1 = create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.1), dist.clone()), + ("y", ParamValue::Float(0.2), dist.clone()), + ], + 1.0, + ); + + // Trial 2: y, z (y connects x and z transitively) + let trial2 = create_trial( + 1, + vec![ + ("y", ParamValue::Float(0.3), dist.clone()), + ("z", ParamValue::Float(0.4), dist), + ], + 0.5, + ); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2]); + + // All should be in one group due to transitive connection via y + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].len(), 3); + assert!(groups[0].contains("x")); + assert!(groups[0].contains("y")); + assert!(groups[0].contains("z")); + } + + #[test] + fn test_group_chain_connection() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Create a chain: a-b, b-c, c-d + // Trial 1: a, b + let trial1 = create_trial( + 0, + vec![ + ("a", ParamValue::Float(0.1), dist.clone()), + ("b", ParamValue::Float(0.2), dist.clone()), + ], + 1.0, + ); + + // Trial 2: b, c + let trial2 = create_trial( + 1, + vec![ + ("b", ParamValue::Float(0.3), dist.clone()), + ("c", ParamValue::Float(0.4), dist.clone()), + ], + 0.5, + ); + + // Trial 3: c, d + let trial3 = create_trial( + 2, + vec![ + ("c", ParamValue::Float(0.5), dist.clone()), + ("d", ParamValue::Float(0.6), dist), + ], + 0.3, + ); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]); + + // All should be connected in one group + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].len(), 4); + assert!(groups[0].contains("a")); + assert!(groups[0].contains("b")); + assert!(groups[0].contains("c")); + assert!(groups[0].contains("d")); + } + + #[test] + fn test_group_multiple_isolated_params() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Each trial has exactly one parameter (all isolated) + let trial1 = create_trial(0, vec![("x", ParamValue::Float(0.1), dist.clone())], 1.0); + let trial2 = create_trial(1, vec![("y", ParamValue::Float(0.2), dist.clone())], 0.5); + let trial3 = create_trial(2, vec![("z", ParamValue::Float(0.3), dist)], 0.3); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]); + + // Each param should be its own group + assert_eq!(groups.len(), 3); + for group in &groups { + assert_eq!(group.len(), 1); + } + + // Verify all params are covered + let all_params: HashSet = groups.iter().flatten().cloned().collect(); + assert!(all_params.contains("x")); + assert!(all_params.contains("y")); + assert!(all_params.contains("z")); + } + + #[test] + fn test_group_complex_scenario() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Group 1: x, y, z connected + // Group 2: a, b connected + // w is isolated + + // Trial 1: x, y + let trial1 = create_trial( + 0, + vec![ + ("x", ParamValue::Float(0.1), dist.clone()), + ("y", ParamValue::Float(0.2), dist.clone()), + ], + 1.0, + ); + + // Trial 2: y, z + let trial2 = create_trial( + 1, + vec![ + ("y", ParamValue::Float(0.3), dist.clone()), + ("z", ParamValue::Float(0.4), dist.clone()), + ], + 0.5, + ); + + // Trial 3: a, b + let trial3 = create_trial( + 2, + vec![ + ("a", ParamValue::Float(0.5), dist.clone()), + ("b", ParamValue::Float(0.6), dist.clone()), + ], + 0.3, + ); + + // Trial 4: w (isolated) + let trial4 = create_trial(3, vec![("w", ParamValue::Float(0.7), dist)], 0.2); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3, trial4]); + + // Should have 3 groups: {x,y,z}, {a,b}, {w} + assert_eq!(groups.len(), 3); + + let group_xyz = groups.iter().find(|g| g.contains("x")); + let group_ab = groups.iter().find(|g| g.contains("a")); + let group_w = groups.iter().find(|g| g.contains("w")); + + assert!(group_xyz.is_some()); + assert!(group_ab.is_some()); + assert!(group_w.is_some()); + + let group_xyz = group_xyz.expect("group with x should exist"); + assert_eq!(group_xyz.len(), 3); + assert!(group_xyz.contains("x")); + assert!(group_xyz.contains("y")); + assert!(group_xyz.contains("z")); + + let group_ab = group_ab.expect("group with a should exist"); + assert_eq!(group_ab.len(), 2); + assert!(group_ab.contains("a")); + assert!(group_ab.contains("b")); + + let group_w = group_w.expect("group with w should exist"); + assert_eq!(group_w.len(), 1); + assert!(group_w.contains("w")); + } + + #[test] + fn test_group_all_params_same_trial() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // All params in single trial + let trial = create_trial( + 0, + vec![ + ("a", ParamValue::Float(0.1), dist.clone()), + ("b", ParamValue::Float(0.2), dist.clone()), + ("c", ParamValue::Float(0.3), dist.clone()), + ("d", ParamValue::Float(0.4), dist), + ], + 1.0, + ); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial]); + + // All in one group + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].len(), 4); + } + + #[test] + fn test_group_with_mixed_distribution_types() { + let dist_float = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist_int = Distribution::Int(IntDistribution { + low: 1, + high: 10, + log_scale: false, + step: None, + }); + let dist_cat = Distribution::Categorical(CategoricalDistribution { n_choices: 3 }); + + // Group 1: learning_rate (float), n_layers (int) + // Group 2: optimizer (categorical) + let trial1 = create_trial( + 0, + vec![ + ("learning_rate", ParamValue::Float(0.01), dist_float.clone()), + ("n_layers", ParamValue::Int(3), dist_int.clone()), + ], + 1.0, + ); + + let trial2 = create_trial( + 1, + vec![("optimizer", ParamValue::Categorical(1), dist_cat)], + 0.5, + ); + + let trial3 = create_trial( + 2, + vec![ + ("learning_rate", ParamValue::Float(0.001), dist_float), + ("n_layers", ParamValue::Int(5), dist_int), + ], + 0.8, + ); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]); + + // Should have 2 groups: {learning_rate, n_layers} and {optimizer} + assert_eq!(groups.len(), 2); + + let group_lr = groups.iter().find(|g| g.contains("learning_rate")); + let group_opt = groups.iter().find(|g| g.contains("optimizer")); + + assert!(group_lr.is_some()); + assert!(group_opt.is_some()); + + let group_lr = group_lr.expect("group with learning_rate should exist"); + assert_eq!(group_lr.len(), 2); + assert!(group_lr.contains("learning_rate")); + assert!(group_lr.contains("n_layers")); + + let group_opt = group_opt.expect("group with optimizer should exist"); + assert_eq!(group_opt.len(), 1); + assert!(group_opt.contains("optimizer")); + } + + #[test] + fn test_group_star_topology() { + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Star topology: center connects to all others + // Trial 1: center, a + // Trial 2: center, b + // Trial 3: center, c + let trial1 = create_trial( + 0, + vec![ + ("center", ParamValue::Float(0.1), dist.clone()), + ("a", ParamValue::Float(0.2), dist.clone()), + ], + 1.0, + ); + + let trial2 = create_trial( + 1, + vec![ + ("center", ParamValue::Float(0.3), dist.clone()), + ("b", ParamValue::Float(0.4), dist.clone()), + ], + 0.5, + ); + + let trial3 = create_trial( + 2, + vec![ + ("center", ParamValue::Float(0.5), dist.clone()), + ("c", ParamValue::Float(0.6), dist), + ], + 0.3, + ); + + let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]); + + // All connected via center + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].len(), 4); + assert!(groups[0].contains("center")); + assert!(groups[0].contains("a")); + assert!(groups[0].contains("b")); + assert!(groups[0].contains("c")); + } +} diff --git a/tests/multivariate_tpe_integration.rs b/tests/multivariate_tpe_integration.rs new file mode 100644 index 0000000..c2c6689 --- /dev/null +++ b/tests/multivariate_tpe_integration.rs @@ -0,0 +1,393 @@ +//! Integration tests for the Multivariate TPE sampler. +//! +//! These tests compare the performance of `MultivariateTpeSampler` against +//! the standard `TpeSampler` on problems with and without parameter correlations. + +#![allow( + clippy::cast_sign_loss, + clippy::cast_precision_loss, + clippy::cast_possible_truncation +)] + +use optimizer::sampler::MultivariateTpeSampler; +use optimizer::sampler::tpe::TpeSampler; +use optimizer::{Direction, Error, Study}; + +// ============================================================================= +// Rosenbrock function: f(x,y) = (a-x)^2 + b*(y-x^2)^2 +// This is a classic benchmark with strong parameter correlation. +// Optimal: (x, y) = (a, a^2) with f(x, y) = 0 +// Standard parameters: a = 1, b = 100 +// The "banana-shaped" valley makes this hard for independent samplers. +// ============================================================================= + +/// Computes the Rosenbrock function value. +/// +/// f(x, y) = (a - x)^2 + b * (y - x^2)^2 +/// +/// With standard parameters a = 1, b = 100: +/// - Optimal point: (1, 1) +/// - Optimal value: 0 +fn rosenbrock(x: f64, y: f64) -> f64 { + let a = 1.0; + let b = 100.0; + (a - x).powi(2) + b * (y - x * x).powi(2) +} + +// ============================================================================= +// Test: Multivariate TPE on Rosenbrock function (correlated parameters) +// ============================================================================= + +#[test] +fn test_multivariate_tpe_rosenbrock_finds_good_solution() { + // Multivariate TPE should find a good solution on Rosenbrock + // because it can model the correlation between x and y + let sampler = MultivariateTpeSampler::builder() + .seed(42) + .n_startup_trials(10) + .n_ei_candidates(24) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(100, |trial| { + let x = trial.suggest_float("x", -2.0, 2.0)?; + let y = trial.suggest_float("y", -2.0, 4.0)?; + Ok::<_, Error>(rosenbrock(x, y)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + // Multivariate TPE should find a reasonably good solution + // The global minimum is 0, but getting close is challenging + assert!( + best.value < 10.0, + "Multivariate TPE should find good Rosenbrock solution: best value {} should be < 10.0", + best.value + ); +} + +#[test] +fn test_independent_tpe_rosenbrock() { + // Independent TPE (standard TpeSampler) on Rosenbrock + // This establishes a baseline for comparison + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(10) + .n_ei_candidates(24) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(100, |trial| { + let x = trial.suggest_float("x", -2.0, 2.0)?; + let y = trial.suggest_float("y", -2.0, 4.0)?; + Ok::<_, Error>(rosenbrock(x, y)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + // Independent TPE should still find a decent solution, + // but may not be as good as multivariate TPE on this correlated problem + assert!( + best.value < 50.0, + "Independent TPE should find reasonable Rosenbrock solution: best value {} should be < 50.0", + best.value + ); +} + +#[test] +fn test_multivariate_tpe_outperforms_on_correlated_problem() { + // Run multiple seeds and compare average performance + // Multivariate TPE should generally find better solutions on Rosenbrock + let n_runs = 5; + let n_trials = 80; + + let mut multivariate_best_values = Vec::new(); + let mut independent_best_values = Vec::new(); + + for seed in 0..n_runs { + // Multivariate TPE + let multivariate_sampler = MultivariateTpeSampler::builder() + .seed(seed as u64) + .n_startup_trials(10) + .n_ei_candidates(24) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, multivariate_sampler); + + study + .optimize_with_sampler(n_trials, |trial| { + let x = trial.suggest_float("x", -2.0, 2.0)?; + let y = trial.suggest_float("y", -2.0, 4.0)?; + Ok::<_, Error>(rosenbrock(x, y)) + }) + .unwrap(); + + multivariate_best_values.push(study.best_trial().unwrap().value); + + // Independent TPE + let independent_sampler = TpeSampler::builder() + .seed(seed as u64) + .n_startup_trials(10) + .n_ei_candidates(24) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, independent_sampler); + + study + .optimize_with_sampler(n_trials, |trial| { + let x = trial.suggest_float("x", -2.0, 2.0)?; + let y = trial.suggest_float("y", -2.0, 4.0)?; + Ok::<_, Error>(rosenbrock(x, y)) + }) + .unwrap(); + + independent_best_values.push(study.best_trial().unwrap().value); + } + + let multivariate_mean: f64 = multivariate_best_values.iter().sum::() / n_runs as f64; + let independent_mean: f64 = independent_best_values.iter().sum::() / n_runs as f64; + + // Log results for debugging (these won't show in normal test runs, + // but are useful when running with --nocapture) + eprintln!("Multivariate TPE mean best: {multivariate_mean:.4}"); + eprintln!("Independent TPE mean best: {independent_mean:.4}"); + eprintln!("Multivariate best values: {multivariate_best_values:?}"); + eprintln!("Independent best values: {independent_best_values:?}"); + + // Both methods should find reasonable solutions + assert!( + multivariate_mean < 20.0, + "Multivariate TPE mean {multivariate_mean:.4} should be < 20.0" + ); + assert!( + independent_mean < 100.0, + "Independent TPE mean {independent_mean:.4} should be < 100.0" + ); +} + +// ============================================================================= +// Independent parameter problem: f(x,y) = x^2 + y^2 +// No correlation between parameters - both methods should work equally well. +// ============================================================================= + +/// Simple sphere function with independent parameters. +/// +/// f(x, y) = x^2 + y^2 +/// +/// Optimal point: (0, 0) +/// Optimal value: 0 +fn sphere(x: f64, y: f64) -> f64 { + x * x + y * y +} + +#[test] +fn test_multivariate_tpe_independent_problem() { + // On an independent problem, multivariate TPE should still work well + let sampler = MultivariateTpeSampler::builder() + .seed(42) + .n_startup_trials(10) + .n_ei_candidates(24) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(50, |trial| { + let x = trial.suggest_float("x", -5.0, 5.0)?; + let y = trial.suggest_float("y", -5.0, 5.0)?; + Ok::<_, Error>(sphere(x, y)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + assert!( + best.value < 5.0, + "Multivariate TPE should find good solution on sphere: best value {} should be < 5.0", + best.value + ); +} + +#[test] +fn test_independent_tpe_independent_problem() { + // Baseline: independent TPE on sphere function + let sampler = TpeSampler::builder() + .seed(42) + .n_startup_trials(10) + .n_ei_candidates(24) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(50, |trial| { + let x = trial.suggest_float("x", -5.0, 5.0)?; + let y = trial.suggest_float("y", -5.0, 5.0)?; + Ok::<_, Error>(sphere(x, y)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + assert!( + best.value < 5.0, + "Independent TPE should find good solution on sphere: best value {} should be < 5.0", + best.value + ); +} + +#[test] +fn test_both_samplers_work_on_independent_problem() { + // Run both samplers on the independent sphere function + // and verify they both achieve similar performance + let n_runs = 5; + let n_trials = 50; + + let mut multivariate_results = Vec::new(); + let mut independent_results = Vec::new(); + + for seed in 0..n_runs { + // Multivariate TPE + let sampler = MultivariateTpeSampler::builder() + .seed(seed as u64) + .n_startup_trials(10) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(n_trials, |trial| { + let x = trial.suggest_float("x", -5.0, 5.0)?; + let y = trial.suggest_float("y", -5.0, 5.0)?; + Ok::<_, Error>(sphere(x, y)) + }) + .unwrap(); + + multivariate_results.push(study.best_trial().unwrap().value); + + // Independent TPE + let sampler = TpeSampler::builder() + .seed(seed as u64) + .n_startup_trials(10) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(n_trials, |trial| { + let x = trial.suggest_float("x", -5.0, 5.0)?; + let y = trial.suggest_float("y", -5.0, 5.0)?; + Ok::<_, Error>(sphere(x, y)) + }) + .unwrap(); + + independent_results.push(study.best_trial().unwrap().value); + } + + let multivariate_mean: f64 = multivariate_results.iter().sum::() / n_runs as f64; + let independent_mean: f64 = independent_results.iter().sum::() / n_runs as f64; + + eprintln!("Sphere function results:"); + eprintln!(" Multivariate TPE mean: {multivariate_mean:.4}"); + eprintln!(" Independent TPE mean: {independent_mean:.4}"); + + // Both should find good solutions on this simple problem + assert!( + multivariate_mean < 5.0, + "Multivariate TPE mean {multivariate_mean:.4} should be < 5.0 on sphere" + ); + assert!( + independent_mean < 5.0, + "Independent TPE mean {independent_mean:.4} should be < 5.0 on sphere" + ); +} + +// ============================================================================= +// Test: Multivariate TPE with group decomposition +// ============================================================================= + +#[test] +fn test_multivariate_tpe_with_group_decomposition() { + // Test that group decomposition works correctly + let sampler = MultivariateTpeSampler::builder() + .seed(42) + .n_startup_trials(10) + .group(true) // Enable group decomposition + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(50, |trial| { + let x = trial.suggest_float("x", -5.0, 5.0)?; + let y = trial.suggest_float("y", -5.0, 5.0)?; + Ok::<_, Error>(sphere(x, y)) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + assert!( + best.value < 10.0, + "Multivariate TPE with groups should find good solution: best value {} should be < 10.0", + best.value + ); +} + +// ============================================================================= +// Test: Multivariate TPE with mixed parameter types +// ============================================================================= + +#[test] +fn test_multivariate_tpe_mixed_parameter_types() { + // Test with float, int, and categorical parameters + let sampler = MultivariateTpeSampler::builder() + .seed(42) + .n_startup_trials(10) + .build() + .unwrap(); + + let study: Study = Study::with_sampler(Direction::Minimize, sampler); + + study + .optimize_with_sampler(50, |trial| { + let x = trial.suggest_float("x", -5.0, 5.0)?; + let n = trial.suggest_int("n", 1, 10)?; + let mode = trial.suggest_categorical("mode", &["a", "b", "c"])?; + + // Objective depends on all parameters + let mode_factor = match mode { + "a" => 1.0, + "b" => 0.5, + "c" => 2.0, + _ => unreachable!(), + }; + + Ok::<_, Error>(x * x + (n as f64 - 5.0).powi(2) * mode_factor) + }) + .expect("optimization should succeed"); + + let best = study.best_trial().expect("should have at least one trial"); + + // Should find a reasonable solution + assert!( + best.value < 25.0, + "Multivariate TPE should handle mixed types: best value {} should be < 25.0", + best.value + ); +}