From b482d56e8976157f6bc2b4411d99464095fd4852 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Fri, 30 Jan 2026 19:58:18 +0100 Subject: [PATCH] Implement grid search --- Cargo.toml | 2 +- README.md | 55 +- src/lib.rs | 42 +- src/sampler/grid.rs | 1157 ++++++++++++++++++++++++++++++++++++++++++ src/sampler/mod.rs | 1 + tests/integration.rs | 12 +- 6 files changed, 1257 insertions(+), 12 deletions(-) create mode 100644 src/sampler/grid.rs diff --git a/Cargo.toml b/Cargo.toml index 0813afd..cf86626 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,7 +8,7 @@ authors = ["Manuel Raimann = Study::with_sampler(Direction::Minimize, sampler); study @@ -32,6 +37,50 @@ let best = study.best_trial().unwrap(); println!("Best value: {} at x={:?}", best.value, best.params); ``` +## Samplers + +### Random Search + +```rust +use optimizer::{Direction, Study}; +use optimizer::sampler::random::RandomSampler; + +let study: Study = Study::with_sampler( + Direction::Minimize, + RandomSampler::with_seed(42), +); +``` + +### TPE (Tree-Parzen Estimator) + +```rust +use optimizer::{Direction, Study}; +use optimizer::sampler::tpe::TpeSampler; + +let sampler = TpeSampler::builder() + .gamma(0.15) // Quantile for good/bad split + .n_startup_trials(20) // Random trials before TPE kicks in + .n_ei_candidates(32) // Candidates to evaluate + .seed(42) + .build() + .unwrap(); + +let study: Study = Study::with_sampler(Direction::Minimize, sampler); +``` + +### Grid Search + +```rust +use optimizer::{Direction, Study}; +use optimizer::sampler::grid::GridSearchSampler; + +let sampler = GridSearchSampler::builder() + .n_points_per_param(10) // Number of points per parameter dimension + .build(); + +let study: Study = Study::with_sampler(Direction::Minimize, sampler); +``` + ## Feature Flags - `async` - Enable async optimization methods (requires tokio) diff --git a/src/lib.rs b/src/lib.rs index a0660df..99473ea 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,10 +9,16 @@ #![deny(clippy::pedantic)] #![deny(clippy::std_instead_of_core)] -//! A Tree-Parzen Estimator (TPE) library for black-box optimization. +//! A black-box optimization library with multiple sampling strategies. //! //! This library provides an Optuna-like API for hyperparameter optimization -//! using the Tree-Parzen Estimator algorithm. It supports: +//! with support for multiple sampling algorithms: +//! +//! - **Random Search** - Simple random sampling for baseline comparisons +//! - **TPE (Tree-Parzen Estimator)** - Bayesian optimization for efficient search +//! - **Grid Search** - Exhaustive search over a specified parameter grid +//! +//! Additional features include: //! //! - Float, integer, and categorical parameter types //! - Log-scale and stepped parameter sampling @@ -91,9 +97,22 @@ //! .unwrap(); //! ``` //! -//! # Configuring TPE +//! # Available Samplers //! -//! The [`sampler::tpe::TpeSampler`] can be configured using the builder pattern: +//! ## Random Search +//! +//! The simplest sampling strategy, useful for baselines: +//! +//! ``` +//! use optimizer::sampler::random::RandomSampler; +//! use optimizer::{Direction, Study}; +//! +//! let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42)); +//! ``` +//! +//! ## TPE (Tree-Parzen Estimator) +//! +//! Bayesian optimization that learns from previous trials: //! //! ``` //! use optimizer::sampler::tpe::TpeSampler; @@ -107,6 +126,21 @@ //! .unwrap(); //! ``` //! +//! ## Grid Search +//! +//! Exhaustive search over a discretized parameter space: +//! +//! ``` +//! use optimizer::sampler::grid::GridSearchSampler; +//! use optimizer::{Direction, Study}; +//! +//! let sampler = GridSearchSampler::builder() +//! .n_points_per_param(10) // Points per parameter dimension +//! .build(); +//! +//! let study: Study = Study::with_sampler(Direction::Minimize, sampler); +//! ``` +//! //! # Async and Parallel Optimization //! //! With the `async` feature enabled, you can run trials asynchronously: diff --git a/src/sampler/grid.rs b/src/sampler/grid.rs new file mode 100644 index 0000000..dd6768f --- /dev/null +++ b/src/sampler/grid.rs @@ -0,0 +1,1157 @@ +//! Grid search sampler implementation. +//! +//! `GridSearchSampler` performs exhaustive grid search over the parameter space, +//! systematically evaluating all combinations of discretized parameter values. + +use std::collections::HashMap; + +use parking_lot::Mutex; + +use crate::distribution::{ + CategoricalDistribution, Distribution, FloatDistribution, IntDistribution, +}; +use crate::param::ParamValue; +use crate::sampler::{CompletedTrial, Sampler}; + +/// Generates grid points for an integer distribution. +/// +/// # Behavior +/// +/// - If `step` is `Some(s)`: generates points at `low, low+s, low+2*s, ...` up to `high`. +/// - If `step` is `None` and `log_scale` is `false`: generates `n_points` evenly spaced +/// integers from `low` to `high`. +/// - If `step` is `None` and `log_scale` is `true`: generates `n_points` evenly spaced +/// in log space, rounded to integers. +/// +/// All grid points are clamped to `[low, high]` bounds and deduplicated. +/// +/// # Arguments +/// +/// * `dist` - The integer distribution defining bounds, step, and log scale. +/// * `n_points` - Number of points to generate when auto-discretizing (ignored if step is set). +/// +/// # Returns +/// +/// A vector of unique grid points in ascending order. +#[must_use] +pub fn generate_int_grid_points(dist: &IntDistribution, n_points: usize) -> Vec { + let low = dist.low; + let high = dist.high; + + if low > high { + return vec![]; + } + + if low == high { + return vec![low]; + } + + let points: Vec = if let Some(step) = dist.step { + // Generate points at low, low+step, low+2*step, ... up to high + if step <= 0 { + return vec![low]; + } + let mut result = Vec::new(); + let mut current = low; + while current <= high { + result.push(current); + current = current.saturating_add(step); + // Prevent infinite loop if saturating_add doesn't change the value + if result.last() == Some(¤t) { + break; + } + } + result + } else if dist.log_scale { + // Generate n_points evenly spaced in log space, rounded to integers + // For log scale, low must be positive + if low <= 0 { + // Fall back to linear for non-positive values + generate_linear_int_points(low, high, n_points) + } else { + generate_log_int_points(low, high, n_points) + } + } else { + // Generate n_points evenly spaced integers from low to high + generate_linear_int_points(low, high, n_points) + }; + + // Clamp all points to [low, high] and deduplicate + let mut clamped: Vec = points.into_iter().map(|p| p.clamp(low, high)).collect(); + + // Remove duplicates while preserving order + clamped.sort_unstable(); + clamped.dedup(); + + clamped +} + +/// Generates evenly spaced integers from low to high (linear scale). +#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] +fn generate_linear_int_points(low: i64, high: i64, n_points: usize) -> Vec { + if n_points == 0 { + return vec![]; + } + if n_points == 1 { + return vec![low]; + } + + let range = high - low; + let mut result = Vec::with_capacity(n_points); + + for i in 0..n_points { + // Calculate position as a fraction of the range + let fraction = i as f64 / (n_points - 1) as f64; + let value = low as f64 + fraction * range as f64; + result.push(value.round() as i64); + } + + result +} + +/// Generates evenly spaced integers in log space. +#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] +fn generate_log_int_points(low: i64, high: i64, n_points: usize) -> Vec { + debug_assert!(low > 0, "log scale requires positive low bound"); + + if n_points == 0 { + return vec![]; + } + if n_points == 1 { + return vec![low]; + } + + let log_low = (low as f64).ln(); + let log_high = (high as f64).ln(); + let mut result = Vec::with_capacity(n_points); + + for i in 0..n_points { + let fraction = i as f64 / (n_points - 1) as f64; + let log_value = log_low + fraction * (log_high - log_low); + let value = log_value.exp().round() as i64; + result.push(value); + } + + result +} + +/// Generates grid points for a float distribution. +/// +/// # Behavior +/// +/// - If `step` is `Some(s)`: generates points at `low, low+s, low+2*s, ...` up to `high`. +/// Step size overrides `log_scale` behavior. +/// - If `step` is `None` and `log_scale` is `false`: generates `n_points` evenly spaced +/// floats from `low` to `high`. +/// - If `step` is `None` and `log_scale` is `true`: generates `n_points` evenly spaced +/// in log space. +/// +/// All grid points are within `[low, high]` bounds. +/// +/// # Arguments +/// +/// * `dist` - The float distribution defining bounds, step, and log scale. +/// * `n_points` - Number of points to generate when auto-discretizing (ignored if step is set). +/// +/// # Returns +/// +/// A vector of grid points in ascending order. +#[must_use] +pub fn generate_float_grid_points(dist: &FloatDistribution, n_points: usize) -> Vec { + let low = dist.low; + let high = dist.high; + + if low > high { + return vec![]; + } + + if (low - high).abs() < f64::EPSILON { + return vec![low]; + } + + let points: Vec = if let Some(step) = dist.step { + // Step overrides log_scale - generate points at low, low+step, low+2*step, ... up to high + if step <= 0.0 { + return vec![low]; + } + let mut result = Vec::new(); + let mut current = low; + while current <= high + f64::EPSILON { + result.push(current.clamp(low, high)); + current += step; + // Prevent infinite loop from floating point issues + if result.len() > 1_000_000 { + break; + } + } + result + } else if dist.log_scale { + // Generate n_points evenly spaced in log space + // For log scale, low must be positive + if low <= 0.0 { + // Fall back to linear for non-positive values + generate_linear_float_points(low, high, n_points) + } else { + generate_log_float_points(low, high, n_points) + } + } else { + // Generate n_points evenly spaced floats from low to high + generate_linear_float_points(low, high, n_points) + }; + + // Clamp all points to [low, high] bounds + points.into_iter().map(|p| p.clamp(low, high)).collect() +} + +/// Generates evenly spaced floats from low to high (linear scale). +#[allow(clippy::cast_precision_loss)] +fn generate_linear_float_points(low: f64, high: f64, n_points: usize) -> Vec { + if n_points == 0 { + return vec![]; + } + if n_points == 1 { + return vec![low]; + } + + let range = high - low; + let mut result = Vec::with_capacity(n_points); + + for i in 0..n_points { + let fraction = i as f64 / (n_points - 1) as f64; + let value = low + fraction * range; + result.push(value); + } + + result +} + +/// Generates evenly spaced floats in log space. +#[allow(clippy::cast_precision_loss)] +fn generate_log_float_points(low: f64, high: f64, n_points: usize) -> Vec { + debug_assert!(low > 0.0, "log scale requires positive low bound"); + + if n_points == 0 { + return vec![]; + } + if n_points == 1 { + return vec![low]; + } + + let log_low = low.ln(); + let log_high = high.ln(); + let mut result = Vec::with_capacity(n_points); + + for i in 0..n_points { + let fraction = i as f64 / (n_points - 1) as f64; + let log_value = log_low + fraction * (log_high - log_low); + let value = log_value.exp(); + result.push(value); + } + + result +} + +/// Generates grid points for a categorical distribution. +/// +/// Returns a vector containing all choice indices `[0, 1, 2, ..., n_choices-1]`. +/// All choices are included in the grid for exhaustive evaluation. +/// +/// # Arguments +/// +/// * `dist` - The categorical distribution defining the number of choices. +/// +/// # Returns +/// +/// A vector of all choice indices from 0 to `n_choices - 1`. +/// +/// # Examples +/// +/// ```ignore +/// use optimizer::sampler::grid::generate_categorical_grid_points; +/// +/// // CategoricalDistribution is internal; this shows intended usage +/// let dist = CategoricalDistribution { n_choices: 3 }; +/// let points = generate_categorical_grid_points(&dist); +/// assert_eq!(points, vec![0, 1, 2]); +/// ``` +#[must_use] +pub fn generate_categorical_grid_points(dist: &CategoricalDistribution) -> Vec { + (0..dist.n_choices).collect() +} + +/// Cached grid points for a distribution. +#[derive(Debug, Clone)] +struct CachedGrid { + /// The generated grid points as `ParamValue` instances. + points: Vec, + /// Current index into the grid. + current_index: usize, +} + +/// Internal state for tracking grid position per distribution. +#[derive(Debug, Default)] +struct GridState { + /// Cached grids for each distribution, keyed by distribution identifier. + grids: HashMap, +} + +/// A grid search sampler that exhaustively evaluates all grid points. +/// +/// `GridSearchSampler` divides the parameter space into a grid and systematically +/// samples each point. This is useful when you want to evaluate all combinations +/// of parameter values, especially for discrete or small parameter spaces. +/// +/// # Grid Exhaustion +/// +/// The sampler tracks its position in the grid for each distribution independently. +/// When all grid points for a distribution have been sampled, subsequent calls to +/// `sample()` for that distribution will **panic** with the message: +/// `"GridSearchSampler: all grid points exhausted"`. +/// +/// To avoid panics, use [`is_exhausted()`](Self::is_exhausted) to check if all +/// points have been sampled before calling `sample()`. You can also use +/// [`grid_size()`](Self::grid_size) to determine the total number of grid points +/// that will be sampled. +/// +/// # Thread Safety +/// +/// `GridSearchSampler` is thread-safe (`Send + Sync`) and uses internal locking +/// to ensure safe concurrent access to grid state. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::grid::GridSearchSampler; +/// +/// // Create with default settings (10 points per parameter) +/// let sampler = GridSearchSampler::new(); +/// +/// // Create with custom settings using the builder +/// let sampler = GridSearchSampler::builder().n_points_per_param(20).build(); +/// ``` +pub struct GridSearchSampler { + /// Number of grid points per parameter (used when auto-discretizing). + n_points_per_param: usize, + /// Thread-safe internal state for tracking grid positions. + state: Mutex, +} + +impl GridSearchSampler { + /// Creates a new grid search sampler with default settings. + /// + /// Default settings: + /// - `n_points_per_param`: 10 (each continuous parameter is discretized to 10 points) + #[must_use] + pub fn new() -> Self { + Self { + n_points_per_param: 10, + state: Mutex::new(GridState::default()), + } + } + + /// Creates a builder for configuring a `GridSearchSampler`. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::grid::GridSearchSampler; + /// + /// let sampler = GridSearchSampler::builder().n_points_per_param(20).build(); + /// ``` + #[must_use] + pub fn builder() -> GridSearchSamplerBuilder { + GridSearchSamplerBuilder::new() + } +} + +impl Default for GridSearchSampler { + fn default() -> Self { + Self::new() + } +} + +impl GridSearchSampler { + /// Returns `true` if all grid points for all tracked distributions have been sampled. + /// + /// A distribution is considered exhausted when its `current_index` equals the number + /// of grid points. This method returns `true` only when **all** tracked distributions + /// are exhausted. + /// + /// Note that a newly created sampler with no distributions sampled yet will return + /// `true` (vacuously exhausted). After the first `sample()` call, the distribution + /// is tracked and `is_exhausted()` will return `false` until all its points are used. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::grid::GridSearchSampler; + /// + /// let sampler = GridSearchSampler::new(); + /// // Initially exhausted (no distributions tracked yet) + /// assert!(sampler.is_exhausted()); + /// ``` + #[must_use] + pub fn is_exhausted(&self) -> bool { + let state = self.state.lock(); + + // If no distributions are tracked yet, consider it vacuously exhausted + if state.grids.is_empty() { + return true; + } + + // All distributions must be exhausted + state + .grids + .values() + .all(|grid| grid.current_index >= grid.points.len()) + } + + /// Returns the total number of grid points across all tracked distributions. + /// + /// This method sums the number of grid points for each distribution that has been + /// sampled at least once. Before any `sample()` calls, this returns 0. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::grid::GridSearchSampler; + /// + /// let sampler = GridSearchSampler::new(); + /// // No distributions tracked yet + /// assert_eq!(sampler.grid_size(), 0); + /// ``` + #[must_use] + pub fn grid_size(&self) -> usize { + let state = self.state.lock(); + state.grids.values().map(|grid| grid.points.len()).sum() + } +} + +/// Builder for configuring a [`GridSearchSampler`]. +/// +/// # Examples +/// +/// ``` +/// use optimizer::sampler::grid::GridSearchSamplerBuilder; +/// +/// let sampler = GridSearchSamplerBuilder::new() +/// .n_points_per_param(20) +/// .build(); +/// ``` +#[derive(Debug, Clone)] +pub struct GridSearchSamplerBuilder { + n_points_per_param: usize, +} + +impl GridSearchSamplerBuilder { + /// Creates a new builder with default settings. + /// + /// Default settings: + /// - `n_points_per_param`: 10 + #[must_use] + pub fn new() -> Self { + Self { + n_points_per_param: 10, + } + } + + /// Sets the number of grid points per parameter for auto-discretization. + /// + /// This value is used when a parameter distribution doesn't have an explicit + /// step size. The parameter range is divided into `n` evenly spaced points. + /// + /// # Arguments + /// + /// * `n` - Number of points per parameter. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::grid::GridSearchSamplerBuilder; + /// + /// let sampler = GridSearchSamplerBuilder::new() + /// .n_points_per_param(20) + /// .build(); + /// ``` + #[must_use] + pub fn n_points_per_param(mut self, n: usize) -> Self { + self.n_points_per_param = n; + self + } + + /// Builds the configured [`GridSearchSampler`]. + /// + /// # Examples + /// + /// ``` + /// use optimizer::sampler::grid::GridSearchSamplerBuilder; + /// + /// let sampler = GridSearchSamplerBuilder::new() + /// .n_points_per_param(20) + /// .build(); + /// ``` + #[must_use] + pub fn build(self) -> GridSearchSampler { + GridSearchSampler { + n_points_per_param: self.n_points_per_param, + state: Mutex::new(GridState::default()), + } + } +} + +impl Default for GridSearchSamplerBuilder { + fn default() -> Self { + Self::new() + } +} + +/// Creates a unique identifier string from a distribution. +/// +/// This is used as a key in the internal state to track grid position per distribution. +fn distribution_key(dist: &Distribution) -> String { + match dist { + Distribution::Float(d) => { + format!( + "float:{}:{}:{}:{}", + d.low, + d.high, + d.log_scale, + d.step.map_or("none".to_string(), |s| s.to_string()) + ) + } + Distribution::Int(d) => { + format!( + "int:{}:{}:{}:{}", + d.low, + d.high, + d.log_scale, + d.step.map_or("none".to_string(), |s| s.to_string()) + ) + } + Distribution::Categorical(d) => { + format!("cat:{}", d.n_choices) + } + } +} + +impl Sampler for GridSearchSampler { + fn sample( + &self, + distribution: &Distribution, + _trial_id: u64, + _history: &[CompletedTrial], + ) -> ParamValue { + let mut state = self.state.lock(); + let key = distribution_key(distribution); + + // Get or create the cached grid for this distribution + let cached = state.grids.entry(key).or_insert_with(|| { + let points = match distribution { + Distribution::Float(d) => generate_float_grid_points(d, self.n_points_per_param) + .into_iter() + .map(ParamValue::Float) + .collect(), + Distribution::Int(d) => generate_int_grid_points(d, self.n_points_per_param) + .into_iter() + .map(ParamValue::Int) + .collect(), + Distribution::Categorical(d) => generate_categorical_grid_points(d) + .into_iter() + .map(ParamValue::Categorical) + .collect(), + }; + CachedGrid { + points, + current_index: 0, + } + }); + + // Check if all points have been exhausted + assert!( + cached.current_index < cached.points.len(), + "GridSearchSampler: all grid points exhausted" + ); + + // Get the current grid point and advance the index + let value = cached.points[cached.current_index].clone(); + cached.current_index += 1; + + value + } +} + +#[cfg(test)] +#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] +mod tests { + use super::*; + + // ==================== Int Distribution Tests ==================== + + #[test] + fn test_int_grid_with_step() { + let dist = IntDistribution { + low: 0, + high: 10, + log_scale: false, + step: Some(2), + }; + let points = generate_int_grid_points(&dist, 10); + // Should generate: 0, 2, 4, 6, 8, 10 + assert_eq!(points, vec![0, 2, 4, 6, 8, 10]); + } + + #[test] + fn test_int_grid_with_step_not_exact_multiple() { + let dist = IntDistribution { + low: 0, + high: 9, + log_scale: false, + step: Some(2), + }; + let points = generate_int_grid_points(&dist, 10); + // Should generate: 0, 2, 4, 6, 8 (stops at 8 because next step would be 10 > 9) + assert_eq!(points, vec![0, 2, 4, 6, 8]); + } + + #[test] + fn test_int_grid_without_step_linear() { + let dist = IntDistribution { + low: 0, + high: 100, + log_scale: false, + step: None, + }; + let points = generate_int_grid_points(&dist, 5); + // Should generate 5 evenly spaced points: 0, 25, 50, 75, 100 + assert_eq!(points, vec![0, 25, 50, 75, 100]); + } + + #[test] + fn test_int_grid_without_step_linear_10_points() { + let dist = IntDistribution { + low: 0, + high: 9, + log_scale: false, + step: None, + }; + let points = generate_int_grid_points(&dist, 10); + // Should generate all integers from 0 to 9 + assert_eq!(points, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + } + + #[test] + fn test_int_grid_with_log_scale() { + let dist = IntDistribution { + low: 1, + high: 1000, + log_scale: true, + step: None, + }; + let points = generate_int_grid_points(&dist, 4); + // Log scale: should cover range exponentially + // Points should be roughly: 1, 10, 100, 1000 + assert_eq!(points.len(), 4); + assert_eq!(points[0], 1); + assert_eq!(*points.last().unwrap(), 1000); + // Middle points should be between bounds + for p in &points { + assert!(*p >= 1 && *p <= 1000); + } + } + + #[test] + fn test_int_grid_log_scale_non_positive_fallback() { + // Log scale with non-positive low should fall back to linear + let dist = IntDistribution { + low: 0, + high: 100, + log_scale: true, + step: None, + }; + let points = generate_int_grid_points(&dist, 5); + // Falls back to linear: 0, 25, 50, 75, 100 + assert_eq!(points, vec![0, 25, 50, 75, 100]); + } + + #[test] + fn test_int_grid_single_point() { + let dist = IntDistribution { + low: 5, + high: 5, + log_scale: false, + step: None, + }; + let points = generate_int_grid_points(&dist, 10); + assert_eq!(points, vec![5]); + } + + #[test] + fn test_int_grid_invalid_bounds() { + let dist = IntDistribution { + low: 10, + high: 5, + log_scale: false, + step: None, + }; + let points = generate_int_grid_points(&dist, 10); + assert!(points.is_empty()); + } + + // ==================== Float Distribution Tests ==================== + + #[test] + fn test_float_grid_with_step() { + let dist = FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: Some(0.25), + }; + let points = generate_float_grid_points(&dist, 10); + // Should generate: 0.0, 0.25, 0.5, 0.75, 1.0 + assert_eq!(points.len(), 5); + assert!((points[0] - 0.0).abs() < f64::EPSILON); + assert!((points[1] - 0.25).abs() < f64::EPSILON); + assert!((points[2] - 0.5).abs() < f64::EPSILON); + assert!((points[3] - 0.75).abs() < f64::EPSILON); + assert!((points[4] - 1.0).abs() < f64::EPSILON); + } + + #[test] + fn test_float_grid_step_overrides_log_scale() { + // Step should override log_scale + let dist = FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: true, // This should be ignored + step: Some(0.5), + }; + let points = generate_float_grid_points(&dist, 10); + // Should generate: 0.0, 0.5, 1.0 (step overrides log_scale) + assert_eq!(points.len(), 3); + assert!((points[0] - 0.0).abs() < f64::EPSILON); + assert!((points[1] - 0.5).abs() < f64::EPSILON); + assert!((points[2] - 1.0).abs() < f64::EPSILON); + } + + #[test] + fn test_float_grid_without_step_linear() { + let dist = FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }; + let points = generate_float_grid_points(&dist, 5); + // Should generate 5 evenly spaced points: 0.0, 0.25, 0.5, 0.75, 1.0 + assert_eq!(points.len(), 5); + assert!((points[0] - 0.0).abs() < f64::EPSILON); + assert!((points[1] - 0.25).abs() < f64::EPSILON); + assert!((points[2] - 0.5).abs() < f64::EPSILON); + assert!((points[3] - 0.75).abs() < f64::EPSILON); + assert!((points[4] - 1.0).abs() < f64::EPSILON); + } + + #[test] + fn test_float_grid_with_log_scale() { + let dist = FloatDistribution { + low: 1e-4, + high: 1.0, + log_scale: true, + step: None, + }; + let points = generate_float_grid_points(&dist, 5); + // Log scale: should cover range exponentially + // Points should be roughly: 1e-4, 1e-3, 1e-2, 1e-1, 1.0 + assert_eq!(points.len(), 5); + assert!((points[0] - 1e-4).abs() < 1e-10); + assert!((points[4] - 1.0).abs() < 1e-10); + // Middle points should be between bounds + for p in &points { + assert!(*p >= 1e-4 && *p <= 1.0); + } + // In log scale, ratio between consecutive points should be roughly equal + let ratio1 = points[1] / points[0]; + let ratio2 = points[2] / points[1]; + assert!((ratio1 - ratio2).abs() / ratio1 < 0.01); + } + + #[test] + fn test_float_grid_log_scale_non_positive_fallback() { + // Log scale with non-positive low should fall back to linear + let dist = FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: true, + step: None, + }; + let points = generate_float_grid_points(&dist, 5); + // Falls back to linear: 0.0, 0.25, 0.5, 0.75, 1.0 + assert_eq!(points.len(), 5); + assert!((points[0] - 0.0).abs() < f64::EPSILON); + assert!((points[4] - 1.0).abs() < f64::EPSILON); + } + + #[test] + fn test_float_grid_single_point() { + let dist = FloatDistribution { + low: 0.5, + high: 0.5, + log_scale: false, + step: None, + }; + let points = generate_float_grid_points(&dist, 10); + assert_eq!(points.len(), 1); + assert!((points[0] - 0.5).abs() < f64::EPSILON); + } + + #[test] + fn test_float_grid_invalid_bounds() { + let dist = FloatDistribution { + low: 1.0, + high: 0.0, + log_scale: false, + step: None, + }; + let points = generate_float_grid_points(&dist, 10); + assert!(points.is_empty()); + } + + // ==================== Categorical Distribution Tests ==================== + + #[test] + fn test_categorical_grid() { + let dist = CategoricalDistribution { n_choices: 5 }; + let points = generate_categorical_grid_points(&dist); + assert_eq!(points, vec![0, 1, 2, 3, 4]); + } + + #[test] + fn test_categorical_grid_single_choice() { + let dist = CategoricalDistribution { n_choices: 1 }; + let points = generate_categorical_grid_points(&dist); + assert_eq!(points, vec![0]); + } + + #[test] + fn test_categorical_grid_empty() { + let dist = CategoricalDistribution { n_choices: 0 }; + let points = generate_categorical_grid_points(&dist); + assert!(points.is_empty()); + } + + // ==================== Sampler Exhaustion Tests ==================== + + #[test] + fn test_sampler_exhausts_after_expected_samples() { + let sampler = GridSearchSampler::new(); + let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 3 }); + + // Sample all 3 points + for _ in 0..3 { + let _ = sampler.sample(&dist, 0, &[]); + } + + // Check exhaustion + assert!(sampler.is_exhausted()); + } + + #[test] + fn test_sampler_exhaustion_with_int_distribution() { + let sampler = GridSearchSampler::builder().n_points_per_param(5).build(); + let dist = Distribution::Int(IntDistribution { + low: 0, + high: 100, + log_scale: false, + step: None, + }); + + // Sample all 5 points + for _ in 0..5 { + let _ = sampler.sample(&dist, 0, &[]); + } + + assert!(sampler.is_exhausted()); + assert_eq!(sampler.grid_size(), 5); + } + + #[test] + #[should_panic(expected = "GridSearchSampler: all grid points exhausted")] + fn test_sampler_panics_after_exhaustion() { + let sampler = GridSearchSampler::new(); + let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 2 }); + + // Sample all 2 points + sampler.sample(&dist, 0, &[]); + sampler.sample(&dist, 0, &[]); + + // This should panic + sampler.sample(&dist, 0, &[]); + } + + // ==================== is_exhausted() Tests ==================== + + #[test] + fn test_is_exhausted_before_sampling() { + let sampler = GridSearchSampler::new(); + // Newly created sampler is vacuously exhausted (no distributions tracked) + assert!(sampler.is_exhausted()); + } + + #[test] + fn test_is_exhausted_during_sampling() { + let sampler = GridSearchSampler::new(); + let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 3 }); + + // After first sample, not exhausted + sampler.sample(&dist, 0, &[]); + assert!(!sampler.is_exhausted()); + + // After second sample, still not exhausted + sampler.sample(&dist, 0, &[]); + assert!(!sampler.is_exhausted()); + + // After third sample, exhausted + sampler.sample(&dist, 0, &[]); + assert!(sampler.is_exhausted()); + } + + #[test] + fn test_is_exhausted_multiple_distributions() { + let sampler = GridSearchSampler::new(); + // Use different n_choices so they have different distribution keys + let dist1 = Distribution::Categorical(CategoricalDistribution { n_choices: 2 }); + let dist2 = Distribution::Categorical(CategoricalDistribution { n_choices: 3 }); + + // Exhaust first distribution + sampler.sample(&dist1, 0, &[]); + sampler.sample(&dist1, 0, &[]); + + // Not exhausted yet because dist2 is not exhausted + sampler.sample(&dist2, 0, &[]); + assert!(!sampler.is_exhausted()); + + // Continue sampling dist2 + sampler.sample(&dist2, 0, &[]); + assert!(!sampler.is_exhausted()); + + // Exhaust second distribution + sampler.sample(&dist2, 0, &[]); + assert!(sampler.is_exhausted()); + } + + // ==================== Builder Pattern Tests ==================== + + #[test] + fn test_builder_default() { + let sampler = GridSearchSampler::builder().build(); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Default is 10 points per param + for _ in 0..10 { + let _ = sampler.sample(&dist, 0, &[]); + } + assert!(sampler.is_exhausted()); + } + + #[test] + fn test_builder_custom_n_points() { + let sampler = GridSearchSampler::builder().n_points_per_param(3).build(); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Should have 3 points + for _ in 0..3 { + let _ = sampler.sample(&dist, 0, &[]); + } + assert!(sampler.is_exhausted()); + assert_eq!(sampler.grid_size(), 3); + } + + #[test] + fn test_new_default() { + let sampler = GridSearchSampler::new(); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Default is 10 points per param + for _ in 0..10 { + let _ = sampler.sample(&dist, 0, &[]); + } + assert!(sampler.is_exhausted()); + } + + // ==================== Reproducibility Tests ==================== + + #[test] + fn test_reproducibility_same_grid_order() { + // Two samplers with the same configuration should produce the same grid order + let sampler1 = GridSearchSampler::builder().n_points_per_param(5).build(); + let sampler2 = GridSearchSampler::builder().n_points_per_param(5).build(); + + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Both should produce the same sequence + for _ in 0..5 { + let v1 = sampler1.sample(&dist, 0, &[]); + let v2 = sampler2.sample(&dist, 0, &[]); + assert_eq!(v1, v2); + } + } + + #[test] + fn test_reproducibility_int_distribution() { + let sampler1 = GridSearchSampler::new(); + let sampler2 = GridSearchSampler::new(); + + let dist = Distribution::Int(IntDistribution { + low: 0, + high: 10, + log_scale: false, + step: Some(2), + }); + + // Both should produce: 0, 2, 4, 6, 8, 10 + let expected = vec![0, 2, 4, 6, 8, 10]; + for exp in &expected { + let v1 = sampler1.sample(&dist, 0, &[]); + let v2 = sampler2.sample(&dist, 0, &[]); + assert_eq!(v1, ParamValue::Int(*exp)); + assert_eq!(v2, ParamValue::Int(*exp)); + } + } + + #[test] + fn test_reproducibility_categorical() { + let sampler1 = GridSearchSampler::new(); + let sampler2 = GridSearchSampler::new(); + + let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 }); + + // Both should produce: 0, 1, 2, 3 + for i in 0..4 { + let v1 = sampler1.sample(&dist, 0, &[]); + let v2 = sampler2.sample(&dist, 0, &[]); + assert_eq!(v1, ParamValue::Categorical(i)); + assert_eq!(v2, ParamValue::Categorical(i)); + } + } + + // ==================== Grid Size Tests ==================== + + #[test] + fn test_grid_size_empty() { + let sampler = GridSearchSampler::new(); + assert_eq!(sampler.grid_size(), 0); + } + + #[test] + fn test_grid_size_single_distribution() { + let sampler = GridSearchSampler::builder().n_points_per_param(5).build(); + let dist = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + + // Before sampling + assert_eq!(sampler.grid_size(), 0); + + // After first sample, grid is created + sampler.sample(&dist, 0, &[]); + assert_eq!(sampler.grid_size(), 5); + } + + #[test] + fn test_grid_size_multiple_distributions() { + let sampler = GridSearchSampler::builder().n_points_per_param(3).build(); + let dist1 = Distribution::Float(FloatDistribution { + low: 0.0, + high: 1.0, + log_scale: false, + step: None, + }); + let dist2 = Distribution::Categorical(CategoricalDistribution { n_choices: 5 }); + + // Sample from first distribution + sampler.sample(&dist1, 0, &[]); + assert_eq!(sampler.grid_size(), 3); + + // Sample from second distribution + sampler.sample(&dist2, 0, &[]); + assert_eq!(sampler.grid_size(), 3 + 5); + } + + // ==================== Edge Cases ==================== + + #[test] + fn test_int_step_larger_than_range() { + let dist = IntDistribution { + low: 0, + high: 5, + log_scale: false, + step: Some(10), + }; + let points = generate_int_grid_points(&dist, 10); + // Only the starting point should be generated + assert_eq!(points, vec![0]); + } + + #[test] + fn test_float_step_larger_than_range() { + let dist = FloatDistribution { + low: 0.0, + high: 0.5, + log_scale: false, + step: Some(1.0), + }; + let points = generate_float_grid_points(&dist, 10); + // Only the starting point should be generated + assert_eq!(points.len(), 1); + assert!((points[0] - 0.0).abs() < f64::EPSILON); + } + + #[test] + fn test_n_points_one() { + let dist = IntDistribution { + low: 0, + high: 100, + log_scale: false, + step: None, + }; + let points = generate_int_grid_points(&dist, 1); + // Should return just the low bound + assert_eq!(points, vec![0]); + } + + #[test] + fn test_n_points_zero() { + let dist = IntDistribution { + low: 0, + high: 100, + log_scale: false, + step: None, + }; + let points = generate_int_grid_points(&dist, 0); + assert!(points.is_empty()); + } +} diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index e5e1e0a..b04f05b 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -1,5 +1,6 @@ //! Sampler trait and implementations for parameter sampling. +pub mod grid; pub mod random; pub mod tpe; diff --git a/tests/integration.rs b/tests/integration.rs index dd37f42..e0d5242 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -152,7 +152,7 @@ fn test_random_sampler_uniform_float_distribution() { fn test_random_sampler_uniform_int_distribution() { let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(123)); - let n_samples = 1000; + let n_samples = 5000; let mut counts = [0u32; 10]; // counts for values 1-10 study @@ -165,11 +165,13 @@ fn test_random_sampler_uniform_int_distribution() { .unwrap(); // Each value should appear roughly n_samples / 10 times + // With 5000 samples, expected ~500 per bucket, std dev ~21 + // 20% tolerance allows for ~4.5 std devs which is very safe let expected = n_samples as f64 / 10.0; for (i, &count) in counts.iter().enumerate() { let diff = (count as f64 - expected).abs() / expected; assert!( - diff < 0.3, + diff < 0.2, "value {} appeared {} times, expected ~{}, diff = {:.1}%", i + 1, count, @@ -183,7 +185,7 @@ fn test_random_sampler_uniform_int_distribution() { fn test_random_sampler_uniform_categorical_distribution() { let study: Study = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(456)); - let n_samples = 1000; + let n_samples = 2000; let mut counts = [0u32; 4]; let choices = ["a", "b", "c", "d"]; @@ -197,11 +199,13 @@ fn test_random_sampler_uniform_categorical_distribution() { .unwrap(); // Each category should appear roughly n_samples / 4 times + // With 2000 samples, expected ~500 per bucket, std dev ~19 + // 15% tolerance allows for ~4 std devs which is very safe let expected = n_samples as f64 / 4.0; for (i, &count) in counts.iter().enumerate() { let diff = (count as f64 - expected).abs() / expected; assert!( - diff < 0.25, + diff < 0.15, "category {} appeared {} times, expected ~{}, diff = {:.1}%", i, count,