Skip to main content

hyperopt_core/
context.rs

1use crate::{Distribution, Pruner, Sampler, StudyState, Trial, Value};
2use std::sync::Mutex;
3
4/// The handle passed into the user's objective closure.
5///
6/// This is where define-by-run happens: the search space is discovered by
7/// *calling* `suggest_*` methods here, not declared up front. Each call
8/// (1) asks the active [`Sampler`] for a value given the study history and the
9/// distribution, then (2) records the `(name, distribution, value)` on the
10/// trial. Recording after asking keeps the trial's history accurate for the
11/// storage, pruning, and importance phases that depend on it.
12///
13/// The sampler is held behind a shared [`Mutex`] so the exact same context type
14/// serves both sequential and parallel execution; under sequential runs the
15/// lock is uncontended.
16pub struct TrialContext<'a> {
17    trial: Trial,
18    sampler: &'a Mutex<Box<dyn Sampler>>,
19    study_state: &'a StudyState,
20    pruner: &'a dyn Pruner,
21}
22
23impl<'a> TrialContext<'a> {
24    /// Internal: build a context for one trial. Used by [`crate::Study`].
25    pub(crate) fn new(
26        number: usize,
27        sampler: &'a Mutex<Box<dyn Sampler>>,
28        study_state: &'a StudyState,
29        pruner: &'a dyn Pruner,
30    ) -> Self {
31        TrialContext {
32            trial: Trial::new(number),
33            sampler,
34            study_state,
35            pruner,
36        }
37    }
38
39    /// This trial's number within the study.
40    pub fn number(&self) -> usize {
41        self.trial.number
42    }
43
44    fn suggest(&mut self, name: &str, distribution: Distribution) -> Value {
45        let value = {
46            // Recover from poisoning: a prior trial that panicked mid-suggest
47            // shouldn't take down the rest of the study.
48            let mut sampler = self
49                .sampler
50                .lock()
51                .unwrap_or_else(|poisoned| poisoned.into_inner());
52            sampler.suggest(self.study_state, &self.trial, name, &distribution)
53        };
54        self.trial.record(name, distribution, value.clone());
55        value
56    }
57
58    /// Suggest a continuous value uniformly over `[low, high]`.
59    pub fn suggest_float(&mut self, name: &str, low: f64, high: f64) -> f64 {
60        let v = self.suggest(name, Distribution::Uniform { low, high });
61        coerce_float(&v)
62    }
63
64    /// Suggest a continuous value log-uniformly over `[low, high]`
65    /// (both must be `> 0`). Good for scale parameters like learning rates.
66    pub fn suggest_loguniform(&mut self, name: &str, low: f64, high: f64) -> f64 {
67        let v = self.suggest(name, Distribution::LogUniform { low, high });
68        coerce_float(&v)
69    }
70
71    /// Suggest an integer uniformly over the inclusive range `[low, high]`.
72    pub fn suggest_int(&mut self, name: &str, low: i64, high: i64) -> i64 {
73        let v = self.suggest(name, Distribution::IntUniform { low, high });
74        match v {
75            Value::Int(x) => x,
76            Value::Float(x) => x.round() as i64,
77            Value::Categorical(_) => low,
78        }
79    }
80
81    /// Suggest one of `choices`, returning the chosen label.
82    pub fn suggest_categorical(&mut self, name: &str, choices: &[&str]) -> String {
83        let dist = Distribution::Categorical {
84            choices: choices.iter().map(|s| s.to_string()).collect(),
85        };
86        let v = self.suggest(name, dist);
87        match v {
88            Value::Categorical(s) => s,
89            // Defensive: a sampler returning a numeric value for a categorical
90            // is treated as an index into `choices`.
91            Value::Int(i) => choices
92                .get(i.max(0) as usize)
93                .map(|s| s.to_string())
94                .unwrap_or_default(),
95            Value::Float(f) => choices
96                .get(f as usize)
97                .map(|s| s.to_string())
98                .unwrap_or_default(),
99        }
100    }
101
102    /// Report an intermediate objective value at `step` (e.g. per-epoch
103    /// validation score). Consumed by pruners via [`Self::should_prune`].
104    pub fn report(&mut self, step: usize, value: f64) {
105        // Overwrite an existing report for the same step rather than duplicate.
106        if let Some(slot) = self
107            .trial
108            .intermediate_values
109            .iter_mut()
110            .find(|(s, _)| *s == step)
111        {
112            slot.1 = value;
113        } else {
114            self.trial.intermediate_values.push((step, value));
115        }
116    }
117
118    /// Ask the active pruner whether this trial should stop early, based on the
119    /// intermediate values reported so far versus the rest of the study.
120    pub fn should_prune(&self) -> bool {
121        self.pruner.should_prune(self.study_state, &self.trial)
122    }
123
124    /// Read-only access to the study snapshot this trial sees.
125    pub fn study_state(&self) -> &StudyState {
126        self.study_state
127    }
128
129    /// Consume the context, yielding the trial it built up.
130    pub(crate) fn into_trial(self) -> Trial {
131        self.trial
132    }
133}
134
135fn coerce_float(v: &Value) -> f64 {
136    match v {
137        Value::Float(x) => *x,
138        Value::Int(x) => *x as f64,
139        Value::Categorical(_) => f64::NAN,
140    }
141}