From 86db6361d639039cada0635fe69ade533c3babc1 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Fri, 13 Feb 2026 10:20:32 +0100 Subject: [PATCH] fix(sampler): clamp multivariate TPE candidates to parameter bounds - 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") --- src/sampler/tpe/multivariate/engine.rs | 62 +++++++++++++++++++++----- src/sampler/tpe/multivariate/mod.rs | 37 ++++++++++----- tests/async_tests.rs | 24 +++++++--- tests/sampler/multivariate_tpe.rs | 4 +- tests/stress_tests.rs | 1 + 5 files changed, 98 insertions(+), 30 deletions(-) diff --git a/src/sampler/tpe/multivariate/engine.rs b/src/sampler/tpe/multivariate/engine.rs index 52f5f53..4095329 100644 --- a/src/sampler/tpe/multivariate/engine.rs +++ b/src/sampler/tpe/multivariate/engine.rs @@ -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, ¶m_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 { let mut rng = self.rng.lock(); - // Generate candidates from the good distribution + // Generate candidates from the good distribution, clamped to bounds let candidates: Vec> = (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 { - // Generate candidates from the good distribution + // Generate candidates from the good distribution, clamped to bounds let candidates: Vec> = (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, 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::>()); @@ -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, @@ -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); diff --git a/src/sampler/tpe/multivariate/mod.rs b/src/sampler/tpe/multivariate/mod.rs index 4c452bc..05b2cd0 100644 --- a/src/sampler/tpe/multivariate/mod.rs +++ b/src/sampler/tpe/multivariate/mod.rs @@ -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, rng: &mut fastrand::Rng, ) -> HashMap { - 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); diff --git a/tests/async_tests.rs b/tests/async_tests.rs index a5c271f..f82be17 100644 --- a/tests/async_tests.rs +++ b/tests/async_tests.rs @@ -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 = 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}" ); } diff --git a/tests/sampler/multivariate_tpe.rs b/tests/sampler/multivariate_tpe.rs index a84ca34..78d1718 100644 --- a/tests/sampler/multivariate_tpe.rs +++ b/tests/sampler/multivariate_tpe.rs @@ -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)) diff --git a/tests/stress_tests.rs b/tests/stress_tests.rs index b652b90..48df70d 100644 --- a/tests/stress_tests.rs +++ b/tests/stress_tests.rs @@ -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};