refactor(tpe): split multivariate sampler into focused submodules
- Extract engine.rs (core sampling logic) and trials.rs (trial processing) - Deduplicate TPE sampling functions by delegating to tpe::common - Gate persistence.rs imports with #[cfg(feature = "serde")]
This commit is contained in:
Manuel Raimann
@@ -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<ParamId, Distribution>,
|
||||
history: &[CompletedTrial],
|
||||
) -> HashMap<ParamId, ParamValue> {
|
||||
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<ParamId, ParamValue> = HashMap::new();
|
||||
|
||||
// Sample each group independently
|
||||
for group in &groups {
|
||||
// Build a sub-search space for this group
|
||||
let group_search_space: HashMap<ParamId, Distribution> = 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<CompletedTrial> =
|
||||
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<ParamId> = groups.iter().flatten().copied().collect();
|
||||
let ungrouped_params: HashMap<ParamId, Distribution> = 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<ParamId, Distribution>,
|
||||
history: &[CompletedTrial],
|
||||
rng: &mut fastrand::Rng,
|
||||
) -> HashMap<ParamId, ParamValue> {
|
||||
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<ParamId, ParamValue> = 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<ParamId> = 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<ParamValue> {
|
||||
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<f64> {
|
||||
let mut rng = self.rng.lock();
|
||||
|
||||
// Generate candidates from the good distribution
|
||||
let candidates: Vec<Vec<f64>> = (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<f64> = 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<f64> {
|
||||
// Generate candidates from the good distribution
|
||||
let candidates: Vec<Vec<f64>> = (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<f64> = 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<ParamId, Distribution>,
|
||||
_intersection: &HashMap<ParamId, Distribution>,
|
||||
history: &[CompletedTrial],
|
||||
result: &mut HashMap<ParamId, ParamValue>,
|
||||
) {
|
||||
// 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::<Vec<_>>());
|
||||
|
||||
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<ParamId, Distribution>,
|
||||
_intersection: &HashMap<ParamId, Distribution>,
|
||||
history: &[CompletedTrial],
|
||||
result: &mut HashMap<ParamId, ParamValue>,
|
||||
rng: &mut fastrand::Rng,
|
||||
) {
|
||||
// Identify parameters not in result (and not in intersection)
|
||||
let missing_params: Vec<(&ParamId, &Distribution)> = search_space
|
||||
.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::<Vec<_>>());
|
||||
|
||||
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<ParamId, Distribution>,
|
||||
history: &[CompletedTrial],
|
||||
) -> HashMap<ParamId, ParamValue> {
|
||||
// Split trials for independent sampling
|
||||
let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::<Vec<_>>());
|
||||
|
||||
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<ParamId, Distribution>,
|
||||
history: &[CompletedTrial],
|
||||
rng: &mut fastrand::Rng,
|
||||
) -> HashMap<ParamId, ParamValue> {
|
||||
// Split trials for independent sampling
|
||||
let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::<Vec<_>>());
|
||||
|
||||
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<f64> = 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<f64> = 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<i64> = 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<i64> = 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<usize> = 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<usize> = 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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<CompletedTrial> {
|
||||
// Start with a copy of completed trials
|
||||
let mut result: Vec<CompletedTrial> = 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<ParamId, Distribution>,
|
||||
) -> 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<usize> = (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<Vec<f64>> {
|
||||
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<usize> {
|
||||
trials
|
||||
.iter()
|
||||
.filter_map(|trial| {
|
||||
trial.params.get(¶m_id).and_then(|value| {
|
||||
if let ParamValue::Categorical(idx) = value {
|
||||
Some(*idx)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user