From cc50e1ad4b7cb39688ea3bd4f91e61fed35ebf9a Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Thu, 12 Feb 2026 15:52:59 +0100 Subject: [PATCH] refactor: reduce #[allow] attributes and extract per-distribution helpers - Remove stale #[allow(dead_code)] and #[allow(too_many_lines)] attributes - Replace #[allow(dead_code)] with #[cfg(test)] for test-only methods - Delete unused fill_remaining_independent() and ParamMeta.dist field - Extract sample_float/int/categorical helpers from TpeSampler, MotpeSampler, and SamplingEngine to shorten match-heavy sample() methods - Extract try_fit_kdes() helper to consolidate repeated fallback blocks - Refactor sample_tpe_float/int to take &FloatDistribution/&IntDistribution instead of destructured args, removing #[allow(too_many_arguments)] --- src/kde/multivariate.rs | 9 +- src/sampler/motpe.rs | 234 ++++++++++-------- src/sampler/tpe/common.rs | 34 ++- src/sampler/tpe/multivariate/engine.rs | 330 ++++++++++++------------- src/sampler/tpe/multivariate/mod.rs | 1 - src/sampler/tpe/sampler.rs | 242 +++++++++--------- src/study/export.rs | 2 +- src/visualization.rs | 6 +- 8 files changed, 427 insertions(+), 431 deletions(-) diff --git a/src/kde/multivariate.rs b/src/kde/multivariate.rs index 00bd843..ef1c02e 100644 --- a/src/kde/multivariate.rs +++ b/src/kde/multivariate.rs @@ -30,7 +30,6 @@ use crate::error::{Error, Result}; /// // Get dimensionality /// assert_eq!(kde.n_dims(), 2); /// ``` -#[allow(dead_code)] // Fields and methods will be used in subsequent stories (US-003, US-004) #[derive(Clone, Debug)] pub(crate) struct MultivariateKDE { /// The sample points used to construct the KDE. @@ -43,7 +42,6 @@ pub(crate) struct MultivariateKDE { n_dims: usize, } -#[allow(dead_code)] // Methods will be used in subsequent stories (US-003, US-004) impl MultivariateKDE { /// Creates a new multivariate KDE with automatic bandwidth selection using Scott's rule. /// @@ -100,6 +98,7 @@ impl MultivariateKDE { /// Returns `Error::ZeroDimensions` if samples have zero dimensions. /// Returns `Error::BandwidthDimensionMismatch` if bandwidths length doesn't match dimensions. /// Returns `Error::InvalidBandwidth` if any bandwidth is not positive. + #[cfg(test)] pub(crate) fn with_bandwidths(samples: Vec>, bandwidths: Vec) -> Result { if samples.is_empty() { return Err(Error::EmptySamples); @@ -180,12 +179,13 @@ impl MultivariateKDE { } /// Returns the number of dimensions. + #[cfg(test)] pub(crate) fn n_dims(&self) -> usize { self.n_dims } /// Returns the number of samples. - #[allow(dead_code)] // Will be used in subsequent stories + #[cfg(test)] pub(crate) fn n_samples(&self) -> usize { self.samples.len() } @@ -197,7 +197,7 @@ impl MultivariateKDE { } /// Returns a reference to the samples. - #[allow(dead_code)] // Will be used in subsequent stories + #[cfg(test)] pub(crate) fn samples(&self) -> &[Vec] { &self.samples } @@ -288,6 +288,7 @@ impl MultivariateKDE { /// # Panics /// /// Panics if `x.len() != self.n_dims`. + #[cfg(test)] pub(crate) fn pdf(&self, x: &[f64]) -> f64 { self.log_pdf(x).exp() } diff --git a/src/sampler/motpe.rs b/src/sampler/motpe.rs index a84580d..8049dee 100644 --- a/src/sampler/motpe.rs +++ b/src/sampler/motpe.rs @@ -214,8 +214,130 @@ impl Default for MotpeSampler { } } +impl MotpeSampler { + fn sample_float( + &self, + d: &crate::distribution::FloatDistribution, + good_trials: &[&MultiObjectiveTrial], + bad_trials: &[&MultiObjectiveTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + let good_values: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + if good_values.is_empty() || bad_values.is_empty() { + return ParamValue::Float(rng_util::f64_range(rng, d.low, d.high)); + } + + let value = tpe_common::sample_tpe_float( + d, + good_values, + bad_values, + self.n_ei_candidates, + self.kde_bandwidth, + rng, + ); + ParamValue::Float(value) + } + + fn sample_int( + &self, + d: &crate::distribution::IntDistribution, + good_trials: &[&MultiObjectiveTrial], + bad_trials: &[&MultiObjectiveTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + let good_values: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + if good_values.is_empty() || bad_values.is_empty() { + return common::sample_random(rng, &Distribution::Int(d.clone())); + } + + let value = tpe_common::sample_tpe_int( + d, + good_values, + bad_values, + self.n_ei_candidates, + self.kde_bandwidth, + rng, + ); + ParamValue::Int(value) + } + + #[allow(clippy::unused_self)] + fn sample_categorical( + &self, + d: &crate::distribution::CategoricalDistribution, + good_trials: &[&MultiObjectiveTrial], + bad_trials: &[&MultiObjectiveTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + let good_indices: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + let bad_indices: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + if good_indices.is_empty() || bad_indices.is_empty() { + return common::sample_random(rng, &Distribution::Categorical(d.clone())); + } + + let index = + tpe_common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); + ParamValue::Categorical(index) + } +} + impl MultiObjectiveSampler for MotpeSampler { - #[allow(clippy::too_many_lines)] fn sample( &self, distribution: &Distribution, @@ -247,114 +369,10 @@ impl MultiObjectiveSampler for MotpeSampler { } match distribution { - Distribution::Float(d) => { - let good_values: Vec = good_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Float(f) => Some(*f), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - let bad_values: Vec = bad_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Float(f) => Some(*f), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - if good_values.is_empty() || bad_values.is_empty() { - return common::sample_random(&mut rng, distribution); - } - - 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) - } - Distribution::Int(d) => { - let good_values: Vec = good_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Int(i) => Some(*i), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - let bad_values: Vec = bad_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Int(i) => Some(*i), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - if good_values.is_empty() || bad_values.is_empty() { - return common::sample_random(&mut rng, distribution); - } - - 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) - } + Distribution::Float(d) => self.sample_float(d, &good_trials, &bad_trials, &mut rng), + Distribution::Int(d) => self.sample_int(d, &good_trials, &bad_trials, &mut rng), Distribution::Categorical(d) => { - let good_indices: Vec = good_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Categorical(i) => Some(*i), - _ => None, - }) - .filter(|&i| i < d.n_choices) - .collect(); - - let bad_indices: Vec = bad_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Categorical(i) => Some(*i), - _ => None, - }) - .filter(|&i| i < d.n_choices) - .collect(); - - if good_indices.is_empty() || bad_indices.is_empty() { - return common::sample_random(&mut rng, distribution); - } - - let index = tpe_common::sample_tpe_categorical( - d.n_choices, - &good_indices, - &bad_indices, - &mut rng, - ); - ParamValue::Categorical(index) + self.sample_categorical(d, &good_trials, &bad_trials, &mut rng) } } } diff --git a/src/sampler/tpe/common.rs b/src/sampler/tpe/common.rs index c566ae9..02b41b8 100644 --- a/src/sampler/tpe/common.rs +++ b/src/sampler/tpe/common.rs @@ -1,21 +1,20 @@ //! Shared TPE sampling functions used by both `TpeSampler` and `MotpeSampler`. +use crate::distribution::{FloatDistribution, IntDistribution}; 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, + dist: &FloatDistribution, good_values: Vec, bad_values: Vec, n_ei_candidates: usize, kde_bandwidth: Option, rng: &mut fastrand::Rng, ) -> f64 { + let (low, high, log_scale, step) = (dist.low, dist.high, dist.log_scale, dist.step); + // 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(); @@ -99,32 +98,31 @@ pub(crate) fn sample_tpe_float( } /// Samples using TPE for integer distributions. -#[allow( - clippy::too_many_arguments, - clippy::cast_precision_loss, - clippy::cast_possible_truncation -)] +#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)] pub(crate) fn sample_tpe_int( - low: i64, - high: i64, - log_scale: bool, - step: Option, + dist: &IntDistribution, good_values: Vec, bad_values: Vec, n_ei_candidates: usize, kde_bandwidth: Option, rng: &mut fastrand::Rng, ) -> i64 { + let (low, high, log_scale, step) = (dist.low, dist.high, dist.log_scale, dist.step); + // 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(); + let float_dist = FloatDistribution { + low: low as f64, + high: high as f64, + log_scale, + step: step.map(|s| s as f64), + }; + // Use float TPE sampling let float_value = sample_tpe_float( - low as f64, - high as f64, - log_scale, - step.map(|s| s as f64), + &float_dist, good_floats, bad_floats, n_ei_candidates, diff --git a/src/sampler/tpe/multivariate/engine.rs b/src/sampler/tpe/multivariate/engine.rs index 1002b29..52f5f53 100644 --- a/src/sampler/tpe/multivariate/engine.rs +++ b/src/sampler/tpe/multivariate/engine.rs @@ -103,6 +103,30 @@ impl MultivariateTpeSampler { result } + /// Validates observations and fits multivariate KDEs for good and bad groups. + /// + /// Returns `None` if observations are invalid or KDE construction fails. + fn try_fit_kdes( + good_obs: Vec>, + bad_obs: Vec>, + expected_dims: usize, + ) -> Option<(crate::kde::MultivariateKDE, crate::kde::MultivariateKDE)> { + use crate::kde::MultivariateKDE; + + let valid = !good_obs.is_empty() + && !bad_obs.is_empty() + && good_obs.iter().all(|obs| obs.len() == expected_dims) + && bad_obs.iter().all(|obs| obs.len() == expected_dims); + + if !valid { + return None; + } + + let good_kde = MultivariateKDE::new(good_obs).ok()?; + let bad_kde = MultivariateKDE::new(bad_obs).ok()?; + Some((good_kde, bad_kde)) + } + /// Samples parameters as a single group using multivariate TPE. /// /// This is the core multivariate TPE sampling logic, used both in non-grouped mode @@ -117,14 +141,12 @@ impl MultivariateTpeSampler { /// # Returns /// /// A `HashMap` mapping parameter names to their sampled values. - #[allow(clippy::too_many_lines)] pub(crate) fn sample_single_group( &self, search_space: &HashMap, history: &[CompletedTrial], rng: &mut fastrand::Rng, ) -> HashMap { - use crate::kde::MultivariateKDE; use crate::sampler::tpe::IntersectionSearchSpace; use crate::sampler::tpe::common; @@ -165,7 +187,6 @@ impl MultivariateTpeSampler { .collect(); if param_order.is_empty() { - // Only categorical parameters in intersection - fill remaining with independent TPE self.fill_remaining_independent_with_rng( search_space, &intersection, @@ -178,43 +199,12 @@ impl MultivariateTpeSampler { param_order.sort_by_key(|id| format!("{id}")); - // Extract observations and validate + // Extract observations, validate, and fit KDEs let good_obs = self.extract_observations(&good, ¶m_order); let bad_obs = self.extract_observations(&bad, ¶m_order); - let expected_dims = param_order.len(); - let valid = !good_obs.is_empty() - && !bad_obs.is_empty() - && good_obs.iter().all(|obs| obs.len() == expected_dims) - && bad_obs.iter().all(|obs| obs.len() == expected_dims); - - if !valid { - // Observations invalid - fill remaining with independent TPE - self.fill_remaining_independent_with_rng( - search_space, - &intersection, - history, - &mut result, - rng, - ); - return result; - } - - // Fit KDEs using let...else pattern - let Ok(good_kde) = MultivariateKDE::new(good_obs) else { - // KDE construction failed - fill remaining with independent TPE - self.fill_remaining_independent_with_rng( - search_space, - &intersection, - history, - &mut result, - rng, - ); - return result; - }; - - let Ok(bad_kde) = MultivariateKDE::new(bad_obs) else { - // KDE construction failed - fill remaining with independent TPE + let Some((good_kde, bad_kde)) = Self::try_fit_kdes(good_obs, bad_obs, param_order.len()) + else { self.fill_remaining_independent_with_rng( search_space, &intersection, @@ -289,7 +279,7 @@ impl MultivariateTpeSampler { /// 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))`. #[must_use] - #[allow(dead_code)] // Used by tests + #[cfg(test)] pub(crate) fn select_candidate( &self, good_kde: &crate::kde::MultivariateKDE, @@ -366,43 +356,6 @@ impl MultivariateTpeSampler { candidates.into_iter().nth(best_idx).unwrap_or_default() } - /// Fills remaining parameters in result using independent TPE sampling. - /// - /// This method is used to sample parameters that are not in the intersection - /// search space. It uses independent univariate TPE sampling for each parameter, - /// similar to the standard [`TpeSampler`]. - /// - /// When there isn't enough history for a parameter, falls back to uniform sampling. - #[allow(dead_code)] - pub(crate) fn fill_remaining_independent( - &self, - search_space: &HashMap, - _intersection: &HashMap, - history: &[CompletedTrial], - result: &mut HashMap, - ) { - // Identify parameters not in result (and not in intersection) - let missing_params: Vec<(&ParamId, &Distribution)> = search_space - .iter() - .filter(|(id, _)| !result.contains_key(id)) - .collect(); - - if missing_params.is_empty() { - return; - } - - // Split trials for independent sampling - let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); - - let mut rng = self.rng.lock(); - - for (param_id, dist) in missing_params { - let value = - self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, &mut rng); - result.insert(*param_id, value); - } - } - /// Fills remaining parameters using independent TPE sampling with an external RNG. /// /// This variant accepts an external RNG, used when the caller already holds the lock. @@ -437,7 +390,7 @@ impl MultivariateTpeSampler { /// Samples all parameters using independent TPE sampling. /// /// This is used as a complete fallback when no intersection search space exists. - #[allow(dead_code)] // Used by tests + #[cfg(test)] pub(crate) fn sample_all_independent( &self, search_space: &HashMap, @@ -485,7 +438,6 @@ impl MultivariateTpeSampler { /// /// This method extracts values for the given parameter from good and bad trials, /// fits univariate KDEs, and samples using the TPE acquisition function. - #[allow(clippy::too_many_lines)] pub(crate) fn sample_independent_tpe( &self, param_id: ParamId, @@ -494,114 +446,136 @@ impl MultivariateTpeSampler { bad_trials: &[&CompletedTrial], rng: &mut fastrand::Rng, ) -> ParamValue { - use crate::sampler::tpe::common; - match distribution { Distribution::Float(d) => { - let good_values: Vec = good_trials - .iter() - .filter_map(|t| t.params.get(¶m_id)) - .filter_map(|v| match v { - ParamValue::Float(f) => Some(*f), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - let bad_values: Vec = bad_trials - .iter() - .filter_map(|t| t.params.get(¶m_id)) - .filter_map(|v| match v { - ParamValue::Float(f) => Some(*f), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - if good_values.is_empty() || bad_values.is_empty() { - return crate::sampler::common::sample_random(rng, distribution); - } - - let value = common::sample_tpe_float( - d.low, - d.high, - d.log_scale, - d.step, - good_values, - bad_values, - self.n_ei_candidates, - None, - rng, - ); - ParamValue::Float(value) + self.sample_independent_float(param_id, d, good_trials, bad_trials, rng) } Distribution::Int(d) => { - let good_values: Vec = good_trials - .iter() - .filter_map(|t| t.params.get(¶m_id)) - .filter_map(|v| match v { - ParamValue::Int(i) => Some(*i), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - let bad_values: Vec = bad_trials - .iter() - .filter_map(|t| t.params.get(¶m_id)) - .filter_map(|v| match v { - ParamValue::Int(i) => Some(*i), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - if good_values.is_empty() || bad_values.is_empty() { - return crate::sampler::common::sample_random(rng, distribution); - } - - let value = common::sample_tpe_int( - d.low, - d.high, - d.log_scale, - d.step, - good_values, - bad_values, - self.n_ei_candidates, - None, - rng, - ); - ParamValue::Int(value) + self.sample_independent_int(param_id, d, good_trials, bad_trials, rng) } Distribution::Categorical(d) => { - let good_indices: Vec = good_trials - .iter() - .filter_map(|t| t.params.get(¶m_id)) - .filter_map(|v| match v { - ParamValue::Categorical(i) => Some(*i), - _ => None, - }) - .filter(|&i| i < d.n_choices) - .collect(); - - let bad_indices: Vec = bad_trials - .iter() - .filter_map(|t| t.params.get(¶m_id)) - .filter_map(|v| match v { - ParamValue::Categorical(i) => Some(*i), - _ => None, - }) - .filter(|&i| i < d.n_choices) - .collect(); - - if good_indices.is_empty() || bad_indices.is_empty() { - return crate::sampler::common::sample_random(rng, distribution); - } - - let idx = - common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); - ParamValue::Categorical(idx) + self.sample_independent_categorical(param_id, d, good_trials, bad_trials, rng) } } } + + fn sample_independent_float( + &self, + param_id: ParamId, + d: &crate::distribution::FloatDistribution, + good_trials: &[&CompletedTrial], + bad_trials: &[&CompletedTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + use crate::sampler::tpe::common; + + let good_values: Vec = good_trials + .iter() + .filter_map(|t| t.params.get(¶m_id)) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .filter_map(|t| t.params.get(¶m_id)) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + if good_values.is_empty() || bad_values.is_empty() { + return crate::sampler::common::sample_random(rng, &Distribution::Float(d.clone())); + } + + let value = + common::sample_tpe_float(d, good_values, bad_values, self.n_ei_candidates, None, rng); + ParamValue::Float(value) + } + + fn sample_independent_int( + &self, + param_id: ParamId, + d: &crate::distribution::IntDistribution, + good_trials: &[&CompletedTrial], + bad_trials: &[&CompletedTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + use crate::sampler::tpe::common; + + let good_values: Vec = good_trials + .iter() + .filter_map(|t| t.params.get(¶m_id)) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .filter_map(|t| t.params.get(¶m_id)) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + if good_values.is_empty() || bad_values.is_empty() { + return crate::sampler::common::sample_random(rng, &Distribution::Int(d.clone())); + } + + let value = + common::sample_tpe_int(d, good_values, bad_values, self.n_ei_candidates, None, rng); + ParamValue::Int(value) + } + + #[allow(clippy::unused_self)] + fn sample_independent_categorical( + &self, + param_id: ParamId, + d: &crate::distribution::CategoricalDistribution, + good_trials: &[&CompletedTrial], + bad_trials: &[&CompletedTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + use crate::sampler::tpe::common; + + let good_indices: Vec = good_trials + .iter() + .filter_map(|t| t.params.get(¶m_id)) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + let bad_indices: Vec = bad_trials + .iter() + .filter_map(|t| t.params.get(¶m_id)) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + if good_indices.is_empty() || bad_indices.is_empty() { + return crate::sampler::common::sample_random( + rng, + &Distribution::Categorical(d.clone()), + ); + } + + let idx = common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); + ParamValue::Categorical(idx) + } } diff --git a/src/sampler/tpe/multivariate/mod.rs b/src/sampler/tpe/multivariate/mod.rs index 4900e13..f3738f9 100644 --- a/src/sampler/tpe/multivariate/mod.rs +++ b/src/sampler/tpe/multivariate/mod.rs @@ -389,7 +389,6 @@ impl MultivariateTpeSampler { /// let params = sampler.sample_joint(&search_space, &history); /// ``` #[must_use] - #[allow(clippy::too_many_lines)] pub fn sample_joint( &self, search_space: &HashMap, diff --git a/src/sampler/tpe/sampler.rs b/src/sampler/tpe/sampler.rs index 5c9af1d..32f3f1a 100644 --- a/src/sampler/tpe/sampler.rs +++ b/src/sampler/tpe/sampler.rs @@ -642,8 +642,130 @@ impl Default for TpeSamplerBuilder { } } +impl TpeSampler { + fn sample_float( + &self, + d: &crate::distribution::FloatDistribution, + good_trials: &[&CompletedTrial], + bad_trials: &[&CompletedTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + let good_values: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Float(f) => Some(*f), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + if good_values.is_empty() || bad_values.is_empty() { + return ParamValue::Float(rng_util::f64_range(rng, d.low, d.high)); + } + + let value = tpe_common::sample_tpe_float( + d, + good_values, + bad_values, + self.n_ei_candidates, + self.kde_bandwidth, + rng, + ); + ParamValue::Float(value) + } + + fn sample_int( + &self, + d: &crate::distribution::IntDistribution, + good_trials: &[&CompletedTrial], + bad_trials: &[&CompletedTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + let good_values: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + let bad_values: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Int(i) => Some(*i), + _ => None, + }) + .filter(|&v| v >= d.low && v <= d.high) + .collect(); + + if good_values.is_empty() || bad_values.is_empty() { + return common::sample_random(rng, &Distribution::Int(d.clone())); + } + + let value = tpe_common::sample_tpe_int( + d, + good_values, + bad_values, + self.n_ei_candidates, + self.kde_bandwidth, + rng, + ); + ParamValue::Int(value) + } + + #[allow(clippy::unused_self)] + fn sample_categorical( + &self, + d: &crate::distribution::CategoricalDistribution, + good_trials: &[&CompletedTrial], + bad_trials: &[&CompletedTrial], + rng: &mut fastrand::Rng, + ) -> ParamValue { + let good_indices: Vec = good_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + let bad_indices: Vec = bad_trials + .iter() + .flat_map(|t| t.params.values()) + .filter_map(|v| match v { + ParamValue::Categorical(i) => Some(*i), + _ => None, + }) + .filter(|&i| i < d.n_choices) + .collect(); + + if good_indices.is_empty() || bad_indices.is_empty() { + return common::sample_random(rng, &Distribution::Categorical(d.clone())); + } + + let index = + tpe_common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); + ParamValue::Categorical(index) + } +} + impl Sampler for TpeSampler { - #[allow(clippy::too_many_lines)] fn sample( &self, distribution: &Distribution, @@ -670,123 +792,11 @@ impl Sampler for TpeSampler { return common::sample_random(&mut rng, distribution); } - // Extract parameter values for this distribution - // Since we don't have the parameter name here, we need to look at all - // trials and find matching distributions - // Note: This is a simplification - in practice, we'd need the param name - // For now, we'll collect values from trials that have this exact distribution type - match distribution { - Distribution::Float(d) => { - // Collect float values from trials - let good_values: Vec = good_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Float(f) => Some(*f), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - let bad_values: Vec = bad_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Float(f) => Some(*f), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - // Need values in both groups for TPE - if good_values.is_empty() || bad_values.is_empty() { - return common::sample_random(&mut rng, distribution); - } - - 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) - } - Distribution::Int(d) => { - let good_values: Vec = good_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Int(i) => Some(*i), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - let bad_values: Vec = bad_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Int(i) => Some(*i), - _ => None, - }) - .filter(|&v| v >= d.low && v <= d.high) - .collect(); - - if good_values.is_empty() || bad_values.is_empty() { - return common::sample_random(&mut rng, distribution); - } - - 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) - } + Distribution::Float(d) => self.sample_float(d, &good_trials, &bad_trials, &mut rng), + Distribution::Int(d) => self.sample_int(d, &good_trials, &bad_trials, &mut rng), Distribution::Categorical(d) => { - let good_indices: Vec = good_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Categorical(i) => Some(*i), - _ => None, - }) - .filter(|&i| i < d.n_choices) - .collect(); - - let bad_indices: Vec = bad_trials - .iter() - .flat_map(|t| t.params.values()) - .filter_map(|v| match v { - ParamValue::Categorical(i) => Some(*i), - _ => None, - }) - .filter(|&i| i < d.n_choices) - .collect(); - - if good_indices.is_empty() || bad_indices.is_empty() { - return common::sample_random(&mut rng, distribution); - } - - let index = tpe_common::sample_tpe_categorical( - d.n_choices, - &good_indices, - &bad_indices, - &mut rng, - ); - ParamValue::Categorical(index) + self.sample_categorical(d, &good_trials, &bad_trials, &mut rng) } } } diff --git a/src/study/export.rs b/src/study/export.rs index 0db9ef9..beff14b 100644 --- a/src/study/export.rs +++ b/src/study/export.rs @@ -224,7 +224,7 @@ where impl Study { /// Export trials as a pretty-printed JSON array to a file. /// - /// Each element in the array is a serialized [`crate::CompletedTrial`]. + /// Each element in the array is a serialized [`CompletedTrial`](crate::sampler::CompletedTrial). /// Requires the `serde` feature. /// /// # Errors diff --git a/src/visualization.rs b/src/visualization.rs index 3c9558d..00eac89 100644 --- a/src/visualization.rs +++ b/src/visualization.rs @@ -41,7 +41,6 @@ use core::fmt::Write as _; use std::collections::BTreeMap; use std::path::Path; -use crate::distribution::Distribution; use crate::param::ParamValue; use crate::parameter::ParamId; use crate::sampler::CompletedTrial; @@ -156,8 +155,6 @@ fn build_html( /// Metadata about each parameter seen across trials. struct ParamMeta { label: String, - #[allow(dead_code)] - dist: Option, } /// Collect parameter labels and distributions across all trials. @@ -171,8 +168,7 @@ fn collect_param_info(trials: &[CompletedTrial]) -> BTreeMap