Manuel Raimann
1128 lines
36 KiB
Rust
1128 lines
36 KiB
Rust
//! Search space intersection utilities for multivariate TPE sampling.
|
|
//!
|
|
//! This module provides utilities for computing the intersection of search spaces
|
|
//! across completed trials, which is necessary for multivariate TPE sampling where
|
|
//! we need to identify parameters that appear in all trials.
|
|
//!
|
|
//! It also provides utilities for decomposing the search space into independent
|
|
//! groups based on parameter co-occurrence patterns, which enables more efficient
|
|
//! multivariate modeling when parameters naturally partition into independent subsets.
|
|
|
|
use std::collections::{HashMap, HashSet, VecDeque};
|
|
|
|
use crate::distribution::Distribution;
|
|
use crate::parameter::ParamId;
|
|
use crate::sampler::CompletedTrial;
|
|
|
|
/// Computes the intersection of search spaces across completed trials.
|
|
///
|
|
/// In multivariate TPE sampling, we need to model the joint distribution of parameters.
|
|
/// However, trials may have different sets of parameters due to dynamic search spaces
|
|
/// (e.g., conditional parameters). The intersection search space contains only those
|
|
/// parameters that appear in ALL completed trials, allowing us to fit a joint KDE.
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```ignore
|
|
/// use std::collections::HashMap;
|
|
/// use optimizer::distribution::{Distribution, FloatDistribution};
|
|
/// use optimizer::param::ParamValue;
|
|
/// use optimizer::sampler::CompletedTrial;
|
|
/// use optimizer::sampler::tpe::IntersectionSearchSpace;
|
|
///
|
|
/// // Create two trials with overlapping parameters
|
|
/// let dist_x = Distribution::Float(FloatDistribution {
|
|
/// low: 0.0, high: 1.0, log_scale: false, step: None,
|
|
/// });
|
|
/// let dist_y = Distribution::Float(FloatDistribution {
|
|
/// low: 0.0, high: 1.0, log_scale: false, step: None,
|
|
/// });
|
|
///
|
|
/// let mut params1 = HashMap::new();
|
|
/// params1.insert("x".to_string(), ParamValue::Float(0.5));
|
|
/// params1.insert("y".to_string(), ParamValue::Float(0.3));
|
|
/// let mut dists1 = HashMap::new();
|
|
/// dists1.insert("x".to_string(), dist_x.clone());
|
|
/// dists1.insert("y".to_string(), dist_y.clone());
|
|
/// let trial1 = CompletedTrial::new(0, params1, dists1, 1.0);
|
|
///
|
|
/// let mut params2 = HashMap::new();
|
|
/// params2.insert("x".to_string(), ParamValue::Float(0.7));
|
|
/// // Note: trial2 doesn't have "y"
|
|
/// let mut dists2 = HashMap::new();
|
|
/// dists2.insert("x".to_string(), dist_x.clone());
|
|
/// let trial2 = CompletedTrial::new(1, params2, dists2, 0.5);
|
|
///
|
|
/// let trials = vec![trial1, trial2];
|
|
/// let intersection = IntersectionSearchSpace::calculate(&trials);
|
|
///
|
|
/// // Only "x" appears in both trials
|
|
/// assert!(intersection.contains_key("x"));
|
|
/// assert!(!intersection.contains_key("y"));
|
|
/// ```
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct IntersectionSearchSpace;
|
|
|
|
impl IntersectionSearchSpace {
|
|
/// Calculates the intersection of search spaces across all completed trials.
|
|
///
|
|
/// This method identifies parameters that appear in ALL completed trials and returns
|
|
/// a mapping from parameter names to their distributions. If any parameter has
|
|
/// different distributions across trials (e.g., different bounds), the first
|
|
/// encountered distribution is used.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `trials` - A slice of completed trials to compute the intersection from.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// A `HashMap` mapping parameter names to their distributions, containing only
|
|
/// parameters that appear in every trial. Returns an empty map if `trials` is empty.
|
|
///
|
|
/// # Notes
|
|
///
|
|
/// - Parameters must appear in ALL trials to be included in the intersection.
|
|
/// - If a parameter has different distributions in different trials, the distribution
|
|
/// from the first trial containing that parameter is used.
|
|
/// - For dynamic search spaces where not all parameters are sampled in every trial,
|
|
/// this helps identify the "stable" set of parameters that can be modeled jointly.
|
|
#[must_use]
|
|
pub fn calculate(trials: &[CompletedTrial]) -> HashMap<ParamId, Distribution> {
|
|
if trials.is_empty() {
|
|
return HashMap::new();
|
|
}
|
|
|
|
// Get parameter ids from the first trial as the initial candidate set
|
|
let first_trial = &trials[0];
|
|
let mut candidate_params: HashSet<ParamId> =
|
|
first_trial.distributions.keys().copied().collect();
|
|
|
|
// Intersect with parameter sets from all other trials
|
|
for trial in trials.iter().skip(1) {
|
|
let trial_params: HashSet<ParamId> = trial.distributions.keys().copied().collect();
|
|
candidate_params.retain(|param| trial_params.contains(param));
|
|
}
|
|
|
|
// Build the result map using distributions from the first trial
|
|
// that contains each parameter
|
|
let mut result = HashMap::new();
|
|
for param_id in candidate_params {
|
|
// Find the first trial that has this parameter and use its distribution
|
|
for trial in trials {
|
|
if let Some(dist) = trial.distributions.get(¶m_id) {
|
|
result.insert(param_id, dist.clone());
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
result
|
|
}
|
|
}
|
|
|
|
/// Decomposes the search space into independent parameter groups based on co-occurrence.
|
|
///
|
|
/// In multivariate TPE, we model the joint distribution of parameters. However, when
|
|
/// certain parameters never appear together (e.g., they are in different branches of
|
|
/// a conditional), modeling them jointly is wasteful and can hurt performance.
|
|
///
|
|
/// `GroupDecomposedSearchSpace` analyzes trial history to identify groups of parameters
|
|
/// that always appear together. By building a co-occurrence graph (where edges connect
|
|
/// parameters that appear in the same trial) and finding its connected components, we
|
|
/// can partition parameters into independent groups that can be sampled separately.
|
|
///
|
|
/// # Example
|
|
///
|
|
/// ```ignore
|
|
/// use std::collections::HashMap;
|
|
/// use optimizer::distribution::{Distribution, FloatDistribution};
|
|
/// use optimizer::param::ParamValue;
|
|
/// use optimizer::sampler::CompletedTrial;
|
|
/// use optimizer::sampler::tpe::GroupDecomposedSearchSpace;
|
|
///
|
|
/// // Trials with two independent parameter groups:
|
|
/// // Group 1: "x", "y" always appear together
|
|
/// // Group 2: "a", "b" always appear together
|
|
/// // But "x"/"y" never appear with "a"/"b"
|
|
///
|
|
/// let dist = Distribution::Float(FloatDistribution {
|
|
/// low: 0.0, high: 1.0, log_scale: false, step: None,
|
|
/// });
|
|
///
|
|
/// let mut params1 = HashMap::new();
|
|
/// params1.insert("x".to_string(), ParamValue::Float(0.1));
|
|
/// params1.insert("y".to_string(), ParamValue::Float(0.2));
|
|
/// let mut dists1 = HashMap::new();
|
|
/// dists1.insert("x".to_string(), dist.clone());
|
|
/// dists1.insert("y".to_string(), dist.clone());
|
|
/// let trial1 = CompletedTrial::new(0, params1, dists1, 1.0);
|
|
///
|
|
/// let mut params2 = HashMap::new();
|
|
/// params2.insert("a".to_string(), ParamValue::Float(0.3));
|
|
/// params2.insert("b".to_string(), ParamValue::Float(0.4));
|
|
/// let mut dists2 = HashMap::new();
|
|
/// dists2.insert("a".to_string(), dist.clone());
|
|
/// dists2.insert("b".to_string(), dist.clone());
|
|
/// let trial2 = CompletedTrial::new(1, params2, dists2, 0.5);
|
|
///
|
|
/// let trials = vec![trial1, trial2];
|
|
/// let groups = GroupDecomposedSearchSpace::calculate(&trials);
|
|
///
|
|
/// // Should produce two groups: {"x", "y"} and {"a", "b"}
|
|
/// assert_eq!(groups.len(), 2);
|
|
/// ```
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct GroupDecomposedSearchSpace;
|
|
|
|
impl GroupDecomposedSearchSpace {
|
|
/// Calculates parameter groups based on co-occurrence in completed trials.
|
|
///
|
|
/// This method builds a co-occurrence graph where:
|
|
/// - Each parameter is a node
|
|
/// - An edge exists between two parameters if they appear in the same trial
|
|
///
|
|
/// Then finds connected components to partition parameters into independent groups.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `trials` - A slice of completed trials to analyze.
|
|
///
|
|
/// # Returns
|
|
///
|
|
/// A `Vec<HashSet<ParamId>>` where each `HashSet` represents an independent group
|
|
/// of parameters that co-occur. Parameters within the same group have appeared
|
|
/// together in at least one trial (directly or transitively). Parameters in
|
|
/// different groups have never appeared in the same trial.
|
|
///
|
|
/// Returns an empty vector if `trials` is empty.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// This function does not panic under normal circumstances. The internal
|
|
/// `.expect()` calls are guarded by the algorithm's invariants.
|
|
///
|
|
/// # Notes
|
|
///
|
|
/// - Single-parameter groups are included in the output.
|
|
/// - The order of groups in the output vector is not guaranteed.
|
|
/// - This is useful for `MultivariateTpeSampler` when `group=true` to sample
|
|
/// independent parameter groups separately.
|
|
#[must_use]
|
|
pub fn calculate(trials: &[CompletedTrial]) -> Vec<HashSet<ParamId>> {
|
|
if trials.is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
// Collect all unique parameter ids
|
|
let mut all_params: HashSet<ParamId> = HashSet::new();
|
|
for trial in trials {
|
|
for ¶m_id in trial.distributions.keys() {
|
|
all_params.insert(param_id);
|
|
}
|
|
}
|
|
|
|
if all_params.is_empty() {
|
|
return Vec::new();
|
|
}
|
|
|
|
// Build adjacency list for co-occurrence graph
|
|
// Two parameters are adjacent if they appear in the same trial
|
|
let mut adjacency: HashMap<ParamId, HashSet<ParamId>> = HashMap::new();
|
|
for ¶m in &all_params {
|
|
adjacency.insert(param, HashSet::new());
|
|
}
|
|
|
|
for trial in trials {
|
|
let trial_params: Vec<ParamId> = trial.distributions.keys().copied().collect();
|
|
// Connect all pairs of parameters in this trial
|
|
for (i, ¶m1) in trial_params.iter().enumerate() {
|
|
for ¶m2 in trial_params.iter().skip(i + 1) {
|
|
adjacency
|
|
.get_mut(¶m1)
|
|
.expect("param should exist in adjacency map")
|
|
.insert(param2);
|
|
adjacency
|
|
.get_mut(¶m2)
|
|
.expect("param should exist in adjacency map")
|
|
.insert(param1);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Find connected components using BFS
|
|
let mut visited: HashSet<ParamId> = HashSet::new();
|
|
let mut groups: Vec<HashSet<ParamId>> = Vec::new();
|
|
|
|
for ¶m in &all_params {
|
|
if visited.contains(¶m) {
|
|
continue;
|
|
}
|
|
|
|
// BFS to find all parameters in this component
|
|
let mut component: HashSet<ParamId> = HashSet::new();
|
|
let mut queue: VecDeque<ParamId> = VecDeque::new();
|
|
queue.push_back(param);
|
|
visited.insert(param);
|
|
|
|
while let Some(current) = queue.pop_front() {
|
|
component.insert(current);
|
|
|
|
if let Some(neighbors) = adjacency.get(¤t) {
|
|
for &neighbor in neighbors {
|
|
if !visited.contains(&neighbor) {
|
|
visited.insert(neighbor);
|
|
queue.push_back(neighbor);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
groups.push(component);
|
|
}
|
|
|
|
groups
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
|
|
use crate::param::ParamValue;
|
|
use crate::parameter::ParamId;
|
|
|
|
fn create_trial(
|
|
id: u64,
|
|
params: Vec<(ParamId, ParamValue, Distribution)>,
|
|
value: f64,
|
|
) -> CompletedTrial {
|
|
let mut param_map = HashMap::new();
|
|
let mut dist_map = HashMap::new();
|
|
for (param_id, pv, dist) in params {
|
|
param_map.insert(param_id, pv);
|
|
dist_map.insert(param_id, dist);
|
|
}
|
|
CompletedTrial::new(id, param_map, dist_map, HashMap::new(), value)
|
|
}
|
|
|
|
#[test]
|
|
fn test_empty_trials() {
|
|
let trials: Vec<CompletedTrial> = vec![];
|
|
let result = IntersectionSearchSpace::calculate(&trials);
|
|
assert!(result.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_single_trial() {
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let dist_x = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_y = Distribution::Int(IntDistribution {
|
|
low: 1,
|
|
high: 10,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let trial = create_trial(
|
|
0,
|
|
vec![
|
|
(x_id, ParamValue::Float(0.5), dist_x.clone()),
|
|
(y_id, ParamValue::Int(5), dist_y.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
let result = IntersectionSearchSpace::calculate(&[trial]);
|
|
assert_eq!(result.len(), 2);
|
|
assert!(result.contains_key(&x_id));
|
|
assert!(result.contains_key(&y_id));
|
|
assert_eq!(result.get(&x_id), Some(&dist_x));
|
|
assert_eq!(result.get(&y_id), Some(&dist_y));
|
|
}
|
|
|
|
#[test]
|
|
fn test_all_trials_same_params() {
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let dist_x = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_y = Distribution::Float(FloatDistribution {
|
|
low: -1.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let trials: Vec<CompletedTrial> = (0..5)
|
|
.map(|i| {
|
|
#[allow(clippy::cast_precision_loss)]
|
|
let val = i as f64 * 0.1;
|
|
create_trial(
|
|
i,
|
|
vec![
|
|
(x_id, ParamValue::Float(val), dist_x.clone()),
|
|
(y_id, ParamValue::Float(val - 0.5), dist_y.clone()),
|
|
],
|
|
val * val,
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
let result = IntersectionSearchSpace::calculate(&trials);
|
|
assert_eq!(result.len(), 2);
|
|
assert!(result.contains_key(&x_id));
|
|
assert!(result.contains_key(&y_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_partial_overlap() {
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let z_id = ParamId::new();
|
|
let dist_x = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_y = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_z = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
// Trial 1: x, y
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(x_id, ParamValue::Float(0.5), dist_x.clone()),
|
|
(y_id, ParamValue::Float(0.3), dist_y.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
// Trial 2: x, z (no y)
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![
|
|
(x_id, ParamValue::Float(0.7), dist_x.clone()),
|
|
(z_id, ParamValue::Float(0.2), dist_z.clone()),
|
|
],
|
|
0.5,
|
|
);
|
|
|
|
// Trial 3: x, y, z (has all)
|
|
let trial3 = create_trial(
|
|
2,
|
|
vec![
|
|
(x_id, ParamValue::Float(0.6), dist_x.clone()),
|
|
(y_id, ParamValue::Float(0.4), dist_y.clone()),
|
|
(z_id, ParamValue::Float(0.1), dist_z.clone()),
|
|
],
|
|
0.8,
|
|
);
|
|
|
|
let result = IntersectionSearchSpace::calculate(&[trial1, trial2, trial3]);
|
|
|
|
// Only x appears in all three trials
|
|
assert_eq!(result.len(), 1);
|
|
assert!(result.contains_key(&x_id));
|
|
assert!(!result.contains_key(&y_id));
|
|
assert!(!result.contains_key(&z_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_no_common_params() {
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let dist_x = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_y = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
// Trial 1: only x
|
|
let trial1 = create_trial(0, vec![(x_id, ParamValue::Float(0.5), dist_x.clone())], 1.0);
|
|
|
|
// Trial 2: only y
|
|
let trial2 = create_trial(1, vec![(y_id, ParamValue::Float(0.3), dist_y.clone())], 0.5);
|
|
|
|
let result = IntersectionSearchSpace::calculate(&[trial1, trial2]);
|
|
|
|
// No common parameters
|
|
assert!(result.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_mixed_distribution_types() {
|
|
let lr_id = ParamId::new();
|
|
let n_layers_id = ParamId::new();
|
|
let optimizer_id = ParamId::new();
|
|
let dist_float = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_int = Distribution::Int(IntDistribution {
|
|
low: 1,
|
|
high: 100,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_cat = Distribution::Categorical(CategoricalDistribution { n_choices: 3 });
|
|
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(lr_id, ParamValue::Float(0.01), dist_float.clone()),
|
|
(n_layers_id, ParamValue::Int(3), dist_int.clone()),
|
|
(optimizer_id, ParamValue::Categorical(0), dist_cat.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![
|
|
(lr_id, ParamValue::Float(0.001), dist_float.clone()),
|
|
(n_layers_id, ParamValue::Int(5), dist_int.clone()),
|
|
(optimizer_id, ParamValue::Categorical(1), dist_cat.clone()),
|
|
],
|
|
0.8,
|
|
);
|
|
|
|
let result = IntersectionSearchSpace::calculate(&[trial1, trial2]);
|
|
|
|
assert_eq!(result.len(), 3);
|
|
assert!(matches!(result.get(&lr_id), Some(Distribution::Float(_))));
|
|
assert!(matches!(
|
|
result.get(&n_layers_id),
|
|
Some(Distribution::Int(_))
|
|
));
|
|
assert!(matches!(
|
|
result.get(&optimizer_id),
|
|
Some(Distribution::Categorical(_))
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn test_distribution_from_first_trial() {
|
|
let x_id = ParamId::new();
|
|
// Test that when distributions differ, the first trial's distribution is used
|
|
let dist_x_v1 = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_x_v2 = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 10.0, // Different upper bound
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![(x_id, ParamValue::Float(0.5), dist_x_v1.clone())],
|
|
1.0,
|
|
);
|
|
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![(x_id, ParamValue::Float(5.0), dist_x_v2.clone())],
|
|
0.5,
|
|
);
|
|
|
|
let result = IntersectionSearchSpace::calculate(&[trial1, trial2]);
|
|
|
|
assert_eq!(result.len(), 1);
|
|
// Should use the distribution from the first trial
|
|
assert_eq!(result.get(&x_id), Some(&dist_x_v1));
|
|
}
|
|
|
|
#[test]
|
|
fn test_many_trials_with_conditional_params() {
|
|
let lr_id = ParamId::new();
|
|
let use_dropout_id = ParamId::new();
|
|
let dropout_rate_id = ParamId::new();
|
|
// Simulate a scenario with conditional parameters
|
|
let dist_lr = Distribution::Float(FloatDistribution {
|
|
low: 1e-5,
|
|
high: 1e-1,
|
|
log_scale: true,
|
|
step: None,
|
|
});
|
|
let dist_dropout = Distribution::Categorical(CategoricalDistribution { n_choices: 2 });
|
|
let dist_dropout_rate = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 0.5,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
// Trial 1: use_dropout=true, has dropout_rate
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(lr_id, ParamValue::Float(0.01), dist_lr.clone()),
|
|
(
|
|
use_dropout_id,
|
|
ParamValue::Categorical(1),
|
|
dist_dropout.clone(),
|
|
),
|
|
(
|
|
dropout_rate_id,
|
|
ParamValue::Float(0.2),
|
|
dist_dropout_rate.clone(),
|
|
),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
// Trial 2: use_dropout=false, no dropout_rate
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![
|
|
(lr_id, ParamValue::Float(0.001), dist_lr.clone()),
|
|
(
|
|
use_dropout_id,
|
|
ParamValue::Categorical(0),
|
|
dist_dropout.clone(),
|
|
),
|
|
],
|
|
0.8,
|
|
);
|
|
|
|
// Trial 3: use_dropout=true, has dropout_rate
|
|
let trial3 = create_trial(
|
|
2,
|
|
vec![
|
|
(lr_id, ParamValue::Float(0.005), dist_lr.clone()),
|
|
(
|
|
use_dropout_id,
|
|
ParamValue::Categorical(1),
|
|
dist_dropout.clone(),
|
|
),
|
|
(
|
|
dropout_rate_id,
|
|
ParamValue::Float(0.3),
|
|
dist_dropout_rate.clone(),
|
|
),
|
|
],
|
|
0.9,
|
|
);
|
|
|
|
let result = IntersectionSearchSpace::calculate(&[trial1, trial2, trial3]);
|
|
|
|
// Only lr and use_dropout appear in all trials
|
|
assert_eq!(result.len(), 2);
|
|
assert!(result.contains_key(&lr_id));
|
|
assert!(result.contains_key(&use_dropout_id));
|
|
assert!(!result.contains_key(&dropout_rate_id)); // Not in trial2
|
|
}
|
|
|
|
// ==================== GroupDecomposedSearchSpace Tests ====================
|
|
|
|
#[test]
|
|
fn test_group_empty_trials() {
|
|
let trials: Vec<CompletedTrial> = vec![];
|
|
let groups = GroupDecomposedSearchSpace::calculate(&trials);
|
|
assert!(groups.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_single_trial_single_param() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let x_id = ParamId::new();
|
|
let trial = create_trial(0, vec![(x_id, ParamValue::Float(0.5), dist)], 1.0);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial]);
|
|
|
|
assert_eq!(groups.len(), 1);
|
|
assert!(groups[0].contains(&x_id));
|
|
assert_eq!(groups[0].len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_single_trial_multiple_params() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let z_id = ParamId::new();
|
|
let trial = create_trial(
|
|
0,
|
|
vec![
|
|
(x_id, ParamValue::Float(0.5), dist.clone()),
|
|
(y_id, ParamValue::Float(0.3), dist.clone()),
|
|
(z_id, ParamValue::Float(0.7), dist),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial]);
|
|
|
|
// All params appear together, so one group
|
|
assert_eq!(groups.len(), 1);
|
|
assert_eq!(groups[0].len(), 3);
|
|
assert!(groups[0].contains(&x_id));
|
|
assert!(groups[0].contains(&y_id));
|
|
assert!(groups[0].contains(&z_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_two_independent_groups() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let a_id = ParamId::new();
|
|
let b_id = ParamId::new();
|
|
|
|
// Trial 1: x, y (group 1)
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(x_id, ParamValue::Float(0.1), dist.clone()),
|
|
(y_id, ParamValue::Float(0.2), dist.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
// Trial 2: a, b (group 2 - never appears with x or y)
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![
|
|
(a_id, ParamValue::Float(0.3), dist.clone()),
|
|
(b_id, ParamValue::Float(0.4), dist),
|
|
],
|
|
0.5,
|
|
);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2]);
|
|
|
|
// Should have 2 independent groups
|
|
assert_eq!(groups.len(), 2);
|
|
|
|
// Find which group has x/y and which has a/b
|
|
let group_xy = groups.iter().find(|g| g.contains(&x_id));
|
|
let group_ab = groups.iter().find(|g| g.contains(&a_id));
|
|
|
|
assert!(group_xy.is_some());
|
|
assert!(group_ab.is_some());
|
|
|
|
let group_xy = group_xy.expect("group with x should exist");
|
|
let group_ab = group_ab.expect("group with a should exist");
|
|
|
|
assert_eq!(group_xy.len(), 2);
|
|
assert!(group_xy.contains(&x_id));
|
|
assert!(group_xy.contains(&y_id));
|
|
|
|
assert_eq!(group_ab.len(), 2);
|
|
assert!(group_ab.contains(&a_id));
|
|
assert!(group_ab.contains(&b_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_transitive_connection() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let z_id = ParamId::new();
|
|
|
|
// Trial 1: x, y
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(x_id, ParamValue::Float(0.1), dist.clone()),
|
|
(y_id, ParamValue::Float(0.2), dist.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
// Trial 2: y, z (y connects x and z transitively)
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![
|
|
(y_id, ParamValue::Float(0.3), dist.clone()),
|
|
(z_id, ParamValue::Float(0.4), dist),
|
|
],
|
|
0.5,
|
|
);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2]);
|
|
|
|
// All should be in one group due to transitive connection via y
|
|
assert_eq!(groups.len(), 1);
|
|
assert_eq!(groups[0].len(), 3);
|
|
assert!(groups[0].contains(&x_id));
|
|
assert!(groups[0].contains(&y_id));
|
|
assert!(groups[0].contains(&z_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_chain_connection() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let a_id = ParamId::new();
|
|
let b_id = ParamId::new();
|
|
let c_id = ParamId::new();
|
|
let d_id = ParamId::new();
|
|
|
|
// Create a chain: a-b, b-c, c-d
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(a_id, ParamValue::Float(0.1), dist.clone()),
|
|
(b_id, ParamValue::Float(0.2), dist.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![
|
|
(b_id, ParamValue::Float(0.3), dist.clone()),
|
|
(c_id, ParamValue::Float(0.4), dist.clone()),
|
|
],
|
|
0.5,
|
|
);
|
|
|
|
let trial3 = create_trial(
|
|
2,
|
|
vec![
|
|
(c_id, ParamValue::Float(0.5), dist.clone()),
|
|
(d_id, ParamValue::Float(0.6), dist),
|
|
],
|
|
0.3,
|
|
);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]);
|
|
|
|
// All should be connected in one group
|
|
assert_eq!(groups.len(), 1);
|
|
assert_eq!(groups[0].len(), 4);
|
|
assert!(groups[0].contains(&a_id));
|
|
assert!(groups[0].contains(&b_id));
|
|
assert!(groups[0].contains(&c_id));
|
|
assert!(groups[0].contains(&d_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_multiple_isolated_params() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let z_id = ParamId::new();
|
|
|
|
// Each trial has exactly one parameter (all isolated)
|
|
let trial1 = create_trial(0, vec![(x_id, ParamValue::Float(0.1), dist.clone())], 1.0);
|
|
let trial2 = create_trial(1, vec![(y_id, ParamValue::Float(0.2), dist.clone())], 0.5);
|
|
let trial3 = create_trial(2, vec![(z_id, ParamValue::Float(0.3), dist)], 0.3);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]);
|
|
|
|
// Each param should be its own group
|
|
assert_eq!(groups.len(), 3);
|
|
for group in &groups {
|
|
assert_eq!(group.len(), 1);
|
|
}
|
|
|
|
// Verify all params are covered
|
|
let all_params: HashSet<ParamId> = groups.iter().flatten().copied().collect();
|
|
assert!(all_params.contains(&x_id));
|
|
assert!(all_params.contains(&y_id));
|
|
assert!(all_params.contains(&z_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_complex_scenario() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let x_id = ParamId::new();
|
|
let y_id = ParamId::new();
|
|
let z_id = ParamId::new();
|
|
let a_id = ParamId::new();
|
|
let b_id = ParamId::new();
|
|
let w_id = ParamId::new();
|
|
|
|
// Group 1: x, y, z connected
|
|
// Group 2: a, b connected
|
|
// w is isolated
|
|
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(x_id, ParamValue::Float(0.1), dist.clone()),
|
|
(y_id, ParamValue::Float(0.2), dist.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![
|
|
(y_id, ParamValue::Float(0.3), dist.clone()),
|
|
(z_id, ParamValue::Float(0.4), dist.clone()),
|
|
],
|
|
0.5,
|
|
);
|
|
|
|
let trial3 = create_trial(
|
|
2,
|
|
vec![
|
|
(a_id, ParamValue::Float(0.5), dist.clone()),
|
|
(b_id, ParamValue::Float(0.6), dist.clone()),
|
|
],
|
|
0.3,
|
|
);
|
|
|
|
let trial4 = create_trial(3, vec![(w_id, ParamValue::Float(0.7), dist)], 0.2);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3, trial4]);
|
|
|
|
// Should have 3 groups: {x,y,z}, {a,b}, {w}
|
|
assert_eq!(groups.len(), 3);
|
|
|
|
let group_xyz = groups.iter().find(|g| g.contains(&x_id));
|
|
let group_ab = groups.iter().find(|g| g.contains(&a_id));
|
|
let group_w = groups.iter().find(|g| g.contains(&w_id));
|
|
|
|
assert!(group_xyz.is_some());
|
|
assert!(group_ab.is_some());
|
|
assert!(group_w.is_some());
|
|
|
|
let group_xyz = group_xyz.expect("group with x should exist");
|
|
assert_eq!(group_xyz.len(), 3);
|
|
assert!(group_xyz.contains(&x_id));
|
|
assert!(group_xyz.contains(&y_id));
|
|
assert!(group_xyz.contains(&z_id));
|
|
|
|
let group_ab = group_ab.expect("group with a should exist");
|
|
assert_eq!(group_ab.len(), 2);
|
|
assert!(group_ab.contains(&a_id));
|
|
assert!(group_ab.contains(&b_id));
|
|
|
|
let group_w = group_w.expect("group with w should exist");
|
|
assert_eq!(group_w.len(), 1);
|
|
assert!(group_w.contains(&w_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_all_params_same_trial() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let a_id = ParamId::new();
|
|
let b_id = ParamId::new();
|
|
let c_id = ParamId::new();
|
|
let d_id = ParamId::new();
|
|
|
|
// All params in single trial
|
|
let trial = create_trial(
|
|
0,
|
|
vec![
|
|
(a_id, ParamValue::Float(0.1), dist.clone()),
|
|
(b_id, ParamValue::Float(0.2), dist.clone()),
|
|
(c_id, ParamValue::Float(0.3), dist.clone()),
|
|
(d_id, ParamValue::Float(0.4), dist),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial]);
|
|
|
|
// All in one group
|
|
assert_eq!(groups.len(), 1);
|
|
assert_eq!(groups[0].len(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_with_mixed_distribution_types() {
|
|
let lr_id = ParamId::new();
|
|
let n_layers_id = ParamId::new();
|
|
let optimizer_id = ParamId::new();
|
|
let dist_float = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_int = Distribution::Int(IntDistribution {
|
|
low: 1,
|
|
high: 10,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
let dist_cat = Distribution::Categorical(CategoricalDistribution { n_choices: 3 });
|
|
|
|
// Group 1: learning_rate (float), n_layers (int)
|
|
// Group 2: optimizer (categorical)
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(lr_id, ParamValue::Float(0.01), dist_float.clone()),
|
|
(n_layers_id, ParamValue::Int(3), dist_int.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![(optimizer_id, ParamValue::Categorical(1), dist_cat)],
|
|
0.5,
|
|
);
|
|
|
|
let trial3 = create_trial(
|
|
2,
|
|
vec![
|
|
(lr_id, ParamValue::Float(0.001), dist_float),
|
|
(n_layers_id, ParamValue::Int(5), dist_int),
|
|
],
|
|
0.8,
|
|
);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]);
|
|
|
|
// Should have 2 groups: {learning_rate, n_layers} and {optimizer}
|
|
assert_eq!(groups.len(), 2);
|
|
|
|
let group_lr = groups.iter().find(|g| g.contains(&lr_id));
|
|
let group_opt = groups.iter().find(|g| g.contains(&optimizer_id));
|
|
|
|
assert!(group_lr.is_some());
|
|
assert!(group_opt.is_some());
|
|
|
|
let group_lr = group_lr.expect("group with learning_rate should exist");
|
|
assert_eq!(group_lr.len(), 2);
|
|
assert!(group_lr.contains(&lr_id));
|
|
assert!(group_lr.contains(&n_layers_id));
|
|
|
|
let group_opt = group_opt.expect("group with optimizer should exist");
|
|
assert_eq!(group_opt.len(), 1);
|
|
assert!(group_opt.contains(&optimizer_id));
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_star_topology() {
|
|
let dist = Distribution::Float(FloatDistribution {
|
|
low: 0.0,
|
|
high: 1.0,
|
|
log_scale: false,
|
|
step: None,
|
|
});
|
|
|
|
let center_id = ParamId::new();
|
|
let a_id = ParamId::new();
|
|
let b_id = ParamId::new();
|
|
let c_id = ParamId::new();
|
|
|
|
// Star topology: center connects to all others
|
|
let trial1 = create_trial(
|
|
0,
|
|
vec![
|
|
(center_id, ParamValue::Float(0.1), dist.clone()),
|
|
(a_id, ParamValue::Float(0.2), dist.clone()),
|
|
],
|
|
1.0,
|
|
);
|
|
|
|
let trial2 = create_trial(
|
|
1,
|
|
vec![
|
|
(center_id, ParamValue::Float(0.3), dist.clone()),
|
|
(b_id, ParamValue::Float(0.4), dist.clone()),
|
|
],
|
|
0.5,
|
|
);
|
|
|
|
let trial3 = create_trial(
|
|
2,
|
|
vec![
|
|
(center_id, ParamValue::Float(0.5), dist.clone()),
|
|
(c_id, ParamValue::Float(0.6), dist),
|
|
],
|
|
0.3,
|
|
);
|
|
|
|
let groups = GroupDecomposedSearchSpace::calculate(&[trial1, trial2, trial3]);
|
|
|
|
// All connected via center
|
|
assert_eq!(groups.len(), 1);
|
|
assert_eq!(groups[0].len(), 4);
|
|
assert!(groups[0].contains(¢er_id));
|
|
assert!(groups[0].contains(&a_id));
|
|
assert!(groups[0].contains(&b_id));
|
|
assert!(groups[0].contains(&c_id));
|
|
}
|
|
}
|