hyperopt_samplers/
tpe_sampler.rs

1use hyperopt_core::{Direction, Distribution, Sampler, StudyState, Trial, Value};
2use rand::rngs::StdRng;
3use rand::{RngExt, SeedableRng};
4use tpe::density_estimation::DefaultEstimatorBuilder;
5use tpe::{categorical_range, histogram_estimator, parzen_estimator, range, TpeOptimizer};
6
7use crate::random::sample_value;
8
9/// Tree-structured Parzen Estimator sampler.
10///
11/// This **wraps the [`tpe`](https://docs.rs/tpe) crate** rather than
12/// reimplementing TPE from scratch — a real integration point, not a stubbed
13/// gap. The work this type does is translating between `hyperopt-core`'s
14/// `Distribution`/`Trial` history and `tpe`'s per-parameter `TpeOptimizer`
15/// (one optimizer handles exactly one hyperparameter), including:
16///
17/// - mapping each [`Distribution`] variant onto a `tpe` range + estimator
18///   (Parzen for numeric, histogram for categorical),
19/// - encoding values into `tpe`'s coordinate space (log for `LogUniform`,
20///   index for `Categorical`, `[low, high+1)` for `IntUniform`) and decoding
21///   the sampled result back,
22/// - honouring the study [`Direction`] by negating objective values under
23///   `Maximize`, since `tpe` always minimizes.
24///
25/// The sampler is effectively **stateless over trial history**: each
26/// `suggest` rebuilds a `TpeOptimizer` from the study snapshot and replays the
27/// relevant observations. This makes it correct under both parallel execution
28/// (it simply works from whatever snapshot it's given) and study reload from
29/// storage, at an `O(n)` per-suggestion cost that is negligible next to a real
30/// objective evaluation.
31///
32/// ### Limitations (documented, not silent)
33/// - The first `n_startup_trials` completed trials are sampled at random to
34///   seed the estimators — matching Optuna's warmup and avoiding a biased model
35///   from too few points.
36/// - Every `hyperopt-core` [`Distribution`] variant maps onto `tpe`. There is
37///   no variant `tpe` cannot represent here; should a future distribution not
38///   map cleanly, this sampler falls back to a random draw for it rather than
39///   producing an out-of-range value.
40pub struct TpeSampler {
41    rng: StdRng,
42    n_startup_trials: usize,
43}
44
45impl TpeSampler {
46    /// A TPE sampler with default warmup (`n_startup_trials = 10`), seeded from
47    /// OS entropy.
48    pub fn new() -> Self {
49        let mut seeder = rand::rng();
50        TpeSampler {
51            rng: StdRng::seed_from_u64(seeder.random()),
52            n_startup_trials: 10,
53        }
54    }
55
56    /// A TPE sampler with a fixed seed — reproducible for tests/benchmarks.
57    pub fn seeded(seed: u64) -> Self {
58        TpeSampler {
59            rng: StdRng::seed_from_u64(seed),
60            n_startup_trials: 10,
61        }
62    }
63
64    /// Number of initial trials to sample randomly before building TPE models.
65    pub fn n_startup_trials(mut self, n: usize) -> Self {
66        self.n_startup_trials = n;
67        self
68    }
69
70    /// Collect `(tpe-coordinate, value-to-minimize)` observations for `name`
71    /// from the completed trials, honouring the study direction.
72    fn observations(
73        study_state: &StudyState,
74        name: &str,
75        distribution: &Distribution,
76    ) -> Vec<(f64, f64)> {
77        let direction = study_state.direction();
78        let mut obs = Vec::new();
79        for t in study_state.completed_trials() {
80            let (Some(record), Some(value)) = (
81                t.params.iter().find(|p| p.name == name),
82                t.value,
83            ) else {
84                continue;
85            };
86            if let Some(coord) = encode(distribution, &record.value) {
87                let objective = match direction {
88                    Direction::Minimize => value,
89                    Direction::Maximize => -value,
90                };
91                obs.push((coord, objective));
92            }
93        }
94        obs
95    }
96}
97
98impl Default for TpeSampler {
99    fn default() -> Self {
100        Self::new()
101    }
102}
103
104impl Sampler for TpeSampler {
105    fn suggest(
106        &mut self,
107        study_state: &StudyState,
108        _trial: &Trial,
109        param_name: &str,
110        distribution: &Distribution,
111    ) -> Value {
112        let obs = Self::observations(study_state, param_name, distribution);
113
114        // Warmup / not enough data: sample randomly.
115        if obs.len() < self.n_startup_trials {
116            return sample_value(&mut self.rng, distribution);
117        }
118
119        // Build the tpe optimizer for this distribution; fall back to random if
120        // the range is degenerate.
121        let Some((estimator, tpe_range)) = build_estimator(distribution) else {
122            return sample_value(&mut self.rng, distribution);
123        };
124        let mut optim: TpeOptimizer<DefaultEstimatorBuilder> =
125            TpeOptimizer::new(estimator, tpe_range);
126
127        for (coord, value) in &obs {
128            // A coord outside the range (shouldn't happen given encode) or a
129            // NaN value is simply skipped rather than aborting the suggestion.
130            let _ = optim.tell(*coord, *value);
131        }
132
133        match optim.ask(&mut self.rng) {
134            Ok(raw) => decode(distribution, raw),
135            Err(_) => sample_value(&mut self.rng, distribution),
136        }
137    }
138}
139
140/// Build the `tpe` estimator + range for a distribution, or `None` if the
141/// range is degenerate/unrepresentable.
142fn build_estimator(
143    distribution: &Distribution,
144) -> Option<(DefaultEstimatorBuilder, tpe::range::Range)> {
145    match distribution {
146        Distribution::Uniform { low, high } => {
147            Some((parzen_estimator(), range(*low, *high).ok()?))
148        }
149        Distribution::LogUniform { low, high } => {
150            if *low <= 0.0 {
151                return None;
152            }
153            Some((parzen_estimator(), range(low.ln(), high.ln()).ok()?))
154        }
155        Distribution::IntUniform { low, high } => {
156            // Continuous [low, high+1); decode floors back to an integer.
157            Some((parzen_estimator(), range(*low as f64, *high as f64 + 1.0).ok()?))
158        }
159        Distribution::Categorical { choices } => {
160            Some((histogram_estimator(), categorical_range(choices.len()).ok()?))
161        }
162    }
163}
164
165/// Encode a recorded [`Value`] into `tpe`'s coordinate for `distribution`.
166fn encode(distribution: &Distribution, value: &Value) -> Option<f64> {
167    match distribution {
168        Distribution::Uniform { .. } => value.as_float(),
169        Distribution::LogUniform { .. } => value.as_float().map(f64::ln),
170        Distribution::IntUniform { .. } => value.as_int().map(|i| i as f64),
171        Distribution::Categorical { choices } => {
172            let label = value.as_categorical()?;
173            choices.iter().position(|c| c == label).map(|i| i as f64)
174        }
175    }
176}
177
178/// Decode a value sampled by `tpe` back into a [`Value`] for `distribution`.
179fn decode(distribution: &Distribution, raw: f64) -> Value {
180    match distribution {
181        Distribution::Uniform { .. } => Value::Float(raw),
182        Distribution::LogUniform { .. } => Value::Float(raw.exp()),
183        Distribution::IntUniform { low, high } => {
184            let i = (raw.floor() as i64).clamp(*low, *high);
185            Value::Int(i)
186        }
187        Distribution::Categorical { choices } => {
188            if choices.is_empty() {
189                return Value::Categorical(String::new());
190            }
191            let idx = (raw.floor() as usize).min(choices.len() - 1);
192            Value::Categorical(choices[idx].clone())
193        }
194    }
195}
hyperopt_samplers/tpe_sampler.rs

hyperopt_samplers/
tpe_sampler.rs
