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

refactor(sampler): extract shared utilities to reduce duplication

Browse Source 
- 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
This commit is contained in:
Manuel Raimann
2026-02-12 15:25:27 +01:00
parent 4781107ede
commit d80d2bb1fb
12 changed files with 385 additions and 945 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/sampler/cma_es.rs
+3 -98
View File
@@ -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<f64>, 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)
}
}
}
src/sampler/common.rs
+116
View File
@@ -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)),
}
}
src/sampler/de.rs
+2 -89
View File
@@ -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)
src/sampler/genetic.rs
+1 -36
View File
@@ -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
src/sampler/gp.rs
+2 -111
View File
@@ -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
// ---------------------------------------------------------------------------
src/sampler/mod.rs
+1
View File
@@ -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")]
src/sampler/motpe.rs
+14 -251
View File
@@ -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<f64>,
good_values: Vec<f64>,
bad_values: Vec<f64>,
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<i64>,
good_values: Vec<i64>,
bad_values: Vec<i64>,
rng: &mut fastrand::Rng,
) -> i64 {
let good_floats: Vec<f64> = good_values.into_iter().map(|v| v as f64).collect();
let bad_floats: Vec<f64> = 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,
src/sampler/random.rs
+1 -44
View File
@@ -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)
}
}
src/sampler/tpe/common.rs
+218
View File
@@ -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<f64>,
good_values: Vec<f64>,
bad_values: Vec<f64>,
n_ei_candidates: usize,
kde_bandwidth: Option<f64>,
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<i64>,
good_values: Vec<i64>,
bad_values: Vec<i64>,
n_ei_candidates: usize,
kde_bandwidth: Option<f64>,
rng: &mut fastrand::Rng,
) -> i64 {
// Convert to floats for KDE
let good_floats: Vec<f64> = good_values.into_iter().map(|v| v as f64).collect();
let bad_floats: Vec<f64> = 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
}
src/sampler/tpe/mod.rs
+1
View File
@@ -59,6 +59,7 @@
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! ```
pub(crate) mod common;
mod gamma;
mod multivariate;
mod sampler;
src/sampler/tpe/multivariate.rs
+6 -51
View File
@@ -810,55 +810,10 @@ impl MultivariateTpeSampler {
) -> HashMap<ParamId, ParamValue> {
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)
})
}
}
src/sampler/tpe/sampler.rs
+20 -265
View File
@@ -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<f64>,
good_values: Vec<f64>,
bad_values: Vec<f64>,
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<i64>,
good_values: Vec<i64>,
bad_values: Vec<i64>,
rng: &mut fastrand::Rng,
) -> i64 {
// Convert to floats for KDE
let good_floats: Vec<f64> = good_values.into_iter().map(|v| v as f64).collect();
let bad_floats: Vec<f64> = 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)
}
}