#![expect(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
mod history;
mod math;
mod mixture;
mod workspace;
use std::{num::NonZeroUsize, sync::Arc};
use hashbrown::HashMap;
use rand::{Rng, SeedableRng, distr::Uniform, rngs::StdRng};
use smallvec::SmallVec;
use self::{
history::{BoundedHistory, FullHistory, RankKey},
mixture::ProductMixture,
workspace::AcquisitionWorkspace,
};
use crate::{
Direction, Distribution, ParamValue, ParzenError, SearchSpace, TrialId,
search_space::{GroupId, ParamId},
storage::TrialStorage,
};
pub enum GammaStrategy {
Optuna,
Hyperopt,
Custom(Arc<dyn Fn(usize) -> usize + Send + Sync>),
}
impl std::fmt::Debug for GammaStrategy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Optuna => f.write_str("Optuna"),
Self::Hyperopt => f.write_str("Hyperopt"),
Self::Custom(_) => f.write_str("Custom(..)"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WeightStrategy {
Uniform,
Optuna,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModelStrategy {
Independent,
Grouped { max_group_size: usize },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HistoryPolicy {
Full,
Bounded {
max_good_trials: NonZeroUsize,
max_bad_trials: NonZeroUsize,
recent_bad_trials: usize,
},
}
#[derive(Debug)]
pub struct TpeSamplerConfig {
seed: u64,
startup_trials: usize,
ei_candidates: NonZeroUsize,
prior_weight: f64,
gamma: GammaStrategy,
weights: WeightStrategy,
model: ModelStrategy,
history: HistoryPolicy,
}
impl TpeSamplerConfig {
#[must_use]
pub fn performance(seed: u64) -> Self {
Self {
seed,
startup_trials: 10,
ei_candidates: NonZeroUsize::new(24).unwrap_or(NonZeroUsize::MIN),
prior_weight: 1.0,
gamma: GammaStrategy::Optuna,
weights: WeightStrategy::Uniform,
model: ModelStrategy::Independent,
history: HistoryPolicy::Bounded {
max_good_trials: NonZeroUsize::new(25).unwrap_or(NonZeroUsize::MIN),
max_bad_trials: NonZeroUsize::new(512).unwrap_or(NonZeroUsize::MIN),
recent_bad_trials: 64,
},
}
}
#[must_use]
pub fn optuna_compatible(seed: u64) -> Self {
Self {
weights: WeightStrategy::Optuna,
history: HistoryPolicy::Full,
..Self::performance(seed)
}
}
#[must_use]
pub const fn startup_trials(mut self, value: usize) -> Self {
self.startup_trials = value;
self
}
#[must_use]
pub const fn ei_candidates(mut self, value: NonZeroUsize) -> Self {
self.ei_candidates = value;
self
}
#[must_use]
pub const fn prior_weight(mut self, value: f64) -> Self {
self.prior_weight = value;
self
}
#[must_use]
pub fn gamma(mut self, value: GammaStrategy) -> Self {
self.gamma = value;
self
}
#[must_use]
pub const fn weights(mut self, value: WeightStrategy) -> Self {
self.weights = value;
self
}
#[must_use]
pub const fn model(mut self, value: ModelStrategy) -> Self {
self.model = value;
self
}
#[must_use]
pub const fn history(mut self, value: HistoryPolicy) -> Self {
self.history = value;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum EstimatorKey {
Param(ParamId),
Group(GroupId),
}
struct Tracker {
key: EstimatorKey,
params: SmallVec<[ParamId; 8]>,
history: BoundedHistory,
}
enum Histories {
Uninitialized,
Bounded(Vec<Tracker>),
Full(FullHistory),
}
struct ModelCache {
generation: u64,
good: ProductMixture,
bad: ProductMixture,
}
pub struct TpeSampler {
rng: StdRng,
config: TpeSamplerConfig,
histories: Histories,
caches: HashMap<EstimatorKey, ModelCache>,
workspace: AcquisitionWorkspace,
}
impl TpeSampler {
pub fn new(config: TpeSamplerConfig) -> Result<Self, ParzenError> {
if !config.prior_weight.is_finite() || config.prior_weight <= 0.0 {
return Err(ParzenError::InvalidConfig(
"prior weight must be finite and positive".into(),
));
}
if let ModelStrategy::Grouped { max_group_size } = config.model {
if !(2..=8).contains(&max_group_size) {
return Err(ParzenError::InvalidConfig(
"maximum group size must be between two and eight".into(),
));
}
}
if let HistoryPolicy::Bounded {
max_good_trials,
max_bad_trials,
recent_bad_trials,
} = config.history
{
let required = recent_bad_trials + max_good_trials.get().saturating_sub(1);
if recent_bad_trials > max_bad_trials.get() || max_bad_trials.get() < required {
return Err(ParzenError::InvalidConfig(
"bounded bad history must fit recent trials and non-good top entries".into(),
));
}
}
Ok(Self {
rng: StdRng::seed_from_u64(config.seed),
config,
histories: Histories::Uninitialized,
caches: HashMap::new(),
workspace: AcquisitionWorkspace::default(),
})
}
pub(crate) const fn model_strategy(&self) -> ModelStrategy {
self.config.model
}
pub(crate) fn initialize(&mut self, space: &SearchSpace) {
self.histories = match self.config.history {
HistoryPolicy::Full => Histories::Full(FullHistory::default()),
HistoryPolicy::Bounded {
max_good_trials,
max_bad_trials,
recent_bad_trials,
} => {
let definitions = estimator_definitions(self.config.model, space);
Histories::Bounded(
definitions
.into_iter()
.enumerate()
.map(|(index, (key, params))| Tracker {
key,
params,
history: BoundedHistory::new(
max_good_trials.get(),
max_bad_trials.get(),
recent_bad_trials,
self.config.seed ^ index as u64,
),
})
.collect(),
)
}
};
}
pub(crate) fn on_trial_added(
&mut self,
id: TrialId,
storage: &TrialStorage,
space: &SearchSpace,
direction: Direction,
) {
let rank = RankKey::new(id, storage.header(id).value, direction, self.config.seed);
match &mut self.histories {
Histories::Full(history) => history.insert(rank),
Histories::Bounded(trackers) => {
for tracker in trackers {
if tracker.params.iter().all(|param| {
storage
.typed_value(
id,
*param,
&space.parameters[param.0 as usize].distribution,
)
.is_some()
}) {
tracker.history.insert(rank);
}
}
}
Histories::Uninitialized => {}
}
}
pub(crate) fn sample_param(
&mut self,
param: ParamId,
space: &SearchSpace,
storage: &TrialStorage,
) -> Result<ParamValue, ParzenError> {
let values = self.sample_estimator(EstimatorKey::Param(param), &[param], space, storage)?;
values.first().map(|(_, value)| *value).ok_or_else(|| {
ParzenError::InternalModel("parameter estimator returned no value".into())
})
}
pub(crate) fn sample_group(
&mut self,
group: GroupId,
space: &SearchSpace,
storage: &TrialStorage,
) -> Result<SmallVec<[(ParamId, ParamValue); 8]>, ParzenError> {
self.sample_estimator(
EstimatorKey::Group(group),
&space.groups[group.0 as usize],
space,
storage,
)
}
fn sample_estimator(
&mut self,
key: EstimatorKey,
params: &[ParamId],
space: &SearchSpace,
storage: &TrialStorage,
) -> Result<SmallVec<[(ParamId, ParamValue); 8]>, ParzenError> {
let generation = self.generation_for(key)?;
if self
.caches
.get(&key)
.is_some_and(|cache| cache.generation == generation)
{
return self.acquire(key);
}
let (seen, generation, applicable) =
self.applicable_history(key, params, space, storage)?;
let all_categorical = params.iter().all(|param| {
matches!(
space.parameters[param.0 as usize].distribution,
Distribution::Categorical(_)
)
});
let flat = applicable.first().is_some_and(|first| {
applicable
.iter()
.all(|trial| storage.header(*trial).value == storage.header(*first).value)
});
if seen < self.config.startup_trials || seen == 0 || (all_categorical && flat) {
if all_categorical {
return self.sample_unseen_categorical(params, &applicable, space, storage);
}
return params
.iter()
.map(|param| {
Ok((
*param,
sample_prior(
&mut self.rng,
&space.parameters[param.0 as usize].distribution,
)?,
))
})
.collect();
}
let good_count = self.good_count(seen)?;
let (good_trials, bad_trials) = match &self.histories {
Histories::Bounded(trackers) => trackers
.iter()
.find(|tracker| tracker.key == key)
.ok_or_else(|| ParzenError::InternalModel("bounded tracker is missing".into()))?
.history
.split(good_count),
Histories::Full(_) => {
let count = good_count.min(applicable.len().saturating_sub(1)).max(1);
(applicable[..count].to_vec(), applicable[count..].to_vec())
}
Histories::Uninitialized => {
return Err(ParzenError::InternalModel(
"sampler is not initialized".into(),
));
}
};
let good = ProductMixture::build(
params,
&good_trials,
storage,
space,
self.config.prior_weight,
self.config.weights,
)?;
let bad = ProductMixture::build(
params,
&bad_trials,
storage,
space,
self.config.prior_weight,
self.config.weights,
)?;
self.caches.insert(
key,
ModelCache {
generation,
good,
bad,
},
);
self.acquire(key)
}
fn acquire(
&mut self,
key: EstimatorKey,
) -> Result<SmallVec<[(ParamId, ParamValue); 8]>, ParzenError> {
let cache = self
.caches
.get(&key)
.ok_or_else(|| ParzenError::InternalModel("model cache insertion failed".into()))?;
let mut best_score = f64::NEG_INFINITY;
let mut best = None;
for _ in 0..self.config.ei_candidates.get() {
let candidate = cache.good.sample(&mut self.rng)?;
let score = cache
.good
.log_pdf(&candidate, &mut self.workspace.good_scores)?
- cache
.bad
.log_pdf(&candidate, &mut self.workspace.bad_scores)?;
if score.is_finite() && score > best_score {
best_score = score;
best = Some(candidate);
}
}
best.ok_or_else(|| {
ParzenError::InternalModel("acquisition produced no finite candidate".into())
})
}
fn generation_for(&self, key: EstimatorKey) -> Result<u64, ParzenError> {
match &self.histories {
Histories::Bounded(trackers) => trackers
.iter()
.find(|tracker| tracker.key == key)
.map(|tracker| tracker.history.generation())
.ok_or_else(|| ParzenError::InternalModel("bounded tracker is missing".into())),
Histories::Full(history) => Ok(history.generation()),
Histories::Uninitialized => Err(ParzenError::InternalModel(
"sampler is not initialized".into(),
)),
}
}
fn applicable_history(
&self,
key: EstimatorKey,
params: &[ParamId],
space: &SearchSpace,
storage: &TrialStorage,
) -> Result<(usize, u64, Vec<TrialId>), ParzenError> {
match &self.histories {
Histories::Bounded(trackers) => {
let tracker = trackers
.iter()
.find(|tracker| tracker.key == key)
.ok_or_else(|| {
ParzenError::InternalModel("bounded tracker is missing".into())
})?;
let (good, bad) = tracker.history.split(tracker.history.seen().min(1));
let mut retained = good;
retained.extend(bad);
Ok((
tracker.history.seen(),
tracker.history.generation(),
retained,
))
}
Histories::Full(history) => {
let applicable: Vec<TrialId> = history
.iter()
.filter(|trial| {
params.iter().all(|param| {
storage
.typed_value(
*trial,
*param,
&space.parameters[param.0 as usize].distribution,
)
.is_some()
})
})
.collect();
Ok((applicable.len(), history.generation(), applicable))
}
Histories::Uninitialized => Err(ParzenError::InternalModel(
"sampler is not initialized".into(),
)),
}
}
fn good_count(&self, seen: usize) -> Result<usize, ParzenError> {
if seen <= 1 {
return Ok(1);
}
let requested = match &self.config.gamma {
GammaStrategy::Optuna => ((seen as f64 * 0.1).ceil() as usize).min(25),
GammaStrategy::Hyperopt => {
((seen as f64).sqrt().mul_add(0.25, 0.0).ceil() as usize).min(25)
}
GammaStrategy::Custom(function) => function(seen),
};
if let HistoryPolicy::Bounded {
max_good_trials, ..
} = self.config.history
{
if requested > max_good_trials.get() {
return Err(ParzenError::GammaExceedsHistoryLimit {
requested,
limit: max_good_trials.get(),
});
}
}
Ok(requested.clamp(1, seen - 1))
}
fn sample_unseen_categorical(
&mut self,
params: &[ParamId],
trials: &[TrialId],
space: &SearchSpace,
storage: &TrialStorage,
) -> Result<SmallVec<[(ParamId, ParamValue); 8]>, ParzenError> {
let counts: SmallVec<[u32; 8]> = params
.iter()
.map(
|param| match space.parameters[param.0 as usize].distribution {
Distribution::Categorical(dist) => Ok(dist.num_choices()),
_ => Err(ParzenError::InternalModel(
"non-categorical parameter in categorical startup".into(),
)),
},
)
.collect::<Result<_, _>>()?;
let product = counts
.iter()
.try_fold(1_u64, |total, count| total.checked_mul(u64::from(*count)));
let Some(product) = product.filter(|product| *product <= 1_000_000) else {
return params
.iter()
.map(|param| {
Ok((
*param,
sample_prior(
&mut self.rng,
&space.parameters[param.0 as usize].distribution,
)?,
))
})
.collect();
};
let mut seen = hashbrown::HashSet::with_capacity(trials.len());
for trial in trials {
let mut code = 0_u64;
for (param, count) in params.iter().zip(&counts) {
let value = storage
.typed_value(
*trial,
*param,
&space.parameters[param.0 as usize].distribution,
)
.and_then(ParamValue::as_categorical)
.ok_or_else(|| {
ParzenError::InternalModel("categorical history is incomplete".into())
})?;
code = code * u64::from(*count) + u64::from(value);
}
seen.insert(code);
}
let start = if product == 1 {
0
} else {
self.rng.random_range(0..product)
};
let code = (0..product)
.map(|offset| (start + offset) % product)
.find(|code| !seen.contains(code))
.unwrap_or(start);
Ok(decode_categorical(code, params, &counts))
}
pub(crate) fn retained_history_len(&self) -> usize {
match &self.histories {
Histories::Bounded(trackers) => trackers.iter().map(|t| t.history.retained()).sum(),
Histories::Full(history) => history.len(),
Histories::Uninitialized => 0,
}
}
}
fn estimator_definitions(
strategy: ModelStrategy,
space: &SearchSpace,
) -> Vec<(EstimatorKey, SmallVec<[ParamId; 8]>)> {
let mut definitions = Vec::new();
for (index, def) in space.parameters.iter().enumerate() {
let param = ParamId(index as u32);
if matches!(strategy, ModelStrategy::Grouped { .. }) && def.group.is_some() {
continue;
}
definitions.push((EstimatorKey::Param(param), smallvec::smallvec![param]));
}
if matches!(strategy, ModelStrategy::Grouped { .. }) {
definitions.extend(space.groups.iter().enumerate().map(|(index, params)| {
(
EstimatorKey::Group(GroupId(index as u32)),
params.iter().copied().collect(),
)
}));
}
definitions
}
fn sample_prior(rng: &mut StdRng, distribution: &Distribution) -> Result<ParamValue, ParzenError> {
match distribution {
Distribution::Categorical(dist) => Ok(ParamValue::Categorical(
rng.random_range(0..dist.num_choices()),
)),
Distribution::Float(dist) => {
if let Some(max_index) = dist.max_step_index() {
let index = rng.random_range(0..=max_index);
return Ok(ParamValue::Float(dist.grid_value(index)));
}
let uniform =
Uniform::new_inclusive(dist.transform(dist.low()), dist.transform(dist.high()))
.map_err(|_| {
ParzenError::InternalModel("float prior range is invalid".into())
})?;
Ok(ParamValue::Float(dist.untransform(rng.sample(uniform))))
}
Distribution::Int(dist) => {
if dist.low() == dist.high() {
return Ok(ParamValue::Int(dist.low()));
}
if dist.scale() == crate::IntScale::Linear {
let max_index = dist.max_step_index();
let index = if max_index == u64::MAX {
rng.random::<u64>()
} else {
rng.random_range(0..=max_index)
};
return Ok(ParamValue::Int(dist.grid_value(index)));
}
let low = (dist.low() as f64 - 0.5).ln();
let high = (dist.high() as f64 + 0.5).ln();
let uniform = Uniform::new_inclusive(low, high).map_err(|_| {
ParzenError::InternalModel("log-integer prior range is invalid".into())
})?;
Ok(ParamValue::Int(dist.untransform(rng.sample(uniform))))
}
}
}
fn decode_categorical(
mut code: u64,
params: &[ParamId],
counts: &[u32],
) -> SmallVec<[(ParamId, ParamValue); 8]> {
let mut values = smallvec::smallvec![(ParamId(0), ParamValue::Categorical(0)); params.len()];
for index in (0..params.len()).rev() {
let choices = u64::from(counts[index]);
values[index] = (
params[index],
ParamValue::Categorical((code % choices) as u32),
);
code /= choices;
}
values
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn gamma_strategies_match_named_formulas() {
let sampler = TpeSampler::new(TpeSamplerConfig::performance(1)).unwrap();
assert_eq!(sampler.good_count(100).unwrap(), 10);
assert_eq!(((100_f64.sqrt() * 0.25).ceil() as usize).min(25), 3);
}
}