Skip to main content

hyperopt_core/
study_state.rs

1use crate::{Direction, Trial, TrialState};
2
3/// A read-only snapshot of a study's history, handed to samplers and pruners.
4///
5/// Both [`crate::Sampler::suggest`] and [`crate::Pruner::should_prune`] receive
6/// a `StudyState` rather than a single trial so that adaptive samplers (TPE)
7/// and cross-trial pruners (median, successive-halving) have everything they
8/// need through one channel. Under parallel execution this snapshot is
9/// deliberately *slightly stale* — see [`crate::Study::optimize_parallel`].
10#[derive(Debug, Clone)]
11pub struct StudyState {
12    direction: Direction,
13    trials: Vec<Trial>,
14}
15
16impl StudyState {
17    /// Builds a snapshot from the study direction and its trials so far.
18    pub fn new(direction: Direction, trials: Vec<Trial>) -> Self {
19        StudyState { direction, trials }
20    }
21
22    /// The optimization direction of the owning study.
23    pub fn direction(&self) -> Direction {
24        self.direction
25    }
26
27    /// All trials in the snapshot, in trial-number order.
28    pub fn trials(&self) -> &[Trial] {
29        &self.trials
30    }
31
32    /// Trials that completed successfully and carry a final value.
33    pub fn completed_trials(&self) -> impl Iterator<Item = &Trial> {
34        self.trials
35            .iter()
36            .filter(|t| t.state == TrialState::Complete && t.value.is_some())
37    }
38
39    /// Number of completed trials — used by samplers/pruners for warmup gates.
40    pub fn n_completed(&self) -> usize {
41        self.completed_trials().count()
42    }
43
44    /// Intermediate values reported at exactly `step` across all trials that
45    /// reached it (completed or pruned). Used by [`MedianPruner`](crate).
46    pub fn intermediate_values_at(&self, step: usize) -> Vec<f64> {
47        self.trials
48            .iter()
49            .filter(|t| matches!(t.state, TrialState::Complete | TrialState::Pruned))
50            .filter_map(|t| t.value_at_step(step))
51            .collect()
52    }
53
54    /// Values of all trials that reached resource `>= step`, taken at the first
55    /// step at or after `step`. Used by rung-based pruners.
56    pub fn values_at_or_after(&self, step: usize) -> Vec<f64> {
57        self.trials
58            .iter()
59            .filter_map(|t| t.value_at_or_after(step))
60            .collect()
61    }
62
63    /// The best completed trial under the study direction, if any.
64    pub fn best_trial(&self) -> Option<&Trial> {
65        let mut best: Option<&Trial> = None;
66        for t in self.completed_trials() {
67            let v = t.value.unwrap();
68            match best {
69                None => best = Some(t),
70                Some(b) => {
71                    let bv = b.value.unwrap();
72                    let better = match self.direction {
73                        Direction::Minimize => v < bv,
74                        Direction::Maximize => v > bv,
75                    };
76                    if better {
77                        best = Some(t);
78                    }
79                }
80            }
81        }
82        best
83    }
84}