1use crate::{Distribution, Pruner, Sampler, StudyState, Trial, Value};
2use std::sync::Mutex;
3
4pub 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 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 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 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 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 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 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 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 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 pub fn report(&mut self, step: usize, value: f64) {
105 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 pub fn should_prune(&self) -> bool {
121 self.pruner.should_prune(self.study_state, &self.trial)
122 }
123
124 pub fn study_state(&self) -> &StudyState {
126 self.study_state
127 }
128
129 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}