gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: reduce #[allow] attributes and extract per-distribution helpers

Browse Source 
- Remove stale #[allow(dead_code)] and #[allow(too_many_lines)] attributes
- Replace #[allow(dead_code)] with #[cfg(test)] for test-only methods
- Delete unused fill_remaining_independent() and ParamMeta.dist field
- Extract sample_float/int/categorical helpers from TpeSampler, MotpeSampler,
  and SamplingEngine to shorten match-heavy sample() methods
- Extract try_fit_kdes() helper to consolidate repeated fallback blocks
- Refactor sample_tpe_float/int to take &FloatDistribution/&IntDistribution
  instead of destructured args, removing #[allow(too_many_arguments)]
This commit is contained in:
Manuel Raimann
2026-02-12 15:52:59 +01:00
parent 38a20bbc42
commit cc50e1ad4b
8 changed files with 427 additions and 431 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kde/multivariate.rs
+5 -4
View File
@@ -30,7 +30,6 @@ use crate::error::{Error, Result};
/// // Get dimensionality
/// assert_eq!(kde.n_dims(), 2);
/// ```
#[allow(dead_code)] // Fields and methods will be used in subsequent stories (US-003, US-004)
#[derive(Clone, Debug)]
pub(crate) struct MultivariateKDE {
/// The sample points used to construct the KDE.
@@ -43,7 +42,6 @@ pub(crate) struct MultivariateKDE {
n_dims: usize,
}
#[allow(dead_code)] // Methods will be used in subsequent stories (US-003, US-004)
impl MultivariateKDE {
/// Creates a new multivariate KDE with automatic bandwidth selection using Scott's rule.
///
@@ -100,6 +98,7 @@ impl MultivariateKDE {
/// Returns `Error::ZeroDimensions` if samples have zero dimensions.
/// Returns `Error::BandwidthDimensionMismatch` if bandwidths length doesn't match dimensions.
/// Returns `Error::InvalidBandwidth` if any bandwidth is not positive.
#[cfg(test)]
pub(crate) fn with_bandwidths(samples: Vec<Vec<f64>>, bandwidths: Vec<f64>) -> Result<Self> {
if samples.is_empty() {
return Err(Error::EmptySamples);
@@ -180,12 +179,13 @@ impl MultivariateKDE {
}
/// Returns the number of dimensions.
#[cfg(test)]
pub(crate) fn n_dims(&self) -> usize {
self.n_dims
}
/// Returns the number of samples.
#[allow(dead_code)] // Will be used in subsequent stories
#[cfg(test)]
pub(crate) fn n_samples(&self) -> usize {
self.samples.len()
}
@@ -197,7 +197,7 @@ impl MultivariateKDE {
}
/// Returns a reference to the samples.
#[allow(dead_code)] // Will be used in subsequent stories
#[cfg(test)]
pub(crate) fn samples(&self) -> &[Vec<f64>] {
&self.samples
}
@@ -288,6 +288,7 @@ impl MultivariateKDE {
/// # Panics
///
/// Panics if `x.len() != self.n_dims`.
#[cfg(test)]
pub(crate) fn pdf(&self, x: &[f64]) -> f64 {
self.log_pdf(x).exp()
}
src/sampler/motpe.rs
+126 -108
View File
@@ -214,8 +214,130 @@ impl Default for MotpeSampler {
}
}
impl MotpeSampler {
fn sample_float(
&self,
d: &crate::distribution::FloatDistribution,
good_trials: &[&MultiObjectiveTrial],
bad_trials: &[&MultiObjectiveTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let good_values: Vec<f64> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<f64> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return ParamValue::Float(rng_util::f64_range(rng, d.low, d.high));
}
let value = tpe_common::sample_tpe_float(
d,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
rng,
);
ParamValue::Float(value)
}
fn sample_int(
&self,
d: &crate::distribution::IntDistribution,
good_trials: &[&MultiObjectiveTrial],
bad_trials: &[&MultiObjectiveTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let good_values: Vec<i64> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<i64> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return common::sample_random(rng, &Distribution::Int(d.clone()));
}
let value = tpe_common::sample_tpe_int(
d,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
rng,
);
ParamValue::Int(value)
}
#[allow(clippy::unused_self)]
fn sample_categorical(
&self,
d: &crate::distribution::CategoricalDistribution,
good_trials: &[&MultiObjectiveTrial],
bad_trials: &[&MultiObjectiveTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let good_indices: Vec<usize> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
let bad_indices: Vec<usize> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
if good_indices.is_empty() || bad_indices.is_empty() {
return common::sample_random(rng, &Distribution::Categorical(d.clone()));
}
let index =
tpe_common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng);
ParamValue::Categorical(index)
}
}
impl MultiObjectiveSampler for MotpeSampler {
#[allow(clippy::too_many_lines)]
fn sample(
&self,
distribution: &Distribution,
@@ -247,114 +369,10 @@ impl MultiObjectiveSampler for MotpeSampler {
}
match distribution {
Distribution::Float(d) => {
let good_values: Vec<f64> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<f64> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return common::sample_random(&mut rng, distribution);
}
let value = tpe_common::sample_tpe_float(
d.low,
d.high,
d.log_scale,
d.step,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
&mut rng,
);
ParamValue::Float(value)
}
Distribution::Int(d) => {
let good_values: Vec<i64> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<i64> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return common::sample_random(&mut rng, distribution);
}
let value = tpe_common::sample_tpe_int(
d.low,
d.high,
d.log_scale,
d.step,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
&mut rng,
);
ParamValue::Int(value)
}
Distribution::Float(d) => self.sample_float(d, &good_trials, &bad_trials, &mut rng),
Distribution::Int(d) => self.sample_int(d, &good_trials, &bad_trials, &mut rng),
Distribution::Categorical(d) => {
let good_indices: Vec<usize> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
let bad_indices: Vec<usize> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
if good_indices.is_empty() || bad_indices.is_empty() {
return common::sample_random(&mut rng, distribution);
}
let index = tpe_common::sample_tpe_categorical(
d.n_choices,
&good_indices,
&bad_indices,
&mut rng,
);
ParamValue::Categorical(index)
self.sample_categorical(d, &good_trials, &bad_trials, &mut rng)
}
}
}
src/sampler/tpe/common.rs
+16 -18
View File
@@ -1,21 +1,20 @@
//! Shared TPE sampling functions used by both `TpeSampler` and `MotpeSampler`.
use crate::distribution::{FloatDistribution, IntDistribution};
use crate::kde::KernelDensityEstimator;
use crate::rng_util;
/// Samples using TPE for float distributions.
#[allow(clippy::too_many_arguments)]
pub(crate) fn sample_tpe_float(
low: f64,
high: f64,
log_scale: bool,
step: Option<f64>,
dist: &FloatDistribution,
good_values: Vec<f64>,
bad_values: Vec<f64>,
n_ei_candidates: usize,
kde_bandwidth: Option<f64>,
rng: &mut fastrand::Rng,
) -> f64 {
let (low, high, log_scale, step) = (dist.low, dist.high, dist.log_scale, dist.step);
// Transform to internal space (log space if needed)
let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
let i_low = low.ln();
@@ -99,32 +98,31 @@ pub(crate) fn sample_tpe_float(
}
/// Samples using TPE for integer distributions.
#[allow(
clippy::too_many_arguments,
clippy::cast_precision_loss,
clippy::cast_possible_truncation
)]
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
pub(crate) fn sample_tpe_int(
low: i64,
high: i64,
log_scale: bool,
step: Option<i64>,
dist: &IntDistribution,
good_values: Vec<i64>,
bad_values: Vec<i64>,
n_ei_candidates: usize,
kde_bandwidth: Option<f64>,
rng: &mut fastrand::Rng,
) -> i64 {
let (low, high, log_scale, step) = (dist.low, dist.high, dist.log_scale, dist.step);
// Convert to floats for KDE
let good_floats: Vec<f64> = good_values.into_iter().map(|v| v as f64).collect();
let bad_floats: Vec<f64> = bad_values.into_iter().map(|v| v as f64).collect();
let float_dist = FloatDistribution {
low: low as f64,
high: high as f64,
log_scale,
step: step.map(|s| s as f64),
};
// Use float TPE sampling
let float_value = sample_tpe_float(
low as f64,
high as f64,
log_scale,
step.map(|s| s as f64),
&float_dist,
good_floats,
bad_floats,
n_ei_candidates,
src/sampler/tpe/multivariate/engine.rs
+152 -178
View File
@@ -103,6 +103,30 @@ impl MultivariateTpeSampler {
result
}
/// Validates observations and fits multivariate KDEs for good and bad groups.
///
/// Returns `None` if observations are invalid or KDE construction fails.
fn try_fit_kdes(
good_obs: Vec<Vec<f64>>,
bad_obs: Vec<Vec<f64>>,
expected_dims: usize,
) -> Option<(crate::kde::MultivariateKDE, crate::kde::MultivariateKDE)> {
use crate::kde::MultivariateKDE;
let valid = !good_obs.is_empty()
&& !bad_obs.is_empty()
&& good_obs.iter().all(|obs| obs.len() == expected_dims)
&& bad_obs.iter().all(|obs| obs.len() == expected_dims);
if !valid {
return None;
}
let good_kde = MultivariateKDE::new(good_obs).ok()?;
let bad_kde = MultivariateKDE::new(bad_obs).ok()?;
Some((good_kde, bad_kde))
}
/// Samples parameters as a single group using multivariate TPE.
///
/// This is the core multivariate TPE sampling logic, used both in non-grouped mode
@@ -117,14 +141,12 @@ impl MultivariateTpeSampler {
/// # Returns
///
/// A `HashMap` mapping parameter names to their sampled values.
#[allow(clippy::too_many_lines)]
pub(crate) fn sample_single_group(
&self,
search_space: &HashMap<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;
@@ -165,7 +187,6 @@ impl MultivariateTpeSampler {
.collect();
if param_order.is_empty() {
// Only categorical parameters in intersection - fill remaining with independent TPE
self.fill_remaining_independent_with_rng(
search_space,
&intersection,
@@ -178,43 +199,12 @@ impl MultivariateTpeSampler {
param_order.sort_by_key(|id| format!("{id}"));
// Extract observations and validate
// Extract observations, validate, and fit KDEs
let good_obs = self.extract_observations(&good, &param_order);
let bad_obs = self.extract_observations(&bad, &param_order);
let expected_dims = param_order.len();
let valid = !good_obs.is_empty()
&& !bad_obs.is_empty()
&& good_obs.iter().all(|obs| obs.len() == expected_dims)
&& bad_obs.iter().all(|obs| obs.len() == expected_dims);
if !valid {
// Observations invalid - fill remaining with independent TPE
self.fill_remaining_independent_with_rng(
search_space,
&intersection,
history,
&mut result,
rng,
);
return result;
}
// Fit KDEs using let...else pattern
let Ok(good_kde) = MultivariateKDE::new(good_obs) else {
// KDE construction failed - fill remaining with independent TPE
self.fill_remaining_independent_with_rng(
search_space,
&intersection,
history,
&mut result,
rng,
);
return result;
};
let Ok(bad_kde) = MultivariateKDE::new(bad_obs) else {
// KDE construction failed - fill remaining with independent TPE
let Some((good_kde, bad_kde)) = Self::try_fit_kdes(good_obs, bad_obs, param_order.len())
else {
self.fill_remaining_independent_with_rng(
search_space,
&intersection,
@@ -289,7 +279,7 @@ impl MultivariateTpeSampler {
/// from the "good" KDE (l(x)) and selects the one that maximizes the ratio l(x)/g(x),
/// which is equivalent to maximizing `log(l(x)) - log(g(x))`.
#[must_use]
#[allow(dead_code)] // Used by tests
#[cfg(test)]
pub(crate) fn select_candidate(
&self,
good_kde: &crate::kde::MultivariateKDE,
@@ -366,43 +356,6 @@ impl MultivariateTpeSampler {
candidates.into_iter().nth(best_idx).unwrap_or_default()
}
/// Fills remaining parameters in result using independent TPE sampling.
///
/// This method is used to sample parameters that are not in the intersection
/// search space. It uses independent univariate TPE sampling for each parameter,
/// similar to the standard [`TpeSampler`].
///
/// When there isn't enough history for a parameter, falls back to uniform sampling.
#[allow(dead_code)]
pub(crate) fn fill_remaining_independent(
&self,
search_space: &HashMap<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.
@@ -437,7 +390,7 @@ impl MultivariateTpeSampler {
/// Samples all parameters using independent TPE sampling.
///
/// This is used as a complete fallback when no intersection search space exists.
#[allow(dead_code)] // Used by tests
#[cfg(test)]
pub(crate) fn sample_all_independent(
&self,
search_space: &HashMap<ParamId, Distribution>,
@@ -485,7 +438,6 @@ impl MultivariateTpeSampler {
///
/// This method extracts values for the given parameter from good and bad trials,
/// fits univariate KDEs, and samples using the TPE acquisition function.
#[allow(clippy::too_many_lines)]
pub(crate) fn sample_independent_tpe(
&self,
param_id: ParamId,
@@ -494,114 +446,136 @@ impl MultivariateTpeSampler {
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
use crate::sampler::tpe::common;
match distribution {
Distribution::Float(d) => {
let good_values: Vec<f64> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_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(&param_id))
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return crate::sampler::common::sample_random(rng, distribution);
}
let value = common::sample_tpe_float(
d.low,
d.high,
d.log_scale,
d.step,
good_values,
bad_values,
self.n_ei_candidates,
None,
rng,
);
ParamValue::Float(value)
self.sample_independent_float(param_id, d, good_trials, bad_trials, rng)
}
Distribution::Int(d) => {
let good_values: Vec<i64> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_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(&param_id))
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return crate::sampler::common::sample_random(rng, distribution);
}
let value = common::sample_tpe_int(
d.low,
d.high,
d.log_scale,
d.step,
good_values,
bad_values,
self.n_ei_candidates,
None,
rng,
);
ParamValue::Int(value)
self.sample_independent_int(param_id, d, good_trials, bad_trials, rng)
}
Distribution::Categorical(d) => {
let good_indices: Vec<usize> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_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(&param_id))
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
if good_indices.is_empty() || bad_indices.is_empty() {
return crate::sampler::common::sample_random(rng, distribution);
}
let idx =
common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng);
ParamValue::Categorical(idx)
self.sample_independent_categorical(param_id, d, good_trials, bad_trials, rng)
}
}
}
fn sample_independent_float(
&self,
param_id: ParamId,
d: &crate::distribution::FloatDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
use crate::sampler::tpe::common;
let good_values: Vec<f64> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_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(&param_id))
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return crate::sampler::common::sample_random(rng, &Distribution::Float(d.clone()));
}
let value =
common::sample_tpe_float(d, good_values, bad_values, self.n_ei_candidates, None, rng);
ParamValue::Float(value)
}
fn sample_independent_int(
&self,
param_id: ParamId,
d: &crate::distribution::IntDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
use crate::sampler::tpe::common;
let good_values: Vec<i64> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_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(&param_id))
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return crate::sampler::common::sample_random(rng, &Distribution::Int(d.clone()));
}
let value =
common::sample_tpe_int(d, good_values, bad_values, self.n_ei_candidates, None, rng);
ParamValue::Int(value)
}
#[allow(clippy::unused_self)]
fn sample_independent_categorical(
&self,
param_id: ParamId,
d: &crate::distribution::CategoricalDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
use crate::sampler::tpe::common;
let good_indices: Vec<usize> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_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(&param_id))
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
if good_indices.is_empty() || bad_indices.is_empty() {
return crate::sampler::common::sample_random(
rng,
&Distribution::Categorical(d.clone()),
);
}
let idx = common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng);
ParamValue::Categorical(idx)
}
}
src/sampler/tpe/multivariate/mod.rs
-1
View File
@@ -389,7 +389,6 @@ impl MultivariateTpeSampler {
/// let params = sampler.sample_joint(&search_space, &history);
/// ```
#[must_use]
#[allow(clippy::too_many_lines)]
pub fn sample_joint(
&self,
search_space: &HashMap<ParamId, Distribution>,
src/sampler/tpe/sampler.rs
+126 -116
View File
@@ -642,8 +642,130 @@ impl Default for TpeSamplerBuilder {
}
}
impl TpeSampler {
fn sample_float(
&self,
d: &crate::distribution::FloatDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let good_values: Vec<f64> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<f64> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return ParamValue::Float(rng_util::f64_range(rng, d.low, d.high));
}
let value = tpe_common::sample_tpe_float(
d,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
rng,
);
ParamValue::Float(value)
}
fn sample_int(
&self,
d: &crate::distribution::IntDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let good_values: Vec<i64> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<i64> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return common::sample_random(rng, &Distribution::Int(d.clone()));
}
let value = tpe_common::sample_tpe_int(
d,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
rng,
);
ParamValue::Int(value)
}
#[allow(clippy::unused_self)]
fn sample_categorical(
&self,
d: &crate::distribution::CategoricalDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let good_indices: Vec<usize> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
let bad_indices: Vec<usize> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
if good_indices.is_empty() || bad_indices.is_empty() {
return common::sample_random(rng, &Distribution::Categorical(d.clone()));
}
let index =
tpe_common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng);
ParamValue::Categorical(index)
}
}
impl Sampler for TpeSampler {
#[allow(clippy::too_many_lines)]
fn sample(
&self,
distribution: &Distribution,
@@ -670,123 +792,11 @@ impl Sampler for TpeSampler {
return common::sample_random(&mut rng, distribution);
}
// Extract parameter values for this distribution
// Since we don't have the parameter name here, we need to look at all
// trials and find matching distributions
// Note: This is a simplification - in practice, we'd need the param name
// For now, we'll collect values from trials that have this exact distribution type
match distribution {
Distribution::Float(d) => {
// Collect float values from trials
let good_values: Vec<f64> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<f64> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
// Need values in both groups for TPE
if good_values.is_empty() || bad_values.is_empty() {
return common::sample_random(&mut rng, distribution);
}
let value = tpe_common::sample_tpe_float(
d.low,
d.high,
d.log_scale,
d.step,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
&mut rng,
);
ParamValue::Float(value)
}
Distribution::Int(d) => {
let good_values: Vec<i64> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<i64> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return common::sample_random(&mut rng, distribution);
}
let value = tpe_common::sample_tpe_int(
d.low,
d.high,
d.log_scale,
d.step,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
&mut rng,
);
ParamValue::Int(value)
}
Distribution::Float(d) => self.sample_float(d, &good_trials, &bad_trials, &mut rng),
Distribution::Int(d) => self.sample_int(d, &good_trials, &bad_trials, &mut rng),
Distribution::Categorical(d) => {
let good_indices: Vec<usize> = good_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
let bad_indices: Vec<usize> = bad_trials
.iter()
.flat_map(|t| t.params.values())
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
if good_indices.is_empty() || bad_indices.is_empty() {
return common::sample_random(&mut rng, distribution);
}
let index = tpe_common::sample_tpe_categorical(
d.n_choices,
&good_indices,
&bad_indices,
&mut rng,
);
ParamValue::Categorical(index)
self.sample_categorical(d, &good_trials, &bad_trials, &mut rng)
}
}
}
src/study/export.rs
+1 -1
View File
@@ -224,7 +224,7 @@ where
impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
/// Export trials as a pretty-printed JSON array to a file.
///
/// Each element in the array is a serialized [`crate::CompletedTrial`].
/// Each element in the array is a serialized [`CompletedTrial`](crate::sampler::CompletedTrial).
/// Requires the `serde` feature.
///
/// # Errors
src/visualization.rs
+1 -5
View File
@@ -41,7 +41,6 @@ use core::fmt::Write as _;
use std::collections::BTreeMap;
use std::path::Path;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::parameter::ParamId;
use crate::sampler::CompletedTrial;
@@ -156,8 +155,6 @@ fn build_html(
/// Metadata about each parameter seen across trials.
struct ParamMeta {
label: String,
#[allow(dead_code)]
dist: Option<Distribution>,
}
/// Collect parameter labels and distributions across all trials.
@@ -171,8 +168,7 @@ fn collect_param_info(trials: &[CompletedTrial<f64>]) -> BTreeMap<ParamId, Param
.get(&id)
.cloned()
.unwrap_or_else(|| id.to_string());
let dist = trial.distributions.get(&id).cloned();
ParamMeta { label, dist }
ParamMeta { label }
});
}
}