gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement Multivariant TPE

Browse Source
This commit is contained in:
Manuel Raimann
2026-02-02 14:43:56 +01:00
parent 293b992e75
commit 4ce4311c68
12 changed files with 9702 additions and 752 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
+1
View File
@@ -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"
src/error.rs
+26
View File
@@ -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),
src/kde/mod.rs
+13
View File
@@ -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;
src/kde/multivariate.rs
+921
View File
@@ -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<Vec<f64>>,
/// The bandwidth (standard deviation) for each dimension.
/// Uses a diagonal bandwidth matrix (independent bandwidths per dimension).
bandwidths: Vec<f64>,
/// 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<Vec<f64>>) -> Result<Self> {
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<Vec<f64>>, bandwidths: Vec<f64>) -> Result<Self> {
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<f64>], n_dims: usize) -> Vec<f64> {
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<f64>], dim: usize) -> f64 {
let n = samples.len() as f64;
let values: Vec<f64> = samples.iter().map(|s| s[dim]).collect();
let mean = values.iter().sum::<f64>() / n;
let variance = values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / 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<f64>] {
&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<f64> = 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<f64> = 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<f64>` of length `n_dims` representing a sample from the KDE.
pub(crate) fn sample<R: Rng>(&self, rng: &mut R) -> Vec<f64> {
// 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(|(&center_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<Vec<f64>> = (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<f64>> = 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<Vec<f64>> = (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<Vec<f64>> = (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<Vec<f64>> = (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<f64> = Vec::with_capacity(n_samples);
let mut values_y: Vec<f64> = 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::<f64>() / n;
let mean_y: f64 = values_y.iter().sum::<f64>() / n;
let var_x: f64 = values_x.iter().map(|x| (x - mean_x).powi(2)).sum::<f64>() / n;
let var_y: f64 = values_y.iter().map(|y| (y - mean_y).powi(2)).sum::<f64>() / 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"
);
}
}
src/kde.rs → src/kde/univariate.rs
View File
src/sampler/mod.rs
+56
View File
@@ -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<V> CompletedTrial<V> {
}
}
/// 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<String, ParamValue>,
/// The parameter distributions used, keyed by parameter name.
pub distributions: HashMap<String, Distribution>,
}
impl PendingTrial {
/// Creates a new pending trial.
#[must_use]
pub fn new(
id: u64,
params: HashMap<String, ParamValue>,
distributions: HashMap<String, Distribution>,
) -> Self {
Self {
id,
params,
distributions,
}
}
}
/// Trait for pluggable parameter sampling strategies.
///
/// Samplers are responsible for generating parameter values based on
src/sampler/multivariate_tpe.rs
+6522
View File
File diff suppressed because it is too large Load Diff
src/sampler/tpe/gamma.rs
+670
View File
@@ -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<dyn GammaStrategy> {
/// Box::new(self.clone())
/// }
/// }
/// ```
pub trait GammaStrategy: Send + Sync + Debug {
/// Computes the gamma quantile based on the number of completed trials.
///
/// # Arguments
///
/// * `n_trials` - The number of completed trials in the history.
///
/// # Returns
///
/// A gamma value in the range (0.0, 1.0). Values outside this range
/// will be clamped by the sampler.
fn gamma(&self, n_trials: usize) -> f64;
/// Creates a boxed clone of this strategy.
///
/// This method enables cloning of trait objects, which is necessary
/// for the builder pattern and sampler configuration.
fn clone_box(&self) -> Box<dyn GammaStrategy>;
}
impl Clone for Box<dyn GammaStrategy> {
fn clone(&self) -> Self {
self.clone_box()
}
}
/// A fixed gamma strategy that returns a constant value.
///
/// This is the simplest strategy and the default behavior of TPE.
/// The gamma value remains constant regardless of the number of trials.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{FixedGamma, TpeSampler};
///
/// // Use 15% of trials as "good"
/// let sampler = TpeSampler::builder()
/// .gamma_strategy(FixedGamma::new(0.15).unwrap())
/// .build()
/// .unwrap();
/// ```
#[derive(Debug, Clone, Copy)]
pub struct FixedGamma {
gamma: f64,
}
impl FixedGamma {
/// Creates a new fixed gamma strategy.
///
/// # Arguments
///
/// * `gamma` - The constant gamma value to use.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0).
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::FixedGamma;
///
/// let strategy = FixedGamma::new(0.25).unwrap();
/// assert!((strategy.value() - 0.25).abs() < f64::EPSILON);
/// ```
pub fn new(gamma: f64) -> crate::Result<Self> {
if gamma <= 0.0 || gamma >= 1.0 {
return Err(Error::InvalidGamma(gamma));
}
Ok(Self { gamma })
}
/// Returns the fixed gamma value.
#[must_use]
pub fn value(&self) -> f64 {
self.gamma
}
}
impl Default for FixedGamma {
/// Creates a fixed gamma strategy with the default value of 0.25.
fn default() -> Self {
Self { gamma: 0.25 }
}
}
impl GammaStrategy for FixedGamma {
fn gamma(&self, _n_trials: usize) -> f64 {
self.gamma
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A linear gamma strategy that interpolates between min and max values.
///
/// The gamma value increases linearly from `gamma_min` to `gamma_max` as the
/// number of trials grows from 0 to `n_trials_max`. Beyond `n_trials_max`,
/// gamma remains at `gamma_max`.
///
/// This strategy is useful when you want to be more explorative early on
/// (smaller gamma = fewer "good" trials) and more exploitative later
/// (larger gamma = more "good" trials).
///
/// # Formula
///
/// ```text
/// gamma = gamma_min + (gamma_max - gamma_min) * min(n_trials / n_trials_max, 1.0)
/// ```
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, LinearGamma, TpeSampler};
///
/// let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap();
///
/// // At 0 trials: gamma = 0.1
/// assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON);
///
/// // At 50 trials: gamma = 0.25 (midpoint)
/// assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON);
///
/// // At 100+ trials: gamma = 0.4
/// assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON);
/// assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON);
/// ```
#[derive(Debug, Clone, Copy)]
pub struct LinearGamma {
gamma_min: f64,
gamma_max: f64,
n_trials_max: usize,
}
impl LinearGamma {
/// Creates a new linear gamma strategy.
///
/// # Arguments
///
/// * `gamma_min` - The minimum gamma value (at 0 trials).
/// * `gamma_max` - The maximum gamma value (at `n_trials_max` trials).
/// * `n_trials_max` - The number of trials at which gamma reaches its maximum.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_min` is not in (0.0, 1.0)
/// - `gamma_max` is not in (0.0, 1.0)
/// - `gamma_min > gamma_max`
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::LinearGamma;
///
/// // Gamma goes from 0.1 to 0.3 over 50 trials
/// let strategy = LinearGamma::new(0.1, 0.3, 50).unwrap();
/// ```
pub fn new(gamma_min: f64, gamma_max: f64, n_trials_max: usize) -> crate::Result<Self> {
if gamma_min <= 0.0 || gamma_min >= 1.0 {
return Err(Error::InvalidGamma(gamma_min));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
if gamma_min > gamma_max {
return Err(Error::InvalidGamma(gamma_min));
}
Ok(Self {
gamma_min,
gamma_max,
n_trials_max,
})
}
/// Returns the minimum gamma value.
#[must_use]
pub fn gamma_min(&self) -> f64 {
self.gamma_min
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
/// Returns the number of trials at which gamma reaches its maximum.
#[must_use]
pub fn n_trials_max(&self) -> usize {
self.n_trials_max
}
}
impl Default for LinearGamma {
/// Creates a linear gamma strategy with default values:
/// - `gamma_min`: 0.10
/// - `gamma_max`: 0.25
/// - `n_trials_max`: 100
fn default() -> Self {
Self {
gamma_min: 0.10,
gamma_max: 0.25,
n_trials_max: 100,
}
}
}
impl GammaStrategy for LinearGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if self.n_trials_max == 0 {
return self.gamma_max;
}
let t = (n_trials as f64 / self.n_trials_max as f64).min(1.0);
self.gamma_min + (self.gamma_max - self.gamma_min) * t
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A square root gamma strategy inspired by Optuna's default behavior.
///
/// The gamma value is computed based on the inverse square root of the number
/// of trials, providing a balance between exploration and exploitation that
/// naturally adapts as more data becomes available.
///
/// # Formula
///
/// ```text
/// n_good = max(1, floor(gamma_factor / sqrt(n_trials)))
/// gamma = min(gamma_max, n_good / n_trials)
/// ```
///
/// When `n_trials` is 0, returns `gamma_max`.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, SqrtGamma, TpeSampler};
///
/// let strategy = SqrtGamma::default();
///
/// // Gamma decreases as trials increase
/// let g10 = strategy.gamma(10);
/// let g100 = strategy.gamma(100);
/// assert!(g10 > g100, "Gamma should decrease with more trials");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct SqrtGamma {
gamma_factor: f64,
gamma_max: f64,
}
impl SqrtGamma {
/// Creates a new square root gamma strategy.
///
/// # Arguments
///
/// * `gamma_factor` - The factor controlling how quickly gamma decreases.
/// Higher values mean more "good" trials at any given point.
/// * `gamma_max` - The maximum gamma value (used when `n_trials` is small).
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_factor` is not positive
/// - `gamma_max` is not in (0.0, 1.0)
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::SqrtGamma;
///
/// let strategy = SqrtGamma::new(1.0, 0.25).unwrap();
/// ```
pub fn new(gamma_factor: f64, gamma_max: f64) -> crate::Result<Self> {
if gamma_factor <= 0.0 {
return Err(Error::InvalidGamma(gamma_factor));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
Ok(Self {
gamma_factor,
gamma_max,
})
}
/// Returns the gamma factor.
#[must_use]
pub fn gamma_factor(&self) -> f64 {
self.gamma_factor
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
}
impl Default for SqrtGamma {
/// Creates a square root gamma strategy with default values:
/// - `gamma_factor`: 1.0
/// - `gamma_max`: 0.25
fn default() -> Self {
Self {
gamma_factor: 1.0,
gamma_max: 0.25,
}
}
}
impl GammaStrategy for SqrtGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if n_trials == 0 {
return self.gamma_max;
}
let n_good = (self.gamma_factor / (n_trials as f64).sqrt()).max(1.0);
(n_good / n_trials as f64).min(self.gamma_max)
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A Hyperopt-style gamma strategy.
///
/// This strategy computes gamma as `min(gamma_max, (gamma_base + 1) / n_trials)`,
/// which is inspired by the original Hyperopt TPE implementation.
///
/// # Formula
///
/// ```text
/// gamma = min(gamma_max, (gamma_base + 1) / n_trials)
/// ```
///
/// When `n_trials` is 0, returns `gamma_max`.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, HyperoptGamma};
///
/// // With gamma_base=24 and gamma_max=0.5:
/// // - At n=25: gamma = min(0.5, 25/25) = 0.5 (capped)
/// // - At n=100: gamma = min(0.5, 25/100) = 0.25
/// let strategy = HyperoptGamma::new(24.0, 0.5).unwrap();
///
/// // Early trials have higher gamma
/// let g50 = strategy.gamma(50);
/// let g200 = strategy.gamma(200);
/// assert!(g50 > g200, "Gamma should decrease with more trials");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct HyperoptGamma {
gamma_base: f64,
gamma_max: f64,
}
impl HyperoptGamma {
/// Creates a new Hyperopt-style gamma strategy.
///
/// # Arguments
///
/// * `gamma_base` - The base value added to 1 in the numerator.
/// * `gamma_max` - The maximum gamma value.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_base` is negative
/// - `gamma_max` is not in (0.0, 1.0)
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::HyperoptGamma;
///
/// let strategy = HyperoptGamma::new(24.0, 0.25).unwrap();
/// ```
pub fn new(gamma_base: f64, gamma_max: f64) -> crate::Result<Self> {
if gamma_base < 0.0 {
return Err(Error::InvalidGamma(gamma_base));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
Ok(Self {
gamma_base,
gamma_max,
})
}
/// Returns the gamma base value.
#[must_use]
pub fn gamma_base(&self) -> f64 {
self.gamma_base
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
}
impl Default for HyperoptGamma {
/// Creates a Hyperopt-style gamma strategy with default values:
/// - `gamma_base`: 24.0
/// - `gamma_max`: 0.25
fn default() -> Self {
Self {
gamma_base: 24.0,
gamma_max: 0.25,
}
}
}
impl GammaStrategy for HyperoptGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if n_trials == 0 {
return self.gamma_max;
}
((self.gamma_base + 1.0) / n_trials as f64).min(self.gamma_max)
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
#[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<dyn GammaStrategy> = Box::new(FixedGamma::new(0.3).unwrap());
let cloned = fixed.clone();
assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON);
let linear: Box<dyn GammaStrategy> = Box::new(LinearGamma::default());
let cloned = linear.clone();
assert!((cloned.gamma(0) - 0.10).abs() < f64::EPSILON);
}
#[test]
fn test_tpe_with_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<dyn GammaStrategy> {
Box::new(self.clone())
}
}
let sampler = TpeSampler::builder()
.gamma_strategy(DoubleGamma)
.build()
.unwrap();
assert!((sampler.gamma_strategy().gamma(10) - 0.1).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(50) - 0.5).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(100) - 0.5).abs() < f64::EPSILON);
}
}
src/sampler/tpe/mod.rs
+12
View File
@@ -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};
src/sampler/tpe.rs → src/sampler/tpe/sampler.rs
+2 -752
View File
@@ -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<dyn GammaStrategy> {
/// Box::new(self.clone())
/// }
/// }
/// ```
pub trait GammaStrategy: Send + Sync + Debug {
/// Computes the gamma quantile based on the number of completed trials.
///
/// # Arguments
///
/// * `n_trials` - The number of completed trials in the history.
///
/// # Returns
///
/// A gamma value in the range (0.0, 1.0). Values outside this range
/// will be clamped by the sampler.
fn gamma(&self, n_trials: usize) -> f64;
/// Creates a boxed clone of this strategy.
///
/// This method enables cloning of trait objects, which is necessary
/// for the builder pattern and sampler configuration.
fn clone_box(&self) -> Box<dyn GammaStrategy>;
}
impl Clone for Box<dyn GammaStrategy> {
fn clone(&self) -> Self {
self.clone_box()
}
}
/// A fixed gamma strategy that returns a constant value.
///
/// This is the simplest strategy and the default behavior of TPE.
/// The gamma value remains constant regardless of the number of trials.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{FixedGamma, TpeSampler};
///
/// // Use 15% of trials as "good"
/// let sampler = TpeSampler::builder()
/// .gamma_strategy(FixedGamma::new(0.15).unwrap())
/// .build()
/// .unwrap();
/// ```
#[derive(Debug, Clone, Copy)]
pub struct FixedGamma {
gamma: f64,
}
impl FixedGamma {
/// Creates a new fixed gamma strategy.
///
/// # Arguments
///
/// * `gamma` - The constant gamma value to use.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0).
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::FixedGamma;
///
/// let strategy = FixedGamma::new(0.25).unwrap();
/// assert!((strategy.value() - 0.25).abs() < f64::EPSILON);
/// ```
pub fn new(gamma: f64) -> Result<Self> {
if gamma <= 0.0 || gamma >= 1.0 {
return Err(Error::InvalidGamma(gamma));
}
Ok(Self { gamma })
}
/// Returns the fixed gamma value.
#[must_use]
pub fn value(&self) -> f64 {
self.gamma
}
}
impl Default for FixedGamma {
/// Creates a fixed gamma strategy with the default value of 0.25.
fn default() -> Self {
Self { gamma: 0.25 }
}
}
impl GammaStrategy for FixedGamma {
fn gamma(&self, _n_trials: usize) -> f64 {
self.gamma
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A linear gamma strategy that interpolates between min and max values.
///
/// The gamma value increases linearly from `gamma_min` to `gamma_max` as the
/// number of trials grows from 0 to `n_trials_max`. Beyond `n_trials_max`,
/// gamma remains at `gamma_max`.
///
/// This strategy is useful when you want to be more explorative early on
/// (smaller gamma = fewer "good" trials) and more exploitative later
/// (larger gamma = more "good" trials).
///
/// # Formula
///
/// ```text
/// gamma = gamma_min + (gamma_max - gamma_min) * min(n_trials / n_trials_max, 1.0)
/// ```
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, LinearGamma, TpeSampler};
///
/// let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap();
///
/// // At 0 trials: gamma = 0.1
/// assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON);
///
/// // At 50 trials: gamma = 0.25 (midpoint)
/// assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON);
///
/// // At 100+ trials: gamma = 0.4
/// assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON);
/// assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON);
/// ```
#[derive(Debug, Clone, Copy)]
pub struct LinearGamma {
gamma_min: f64,
gamma_max: f64,
n_trials_max: usize,
}
impl LinearGamma {
/// Creates a new linear gamma strategy.
///
/// # Arguments
///
/// * `gamma_min` - The minimum gamma value (at 0 trials).
/// * `gamma_max` - The maximum gamma value (at `n_trials_max` trials).
/// * `n_trials_max` - The number of trials at which gamma reaches its maximum.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_min` is not in (0.0, 1.0)
/// - `gamma_max` is not in (0.0, 1.0)
/// - `gamma_min > gamma_max`
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::LinearGamma;
///
/// // Gamma goes from 0.1 to 0.3 over 50 trials
/// let strategy = LinearGamma::new(0.1, 0.3, 50).unwrap();
/// ```
pub fn new(gamma_min: f64, gamma_max: f64, n_trials_max: usize) -> Result<Self> {
if gamma_min <= 0.0 || gamma_min >= 1.0 {
return Err(Error::InvalidGamma(gamma_min));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
if gamma_min > gamma_max {
return Err(Error::InvalidGamma(gamma_min));
}
Ok(Self {
gamma_min,
gamma_max,
n_trials_max,
})
}
/// Returns the minimum gamma value.
#[must_use]
pub fn gamma_min(&self) -> f64 {
self.gamma_min
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
/// Returns the number of trials at which gamma reaches its maximum.
#[must_use]
pub fn n_trials_max(&self) -> usize {
self.n_trials_max
}
}
impl Default for LinearGamma {
/// Creates a linear gamma strategy with default values:
/// - `gamma_min`: 0.10
/// - `gamma_max`: 0.25
/// - `n_trials_max`: 100
fn default() -> Self {
Self {
gamma_min: 0.10,
gamma_max: 0.25,
n_trials_max: 100,
}
}
}
impl GammaStrategy for LinearGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if self.n_trials_max == 0 {
return self.gamma_max;
}
let t = (n_trials as f64 / self.n_trials_max as f64).min(1.0);
self.gamma_min + (self.gamma_max - self.gamma_min) * t
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A square root gamma strategy inspired by Optuna's default behavior.
///
/// The gamma value is computed based on the inverse square root of the number
/// of trials, providing a balance between exploration and exploitation that
/// naturally adapts as more data becomes available.
///
/// # Formula
///
/// ```text
/// n_good = max(1, floor(gamma_factor / sqrt(n_trials)))
/// gamma = min(gamma_max, n_good / n_trials)
/// ```
///
/// When `n_trials` is 0, returns `gamma_max`.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, SqrtGamma, TpeSampler};
///
/// let strategy = SqrtGamma::default();
///
/// // Gamma decreases as trials increase
/// let g10 = strategy.gamma(10);
/// let g100 = strategy.gamma(100);
/// assert!(g10 > g100, "Gamma should decrease with more trials");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct SqrtGamma {
gamma_factor: f64,
gamma_max: f64,
}
impl SqrtGamma {
/// Creates a new square root gamma strategy.
///
/// # Arguments
///
/// * `gamma_factor` - The factor controlling how quickly gamma decreases.
/// Higher values mean more "good" trials at any given point.
/// * `gamma_max` - The maximum gamma value (used when `n_trials` is small).
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_factor` is not positive
/// - `gamma_max` is not in (0.0, 1.0)
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::SqrtGamma;
///
/// let strategy = SqrtGamma::new(1.0, 0.25).unwrap();
/// ```
pub fn new(gamma_factor: f64, gamma_max: f64) -> Result<Self> {
if gamma_factor <= 0.0 {
return Err(Error::InvalidGamma(gamma_factor));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
Ok(Self {
gamma_factor,
gamma_max,
})
}
/// Returns the gamma factor.
#[must_use]
pub fn gamma_factor(&self) -> f64 {
self.gamma_factor
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
}
impl Default for SqrtGamma {
/// Creates a square root gamma strategy with default values:
/// - `gamma_factor`: 1.0
/// - `gamma_max`: 0.25
fn default() -> Self {
Self {
gamma_factor: 1.0,
gamma_max: 0.25,
}
}
}
impl GammaStrategy for SqrtGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if n_trials == 0 {
return self.gamma_max;
}
let n_good = (self.gamma_factor / (n_trials as f64).sqrt()).max(1.0);
(n_good / n_trials as f64).min(self.gamma_max)
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A Hyperopt-style gamma strategy.
///
/// This strategy computes gamma as `min(gamma_max, (gamma_base + 1) / n_trials)`,
/// which is inspired by the original Hyperopt TPE implementation.
///
/// # Formula
///
/// ```text
/// gamma = min(gamma_max, (gamma_base + 1) / n_trials)
/// ```
///
/// When `n_trials` is 0, returns `gamma_max`.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, HyperoptGamma};
///
/// // With gamma_base=24 and gamma_max=0.5:
/// // - At n=25: gamma = min(0.5, 25/25) = 0.5 (capped)
/// // - At n=100: gamma = min(0.5, 25/100) = 0.25
/// let strategy = HyperoptGamma::new(24.0, 0.5).unwrap();
///
/// // Early trials have higher gamma
/// let g50 = strategy.gamma(50);
/// let g200 = strategy.gamma(200);
/// assert!(g50 > g200, "Gamma should decrease with more trials");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct HyperoptGamma {
gamma_base: f64,
gamma_max: f64,
}
impl HyperoptGamma {
/// Creates a new Hyperopt-style gamma strategy.
///
/// # Arguments
///
/// * `gamma_base` - The base value added to 1 in the numerator.
/// * `gamma_max` - The maximum gamma value.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_base` is negative
/// - `gamma_max` is not in (0.0, 1.0)
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::HyperoptGamma;
///
/// let strategy = HyperoptGamma::new(24.0, 0.25).unwrap();
/// ```
pub fn new(gamma_base: f64, gamma_max: f64) -> Result<Self> {
if gamma_base < 0.0 {
return Err(Error::InvalidGamma(gamma_base));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
Ok(Self {
gamma_base,
gamma_max,
})
}
/// Returns the gamma base value.
#[must_use]
pub fn gamma_base(&self) -> f64 {
self.gamma_base
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
}
impl Default for HyperoptGamma {
/// Creates a Hyperopt-style gamma strategy with default values:
/// - `gamma_base`: 24.0
/// - `gamma_max`: 0.25
fn default() -> Self {
Self {
gamma_base: 24.0,
gamma_max: 0.25,
}
}
}
impl GammaStrategy for HyperoptGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if n_trials == 0 {
return self.gamma_max;
}
((self.gamma_base + 1.0) / n_trials as f64).min(self.gamma_max)
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
// ============================================================================
// TPE Sampler
// ============================================================================
@@ -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<dyn GammaStrategy> = Box::new(FixedGamma::new(0.3).unwrap());
let cloned = fixed.clone();
assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON);
let linear: Box<dyn GammaStrategy> = Box::new(LinearGamma::default());
let cloned = linear.clone();
assert!((cloned.gamma(0) - 0.10).abs() < f64::EPSILON);
}
#[test]
fn test_tpe_with_sqrt_gamma_strategy() {
let sampler = TpeSampler::builder()
.gamma_strategy(SqrtGamma::default())
.n_startup_trials(5)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
create_trial(
i as u64,
f64::from(i),
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
)
})
.collect();
// Should be able to sample with the sqrt gamma strategy
let value = sampler.sample(&dist, 100, &history);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v));
} else {
panic!("Expected Float value");
}
}
#[test]
fn test_tpe_with_linear_gamma_strategy() {
let sampler = TpeSampler::builder()
.gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap())
.n_startup_trials(5)
.seed(42)
.build()
.unwrap();
// Verify the strategy is applied
let g = sampler.gamma_strategy().gamma(25);
assert!((g - 0.2).abs() < f64::EPSILON); // midpoint of 0.1 to 0.3
}
#[test]
fn test_tpe_with_hyperopt_gamma_strategy() {
let sampler = TpeSampler::builder()
.gamma_strategy(HyperoptGamma::default())
.n_startup_trials(5)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
create_trial(
i as u64,
f64::from(i),
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
)
})
.collect();
// Should be able to sample with the hyperopt gamma strategy
let value = sampler.sample(&dist, 100, &history);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v));
} else {
panic!("Expected Float value");
}
}
#[test]
fn test_gamma_overrides_gamma_strategy() {
// When gamma() is called after gamma_strategy(), it should take precedence
let sampler = TpeSampler::builder()
.gamma_strategy(SqrtGamma::default())
.gamma(0.15) // This should override
.build()
.unwrap();
// Should use fixed gamma of 0.15
assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(100) - 0.15).abs() < f64::EPSILON);
}
#[test]
fn test_gamma_strategy_overrides_gamma() {
// When gamma_strategy() is called after gamma(), it should take precedence
let sampler = TpeSampler::builder()
.gamma(0.15)
.gamma_strategy(SqrtGamma::default()) // This should override
.build()
.unwrap();
// Should use SqrtGamma - gamma decreases with trials
let g10 = sampler.gamma_strategy().gamma(10);
let g100 = sampler.gamma_strategy().gamma(100);
assert!(g10 > g100, "SqrtGamma should decrease with more trials");
}
#[test]
fn test_with_strategy_constructor() {
let sampler = TpeSampler::with_strategy(
LinearGamma::new(0.1, 0.4, 100).unwrap(),
15,
32,
None,
Some(42),
)
.unwrap();
assert_eq!(sampler.n_startup_trials, 15);
assert_eq!(sampler.n_ei_candidates, 32);
assert!((sampler.gamma_strategy().gamma(0) - 0.1).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(100) - 0.4).abs() < f64::EPSILON);
}
#[test]
fn test_custom_gamma_strategy() {
#[derive(Debug, Clone)]
struct DoubleGamma;
impl GammaStrategy for DoubleGamma {
fn gamma(&self, n_trials: usize) -> f64 {
// Double the trial count-based calculation, capped at 0.5
(0.01 * n_trials as f64).min(0.5)
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(self.clone())
}
}
let sampler = TpeSampler::builder()
.gamma_strategy(DoubleGamma)
.build()
.unwrap();
assert!((sampler.gamma_strategy().gamma(10) - 0.1).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(50) - 0.5).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(100) - 0.5).abs() < f64::EPSILON);
}
}
src/sampler/tpe/search_space.rs
+1086
View File
File diff suppressed because it is too large Load Diff
tests/multivariate_tpe_integration.rs
+393
View File
@@ -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<f64> = 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<f64> = 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<f64> = 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<f64> = 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::<f64>() / n_runs as f64;
let independent_mean: f64 = independent_best_values.iter().sum::<f64>() / 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<f64> = 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<f64> = 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<f64> = 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<f64> = 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::<f64>() / n_runs as f64;
let independent_mean: f64 = independent_results.iter().sum::<f64>() / 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<f64> = 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<f64> = 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
);
}