From d80d2bb1fbca72284e2538fc540d912aa0197467 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Thu, 12 Feb 2026 15:25:27 +0100 Subject: [PATCH] refactor(sampler): extract shared utilities to reduce duplication - Add sampler/common.rs with distribution helpers (internal_bounds, from_internal, to_internal, sample_random) used by 8 samplers - Add sampler/tpe/common.rs with TPE sampling functions (sample_tpe_float, sample_tpe_int, sample_tpe_categorical) shared by TpeSampler, MultivariateTpeSampler, and MotpeSampler - Remove ~940 lines of near-identical code across sampler modules --- src/sampler/cma_es.rs | 101 +---------- src/sampler/common.rs | 116 +++++++++++++ src/sampler/de.rs | 91 +--------- src/sampler/genetic.rs | 37 +---- src/sampler/gp.rs | 113 +------------ src/sampler/mod.rs | 1 + src/sampler/motpe.rs | 265 ++--------------------------- src/sampler/random.rs | 45 +---- src/sampler/tpe/common.rs | 218 ++++++++++++++++++++++++ src/sampler/tpe/mod.rs | 1 + src/sampler/tpe/multivariate.rs | 57 +------ src/sampler/tpe/sampler.rs | 285 +++----------------------------- 12 files changed, 385 insertions(+), 945 deletions(-) create mode 100644 src/sampler/common.rs create mode 100644 src/sampler/tpe/common.rs diff --git a/src/sampler/cma_es.rs b/src/sampler/cma_es.rs index f243a88..ac8b501 100644 --- a/src/sampler/cma_es.rs +++ b/src/sampler/cma_es.rs @@ -77,6 +77,8 @@ use crate::param::ParamValue; use crate::rng_util; use crate::sampler::{CompletedTrial, Sampler}; +use super::common::{from_internal, internal_bounds, sample_random}; + /// CMA-ES sampler for continuous optimization. /// /// Adapts a multivariate Gaussian to concentrate around promising regions @@ -671,58 +673,6 @@ fn clip_to_bounds(x: &mut DVector, dimensions: &[DimensionInfo]) { } } -/// Convert an internal-space value back to a `ParamValue`. -#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] -fn from_internal(value: f64, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let v = if d.log_scale { value.exp() } else { value }; - let v = if let Some(step) = d.step { - let k = ((v - d.low) / step).round(); - d.low + k * step - } else { - v - }; - ParamValue::Float(v.clamp(d.low, d.high)) - } - Distribution::Int(d) => { - let v = if d.log_scale { value.exp() } else { value }; - let v = if let Some(step) = d.step { - let k = ((v - d.low as f64) / step as f64).round() as i64; - d.low + k * step - } else { - v.round() as i64 - }; - ParamValue::Int(v.clamp(d.low, d.high)) - } - Distribution::Categorical(_) => { - unreachable!("from_internal should not be called for categorical distributions") - } - } -} - -/// Compute internal-space bounds for a distribution. -#[allow(clippy::cast_precision_loss)] -fn internal_bounds(distribution: &Distribution) -> Option<(f64, f64)> { - match distribution { - Distribution::Float(d) => { - if d.log_scale { - Some((d.low.ln(), d.high.ln())) - } else { - Some((d.low, d.high)) - } - } - Distribution::Int(d) => { - if d.log_scale { - Some(((d.low as f64).ln(), (d.high as f64).ln())) - } else { - Some((d.low as f64, d.high as f64)) - } - } - Distribution::Categorical(_) => None, - } -} - /// Sample a value from the standard normal distribution using Box-Muller transform. fn sample_standard_normal(rng: &mut fastrand::Rng) -> f64 { // Box-Muller transform @@ -731,51 +681,6 @@ fn sample_standard_normal(rng: &mut fastrand::Rng) -> f64 { (-2.0 * u1.ln()).sqrt() * u2.cos() } -/// Sample a categorical value randomly. -fn sample_random_categorical(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), - _ => unreachable!("sample_random_categorical called with non-categorical distribution"), - } -} - -/// Sample a random value for any distribution (used during discovery phase). -#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] -fn sample_random(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - let log_low = d.low.ln(); - let log_high = d.high.ln(); - rng_util::f64_range(rng, log_low, log_high).exp() - } else if let Some(step) = d.step { - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - d.low + (k as f64) * step - } else { - rng_util::f64_range(rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - let value = if d.log_scale { - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), - } -} - // --------------------------------------------------------------------------- // Sampler trait implementation // --------------------------------------------------------------------------- @@ -920,7 +825,7 @@ fn sample_active( if let Some(&cat_idx) = candidate.categorical_values.get(&dim_idx) { ParamValue::Categorical(cat_idx) } else { - sample_random_categorical(&mut state.rng, distribution) + sample_random(&mut state.rng, distribution) } } } diff --git a/src/sampler/common.rs b/src/sampler/common.rs new file mode 100644 index 0000000..a2640bd --- /dev/null +++ b/src/sampler/common.rs @@ -0,0 +1,116 @@ +//! Shared distribution-level utilities used across multiple samplers. + +use crate::distribution::Distribution; +use crate::param::ParamValue; +use crate::rng_util; + +/// Compute internal-space bounds for a distribution. +#[allow(clippy::cast_precision_loss)] +pub(crate) fn internal_bounds(distribution: &Distribution) -> Option<(f64, f64)> { + match distribution { + Distribution::Float(d) => { + if d.log_scale { + Some((d.low.ln(), d.high.ln())) + } else { + Some((d.low, d.high)) + } + } + Distribution::Int(d) => { + if d.log_scale { + Some(((d.low as f64).ln(), (d.high as f64).ln())) + } else { + Some((d.low as f64, d.high as f64)) + } + } + Distribution::Categorical(_) => None, + } +} + +/// Convert an internal-space value back to a `ParamValue`. +#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] +pub(crate) fn from_internal(value: f64, distribution: &Distribution) -> ParamValue { + match distribution { + Distribution::Float(d) => { + let v = if d.log_scale { value.exp() } else { value }; + let v = if let Some(step) = d.step { + let k = ((v - d.low) / step).round(); + d.low + k * step + } else { + v + }; + ParamValue::Float(v.clamp(d.low, d.high)) + } + Distribution::Int(d) => { + let v = if d.log_scale { value.exp() } else { value }; + let v = if let Some(step) = d.step { + let k = ((v - d.low as f64) / step as f64).round() as i64; + d.low + k * step + } else { + v.round() as i64 + }; + ParamValue::Int(v.clamp(d.low, d.high)) + } + Distribution::Categorical(_) => { + unreachable!("from_internal should not be called for categorical distributions") + } + } +} + +/// Convert a `ParamValue` to its internal-space representation. +#[allow(clippy::cast_precision_loss, dead_code)] +pub(crate) fn to_internal(value: &ParamValue, distribution: &Distribution) -> f64 { + match (value, distribution) { + (ParamValue::Float(v), Distribution::Float(d)) => { + if d.log_scale { + v.ln() + } else { + *v + } + } + (ParamValue::Int(v), Distribution::Int(d)) => { + if d.log_scale { + (*v as f64).ln() + } else { + *v as f64 + } + } + _ => 0.0, + } +} + +/// Sample a random value for any distribution. +#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] +pub(crate) fn sample_random(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue { + match distribution { + Distribution::Float(d) => { + let value = if d.log_scale { + let log_low = d.low.ln(); + let log_high = d.high.ln(); + rng_util::f64_range(rng, log_low, log_high).exp() + } else if let Some(step) = d.step { + let n_steps = ((d.high - d.low) / step).floor() as i64; + let k = rng.i64(0..=n_steps); + d.low + (k as f64) * step + } else { + rng_util::f64_range(rng, d.low, d.high) + }; + ParamValue::Float(value) + } + Distribution::Int(d) => { + let value = if d.log_scale { + let log_low = (d.low as f64).ln(); + let log_high = (d.high as f64).ln(); + let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; + raw.clamp(d.low, d.high) + } else if let Some(step) = d.step { + let n_steps = (d.high - d.low) / step; + let k = rng.i64(0..=n_steps); + d.low + k * step + } else { + rng.i64(d.low..=d.high) + }; + ParamValue::Int(value) + } + Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), + } +} diff --git a/src/sampler/de.rs b/src/sampler/de.rs index bd6be26..274078a 100644 --- a/src/sampler/de.rs +++ b/src/sampler/de.rs @@ -70,6 +70,8 @@ use crate::param::ParamValue; use crate::rng_util; use crate::sampler::{CompletedTrial, Sampler}; +use super::common::{from_internal, internal_bounds, sample_random}; + /// Differential Evolution mutation strategy. /// /// Controls how mutant vectors are created from the current population. @@ -396,95 +398,6 @@ impl State { // Helpers // --------------------------------------------------------------------------- -/// Compute internal-space bounds for a distribution. -#[allow(clippy::cast_precision_loss)] -fn internal_bounds(distribution: &Distribution) -> Option<(f64, f64)> { - match distribution { - Distribution::Float(d) => { - if d.log_scale { - Some((d.low.ln(), d.high.ln())) - } else { - Some((d.low, d.high)) - } - } - Distribution::Int(d) => { - if d.log_scale { - Some(((d.low as f64).ln(), (d.high as f64).ln())) - } else { - Some((d.low as f64, d.high as f64)) - } - } - Distribution::Categorical(_) => None, - } -} - -/// Convert an internal-space value back to a `ParamValue`. -#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] -fn from_internal(value: f64, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let v = if d.log_scale { value.exp() } else { value }; - let v = if let Some(step) = d.step { - let k = ((v - d.low) / step).round(); - d.low + k * step - } else { - v - }; - ParamValue::Float(v.clamp(d.low, d.high)) - } - Distribution::Int(d) => { - let v = if d.log_scale { value.exp() } else { value }; - let v = if let Some(step) = d.step { - let k = ((v - d.low as f64) / step as f64).round() as i64; - d.low + k * step - } else { - v.round() as i64 - }; - ParamValue::Int(v.clamp(d.low, d.high)) - } - Distribution::Categorical(_) => { - unreachable!("from_internal should not be called for categorical distributions") - } - } -} - -/// Sample a random value for any distribution. -#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] -fn sample_random(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - let log_low = d.low.ln(); - let log_high = d.high.ln(); - rng_util::f64_range(rng, log_low, log_high).exp() - } else if let Some(step) = d.step { - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - d.low + (k as f64) * step - } else { - rng_util::f64_range(rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - let value = if d.log_scale { - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), - } -} - /// Sample a random value in internal space for a continuous dimension. fn sample_random_internal(rng: &mut fastrand::Rng, bounds: (f64, f64)) -> f64 { rng_util::f64_range(rng, bounds.0, bounds.1) diff --git a/src/sampler/genetic.rs b/src/sampler/genetic.rs index d2a539a..d6d6fa2 100644 --- a/src/sampler/genetic.rs +++ b/src/sampler/genetic.rs @@ -406,42 +406,7 @@ pub(crate) fn polynomial_mutation_f64( (x + delta_q * range).clamp(low, high) } -/// Random sampling for a single distribution. -#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] -pub(crate) fn sample_random(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - let log_low = d.low.ln(); - let log_high = d.high.ln(); - rng_util::f64_range(rng, log_low, log_high).exp() - } else if let Some(step) = d.step { - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - d.low + (k as f64) * step - } else { - rng_util::f64_range(rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - let value = if d.log_scale { - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), - } -} +pub(crate) use super::common::sample_random; // --------------------------------------------------------------------------- // Das-Dennis reference point generation diff --git a/src/sampler/gp.rs b/src/sampler/gp.rs index fe5d8f6..05baa3b 100644 --- a/src/sampler/gp.rs +++ b/src/sampler/gp.rs @@ -82,6 +82,8 @@ use crate::param::ParamValue; use crate::rng_util; use crate::sampler::{CompletedTrial, Sampler}; +use super::common::{from_internal, internal_bounds, sample_random, to_internal}; + // --------------------------------------------------------------------------- // Public API // --------------------------------------------------------------------------- @@ -536,58 +538,6 @@ fn optimize_acquisition( // Data preprocessing helpers // --------------------------------------------------------------------------- -/// Compute internal-space bounds for a distribution. -#[allow(clippy::cast_precision_loss)] -fn internal_bounds(distribution: &Distribution) -> Option<(f64, f64)> { - match distribution { - Distribution::Float(d) => { - if d.log_scale { - Some((d.low.ln(), d.high.ln())) - } else { - Some((d.low, d.high)) - } - } - Distribution::Int(d) => { - if d.log_scale { - Some(((d.low as f64).ln(), (d.high as f64).ln())) - } else { - Some((d.low as f64, d.high as f64)) - } - } - Distribution::Categorical(_) => None, - } -} - -/// Convert a value from internal space to a `ParamValue` in original space. -#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] -fn from_internal(value: f64, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let v = if d.log_scale { value.exp() } else { value }; - let v = if let Some(step) = d.step { - let k = ((v - d.low) / step).round(); - d.low + k * step - } else { - v - }; - ParamValue::Float(v.clamp(d.low, d.high)) - } - Distribution::Int(d) => { - let v = if d.log_scale { value.exp() } else { value }; - let v = if let Some(step) = d.step { - let k = ((v - d.low as f64) / step as f64).round() as i64; - d.low + k * step - } else { - v.round() as i64 - }; - ParamValue::Int(v.clamp(d.low, d.high)) - } - Distribution::Categorical(_) => { - unreachable!("from_internal should not be called for categorical distributions") - } - } -} - /// Convert an internal-space value to normalized [0, 1] using bounds. fn to_normalized(value: f64, lo: f64, hi: f64) -> f64 { if (hi - lo).abs() < 1e-15 { @@ -602,65 +552,6 @@ fn from_normalized(value: f64, lo: f64, hi: f64) -> f64 { lo + value * (hi - lo) } -/// Convert a `ParamValue` to its internal-space representation. -#[allow(clippy::cast_precision_loss)] -fn to_internal(value: &ParamValue, distribution: &Distribution) -> f64 { - match (value, distribution) { - (ParamValue::Float(v), Distribution::Float(d)) => { - if d.log_scale { - v.ln() - } else { - *v - } - } - (ParamValue::Int(v), Distribution::Int(d)) => { - if d.log_scale { - (*v as f64).ln() - } else { - *v as f64 - } - } - _ => 0.0, - } -} - -/// Sample a random value for any distribution. -#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)] -fn sample_random(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - let log_low = d.low.ln(); - let log_high = d.high.ln(); - rng_util::f64_range(rng, log_low, log_high).exp() - } else if let Some(step) = d.step { - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - d.low + (k as f64) * step - } else { - rng_util::f64_range(rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - let value = if d.log_scale { - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), - } -} - // --------------------------------------------------------------------------- // Extract training data from history // --------------------------------------------------------------------------- diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index 4e20f4e..a660d11 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -3,6 +3,7 @@ pub mod bohb; #[cfg(feature = "cma-es")] pub mod cma_es; +pub(crate) mod common; pub mod de; pub(crate) mod genetic; #[cfg(feature = "gp")] diff --git a/src/sampler/motpe.rs b/src/sampler/motpe.rs index 883cfe0..a84580d 100644 --- a/src/sampler/motpe.rs +++ b/src/sampler/motpe.rs @@ -61,9 +61,10 @@ use core::sync::atomic::{AtomicU64, Ordering}; use crate::distribution::Distribution; -use crate::kde::KernelDensityEstimator; use crate::multi_objective::{MultiObjectiveSampler, MultiObjectiveTrial}; use crate::param::ParamValue; +use crate::sampler::common; +use crate::sampler::tpe::common as tpe_common; use crate::types::{Direction, TrialState}; use crate::{pareto, rng_util}; @@ -205,248 +206,6 @@ impl MotpeSampler { (good, bad) } - - /// Samples uniformly from a distribution (used during startup phase). - #[allow( - clippy::cast_possible_truncation, - clippy::cast_precision_loss, - clippy::unused_self - )] - fn sample_uniform(distribution: &Distribution, rng: &mut fastrand::Rng) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - let log_low = d.low.ln(); - let log_high = d.high.ln(); - rng_util::f64_range(rng, log_low, log_high).exp() - } else if let Some(step) = d.step { - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - d.low + (k as f64) * step - } else { - rng_util::f64_range(rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - let value = if d.log_scale { - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), - } - } - - /// Samples using TPE for float distributions. - #[allow(clippy::too_many_arguments)] - fn sample_tpe_float( - &self, - low: f64, - high: f64, - log_scale: bool, - step: Option, - good_values: Vec, - bad_values: Vec, - rng: &mut fastrand::Rng, - ) -> f64 { - // Transform to internal space (log space if needed) - let (internal_low, internal_high, good_internal, bad_internal) = if log_scale { - let i_low = low.ln(); - let i_high = high.ln(); - let g = { - let mut v = good_values; - for x in &mut v { - *x = x.ln(); - } - v - }; - let b = { - let mut v = bad_values; - for x in &mut v { - *x = x.ln(); - } - v - }; - (i_low, i_high, g, b) - } else { - (low, high, good_values, bad_values) - }; - - // Fit KDEs to good and bad groups - let l_kde = match self.kde_bandwidth { - Some(bw) => KernelDensityEstimator::with_bandwidth(good_internal, bw), - None => KernelDensityEstimator::new(good_internal), - }; - let g_kde = match self.kde_bandwidth { - Some(bw) => KernelDensityEstimator::with_bandwidth(bad_internal, bw), - None => KernelDensityEstimator::new(bad_internal), - }; - - // If KDE construction fails, fall back to uniform sampling - let (Ok(l_kde), Ok(g_kde)) = (l_kde, g_kde) else { - return rng_util::f64_range(rng, low, high); - }; - - // Generate candidates from l(x) and select the one with best l(x)/g(x) - let mut best_candidate = internal_low; - let mut best_ratio = f64::NEG_INFINITY; - - for _ in 0..self.n_ei_candidates { - let candidate = l_kde.sample(rng).clamp(internal_low, internal_high); - - let l_density = l_kde.pdf(candidate); - let g_density = g_kde.pdf(candidate); - - let ratio = if g_density < f64::EPSILON { - if l_density > f64::EPSILON { - f64::INFINITY - } else { - 0.0 - } - } else { - l_density / g_density - }; - - if ratio > best_ratio { - best_ratio = ratio; - best_candidate = candidate; - } - } - - // Transform back from internal space - let mut value = if log_scale { - best_candidate.exp() - } else { - best_candidate - }; - - // Apply step constraint if present - if let Some(step) = step { - let k = ((value - low) / step).round(); - value = low + k * step; - } - - value.clamp(low, high) - } - - /// Samples using TPE for integer distributions. - #[allow( - clippy::too_many_arguments, - clippy::cast_precision_loss, - clippy::cast_possible_truncation - )] - fn sample_tpe_int( - &self, - low: i64, - high: i64, - log_scale: bool, - step: Option, - good_values: Vec, - bad_values: Vec, - rng: &mut fastrand::Rng, - ) -> i64 { - let good_floats: Vec = good_values.into_iter().map(|v| v as f64).collect(); - let bad_floats: Vec = bad_values.into_iter().map(|v| v as f64).collect(); - - let float_value = self.sample_tpe_float( - low as f64, - high as f64, - log_scale, - step.map(|s| s as f64), - good_floats, - bad_floats, - rng, - ); - - let int_value = float_value.round() as i64; - let int_value = if let Some(step) = step { - let k = ((int_value - low) as f64 / step as f64).round() as i64; - low + k * step - } else { - int_value - }; - - int_value.clamp(low, high) - } - - /// Samples using TPE for categorical distributions. - #[allow(clippy::cast_precision_loss)] - fn sample_tpe_categorical( - n_choices: usize, - good_indices: &[usize], - bad_indices: &[usize], - rng: &mut fastrand::Rng, - ) -> usize { - // Stack-allocate for the common case (<=32 choices), heap for rare large cases - let mut good_buf = [0usize; 32]; - let mut bad_buf = [0usize; 32]; - let mut weight_buf = [0.0f64; 32]; - - let mut good_vec; - let mut bad_vec; - let mut weight_vec; - - let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) = - if n_choices <= 32 { - ( - &mut good_buf[..n_choices], - &mut bad_buf[..n_choices], - &mut weight_buf[..n_choices], - ) - } else { - good_vec = vec![0usize; n_choices]; - bad_vec = vec![0usize; n_choices]; - weight_vec = vec![0.0f64; n_choices]; - (&mut good_vec, &mut bad_vec, &mut weight_vec) - }; - - // Count occurrences in good and bad groups - for &idx in good_indices { - if idx < n_choices { - good_counts[idx] += 1; - } - } - for &idx in bad_indices { - if idx < n_choices { - bad_counts[idx] += 1; - } - } - - // Laplace smoothing - let good_total = good_indices.len() as f64 + n_choices as f64; - let bad_total = bad_indices.len() as f64 + n_choices as f64; - - // Calculate l(x)/g(x) ratio for each category - for i in 0..n_choices { - let l_prob = (good_counts[i] as f64 + 1.0) / good_total; - let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total; - weights[i] = l_prob / g_prob; - } - - // Sample proportionally to weights - let total_weight: f64 = weights.iter().sum(); - let threshold = rng.f64() * total_weight; - - let mut cumulative = 0.0; - for (i, &w) in weights.iter().enumerate() { - cumulative += w; - if cumulative >= threshold { - return i; - } - } - - n_choices - 1 - } } impl Default for MotpeSampler { @@ -477,14 +236,14 @@ impl MultiObjectiveSampler for MotpeSampler { .filter(|t| t.state == TrialState::Complete) .count(); if n_complete < self.n_startup_trials { - return Self::sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } // Split trials into good (Pareto front) and bad (dominated) let (good_trials, bad_trials) = Self::split_trials(history, directions); if good_trials.is_empty() || bad_trials.is_empty() { - return Self::sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } match distribution { @@ -510,16 +269,18 @@ impl MultiObjectiveSampler for MotpeSampler { .collect(); if good_values.is_empty() || bad_values.is_empty() { - return Self::sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } - let value = self.sample_tpe_float( + let value = tpe_common::sample_tpe_float( d.low, d.high, d.log_scale, d.step, good_values, bad_values, + self.n_ei_candidates, + self.kde_bandwidth, &mut rng, ); ParamValue::Float(value) @@ -546,16 +307,18 @@ impl MultiObjectiveSampler for MotpeSampler { .collect(); if good_values.is_empty() || bad_values.is_empty() { - return Self::sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } - let value = self.sample_tpe_int( + let value = tpe_common::sample_tpe_int( d.low, d.high, d.log_scale, d.step, good_values, bad_values, + self.n_ei_candidates, + self.kde_bandwidth, &mut rng, ); ParamValue::Int(value) @@ -582,10 +345,10 @@ impl MultiObjectiveSampler for MotpeSampler { .collect(); if good_indices.is_empty() || bad_indices.is_empty() { - return Self::sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } - let index = Self::sample_tpe_categorical( + let index = tpe_common::sample_tpe_categorical( d.n_choices, &good_indices, &bad_indices, diff --git a/src/sampler/random.rs b/src/sampler/random.rs index 6c750b3..2850376 100644 --- a/src/sampler/random.rs +++ b/src/sampler/random.rs @@ -127,50 +127,7 @@ impl Sampler for RandomSampler { rng_util::distribution_fingerprint(distribution).wrapping_add(seq), )); - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - // Sample uniformly in log space - let log_low = d.low.ln(); - let log_high = d.high.ln(); - let log_value = rng_util::f64_range(&mut rng, log_low, log_high); - log_value.exp() - } else if let Some(step) = d.step { - // Sample from step grid - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - d.low + (k as f64) * step - } else { - // Uniform sampling - rng_util::f64_range(&mut rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - let value = if d.log_scale { - // Sample uniformly in log space, then round - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - let log_value = rng_util::f64_range(&mut rng, log_low, log_high); - let raw = log_value.exp().round() as i64; - // Clamp to bounds since rounding might push outside - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - // Sample from step grid - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - // Uniform sampling - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => { - let index = rng.usize(0..d.n_choices); - ParamValue::Categorical(index) - } - } + super::common::sample_random(&mut rng, distribution) } } diff --git a/src/sampler/tpe/common.rs b/src/sampler/tpe/common.rs new file mode 100644 index 0000000..c566ae9 --- /dev/null +++ b/src/sampler/tpe/common.rs @@ -0,0 +1,218 @@ +//! Shared TPE sampling functions used by both `TpeSampler` and `MotpeSampler`. + +use crate::kde::KernelDensityEstimator; +use crate::rng_util; + +/// Samples using TPE for float distributions. +#[allow(clippy::too_many_arguments)] +pub(crate) fn sample_tpe_float( + low: f64, + high: f64, + log_scale: bool, + step: Option, + good_values: Vec, + bad_values: Vec, + n_ei_candidates: usize, + kde_bandwidth: Option, + rng: &mut fastrand::Rng, +) -> f64 { + // Transform to internal space (log space if needed) + let (internal_low, internal_high, good_internal, bad_internal) = if log_scale { + let i_low = low.ln(); + let i_high = high.ln(); + let g = { + let mut v = good_values; + for x in &mut v { + *x = x.ln(); + } + v + }; + let b = { + let mut v = bad_values; + for x in &mut v { + *x = x.ln(); + } + v + }; + (i_low, i_high, g, b) + } else { + (low, high, good_values, bad_values) + }; + + // Fit KDEs to good and bad groups + let l_kde = match kde_bandwidth { + Some(bw) => KernelDensityEstimator::with_bandwidth(good_internal, bw), + None => KernelDensityEstimator::new(good_internal), + }; + let g_kde = match kde_bandwidth { + Some(bw) => KernelDensityEstimator::with_bandwidth(bad_internal, bw), + None => KernelDensityEstimator::new(bad_internal), + }; + + // If KDE construction fails, fall back to uniform sampling + let (Ok(l_kde), Ok(g_kde)) = (l_kde, g_kde) else { + return rng_util::f64_range(rng, low, high); + }; + + // Generate candidates from l(x) and select the one with best l(x)/g(x) ratio + let mut best_candidate = internal_low; + let mut best_ratio = f64::NEG_INFINITY; + + for _ in 0..n_ei_candidates { + let candidate = l_kde.sample(rng).clamp(internal_low, internal_high); + + let l_density = l_kde.pdf(candidate); + let g_density = g_kde.pdf(candidate); + + // Compute l(x)/g(x) ratio, handling zero density + let ratio = if g_density < f64::EPSILON { + if l_density > f64::EPSILON { + f64::INFINITY + } else { + 0.0 + } + } else { + l_density / g_density + }; + + if ratio > best_ratio { + best_ratio = ratio; + best_candidate = candidate; + } + } + + // Transform back from internal space + let mut value = if log_scale { + best_candidate.exp() + } else { + best_candidate + }; + + // Apply step constraint if present + if let Some(step) = step { + let k = ((value - low) / step).round(); + value = low + k * step; + } + + // Ensure value is within bounds + value.clamp(low, high) +} + +/// Samples using TPE for integer distributions. +#[allow( + clippy::too_many_arguments, + clippy::cast_precision_loss, + clippy::cast_possible_truncation +)] +pub(crate) fn sample_tpe_int( + low: i64, + high: i64, + log_scale: bool, + step: Option, + good_values: Vec, + bad_values: Vec, + n_ei_candidates: usize, + kde_bandwidth: Option, + rng: &mut fastrand::Rng, +) -> i64 { + // Convert to floats for KDE + let good_floats: Vec = good_values.into_iter().map(|v| v as f64).collect(); + let bad_floats: Vec = bad_values.into_iter().map(|v| v as f64).collect(); + + // Use float TPE sampling + let float_value = sample_tpe_float( + low as f64, + high as f64, + log_scale, + step.map(|s| s as f64), + good_floats, + bad_floats, + n_ei_candidates, + kde_bandwidth, + rng, + ); + + // Round to nearest integer + let int_value = float_value.round() as i64; + + // Apply step constraint if present + let int_value = if let Some(step) = step { + let k = ((int_value - low) as f64 / step as f64).round() as i64; + low + k * step + } else { + int_value + }; + + // Ensure value is within bounds + int_value.clamp(low, high) +} + +/// Samples using TPE for categorical distributions. +#[allow(clippy::cast_precision_loss)] +pub(crate) fn sample_tpe_categorical( + n_choices: usize, + good_indices: &[usize], + bad_indices: &[usize], + rng: &mut fastrand::Rng, +) -> usize { + // Stack-allocate for the common case (<=32 choices), heap for rare large cases + let mut good_buf = [0usize; 32]; + let mut bad_buf = [0usize; 32]; + let mut weight_buf = [0.0f64; 32]; + + let mut good_vec; + let mut bad_vec; + let mut weight_vec; + + let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) = + if n_choices <= 32 { + ( + &mut good_buf[..n_choices], + &mut bad_buf[..n_choices], + &mut weight_buf[..n_choices], + ) + } else { + good_vec = vec![0usize; n_choices]; + bad_vec = vec![0usize; n_choices]; + weight_vec = vec![0.0f64; n_choices]; + (&mut good_vec, &mut bad_vec, &mut weight_vec) + }; + + // Count occurrences in good and bad groups + for &idx in good_indices { + if idx < n_choices { + good_counts[idx] += 1; + } + } + for &idx in bad_indices { + if idx < n_choices { + bad_counts[idx] += 1; + } + } + + // Add smoothing (Laplace smoothing) to avoid zero probabilities + let good_total = good_indices.len() as f64 + n_choices as f64; + let bad_total = bad_indices.len() as f64 + n_choices as f64; + + // Calculate l(x)/g(x) ratio for each category + for i in 0..n_choices { + let l_prob = (good_counts[i] as f64 + 1.0) / good_total; + let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total; + weights[i] = l_prob / g_prob; + } + + // Sample proportionally to weights + let total_weight: f64 = weights.iter().sum(); + let threshold = rng.f64() * total_weight; + + let mut cumulative = 0.0; + for (i, &w) in weights.iter().enumerate() { + cumulative += w; + if cumulative >= threshold { + return i; + } + } + + // Fallback to last index (shouldn't happen) + n_choices - 1 +} diff --git a/src/sampler/tpe/mod.rs b/src/sampler/tpe/mod.rs index 2c17e3d..a856de2 100644 --- a/src/sampler/tpe/mod.rs +++ b/src/sampler/tpe/mod.rs @@ -59,6 +59,7 @@ //! let study: Study = Study::with_sampler(Direction::Minimize, sampler); //! ``` +pub(crate) mod common; mod gamma; mod multivariate; mod sampler; diff --git a/src/sampler/tpe/multivariate.rs b/src/sampler/tpe/multivariate.rs index 731fc44..de8b37e 100644 --- a/src/sampler/tpe/multivariate.rs +++ b/src/sampler/tpe/multivariate.rs @@ -810,55 +810,10 @@ impl MultivariateTpeSampler { ) -> HashMap { search_space .iter() - .map(|(id, dist)| (*id, Self::sample_uniform_single(dist, rng))) + .map(|(id, dist)| (*id, crate::sampler::common::sample_random(rng, dist))) .collect() } - /// Samples a single parameter uniformly at random from its distribution. - fn sample_uniform_single(distribution: &Distribution, rng: &mut fastrand::Rng) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - let log_low = d.low.ln(); - let log_high = d.high.ln(); - rng_util::f64_range(rng, log_low, log_high).exp() - } else if let Some(step) = d.step { - #[allow(clippy::cast_possible_truncation)] - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - #[allow(clippy::cast_precision_loss)] - let result = d.low + (k as f64) * step; - result - } else { - rng_util::f64_range(rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - #[allow(clippy::cast_precision_loss)] - let value = if d.log_scale { - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - #[allow(clippy::cast_possible_truncation)] - let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - #[allow(clippy::cast_possible_truncation)] - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => { - let index = rng.usize(0..d.n_choices); - ParamValue::Categorical(index) - } - } - } - /// Samples parameters jointly using multivariate TPE. /// /// This method samples all parameters in the search space jointly, capturing @@ -1025,7 +980,7 @@ impl MultivariateTpeSampler { // Sample ungrouped parameters uniformly (no history for them) let mut rng = self.rng.lock(); for (id, dist) in &ungrouped_params { - let value = Self::sample_uniform_single(dist, &mut rng); + let value = crate::sampler::common::sample_random(&mut rng, dist); result.insert(*id, value); } } @@ -1312,7 +1267,7 @@ impl MultivariateTpeSampler { .collect(); if good_values.is_empty() || bad_values.is_empty() { - return Self::sample_uniform_single(distribution, rng); + return crate::sampler::common::sample_random(rng, distribution); } let value = self.sample_tpe_float( @@ -1348,7 +1303,7 @@ impl MultivariateTpeSampler { .collect(); if good_values.is_empty() || bad_values.is_empty() { - return Self::sample_uniform_single(distribution, rng); + return crate::sampler::common::sample_random(rng, distribution); } let value = self.sample_tpe_int( @@ -1384,7 +1339,7 @@ impl MultivariateTpeSampler { .collect(); if good_indices.is_empty() || bad_indices.is_empty() { - return Self::sample_uniform_single(distribution, rng); + return crate::sampler::common::sample_random(rng, distribution); } let idx = @@ -1765,7 +1720,7 @@ impl Sampler for MultivariateTpeSampler { Self::find_matching_param(distribution, &joint_sample).unwrap_or_else(|| { // Fallback to uniform sampling if no match found let mut rng = self.rng.lock(); - Self::sample_uniform_single(distribution, &mut rng) + crate::sampler::common::sample_random(&mut rng, distribution) }) } } diff --git a/src/sampler/tpe/sampler.rs b/src/sampler/tpe/sampler.rs index fb9efe6..5c9af1d 100644 --- a/src/sampler/tpe/sampler.rs +++ b/src/sampler/tpe/sampler.rs @@ -61,12 +61,14 @@ use std::sync::Arc; use crate::distribution::Distribution; use crate::error::{Error, Result}; -use crate::kde::KernelDensityEstimator; use crate::param::ParamValue; use crate::rng_util; +use crate::sampler::common; use crate::sampler::tpe::gamma::{FixedGamma, GammaStrategy}; use crate::sampler::{CompletedTrial, Sampler}; +use super::common as tpe_common; + // ============================================================================ // Gamma Strategy Trait and Implementations // ============================================================================ @@ -315,261 +317,6 @@ impl TpeSampler { (good, bad) } - - /// Samples uniformly from a distribution (used during startup phase). - #[allow( - clippy::cast_possible_truncation, - clippy::cast_precision_loss, - clippy::unused_self - )] - fn sample_uniform(&self, distribution: &Distribution, rng: &mut fastrand::Rng) -> ParamValue { - match distribution { - Distribution::Float(d) => { - let value = if d.log_scale { - let log_low = d.low.ln(); - let log_high = d.high.ln(); - rng_util::f64_range(rng, log_low, log_high).exp() - } else if let Some(step) = d.step { - let n_steps = ((d.high - d.low) / step).floor() as i64; - let k = rng.i64(0..=n_steps); - d.low + (k as f64) * step - } else { - rng_util::f64_range(rng, d.low, d.high) - }; - ParamValue::Float(value) - } - Distribution::Int(d) => { - let value = if d.log_scale { - let log_low = (d.low as f64).ln(); - let log_high = (d.high as f64).ln(); - let raw = rng_util::f64_range(rng, log_low, log_high).exp().round() as i64; - raw.clamp(d.low, d.high) - } else if let Some(step) = d.step { - let n_steps = (d.high - d.low) / step; - let k = rng.i64(0..=n_steps); - d.low + k * step - } else { - rng.i64(d.low..=d.high) - }; - ParamValue::Int(value) - } - Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)), - } - } - - /// Samples using TPE for float distributions. - #[allow(clippy::too_many_arguments)] - fn sample_tpe_float( - &self, - low: f64, - high: f64, - log_scale: bool, - step: Option, - good_values: Vec, - bad_values: Vec, - rng: &mut fastrand::Rng, - ) -> f64 { - // Transform to internal space (log space if needed) - let (internal_low, internal_high, good_internal, bad_internal) = if log_scale { - let i_low = low.ln(); - let i_high = high.ln(); - let g = { - let mut v = good_values; - for x in &mut v { - *x = x.ln(); - } - v - }; - let b = { - let mut v = bad_values; - for x in &mut v { - *x = x.ln(); - } - v - }; - (i_low, i_high, g, b) - } else { - (low, high, good_values, bad_values) - }; - - // Fit KDEs to good and bad groups - let l_kde = match self.kde_bandwidth { - Some(bw) => KernelDensityEstimator::with_bandwidth(good_internal, bw), - None => KernelDensityEstimator::new(good_internal), - }; - let g_kde = match self.kde_bandwidth { - Some(bw) => KernelDensityEstimator::with_bandwidth(bad_internal, bw), - None => KernelDensityEstimator::new(bad_internal), - }; - - // If KDE construction fails, fall back to uniform sampling - let (Ok(l_kde), Ok(g_kde)) = (l_kde, g_kde) else { - return rng_util::f64_range(rng, low, high); - }; - - // Generate candidates from l(x) and select the one with best l(x)/g(x) ratio - let mut best_candidate = internal_low; - let mut best_ratio = f64::NEG_INFINITY; - - for _ in 0..self.n_ei_candidates { - let candidate = l_kde.sample(rng); - - // Clamp to bounds - let candidate = candidate.clamp(internal_low, internal_high); - - let l_density = l_kde.pdf(candidate); - let g_density = g_kde.pdf(candidate); - - // Compute l(x)/g(x) ratio, handling zero density - let ratio = if g_density < f64::EPSILON { - if l_density > f64::EPSILON { - f64::INFINITY - } else { - 0.0 - } - } else { - l_density / g_density - }; - - if ratio > best_ratio { - best_ratio = ratio; - best_candidate = candidate; - } - } - - // Transform back from internal space - let mut value = if log_scale { - best_candidate.exp() - } else { - best_candidate - }; - - // Apply step constraint if present - if let Some(step) = step { - let k = ((value - low) / step).round(); - value = low + k * step; - } - - // Ensure value is within bounds - value.clamp(low, high) - } - - /// Samples using TPE for integer distributions. - #[allow( - clippy::too_many_arguments, - clippy::cast_precision_loss, - clippy::cast_possible_truncation - )] - fn sample_tpe_int( - &self, - low: i64, - high: i64, - log_scale: bool, - step: Option, - good_values: Vec, - bad_values: Vec, - rng: &mut fastrand::Rng, - ) -> i64 { - // Convert to floats for KDE - let good_floats: Vec = good_values.into_iter().map(|v| v as f64).collect(); - let bad_floats: Vec = bad_values.into_iter().map(|v| v as f64).collect(); - - // Use float TPE sampling - let float_value = self.sample_tpe_float( - low as f64, - high as f64, - log_scale, - step.map(|s| s as f64), - good_floats, - bad_floats, - rng, - ); - - // Round to nearest integer - let int_value = float_value.round() as i64; - - // Apply step constraint if present - let int_value = if let Some(step) = step { - let k = ((int_value - low) as f64 / step as f64).round() as i64; - low + k * step - } else { - int_value - }; - - // Ensure value is within bounds - int_value.clamp(low, high) - } - - /// Samples using TPE for categorical distributions. - #[allow(clippy::cast_precision_loss, clippy::unused_self)] - fn sample_tpe_categorical( - &self, - n_choices: usize, - good_indices: &[usize], - bad_indices: &[usize], - rng: &mut fastrand::Rng, - ) -> usize { - // Stack-allocate for the common case (<=32 choices), heap for rare large cases - let mut good_buf = [0usize; 32]; - let mut bad_buf = [0usize; 32]; - let mut weight_buf = [0.0f64; 32]; - - let mut good_vec; - let mut bad_vec; - let mut weight_vec; - - let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) = - if n_choices <= 32 { - ( - &mut good_buf[..n_choices], - &mut bad_buf[..n_choices], - &mut weight_buf[..n_choices], - ) - } else { - good_vec = vec![0usize; n_choices]; - bad_vec = vec![0usize; n_choices]; - weight_vec = vec![0.0f64; n_choices]; - (&mut good_vec, &mut bad_vec, &mut weight_vec) - }; - - // Count occurrences in good and bad groups - for &idx in good_indices { - if idx < n_choices { - good_counts[idx] += 1; - } - } - for &idx in bad_indices { - if idx < n_choices { - bad_counts[idx] += 1; - } - } - - // Add smoothing (Laplace smoothing) to avoid zero probabilities - let good_total = good_indices.len() as f64 + n_choices as f64; - let bad_total = bad_indices.len() as f64 + n_choices as f64; - - // Calculate l(x)/g(x) ratio for each category - for i in 0..n_choices { - let l_prob = (good_counts[i] as f64 + 1.0) / good_total; - let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total; - weights[i] = l_prob / g_prob; - } - - // Sample proportionally to weights - let total_weight: f64 = weights.iter().sum(); - let threshold = rng.f64() * total_weight; - - let mut cumulative = 0.0; - for (i, &w) in weights.iter().enumerate() { - cumulative += w; - if cumulative >= threshold { - return i; - } - } - - // Fallback to last index (shouldn't happen) - n_choices - 1 - } } impl Default for TpeSampler { @@ -912,7 +659,7 @@ impl Sampler for TpeSampler { // Fall back to random sampling during startup phase if history.len() < self.n_startup_trials { - return self.sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } // Split trials into good and bad groups @@ -920,7 +667,7 @@ impl Sampler for TpeSampler { // Need at least 1 trial in each group for TPE if good_trials.is_empty() || bad_trials.is_empty() { - return self.sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } // Extract parameter values for this distribution @@ -954,16 +701,18 @@ impl Sampler for TpeSampler { // Need values in both groups for TPE if good_values.is_empty() || bad_values.is_empty() { - return self.sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } - let value = self.sample_tpe_float( + let value = tpe_common::sample_tpe_float( d.low, d.high, d.log_scale, d.step, good_values, bad_values, + self.n_ei_candidates, + self.kde_bandwidth, &mut rng, ); ParamValue::Float(value) @@ -990,16 +739,18 @@ impl Sampler for TpeSampler { .collect(); if good_values.is_empty() || bad_values.is_empty() { - return self.sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } - let value = self.sample_tpe_int( + let value = tpe_common::sample_tpe_int( d.low, d.high, d.log_scale, d.step, good_values, bad_values, + self.n_ei_candidates, + self.kde_bandwidth, &mut rng, ); ParamValue::Int(value) @@ -1026,11 +777,15 @@ impl Sampler for TpeSampler { .collect(); if good_indices.is_empty() || bad_indices.is_empty() { - return self.sample_uniform(distribution, &mut rng); + return common::sample_random(&mut rng, distribution); } - let index = - self.sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, &mut rng); + let index = tpe_common::sample_tpe_categorical( + d.n_choices, + &good_indices, + &bad_indices, + &mut rng, + ); ParamValue::Categorical(index) } }