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

fix(sampler): clamp multivariate TPE candidates to parameter bounds

Browse Source 
- Clamp KDE candidates to parameter bounds before evaluating l(x)/g(x),
  matching the univariate TPE behavior; without this, candidates scored
  well at out-of-bounds locations but became suboptimal when clamped
- Sort HashMap iterations by ParamId before consuming the seeded RNG to
  eliminate non-deterministic sampling caused by global ParamId counter
- Replace wall-clock timing assertion in async concurrency test with
  atomic max-active counter to avoid CI flakiness
- Gate unused HashSet import behind cfg(feature = "async")
This commit is contained in:
Manuel Raimann
2026-02-13 10:20:32 +01:00
parent 11a8534b38
commit 86db6361d6
5 changed files with 98 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/sampler/tpe/multivariate/engine.rs
+52 -10
View File
@@ -197,7 +197,7 @@ impl MultivariateTpeSampler {
return result;
}
param_order.sort_by_key(|id| format!("{id}"));
param_order.sort();
// Extract observations, validate, and fit KDEs
let good_obs = self.extract_observations(&good, &param_order);
@@ -215,7 +215,20 @@ impl MultivariateTpeSampler {
return result;
};
let selected = self.select_candidate_with_rng(&good_kde, &bad_kde, rng);
// Compute parameter bounds for each dimension so candidates are clamped
let bounds: Vec<(f64, f64)> = param_order
.iter()
.filter_map(|id| {
intersection.get(id).and_then(|dist| match dist {
Distribution::Float(d) => Some((d.low, d.high)),
#[allow(clippy::cast_precision_loss)]
Distribution::Int(d) => Some((d.low as f64, d.high as f64)),
Distribution::Categorical(_) => None,
})
})
.collect();
let selected = self.select_candidate_with_rng(&good_kde, &bad_kde, &bounds, rng);
// Map selected values to parameter ids
for (idx, param_id) in param_order.iter().enumerate() {
@@ -273,23 +286,38 @@ impl MultivariateTpeSampler {
}
}
/// Clamps each dimension of a candidate to the corresponding parameter bounds.
fn clamp_candidate(candidate: &mut [f64], bounds: &[(f64, f64)]) {
for (val, &(lo, hi)) in candidate.iter_mut().zip(bounds.iter()) {
*val = val.clamp(lo, hi);
}
}
/// Selects the best candidate from a set of samples using the joint acquisition function.
///
/// This method implements the core TPE selection criterion: it generates candidates
/// from the "good" KDE (l(x)) and selects the one that maximizes the ratio l(x)/g(x),
/// which is equivalent to maximizing `log(l(x)) - log(g(x))`.
///
/// Candidates are clamped to parameter bounds before evaluation so the acquisition
/// function scores the values that will actually be used.
#[must_use]
#[cfg(test)]
pub(crate) fn select_candidate(
&self,
good_kde: &crate::kde::MultivariateKDE,
bad_kde: &crate::kde::MultivariateKDE,
bounds: &[(f64, f64)],
) -> Vec<f64> {
let mut rng = self.rng.lock();
// Generate candidates from the good distribution
// Generate candidates from the good distribution, clamped to bounds
let candidates: Vec<Vec<f64>> = (0..self.n_ei_candidates)
.map(|_| good_kde.sample(&mut rng))
.map(|_| {
let mut c = good_kde.sample(&mut rng);
Self::clamp_candidate(&mut c, bounds);
c
})
.collect();
// Compute log(l(x)) - log(g(x)) for each candidate
@@ -321,15 +349,21 @@ impl MultivariateTpeSampler {
/// Selects the best candidate using an external RNG.
///
/// This variant accepts an external RNG, used when the caller already holds the lock.
/// Candidates are clamped to parameter bounds before evaluation.
pub(crate) fn select_candidate_with_rng(
&self,
good_kde: &crate::kde::MultivariateKDE,
bad_kde: &crate::kde::MultivariateKDE,
bounds: &[(f64, f64)],
rng: &mut fastrand::Rng,
) -> Vec<f64> {
// Generate candidates from the good distribution
// Generate candidates from the good distribution, clamped to bounds
let candidates: Vec<Vec<f64>> = (0..self.n_ei_candidates)
.map(|_| good_kde.sample(rng))
.map(|_| {
let mut c = good_kde.sample(rng);
Self::clamp_candidate(&mut c, bounds);
c
})
.collect();
// Compute log(l(x)) - log(g(x)) for each candidate
@@ -367,8 +401,8 @@ impl MultivariateTpeSampler {
result: &mut HashMap<ParamId, ParamValue>,
rng: &mut fastrand::Rng,
) {
// Identify parameters not in result (and not in intersection)
let missing_params: Vec<(&ParamId, &Distribution)> = search_space
// Identify parameters not in result, sorted for deterministic RNG consumption
let mut missing_params: Vec<(&ParamId, &Distribution)> = search_space
.iter()
.filter(|(id, _)| !result.contains_key(id))
.collect();
@@ -377,6 +411,8 @@ impl MultivariateTpeSampler {
return;
}
missing_params.sort_by_key(|(id, _)| *id);
// Split trials for independent sampling
let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::<Vec<_>>());
@@ -390,6 +426,7 @@ impl MultivariateTpeSampler {
/// Samples all parameters using independent TPE sampling.
///
/// This is used as a complete fallback when no intersection search space exists.
/// Parameters are sorted by `ParamId` for deterministic RNG consumption order.
#[cfg(test)]
pub(crate) fn sample_all_independent(
&self,
@@ -402,7 +439,9 @@ impl MultivariateTpeSampler {
let mut rng = self.rng.lock();
let mut result = HashMap::new();
for (param_id, dist) in search_space {
let mut sorted: Vec<_> = search_space.iter().collect();
sorted.sort_by_key(|(id, _)| *id);
for (param_id, dist) in sorted {
let value =
self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, &mut rng);
result.insert(*param_id, value);
@@ -414,6 +453,7 @@ impl MultivariateTpeSampler {
/// Samples all parameters using independent TPE sampling with an external RNG.
///
/// This variant accepts an external RNG, used when the caller already holds the lock.
/// Parameters are sorted by `ParamId` for deterministic RNG consumption order.
pub(crate) fn sample_all_independent_with_rng(
&self,
search_space: &HashMap<ParamId, Distribution>,
@@ -425,7 +465,9 @@ impl MultivariateTpeSampler {
let mut result = HashMap::new();
for (param_id, dist) in search_space {
let mut sorted: Vec<_> = search_space.iter().collect();
sorted.sort_by_key(|(id, _)| *id);
for (param_id, dist) in sorted {
let value =
self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng);
result.insert(*param_id, value);
src/sampler/tpe/multivariate/mod.rs
+25 -12
View File
@@ -322,14 +322,18 @@ impl MultivariateTpeSampler {
/// Samples all parameters uniformly at random.
///
/// This is a fallback method used when multivariate TPE cannot be applied.
/// Parameters are sorted by `ParamId` to ensure deterministic RNG consumption
/// order when using a seeded sampler.
#[allow(clippy::unused_self)]
fn sample_all_uniform(
&self,
search_space: &HashMap<ParamId, Distribution>,
rng: &mut fastrand::Rng,
) -> HashMap<ParamId, ParamValue> {
search_space
.iter()
let mut sorted: Vec<_> = search_space.iter().collect();
sorted.sort_by_key(|(id, _)| *id);
sorted
.into_iter()
.map(|(id, dist)| (*id, crate::sampler::common::sample_random(rng, dist)))
.collect()
}
@@ -2781,8 +2785,9 @@ mod tests {
let good_kde = MultivariateKDE::new(good_samples).unwrap();
let bad_kde = MultivariateKDE::new(bad_samples).unwrap();
let bounds = &[(0.0, 1.0), (0.0, 1.0)];
let selected = sampler.select_candidate(&good_kde, &bad_kde);
let selected = sampler.select_candidate(&good_kde, &bad_kde, bounds);
// The selected candidate should have 2 dimensions
assert_eq!(selected.len(), 2);
@@ -2822,8 +2827,9 @@ mod tests {
let good_kde = MultivariateKDE::new(good_samples).unwrap();
let bad_kde = MultivariateKDE::new(bad_samples).unwrap();
let bounds = &[(-1.0, 2.0), (-1.0, 2.0), (-1.0, 2.0)];
let selected = sampler.select_candidate(&good_kde, &bad_kde);
let selected = sampler.select_candidate(&good_kde, &bad_kde, bounds);
assert_eq!(selected.len(), 3);
}
@@ -2841,8 +2847,9 @@ mod tests {
let good_kde = MultivariateKDE::new(good_samples).unwrap();
let bad_kde = MultivariateKDE::new(bad_samples).unwrap();
let bounds = &[(0.0, 10.0)];
let selected = sampler.select_candidate(&good_kde, &bad_kde);
let selected = sampler.select_candidate(&good_kde, &bad_kde, bounds);
assert_eq!(selected.len(), 1);
// Selected value should be closer to the good region
@@ -2873,14 +2880,15 @@ mod tests {
let good_kde = MultivariateKDE::new(good_samples.clone()).unwrap();
let bad_kde = MultivariateKDE::new(bad_samples.clone()).unwrap();
let bounds = &[(0.0, 10.0), (0.0, 10.0)];
let selected1 = sampler1.select_candidate(&good_kde, &bad_kde);
let selected1 = sampler1.select_candidate(&good_kde, &bad_kde, bounds);
// Need to recreate KDEs for second sampler since we consumed them
let good_kde2 = MultivariateKDE::new(good_samples).unwrap();
let bad_kde2 = MultivariateKDE::new(bad_samples).unwrap();
let selected2 = sampler2.select_candidate(&good_kde2, &bad_kde2);
let selected2 = sampler2.select_candidate(&good_kde2, &bad_kde2, bounds);
// With same seed, should get same result
assert!(
@@ -2911,8 +2919,9 @@ mod tests {
let good_kde = MultivariateKDE::new(good_samples).unwrap();
let bad_kde = MultivariateKDE::new(bad_samples).unwrap();
let bounds = &[(-10.0, 10.0), (-10.0, 10.0)];
let selected = sampler.select_candidate(&good_kde, &bad_kde);
let selected = sampler.select_candidate(&good_kde, &bad_kde, bounds);
assert_eq!(selected.len(), 2);
}
@@ -2942,8 +2951,9 @@ mod tests {
let good_kde = MultivariateKDE::new(good_samples).unwrap();
let bad_kde = MultivariateKDE::new(bad_samples).unwrap();
let bounds = &[(-5.0, 15.0), (-5.0, 15.0)];
let selected = sampler.select_candidate(&good_kde, &bad_kde);
let selected = sampler.select_candidate(&good_kde, &bad_kde, bounds);
// With more candidates, should definitely find a point in the good region
assert!(
@@ -2981,8 +2991,9 @@ mod tests {
let good_kde = MultivariateKDE::new(good_samples).unwrap();
let bad_kde = MultivariateKDE::new(bad_samples).unwrap();
let bounds = &[(-5.0, 5.0), (-5.0, 5.0)];
let selected = sampler.select_candidate(&good_kde, &bad_kde);
let selected = sampler.select_candidate(&good_kde, &bad_kde, bounds);
// Should still return a valid point
assert_eq!(selected.len(), 2);
@@ -3017,8 +3028,9 @@ mod tests {
let good_kde = MultivariateKDE::new(good_samples).unwrap();
let bad_kde = MultivariateKDE::new(bad_samples).unwrap();
let bounds = &[(-5.0, 15.0); 5];
let selected = sampler.select_candidate(&good_kde, &bad_kde);
let selected = sampler.select_candidate(&good_kde, &bad_kde, bounds);
assert_eq!(selected.len(), 5);
@@ -3091,7 +3103,8 @@ mod tests {
let bad_kde = MultivariateKDE::new(bad_obs).unwrap();
// Select candidate
let selected = sampler.select_candidate(&good_kde, &bad_kde);
let bounds = &[(0.0, 10.0), (0.0, 10.0)];
let selected = sampler.select_candidate(&good_kde, &bad_kde, bounds);
assert_eq!(selected.len(), 2);
tests/async_tests.rs
+18 -6
View File
@@ -197,27 +197,39 @@ async fn test_optimize_parallel_single_concurrency() {
#[tokio::test]
async fn test_parallel_executes_concurrently() {
use std::sync::atomic::{AtomicUsize, Ordering};
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(0.0, 10.0);
let active = Arc::new(AtomicUsize::new(0));
let max_active = Arc::new(AtomicUsize::new(0));
let active_c = Arc::clone(&active);
let max_active_c = Arc::clone(&max_active);
let start = tokio::time::Instant::now();
study
.optimize_parallel(4, 4, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
std::thread::sleep(std::time::Duration::from_millis(100));
let current = active_c.fetch_add(1, Ordering::SeqCst) + 1;
max_active_c.fetch_max(current, Ordering::SeqCst);
std::thread::sleep(std::time::Duration::from_millis(50));
active_c.fetch_sub(1, Ordering::SeqCst);
Ok::<_, Error>(x)
})
.await
.expect("parallel optimization should succeed");
let elapsed = start.elapsed();
assert_eq!(study.n_trials(), 4);
// Sequential would take ~400ms; parallel with concurrency=4 should be ~100ms
let max = max_active.load(Ordering::SeqCst);
// With 4 trials and concurrency=4, all should run concurrently
assert!(
elapsed < std::time::Duration::from_millis(350),
"expected parallel execution under 350ms, took {elapsed:?}"
max >= 2,
"expected at least 2 concurrent workers, but max was {max}"
);
}
tests/sampler/multivariate_tpe.rs
+2 -2
View File
@@ -45,7 +45,7 @@ fn test_multivariate_tpe_rosenbrock_finds_good_solution() {
let sampler = MultivariateTpeSampler::builder()
.seed(42)
.n_startup_trials(10)
.n_ei_candidates(24)
.n_ei_candidates(48)
.build()
.unwrap();
@@ -55,7 +55,7 @@ fn test_multivariate_tpe_rosenbrock_finds_good_solution() {
let y_param = FloatParam::new(-2.0, 4.0);
study
.optimize(100, |trial: &mut optimizer::Trial| {
.optimize(200, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(rosenbrock(x, y))
tests/stress_tests.rs
+1
View File
@@ -3,6 +3,7 @@
//! All tests are `#[ignore]`-gated so they don't run in normal CI.
//! Run with: `cargo test --features async -- --ignored`
#[cfg(feature = "async")]
use std::collections::HashSet;
use optimizer::parameter::{FloatParam, Parameter};