diff --git a/src/sampler/tpe/multivariate/engine.rs b/src/sampler/tpe/multivariate/engine.rs new file mode 100644 index 0000000..1002b29 --- /dev/null +++ b/src/sampler/tpe/multivariate/engine.rs @@ -0,0 +1,607 @@ +//! Core multivariate TPE sampling logic. +//! +//! Contains the main sampling engine: group decomposition, single-group multivariate +//! TPE, candidate selection, independent fallbacks, and value conversion. + +use std::collections::HashMap; + +use crate::distribution::Distribution; +use crate::param::ParamValue; +use crate::parameter::ParamId; +use crate::sampler::CompletedTrial; + +use super::MultivariateTpeSampler; + +impl MultivariateTpeSampler { + /// Samples parameters by decomposing the search space into independent groups. + /// + /// When `group=true`, this method analyzes the trial history to identify groups of + /// parameters that always appear together, then samples each group independently + /// using multivariate TPE. This is more efficient when parameters naturally partition + /// into independent subsets (e.g., due to conditional search spaces). + /// + /// # Arguments + /// + /// * `search_space` - The full search space containing all parameters to sample. + /// * `history` - Completed trials from the optimization history. + /// + /// # Returns + /// + /// A `HashMap` mapping parameter names to their sampled values. + pub(crate) fn sample_with_groups( + &self, + search_space: &HashMap, + history: &[CompletedTrial], + ) -> HashMap { + use std::collections::HashSet; + + use crate::sampler::tpe::GroupDecomposedSearchSpace; + + // Decompose the search space into independent parameter groups + let groups = GroupDecomposedSearchSpace::calculate(history); + + let mut result: HashMap = HashMap::new(); + + // Sample each group independently + for group in &groups { + // Build a sub-search space for this group + let group_search_space: HashMap = search_space + .iter() + .filter(|(id, _)| group.contains(id)) + .map(|(id, dist)| (*id, dist.clone())) + .collect(); + + if group_search_space.is_empty() { + continue; + } + + // Filter history to trials that have at least one parameter in this group + let group_history: Vec<&CompletedTrial> = history + .iter() + .filter(|trial| { + trial + .distributions + .keys() + .any(|param_id| group.contains(param_id)) + }) + .collect(); + + // Build completed trials from references for the group + // We need to create a temporary slice for sample_group_internal + let group_history_owned: Vec = + group_history.iter().map(|t| (*t).clone()).collect(); + + // Sample this group using multivariate TPE + let mut rng = self.rng.lock(); + let group_result = + self.sample_single_group(&group_search_space, &group_history_owned, &mut rng); + drop(rng); + + // Merge group results into the main result + for (id, value) in group_result { + result.insert(id, value); + } + } + + // Handle parameters not in any group (sample independently) + let grouped_params: HashSet = groups.iter().flatten().copied().collect(); + let ungrouped_params: HashMap = search_space + .iter() + .filter(|(id, _)| !grouped_params.contains(id) && !result.contains_key(id)) + .map(|(id, dist)| (*id, dist.clone())) + .collect(); + + if !ungrouped_params.is_empty() { + // Sample ungrouped parameters uniformly (no history for them) + let mut rng = self.rng.lock(); + for (id, dist) in &ungrouped_params { + let value = crate::sampler::common::sample_random(&mut rng, dist); + result.insert(*id, value); + } + } + + result + } + + /// Samples parameters as a single group using multivariate TPE. + /// + /// This is the core multivariate TPE sampling logic, used both in non-grouped mode + /// and for sampling individual groups in grouped mode. + /// + /// # Arguments + /// + /// * `search_space` - The search space for this group of parameters. + /// * `history` - Completed trials to use for model fitting. + /// * `rng` - Random number generator (caller must hold lock). + /// + /// # 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; + + // Early returns for cases requiring random sampling + if history.len() < self.n_startup_trials { + return self.sample_all_uniform(search_space, rng); + } + + let intersection = IntersectionSearchSpace::calculate(history); + if intersection.is_empty() { + return self.sample_all_independent_with_rng(search_space, history, rng); + } + + let filtered = self.filter_trials(history, &intersection); + if filtered.len() < 2 { + return self.sample_all_independent_with_rng(search_space, history, rng); + } + + let (good, bad) = self.split_trials(&filtered); + + // Sample categorical parameters using TPE with l(x)/g(x) ratio + let mut result: HashMap = HashMap::new(); + for (param_id, dist) in &intersection { + if let Distribution::Categorical(d) = dist { + let good_indices = Self::extract_categorical_indices(&good, *param_id); + let bad_indices = Self::extract_categorical_indices(&bad, *param_id); + let idx = + common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); + result.insert(*param_id, ParamValue::Categorical(idx)); + } + } + + // Collect continuous parameters + let mut param_order: Vec = intersection + .iter() + .filter(|(_, dist)| !matches!(dist, Distribution::Categorical(_))) + .map(|(id, _)| *id) + .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, + history, + &mut result, + rng, + ); + return result; + } + + param_order.sort_by_key(|id| format!("{id}")); + + // Extract observations and validate + 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 + self.fill_remaining_independent_with_rng( + search_space, + &intersection, + history, + &mut result, + rng, + ); + return result; + }; + + let selected = self.select_candidate_with_rng(&good_kde, &bad_kde, rng); + + // Map selected values to parameter ids + for (idx, param_id) in param_order.iter().enumerate() { + if let Some(dist) = intersection.get(param_id) { + let value = selected[idx]; + let param_value = self.convert_to_param_value(value, dist); + if let Some(pv) = param_value { + result.insert(*param_id, pv); + } + } + } + + // Fill remaining parameters using independent TPE sampling + self.fill_remaining_independent_with_rng( + search_space, + &intersection, + history, + &mut result, + rng, + ); + result + } + + /// Converts a raw f64 value to a `ParamValue` based on the distribution. + #[allow(clippy::unused_self)] + pub(crate) fn convert_to_param_value( + &self, + value: f64, + dist: &Distribution, + ) -> Option { + match dist { + Distribution::Float(d) => { + let clamped = value.clamp(d.low, d.high); + let stepped = if let Some(step) = d.step { + let steps = ((clamped - d.low) / step).round(); + (d.low + steps * step).clamp(d.low, d.high) + } else { + clamped + }; + Some(ParamValue::Float(stepped)) + } + Distribution::Int(d) => { + #[allow(clippy::cast_possible_truncation)] + let int_value = value.round() as i64; + let clamped = int_value.clamp(d.low, d.high); + let stepped = if let Some(step) = d.step { + let steps = (clamped - d.low) / step; + (d.low + steps * step).clamp(d.low, d.high) + } else { + clamped + }; + Some(ParamValue::Int(stepped)) + } + Distribution::Categorical(_) => None, + } + } + + /// 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))`. + #[must_use] + #[allow(dead_code)] // Used by tests + pub(crate) fn select_candidate( + &self, + good_kde: &crate::kde::MultivariateKDE, + bad_kde: &crate::kde::MultivariateKDE, + ) -> Vec { + let mut rng = self.rng.lock(); + + // Generate candidates from the good distribution + let candidates: Vec> = (0..self.n_ei_candidates) + .map(|_| good_kde.sample(&mut rng)) + .collect(); + + // Compute log(l(x)) - log(g(x)) for each candidate + // This is equivalent to log(l(x)/g(x)) which we want to maximize + let log_ratios: Vec = candidates + .iter() + .map(|candidate| { + let log_l = good_kde.log_pdf(candidate); + let log_g = bad_kde.log_pdf(candidate); + log_l - log_g + }) + .collect(); + + // Find the candidate with the maximum log ratio + let mut best_idx = 0; + let mut best_ratio = f64::NEG_INFINITY; + + for (idx, &ratio) in log_ratios.iter().enumerate() { + // Handle NaN by treating it as worse than any finite value + if ratio > best_ratio || (best_ratio.is_nan() && !ratio.is_nan()) { + best_ratio = ratio; + best_idx = idx; + } + } + + candidates.into_iter().nth(best_idx).unwrap_or_default() + } + + /// Selects the best candidate using an external RNG. + /// + /// This variant accepts an external RNG, used when the caller already holds the lock. + pub(crate) fn select_candidate_with_rng( + &self, + good_kde: &crate::kde::MultivariateKDE, + bad_kde: &crate::kde::MultivariateKDE, + rng: &mut fastrand::Rng, + ) -> Vec { + // Generate candidates from the good distribution + let candidates: Vec> = (0..self.n_ei_candidates) + .map(|_| good_kde.sample(rng)) + .collect(); + + // Compute log(l(x)) - log(g(x)) for each candidate + let log_ratios: Vec = candidates + .iter() + .map(|candidate| { + let log_l = good_kde.log_pdf(candidate); + let log_g = bad_kde.log_pdf(candidate); + log_l - log_g + }) + .collect(); + + // Find the candidate with the maximum log ratio + let mut best_idx = 0; + let mut best_ratio = f64::NEG_INFINITY; + + for (idx, &ratio) in log_ratios.iter().enumerate() { + if ratio > best_ratio || (best_ratio.is_nan() && !ratio.is_nan()) { + best_ratio = ratio; + best_idx = idx; + } + } + + 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. + pub(crate) fn fill_remaining_independent_with_rng( + &self, + search_space: &HashMap, + _intersection: &HashMap, + history: &[CompletedTrial], + result: &mut HashMap, + rng: &mut fastrand::Rng, + ) { + // 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::>()); + + for (param_id, dist) in missing_params { + let value = + self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng); + result.insert(*param_id, value); + } + } + + /// 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 + pub(crate) fn sample_all_independent( + &self, + search_space: &HashMap, + history: &[CompletedTrial], + ) -> HashMap { + // Split trials for independent sampling + let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); + + let mut rng = self.rng.lock(); + let mut result = HashMap::new(); + + for (param_id, dist) in search_space { + let value = + self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, &mut rng); + result.insert(*param_id, value); + } + + result + } + + /// 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. + pub(crate) fn sample_all_independent_with_rng( + &self, + search_space: &HashMap, + history: &[CompletedTrial], + rng: &mut fastrand::Rng, + ) -> HashMap { + // Split trials for independent sampling + let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); + + let mut result = HashMap::new(); + + for (param_id, dist) in search_space { + let value = + self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng); + result.insert(*param_id, value); + } + + result + } + + /// Samples a single parameter using independent TPE. + /// + /// 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, + distribution: &Distribution, + good_trials: &[&CompletedTrial], + 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) + } + 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) + } + 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) + } + } + } +} diff --git a/src/sampler/tpe/multivariate.rs b/src/sampler/tpe/multivariate/mod.rs similarity index 80% rename from src/sampler/tpe/multivariate.rs rename to src/sampler/tpe/multivariate/mod.rs index de8b37e..4900e13 100644 --- a/src/sampler/tpe/multivariate.rs +++ b/src/sampler/tpe/multivariate/mod.rs @@ -112,6 +112,9 @@ //! let sampler = MultivariateTpeSampler::builder().build().unwrap(); //! ``` +mod engine; +mod trials; + use std::collections::HashMap; use std::sync::Arc; @@ -122,8 +125,7 @@ use crate::distribution::Distribution; use crate::error::Result; use crate::param::ParamValue; use crate::parameter::ParamId; -use crate::rng_util; -use crate::sampler::{CompletedTrial, PendingTrial, Sampler}; +use crate::sampler::{CompletedTrial, Sampler}; /// Strategy for imputing objective values for pending/running trials during parallel optimization. /// @@ -311,494 +313,6 @@ impl MultivariateTpeSampler { &self.constant_liar } - /// Imputes objective values for pending trials based on the constant liar strategy. - /// - /// In parallel optimization, multiple trials may be running simultaneously. This method - /// assigns "lie" values to pending trials so they can be included in the model fitting, - /// which helps avoid redundant exploration of the same region. - /// - /// # Arguments - /// - /// * `pending_trials` - Trials that are currently running and have no objective value yet. - /// * `completed_trials` - Trials that have completed and have objective values. - /// - /// # Returns - /// - /// A vector of `CompletedTrial` objects containing both the original completed trials - /// and the pending trials with imputed values. If the strategy is `None`, returns - /// only the completed trials (pending trials are ignored). - /// - /// # Imputation Strategies - /// - /// - `None`: Pending trials are ignored (returns only completed trials) - /// - `Mean`: Pending trials get the mean of completed objective values - /// - `Best`: Pending trials get the minimum completed objective value (for minimization) - /// - `Worst`: Pending trials get the maximum completed objective value (for minimization) - /// - `Custom(v)`: Pending trials get the specified value `v` - /// - /// # Examples - /// - /// ```ignore - /// use std::collections::HashMap; - /// use optimizer::sampler::{ - /// ConstantLiarStrategy, MultivariateTpeSampler, CompletedTrial, PendingTrial, - /// }; - /// use optimizer::param::ParamValue; - /// use optimizer::parameter::ParamId; - /// use optimizer::distribution::{Distribution, FloatDistribution}; - /// - /// // Create a sampler with mean imputation - /// let sampler = MultivariateTpeSampler::builder() - /// .constant_liar(ConstantLiarStrategy::Mean) - /// .build() - /// .unwrap(); - /// - /// // Create some completed trials - /// let dist = Distribution::Float(FloatDistribution { - /// low: 0.0, high: 1.0, log_scale: false, step: None, - /// }); - /// let x_id = ParamId::new(); - /// let completed = vec![ - /// CompletedTrial::new( - /// 0, - /// [(x_id, ParamValue::Float(0.2))].into_iter().collect(), - /// [(x_id, dist.clone())].into_iter().collect(), - /// HashMap::new(), - /// 1.0, - /// ), - /// CompletedTrial::new( - /// 1, - /// [(x_id, ParamValue::Float(0.8))].into_iter().collect(), - /// [(x_id, dist.clone())].into_iter().collect(), - /// HashMap::new(), - /// 3.0, - /// ), - /// ]; - /// - /// // Create a pending trial - /// let pending = vec![ - /// PendingTrial::new( - /// 2, - /// [(x_id, ParamValue::Float(0.5))].into_iter().collect(), - /// [(x_id, dist.clone())].into_iter().collect(), - /// HashMap::new(), - /// ), - /// ]; - /// - /// // Impute values - /// let augmented = sampler.impute_pending_trials(&pending, &completed); - /// - /// // Should have 3 trials total - /// assert_eq!(augmented.len(), 3); - /// - /// // The pending trial should have the mean value (1.0 + 3.0) / 2 = 2.0 - /// let imputed = augmented.iter().find(|t| t.id == 2).unwrap(); - /// assert!((imputed.value - 2.0).abs() < f64::EPSILON); - /// ``` - #[must_use] - pub fn impute_pending_trials( - &self, - pending_trials: &[PendingTrial], - completed_trials: &[CompletedTrial], - ) -> Vec { - // Start with a copy of completed trials - let mut result: Vec = completed_trials.to_vec(); - - // If strategy is None or no pending trials, just return completed trials - if matches!(self.constant_liar, ConstantLiarStrategy::None) || pending_trials.is_empty() { - return result; - } - - // Compute the imputation value based on strategy - let imputed_value = self.compute_imputation_value(completed_trials); - - // Convert pending trials to completed trials with imputed values - for pending in pending_trials { - result.push(CompletedTrial::new( - pending.id, - pending.params.clone(), - pending.distributions.clone(), - HashMap::new(), - imputed_value, - )); - } - - result - } - - /// Computes the imputation value based on the constant liar strategy. - /// - /// This is a helper method used by [`impute_pending_trials`](Self::impute_pending_trials). - /// - /// # Arguments - /// - /// * `completed_trials` - The completed trials to compute the imputation value from. - /// - /// # Returns - /// - /// The imputed value based on the strategy. Returns 0.0 if there are no completed - /// trials (except for `Custom` strategy which returns its specified value). - #[allow(clippy::cast_precision_loss)] - fn compute_imputation_value(&self, completed_trials: &[CompletedTrial]) -> f64 { - match self.constant_liar { - ConstantLiarStrategy::None => 0.0, // This case is handled before calling this method - ConstantLiarStrategy::Mean => { - if completed_trials.is_empty() { - 0.0 - } else { - let sum: f64 = completed_trials.iter().map(|t| t.value).sum(); - sum / completed_trials.len() as f64 - } - } - ConstantLiarStrategy::Best => { - // Best means minimum for minimization problems - completed_trials - .iter() - .map(|t| t.value) - .fold(f64::INFINITY, f64::min) - } - ConstantLiarStrategy::Worst => { - // Worst means maximum for minimization problems - completed_trials - .iter() - .map(|t| t.value) - .fold(f64::NEG_INFINITY, f64::max) - } - ConstantLiarStrategy::Custom(v) => v, - } - } - - /// Filters trials to those containing all parameters in the search space. - /// - /// This method is used to identify trials that can be used for multivariate - /// KDE fitting. Only trials that contain ALL parameters in the search space - /// are included, ensuring we can model the joint distribution over all - /// parameters. - /// - /// # Arguments - /// - /// * `history` - All completed trials from the optimization history. - /// * `search_space` - The intersection search space containing parameters that - /// appear in all trials. - /// - /// # Returns - /// - /// A vector of references to trials that contain all parameters in the search space. - /// - /// # Examples - /// - /// ```ignore - /// use std::collections::HashMap; - /// use optimizer::sampler::tpe::MultivariateTpeSampler; - /// use optimizer::sampler::tpe::IntersectionSearchSpace; - /// - /// let sampler = MultivariateTpeSampler::new(); - /// let trials = vec![/* ... completed trials ... */]; - /// let search_space = IntersectionSearchSpace::calculate(&trials); - /// let filtered = sampler.filter_trials(&trials, &search_space); - /// ``` - #[must_use] - pub fn filter_trials<'a>( - &self, - history: &'a [CompletedTrial], - search_space: &HashMap, - ) -> Vec<&'a CompletedTrial> { - history - .iter() - .filter(|trial| { - // Include trial only if it has ALL parameters in the search space - search_space - .keys() - .all(|param_id| trial.params.contains_key(param_id)) - }) - .collect() - } - - /// Splits filtered trials into good and bad groups based on the gamma quantile. - /// - /// The gamma value is computed dynamically using the configured [`GammaStrategy`]. - /// Trials are sorted by objective value (ascending for minimization), and the - /// gamma quantile determines the split point. - /// - /// # Arguments - /// - /// * `trials` - Filtered trials to split (typically from [`filter_trials`](Self::filter_trials)). - /// - /// # Returns - /// - /// A tuple `(good_trials, bad_trials)` where: - /// - `good_trials` contains trials with values below the gamma quantile - /// - `bad_trials` contains trials with values at or above the gamma quantile - /// - /// Both vectors are guaranteed to be non-empty when the input has at least 2 trials. - /// If the input has fewer than 2 trials, both vectors may be empty or one may - /// contain the single trial. - /// - /// # Examples - /// - /// ```ignore - /// use std::collections::HashMap; - /// use optimizer::sampler::tpe::MultivariateTpeSampler; - /// use optimizer::sampler::tpe::IntersectionSearchSpace; - /// - /// let sampler = MultivariateTpeSampler::new(); - /// let trials = vec![/* ... completed trials ... */]; - /// let search_space = IntersectionSearchSpace::calculate(&trials); - /// let filtered = sampler.filter_trials(&trials, &search_space); - /// let (good, bad) = sampler.split_trials(&filtered); - /// - /// // good contains trials with lowest objective values - /// // bad contains trials with higher objective values - /// ``` - #[allow( - clippy::cast_precision_loss, - clippy::cast_possible_truncation, - clippy::cast_sign_loss - )] - #[must_use] - pub fn split_trials<'a>( - &self, - trials: &[&'a CompletedTrial], - ) -> (Vec<&'a CompletedTrial>, Vec<&'a CompletedTrial>) { - if trials.is_empty() { - return (vec![], vec![]); - } - - // Sort trials by objective value (ascending for minimization) - let mut sorted_indices: Vec = (0..trials.len()).collect(); - sorted_indices.sort_by(|&a, &b| { - trials[a] - .value - .partial_cmp(&trials[b].value) - .unwrap_or(core::cmp::Ordering::Equal) - }); - - // Compute gamma using the strategy and clamp to valid range - let gamma = self - .gamma_strategy - .gamma(trials.len()) - .clamp(f64::EPSILON, 1.0 - f64::EPSILON); - - // Calculate the split point (gamma quantile) - // Ensure at least 1 trial in each group if possible - let n_good = ((trials.len() as f64 * gamma).ceil() as usize) - .max(1) - .min(trials.len().saturating_sub(1)); - - // Handle edge case: if we have only 1 trial, put it in good - if trials.len() == 1 { - return (vec![trials[0]], vec![]); - } - - let good: Vec<_> = sorted_indices[..n_good] - .iter() - .map(|&i| trials[i]) - .collect(); - let bad: Vec<_> = sorted_indices[n_good..] - .iter() - .map(|&i| trials[i]) - .collect(); - - (good, bad) - } - - /// Extracts parameter values from trials as a numeric observation matrix. - /// - /// This method converts trial parameter values into a matrix format suitable - /// for multivariate KDE fitting. Each row in the output represents one trial's - /// parameter values in the specified order. - /// - /// # Arguments - /// - /// * `trials` - Trials to extract observations from. - /// * `param_order` - The order of parameters in the output vectors. This ensures - /// consistent column ordering across the observation matrix. - /// - /// # Returns - /// - /// A `Vec>` where: - /// - Outer vec has one entry per trial - /// - Inner vec has one entry per parameter (in `param_order` order) - /// - Float values are used directly - /// - Int values are converted to f64 - /// - Categorical parameters are skipped (not included in output) - /// - /// # Panics - /// - /// This method does not panic. If a parameter is missing from a trial or has - /// an unsupported type (Categorical), it is simply skipped. - /// - /// # Examples - /// - /// ```ignore - /// use std::collections::HashMap; - /// use optimizer::sampler::tpe::MultivariateTpeSampler; - /// use optimizer::parameter::ParamId; - /// - /// let sampler = MultivariateTpeSampler::new(); - /// let trials = vec![/* ... completed trials ... */]; - /// let filtered = sampler.filter_trials(&trials, &search_space); - /// - /// // Extract observations for x and y in that order - /// let x_id = ParamId::new(); - /// let y_id = ParamId::new(); - /// let param_order = vec![x_id, y_id]; - /// let observations = sampler.extract_observations(&filtered, ¶m_order); - /// - /// // observations[i][0] is the x value for trial i - /// // observations[i][1] is the y value for trial i - /// ``` - #[must_use] - #[allow(clippy::cast_precision_loss)] - pub fn extract_observations( - &self, - trials: &[&CompletedTrial], - param_order: &[ParamId], - ) -> Vec> { - trials - .iter() - .map(|trial| { - param_order - .iter() - .filter_map(|param_id| { - trial.params.get(param_id).and_then(|value| match value { - crate::param::ParamValue::Float(f) => Some(*f), - crate::param::ParamValue::Int(i) => Some(*i as f64), - crate::param::ParamValue::Categorical(_) => None, // Skip categorical - }) - }) - .collect() - }) - .collect() - } - - /// 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))`. - /// - /// # Arguments - /// - /// * `good_kde` - The KDE fitted on the "good" trials (low objective values). - /// * `bad_kde` - The KDE fitted on the "bad" trials (high objective values). - /// - /// # Returns - /// - /// A `Vec` representing the selected candidate point in the parameter space. - /// The point is chosen to maximize the expected improvement proxy l(x)/g(x). - /// - /// # Algorithm - /// - /// 1. Generate `n_ei_candidates` samples from `good_kde` - /// 2. For each candidate, compute `log(l(x)) - log(g(x))` - /// 3. Select the candidate with the highest ratio - /// - /// # Edge Cases - /// - /// - If `g(x)` is very small (near zero), the log-space computation handles this - /// gracefully without division issues. - /// - If all candidates have `-inf` log ratios, the first candidate is returned. - /// - /// # Examples - /// - /// ```ignore - /// use optimizer::sampler::tpe::MultivariateTpeSampler; - /// use optimizer::kde::MultivariateKDE; - /// - /// let sampler = MultivariateTpeSampler::builder() - /// .n_ei_candidates(24) - /// .seed(42) - /// .build() - /// .unwrap(); - /// - /// let good_obs = vec![vec![0.1, 0.2], vec![0.2, 0.3], vec![0.15, 0.25]]; - /// let bad_obs = vec![vec![0.8, 0.9], vec![0.7, 0.85], vec![0.9, 0.95]]; - /// - /// let good_kde = MultivariateKDE::new(good_obs).unwrap(); - /// let bad_kde = MultivariateKDE::new(bad_obs).unwrap(); - /// - /// let selected = sampler.select_candidate(&good_kde, &bad_kde); - /// // selected should be a point likely in the "good" region - /// ``` - #[must_use] - #[allow(dead_code)] // Used by tests - pub(crate) fn select_candidate( - &self, - good_kde: &crate::kde::MultivariateKDE, - bad_kde: &crate::kde::MultivariateKDE, - ) -> Vec { - let mut rng = self.rng.lock(); - - // Generate candidates from the good distribution - let candidates: Vec> = (0..self.n_ei_candidates) - .map(|_| good_kde.sample(&mut rng)) - .collect(); - - // Compute log(l(x)) - log(g(x)) for each candidate - // This is equivalent to log(l(x)/g(x)) which we want to maximize - let log_ratios: Vec = candidates - .iter() - .map(|candidate| { - let log_l = good_kde.log_pdf(candidate); - let log_g = bad_kde.log_pdf(candidate); - log_l - log_g - }) - .collect(); - - // Find the candidate with the maximum log ratio - let mut best_idx = 0; - let mut best_ratio = f64::NEG_INFINITY; - - for (idx, &ratio) in log_ratios.iter().enumerate() { - // Handle NaN by treating it as worse than any finite value - if ratio > best_ratio || (best_ratio.is_nan() && !ratio.is_nan()) { - best_ratio = ratio; - best_idx = idx; - } - } - - candidates.into_iter().nth(best_idx).unwrap_or_default() - } - - /// Selects the best candidate using an external RNG. - /// - /// This variant accepts an external RNG, used when the caller already holds the lock. - fn select_candidate_with_rng( - &self, - good_kde: &crate::kde::MultivariateKDE, - bad_kde: &crate::kde::MultivariateKDE, - rng: &mut fastrand::Rng, - ) -> Vec { - // Generate candidates from the good distribution - let candidates: Vec> = (0..self.n_ei_candidates) - .map(|_| good_kde.sample(rng)) - .collect(); - - // Compute log(l(x)) - log(g(x)) for each candidate - let log_ratios: Vec = candidates - .iter() - .map(|candidate| { - let log_l = good_kde.log_pdf(candidate); - let log_g = bad_kde.log_pdf(candidate); - log_l - log_g - }) - .collect(); - - // Find the candidate with the maximum log ratio - let mut best_idx = 0; - let mut best_ratio = f64::NEG_INFINITY; - - for (idx, &ratio) in log_ratios.iter().enumerate() { - if ratio > best_ratio || (best_ratio.is_nan() && !ratio.is_nan()) { - best_ratio = ratio; - best_idx = idx; - } - } - - candidates.into_iter().nth(best_idx).unwrap_or_default() - } - /// Samples all parameters uniformly at random. /// /// This is a fallback method used when multivariate TPE cannot be applied. @@ -897,760 +411,6 @@ impl MultivariateTpeSampler { // Non-grouped mode: use the original single-group logic self.sample_single_group(search_space, history, &mut rng) } - - /// Samples parameters by decomposing the search space into independent groups. - /// - /// When `group=true`, this method analyzes the trial history to identify groups of - /// parameters that always appear together, then samples each group independently - /// using multivariate TPE. This is more efficient when parameters naturally partition - /// into independent subsets (e.g., due to conditional search spaces). - /// - /// # Arguments - /// - /// * `search_space` - The full search space containing all parameters to sample. - /// * `history` - Completed trials from the optimization history. - /// - /// # Returns - /// - /// A `HashMap` mapping parameter names to their sampled values. - fn sample_with_groups( - &self, - search_space: &HashMap, - history: &[CompletedTrial], - ) -> HashMap { - use std::collections::HashSet; - - use super::GroupDecomposedSearchSpace; - - // Decompose the search space into independent parameter groups - let groups = GroupDecomposedSearchSpace::calculate(history); - - let mut result: HashMap = HashMap::new(); - - // Sample each group independently - for group in &groups { - // Build a sub-search space for this group - let group_search_space: HashMap = search_space - .iter() - .filter(|(id, _)| group.contains(id)) - .map(|(id, dist)| (*id, dist.clone())) - .collect(); - - if group_search_space.is_empty() { - continue; - } - - // Filter history to trials that have at least one parameter in this group - let group_history: Vec<&CompletedTrial> = history - .iter() - .filter(|trial| { - trial - .distributions - .keys() - .any(|param_id| group.contains(param_id)) - }) - .collect(); - - // Build completed trials from references for the group - // We need to create a temporary slice for sample_group_internal - let group_history_owned: Vec = - group_history.iter().map(|t| (*t).clone()).collect(); - - // Sample this group using multivariate TPE - let mut rng = self.rng.lock(); - let group_result = - self.sample_single_group(&group_search_space, &group_history_owned, &mut rng); - drop(rng); - - // Merge group results into the main result - for (id, value) in group_result { - result.insert(id, value); - } - } - - // Handle parameters not in any group (sample independently) - let grouped_params: HashSet = groups.iter().flatten().copied().collect(); - let ungrouped_params: HashMap = search_space - .iter() - .filter(|(id, _)| !grouped_params.contains(id) && !result.contains_key(id)) - .map(|(id, dist)| (*id, dist.clone())) - .collect(); - - if !ungrouped_params.is_empty() { - // Sample ungrouped parameters uniformly (no history for them) - let mut rng = self.rng.lock(); - for (id, dist) in &ungrouped_params { - let value = crate::sampler::common::sample_random(&mut rng, dist); - result.insert(*id, value); - } - } - - result - } - - /// Samples parameters as a single group using multivariate TPE. - /// - /// This is the core multivariate TPE sampling logic, used both in non-grouped mode - /// and for sampling individual groups in grouped mode. - /// - /// # Arguments - /// - /// * `search_space` - The search space for this group of parameters. - /// * `history` - Completed trials to use for model fitting. - /// * `rng` - Random number generator (caller must hold lock). - /// - /// # Returns - /// - /// A `HashMap` mapping parameter names to their sampled values. - #[allow(clippy::too_many_lines)] - fn sample_single_group( - &self, - search_space: &HashMap, - history: &[CompletedTrial], - rng: &mut fastrand::Rng, - ) -> HashMap { - use super::IntersectionSearchSpace; - use crate::kde::MultivariateKDE; - - // Early returns for cases requiring random sampling - if history.len() < self.n_startup_trials { - return self.sample_all_uniform(search_space, rng); - } - - let intersection = IntersectionSearchSpace::calculate(history); - if intersection.is_empty() { - return self.sample_all_independent_with_rng(search_space, history, rng); - } - - let filtered = self.filter_trials(history, &intersection); - if filtered.len() < 2 { - return self.sample_all_independent_with_rng(search_space, history, rng); - } - - let (good, bad) = self.split_trials(&filtered); - - // Sample categorical parameters using TPE with l(x)/g(x) ratio - let mut result: HashMap = HashMap::new(); - for (param_id, dist) in &intersection { - if let Distribution::Categorical(d) = dist { - let good_indices = Self::extract_categorical_indices(&good, *param_id); - let bad_indices = Self::extract_categorical_indices(&bad, *param_id); - let idx = - Self::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); - result.insert(*param_id, ParamValue::Categorical(idx)); - } - } - - // Collect continuous parameters - let mut param_order: Vec = intersection - .iter() - .filter(|(_, dist)| !matches!(dist, Distribution::Categorical(_))) - .map(|(id, _)| *id) - .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, - history, - &mut result, - rng, - ); - return result; - } - - param_order.sort_by_key(|id| format!("{id}")); - - // Extract observations and validate - 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 - self.fill_remaining_independent_with_rng( - search_space, - &intersection, - history, - &mut result, - rng, - ); - return result; - }; - - let selected = self.select_candidate_with_rng(&good_kde, &bad_kde, rng); - - // Map selected values to parameter ids - for (idx, param_id) in param_order.iter().enumerate() { - if let Some(dist) = intersection.get(param_id) { - let value = selected[idx]; - let param_value = self.convert_to_param_value(value, dist); - if let Some(pv) = param_value { - result.insert(*param_id, pv); - } - } - } - - // Fill remaining parameters using independent TPE sampling - self.fill_remaining_independent_with_rng( - search_space, - &intersection, - history, - &mut result, - rng, - ); - result - } - - /// Converts a raw f64 value to a `ParamValue` based on the distribution. - #[allow(clippy::unused_self)] - fn convert_to_param_value(&self, value: f64, dist: &Distribution) -> Option { - match dist { - Distribution::Float(d) => { - let clamped = value.clamp(d.low, d.high); - let stepped = if let Some(step) = d.step { - let steps = ((clamped - d.low) / step).round(); - (d.low + steps * step).clamp(d.low, d.high) - } else { - clamped - }; - Some(ParamValue::Float(stepped)) - } - Distribution::Int(d) => { - #[allow(clippy::cast_possible_truncation)] - let int_value = value.round() as i64; - let clamped = int_value.clamp(d.low, d.high); - let stepped = if let Some(step) = d.step { - let steps = (clamped - d.low) / step; - (d.low + steps * step).clamp(d.low, d.high) - } else { - clamped - }; - Some(ParamValue::Int(stepped)) - } - Distribution::Categorical(_) => None, - } - } - - /// 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)] - 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. - fn fill_remaining_independent_with_rng( - &self, - search_space: &HashMap, - _intersection: &HashMap, - history: &[CompletedTrial], - result: &mut HashMap, - rng: &mut fastrand::Rng, - ) { - // 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::>()); - - for (param_id, dist) in missing_params { - let value = - self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng); - result.insert(*param_id, value); - } - } - - /// Samples a single parameter using independent TPE. - /// - /// 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)] - fn sample_independent_tpe( - &self, - param_id: ParamId, - distribution: &Distribution, - good_trials: &[&CompletedTrial], - bad_trials: &[&CompletedTrial], - rng: &mut fastrand::Rng, - ) -> ParamValue { - 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 = self.sample_tpe_float( - d.low, - d.high, - d.log_scale, - d.step, - good_values, - bad_values, - rng, - ); - ParamValue::Float(value) - } - 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 = self.sample_tpe_int( - d.low, - d.high, - d.log_scale, - d.step, - good_values, - bad_values, - rng, - ); - ParamValue::Int(value) - } - 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 = - Self::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng); - ParamValue::Categorical(idx) - } - } - } - - /// Samples using TPE for float distributions. - #[allow(clippy::too_many_arguments)] - fn sample_tpe_float( - &self, - low: f64, - high: f64, - log_scale: bool, - step: Option, - good_values: Vec, - bad_values: Vec, - rng: &mut fastrand::Rng, - ) -> f64 { - use crate::kde::KernelDensityEstimator; - - // 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 = KernelDensityEstimator::new(good_internal); - let g_kde = KernelDensityEstimator::new(bad_internal); - - // If KDE construction fails, fall back to uniform sampling - let (Ok(l_kde), Ok(g_kde)) = (l_kde, g_kde) else { - return rng_util::f64_range(rng, low, high); - }; - - // Generate candidates from l(x) and select the one with best l(x)/g(x) ratio - let mut best_candidate = internal_low; - let mut best_ratio = f64::NEG_INFINITY; - - for _ in 0..self.n_ei_candidates { - let candidate = l_kde.sample(rng); - - // Clamp to bounds - let candidate = candidate.clamp(internal_low, internal_high); - - let l_density = l_kde.pdf(candidate); - let g_density = g_kde.pdf(candidate); - - // Compute l(x)/g(x) ratio, handling zero density - let ratio = if g_density < f64::EPSILON { - if l_density > f64::EPSILON { - f64::INFINITY - } else { - 0.0 - } - } else { - l_density / g_density - }; - - if ratio > best_ratio { - best_ratio = ratio; - best_candidate = candidate; - } - } - - // Transform back from internal space - let mut value = if log_scale { - best_candidate.exp() - } else { - best_candidate - }; - - // Apply step constraint if present - if let Some(step) = step { - let k = ((value - low) / step).round(); - value = low + k * step; - } - - // Ensure value is within bounds - value.clamp(low, high) - } - - /// Samples using TPE for integer distributions. - #[allow( - clippy::too_many_arguments, - clippy::cast_precision_loss, - clippy::cast_possible_truncation - )] - fn sample_tpe_int( - &self, - low: i64, - high: i64, - log_scale: bool, - step: Option, - good_values: Vec, - bad_values: Vec, - rng: &mut fastrand::Rng, - ) -> i64 { - // Convert to floats for KDE - let good_floats: Vec = good_values.into_iter().map(|v| v as f64).collect(); - let bad_floats: Vec = bad_values.into_iter().map(|v| v as f64).collect(); - - // Use float TPE sampling - let float_value = self.sample_tpe_float( - low as f64, - high as f64, - log_scale, - step.map(|s| s as f64), - good_floats, - bad_floats, - rng, - ); - - // Round to nearest integer - let int_value = float_value.round() as i64; - - // Apply step constraint if present - let int_value = if let Some(step) = step { - let k = ((int_value - low) as f64 / step as f64).round() as i64; - low + k * step - } else { - int_value - }; - - // Ensure value is within bounds - int_value.clamp(low, high) - } - - /// Samples 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 - fn sample_all_independent( - &self, - search_space: &HashMap, - history: &[CompletedTrial], - ) -> HashMap { - // Split trials for independent sampling - let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); - - let mut rng = self.rng.lock(); - let mut result = HashMap::new(); - - for (param_id, dist) in search_space { - let value = - self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, &mut rng); - result.insert(*param_id, value); - } - - result - } - - /// 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. - fn sample_all_independent_with_rng( - &self, - search_space: &HashMap, - history: &[CompletedTrial], - rng: &mut fastrand::Rng, - ) -> HashMap { - // Split trials for independent sampling - let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::>()); - - let mut result = HashMap::new(); - - for (param_id, dist) in search_space { - let value = - self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng); - result.insert(*param_id, value); - } - - result - } - - /// Samples using TPE for categorical distributions. - /// - /// This method computes kernel-weighted category probabilities based on - /// the good and bad trial groups, then samples proportionally to the - /// `l(x)/g(x)` ratio for each category. - /// - /// # Arguments - /// - /// * `n_choices` - The number of categories in the categorical distribution. - /// * `good_indices` - Category indices from the "good" trials. - /// * `bad_indices` - Category indices from the "bad" trials. - /// * `rng` - Random number generator for sampling. - /// - /// # Returns - /// - /// The selected category index. - /// - /// # Algorithm - /// - /// 1. Count occurrences of each category in good and bad groups - /// 2. Apply Laplace smoothing (add 1 to each count) to avoid zero probabilities - /// 3. Compute `l(x)/g(x)` ratio for each category - /// 4. Sample proportionally to the computed weights - #[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; - } - } - - // Add 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 - } - - /// Extracts categorical indices from trials for a specific parameter. - /// - /// # Arguments - /// - /// * `trials` - Trials to extract from. - /// * `param_name` - The name of the categorical parameter. - /// - /// # Returns - /// - /// A vector of category indices from the trials. - fn extract_categorical_indices(trials: &[&CompletedTrial], param_id: ParamId) -> Vec { - trials - .iter() - .filter_map(|trial| { - trial.params.get(¶m_id).and_then(|value| { - if let ParamValue::Categorical(idx) = value { - Some(*idx) - } else { - None - } - }) - }) - .collect() - } } impl Default for MultivariateTpeSampler { @@ -4883,7 +3643,7 @@ mod tests { // Sample many times and check bias toward good category let mut counts = [0usize; 3]; for _ in 0..1000 { - let idx = MultivariateTpeSampler::sample_tpe_categorical( + let idx = crate::sampler::tpe::common::sample_tpe_categorical( 3, &good_indices, &bad_indices, @@ -4917,7 +3677,7 @@ mod tests { let mut sampled_two = false; for _ in 0..1000 { - let idx = MultivariateTpeSampler::sample_tpe_categorical( + let idx = crate::sampler::tpe::common::sample_tpe_categorical( 3, &good_indices, &bad_indices, @@ -4945,7 +3705,7 @@ mod tests { let mut counts = [0usize; 3]; for _ in 0..1000 { - let idx = MultivariateTpeSampler::sample_tpe_categorical( + let idx = crate::sampler::tpe::common::sample_tpe_categorical( 3, &good_indices, &bad_indices, @@ -4970,7 +3730,7 @@ mod tests { // All samples should be valid indices for _ in 0..100 { - let idx = MultivariateTpeSampler::sample_tpe_categorical( + let idx = crate::sampler::tpe::common::sample_tpe_categorical( n_choices, &good_indices, &bad_indices, diff --git a/src/sampler/tpe/multivariate/trials.rs b/src/sampler/tpe/multivariate/trials.rs new file mode 100644 index 0000000..e1d31c9 --- /dev/null +++ b/src/sampler/tpe/multivariate/trials.rs @@ -0,0 +1,220 @@ +//! Trial processing for the multivariate TPE sampler. +//! +//! Contains constant-liar imputation, trial filtering, good/bad splitting, +//! observation extraction, and categorical index extraction. + +use std::collections::HashMap; + +use crate::distribution::Distribution; +use crate::param::ParamValue; +use crate::parameter::ParamId; +use crate::sampler::{CompletedTrial, PendingTrial}; + +use super::{ConstantLiarStrategy, MultivariateTpeSampler}; + +impl MultivariateTpeSampler { + /// Imputes objective values for pending trials based on the constant liar strategy. + /// + /// In parallel optimization, multiple trials may be running simultaneously. This method + /// assigns "lie" values to pending trials so they can be included in the model fitting, + /// which helps avoid redundant exploration of the same region. + /// + /// # Arguments + /// + /// * `pending_trials` - Trials that are currently running and have no objective value yet. + /// * `completed_trials` - Trials that have completed and have objective values. + /// + /// # Returns + /// + /// A vector of `CompletedTrial` objects containing both the original completed trials + /// and the pending trials with imputed values. If the strategy is `None`, returns + /// only the completed trials (pending trials are ignored). + #[must_use] + pub fn impute_pending_trials( + &self, + pending_trials: &[PendingTrial], + completed_trials: &[CompletedTrial], + ) -> Vec { + // Start with a copy of completed trials + let mut result: Vec = completed_trials.to_vec(); + + // If strategy is None or no pending trials, just return completed trials + if matches!(self.constant_liar, ConstantLiarStrategy::None) || pending_trials.is_empty() { + return result; + } + + // Compute the imputation value based on strategy + let imputed_value = self.compute_imputation_value(completed_trials); + + // Convert pending trials to completed trials with imputed values + for pending in pending_trials { + result.push(CompletedTrial::new( + pending.id, + pending.params.clone(), + pending.distributions.clone(), + HashMap::new(), + imputed_value, + )); + } + + result + } + + /// Computes the imputation value based on the constant liar strategy. + /// + /// This is a helper method used by [`impute_pending_trials`](Self::impute_pending_trials). + #[allow(clippy::cast_precision_loss)] + pub(crate) fn compute_imputation_value(&self, completed_trials: &[CompletedTrial]) -> f64 { + match self.constant_liar { + ConstantLiarStrategy::None => 0.0, // This case is handled before calling this method + ConstantLiarStrategy::Mean => { + if completed_trials.is_empty() { + 0.0 + } else { + let sum: f64 = completed_trials.iter().map(|t| t.value).sum(); + sum / completed_trials.len() as f64 + } + } + ConstantLiarStrategy::Best => { + // Best means minimum for minimization problems + completed_trials + .iter() + .map(|t| t.value) + .fold(f64::INFINITY, f64::min) + } + ConstantLiarStrategy::Worst => { + // Worst means maximum for minimization problems + completed_trials + .iter() + .map(|t| t.value) + .fold(f64::NEG_INFINITY, f64::max) + } + ConstantLiarStrategy::Custom(v) => v, + } + } + + /// Filters trials to those containing all parameters in the search space. + /// + /// Only trials that contain ALL parameters in the search space are included, + /// ensuring we can model the joint distribution over all parameters. + #[must_use] + pub fn filter_trials<'a>( + &self, + history: &'a [CompletedTrial], + search_space: &HashMap, + ) -> Vec<&'a CompletedTrial> { + history + .iter() + .filter(|trial| { + // Include trial only if it has ALL parameters in the search space + search_space + .keys() + .all(|param_id| trial.params.contains_key(param_id)) + }) + .collect() + } + + /// Splits filtered trials into good and bad groups based on the gamma quantile. + /// + /// The gamma value is computed dynamically using the configured [`GammaStrategy`]. + /// Trials are sorted by objective value (ascending for minimization), and the + /// gamma quantile determines the split point. + #[allow( + clippy::cast_precision_loss, + clippy::cast_possible_truncation, + clippy::cast_sign_loss + )] + #[must_use] + pub fn split_trials<'a>( + &self, + trials: &[&'a CompletedTrial], + ) -> (Vec<&'a CompletedTrial>, Vec<&'a CompletedTrial>) { + if trials.is_empty() { + return (vec![], vec![]); + } + + // Sort trials by objective value (ascending for minimization) + let mut sorted_indices: Vec = (0..trials.len()).collect(); + sorted_indices.sort_by(|&a, &b| { + trials[a] + .value + .partial_cmp(&trials[b].value) + .unwrap_or(core::cmp::Ordering::Equal) + }); + + // Compute gamma using the strategy and clamp to valid range + let gamma = self + .gamma_strategy + .gamma(trials.len()) + .clamp(f64::EPSILON, 1.0 - f64::EPSILON); + + // Calculate the split point (gamma quantile) + // Ensure at least 1 trial in each group if possible + let n_good = ((trials.len() as f64 * gamma).ceil() as usize) + .max(1) + .min(trials.len().saturating_sub(1)); + + // Handle edge case: if we have only 1 trial, put it in good + if trials.len() == 1 { + return (vec![trials[0]], vec![]); + } + + let good: Vec<_> = sorted_indices[..n_good] + .iter() + .map(|&i| trials[i]) + .collect(); + let bad: Vec<_> = sorted_indices[n_good..] + .iter() + .map(|&i| trials[i]) + .collect(); + + (good, bad) + } + + /// Extracts parameter values from trials as a numeric observation matrix. + /// + /// Each row in the output represents one trial's parameter values in the specified order. + /// Categorical parameters are skipped. + #[must_use] + #[allow(clippy::cast_precision_loss)] + pub fn extract_observations( + &self, + trials: &[&CompletedTrial], + param_order: &[ParamId], + ) -> Vec> { + trials + .iter() + .map(|trial| { + param_order + .iter() + .filter_map(|param_id| { + trial.params.get(param_id).and_then(|value| match value { + crate::param::ParamValue::Float(f) => Some(*f), + crate::param::ParamValue::Int(i) => Some(*i as f64), + crate::param::ParamValue::Categorical(_) => None, // Skip categorical + }) + }) + .collect() + }) + .collect() + } + + /// Extracts categorical indices from trials for a specific parameter. + pub(crate) fn extract_categorical_indices( + trials: &[&CompletedTrial], + param_id: ParamId, + ) -> Vec { + trials + .iter() + .filter_map(|trial| { + trial.params.get(¶m_id).and_then(|value| { + if let ParamValue::Categorical(idx) = value { + Some(*idx) + } else { + None + } + }) + }) + .collect() + } +} diff --git a/src/study/persistence.rs b/src/study/persistence.rs index 78b94cf..6e6ffc9 100644 --- a/src/study/persistence.rs +++ b/src/study/persistence.rs @@ -1,8 +1,12 @@ +#[cfg(feature = "serde")] use std::collections::HashMap; +#[cfg(feature = "serde")] use crate::sampler::CompletedTrial; +#[cfg(feature = "serde")] use crate::types::Direction; +#[cfg(feature = "serde")] use super::Study; /// A serializable snapshot of a study's state.