Initial Implementation

This commit is contained in:
Manuel Raimann
2026-01-30 16:02:42 +01:00
parent 164aafc209
commit 4db6e56466
14 changed files with 4860 additions and 0 deletions
+52
View File
@@ -0,0 +1,52 @@
//! Parameter distribution types.
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
/// Distribution for floating-point parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct FloatDistribution {
/// Lower bound (inclusive).
pub low: f64,
/// Upper bound (inclusive).
pub high: f64,
/// Whether to sample in log space.
pub log_scale: bool,
/// Optional step size for discretization.
pub step: Option<f64>,
}
/// Distribution for integer parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct IntDistribution {
/// Lower bound (inclusive).
pub low: i64,
/// Upper bound (inclusive).
pub high: i64,
/// Whether to sample in log space.
pub log_scale: bool,
/// Optional step size for discretization.
pub step: Option<i64>,
}
/// Distribution for categorical parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct CategoricalDistribution {
/// Number of choices available.
pub n_choices: usize,
}
/// Enum wrapping all parameter distribution types.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum Distribution {
/// A floating-point distribution.
Float(FloatDistribution),
/// An integer distribution.
Int(IntDistribution),
/// A categorical distribution.
Categorical(CategoricalDistribution),
}
+44
View File
@@ -0,0 +1,44 @@
//! Error types for the optimize library.
use thiserror::Error;
/// The error type for TPE operations.
#[derive(Debug, Error)]
pub enum TpeError {
/// Returned when the lower bound is greater than the upper bound.
#[error("invalid bounds: low ({low}) must be less than or equal to high ({high})")]
InvalidBounds {
/// The lower bound value.
low: f64,
/// The upper bound value.
high: f64,
},
/// Returned when log scale is used with non-positive bounds.
#[error("invalid log bounds: low must be positive for log scale")]
InvalidLogBounds,
/// Returned when step size is not positive.
#[error("invalid step: step must be positive")]
InvalidStep,
/// Returned when categorical choices are empty.
#[error("categorical choices cannot be empty")]
EmptyChoices,
/// Returned when a parameter is suggested with a different configuration.
#[error("parameter conflict for '{name}': {reason}")]
ParameterConflict {
/// The name of the conflicting parameter.
name: String,
/// The reason for the conflict.
reason: String,
},
/// Returned when requesting the best trial but no trials have completed.
#[error("no completed trials available")]
NoCompletedTrials,
}
/// A specialized Result type for TPE operations.
pub type Result<T> = std::result::Result<T, TpeError>;
+267
View File
@@ -0,0 +1,267 @@
//! Kernel Density Estimation for continuous parameters.
//!
//! This module provides a Gaussian kernel density estimator used by the TPE
//! sampler to model probability distributions over good and bad trial regions.
use rand::Rng;
/// A Gaussian kernel density estimator for continuous distributions.
///
/// KDE estimates a probability density function from a set of samples by
/// placing Gaussian kernels centered at each sample point. This is used
/// in TPE to model the distributions l(x) (good trials) and g(x) (bad trials).
///
/// # Examples
///
/// ```ignore
/// use crate::kde::KernelDensityEstimator;
///
/// let samples = vec![1.0, 2.0, 3.0, 4.0, 5.0];
/// let kde = KernelDensityEstimator::new(samples);
///
/// // Get probability density at a point
/// let density = kde.pdf(2.5);
/// assert!(density > 0.0);
///
/// // Sample from the estimated distribution
/// let mut rng = rand::rng();
/// let sample = kde.sample(&mut rng);
/// ```
#[derive(Clone, Debug)]
pub struct KernelDensityEstimator {
/// The sample points used to construct the KDE.
samples: Vec<f64>,
/// The bandwidth (standard deviation) of the Gaussian kernels.
bandwidth: f64,
}
impl KernelDensityEstimator {
/// Creates a new KDE with automatic bandwidth selection using Scott's rule.
///
/// Scott's rule sets bandwidth = n^(-1/5) * std_dev, which works well
/// for unimodal distributions close to normal.
///
/// # Panics
///
/// Panics if `samples` is empty.
pub fn new(samples: Vec<f64>) -> Self {
assert!(!samples.is_empty(), "KDE requires at least one sample");
let bandwidth = Self::scotts_rule(&samples);
Self { samples, bandwidth }
}
/// Creates a new KDE with a specified bandwidth.
///
/// Use this when you want explicit control over the smoothing parameter.
///
/// # Panics
///
/// Panics if `samples` is empty or `bandwidth` is not positive.
pub fn with_bandwidth(samples: Vec<f64>, bandwidth: f64) -> Self {
assert!(!samples.is_empty(), "KDE requires at least one sample");
assert!(bandwidth > 0.0, "Bandwidth must be positive");
Self { samples, bandwidth }
}
/// Computes bandwidth using Scott's rule.
///
/// Scott's rule: h = n^(-1/5) * sigma
/// where sigma is the sample standard deviation.
fn scotts_rule(samples: &[f64]) -> f64 {
let n = samples.len() as f64;
let std_dev = Self::sample_std_dev(samples);
// For degenerate case where all samples are identical,
// use a small positive bandwidth
if std_dev < f64::EPSILON {
return 1.0;
}
n.powf(-0.2) * std_dev
}
/// Computes the sample standard deviation.
fn sample_std_dev(samples: &[f64]) -> f64 {
let n = samples.len() as f64;
let mean = samples.iter().sum::<f64>() / n;
let variance = samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n;
variance.sqrt()
}
/// Returns the probability density at point `x`.
///
/// The density is computed as the average of Gaussian kernels centered
/// at each sample point:
///
/// f(x) = (1/n) * sum_i K((x - x_i) / h)
///
/// where K is the standard Gaussian kernel and h is the bandwidth.
pub fn pdf(&self, x: f64) -> f64 {
let n = self.samples.len() as f64;
let inv_bandwidth = 1.0 / self.bandwidth;
let normalization = inv_bandwidth / (2.0 * std::f64::consts::PI).sqrt();
let density: f64 = self
.samples
.iter()
.map(|&xi| {
let z = (x - xi) * inv_bandwidth;
normalization * (-0.5 * z * z).exp()
})
.sum();
density / n
}
/// Samples a value from the estimated density distribution.
///
/// Sampling works by:
/// 1. Uniformly selecting one of the kernel centers (samples)
/// 2. Adding Gaussian noise with the bandwidth as standard deviation
pub fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
// Select a random sample to center the kernel on
let idx = rng.random_range(0..self.samples.len());
let center = self.samples[idx];
// Add Gaussian noise with bandwidth as standard deviation
// Using Box-Muller transform for Gaussian sampling
let u1: f64 = rng.random();
let u2: f64 = rng.random();
let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
center + z * self.bandwidth
}
/// Returns the bandwidth of this KDE.
pub fn bandwidth(&self) -> f64 {
self.bandwidth
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_kde_pdf_basic() {
let samples = vec![0.0, 1.0, 2.0];
let kde = KernelDensityEstimator::new(samples);
// Density should be positive everywhere
assert!(kde.pdf(0.0) > 0.0);
assert!(kde.pdf(1.0) > 0.0);
assert!(kde.pdf(2.0) > 0.0);
// Density should be higher near sample points
let mid_density = kde.pdf(1.0);
let far_density = kde.pdf(10.0);
assert!(mid_density > far_density);
}
#[test]
fn test_kde_pdf_integrates_to_one() {
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
let kde = KernelDensityEstimator::new(samples);
// Numerical integration over a wide range
let n_points = 10000;
let low = -10.0;
let high = 15.0;
let dx = (high - low) / n_points as f64;
let integral: f64 = (0..n_points)
.map(|i| {
let x = low + (i as f64 + 0.5) * dx;
kde.pdf(x) * dx
})
.sum();
// Should be approximately 1.0 (within numerical error)
assert!(
(integral - 1.0).abs() < 0.01,
"Integral = {integral}, expected ~1.0"
);
}
#[test]
fn test_kde_with_bandwidth() {
let samples = vec![0.0, 1.0, 2.0];
let kde = KernelDensityEstimator::with_bandwidth(samples, 0.5);
assert_eq!(kde.bandwidth(), 0.5);
assert!(kde.pdf(1.0) > 0.0);
}
#[test]
fn test_kde_sample_in_reasonable_range() {
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
let kde = KernelDensityEstimator::new(samples);
let mut rng = rand::rng();
// Samples should generally be in a reasonable range around the data
for _ in 0..100 {
let s = kde.sample(&mut rng);
// With high probability, samples should be within a few bandwidths
// of the data range. Use a generous range to avoid flaky tests.
assert!(s > -10.0 && s < 15.0, "Sample {s} outside expected range");
}
}
#[test]
fn test_kde_single_sample() {
let samples = vec![5.0];
let kde = KernelDensityEstimator::new(samples);
// Should have positive density near the sample
assert!(kde.pdf(5.0) > 0.0);
assert!(kde.pdf(4.5) > 0.0);
}
#[test]
fn test_kde_identical_samples() {
let samples = vec![3.0, 3.0, 3.0, 3.0];
let kde = KernelDensityEstimator::new(samples);
// Should handle degenerate case with identical samples
assert!(kde.bandwidth() > 0.0);
assert!(kde.pdf(3.0) > 0.0);
}
#[test]
fn test_scotts_rule_bandwidth() {
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
let kde = KernelDensityEstimator::new(samples);
// n = 10, n^(-1/5) ≈ 0.631
// std_dev ≈ 2.87
// bandwidth ≈ 0.631 * 2.87 ≈ 1.81
let bandwidth = kde.bandwidth();
assert!(
bandwidth > 1.0 && bandwidth < 3.0,
"Bandwidth {bandwidth} outside expected range"
);
}
#[test]
#[should_panic(expected = "KDE requires at least one sample")]
fn test_kde_empty_samples() {
let samples: Vec<f64> = vec![];
KernelDensityEstimator::new(samples);
}
#[test]
#[should_panic(expected = "Bandwidth must be positive")]
fn test_kde_zero_bandwidth() {
let samples = vec![1.0, 2.0, 3.0];
KernelDensityEstimator::with_bandwidth(samples, 0.0);
}
#[test]
#[should_panic(expected = "Bandwidth must be positive")]
fn test_kde_negative_bandwidth() {
let samples = vec![1.0, 2.0, 3.0];
KernelDensityEstimator::with_bandwidth(samples, -1.0);
}
}
+151
View File
@@ -0,0 +1,151 @@
//! A Tree-Parzen Estimator (TPE) library for black-box optimization.
//!
//! This library provides an Optuna-like API for hyperparameter optimization
//! using the Tree-Parzen Estimator algorithm. It supports:
//!
//! - Float, integer, and categorical parameter types
//! - Log-scale and stepped parameter sampling
//! - Synchronous and async optimization
//! - Parallel trial evaluation with bounded concurrency
//! - Serialization for saving/loading study state
//!
//! # Quick Start
//!
//! ```
//! use optimize::{Direction, Study, TpeSampler};
//!
//! // Create a study with TPE sampler
//! let sampler = TpeSampler::builder().seed(42).build();
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//!
//! // Optimize x^2 for 20 trials
//! study
//! .optimize_with_sampler(20, |trial| {
//! let x = trial.suggest_float("x", -10.0, 10.0)?;
//! Ok::<_, optimize::TpeError>(x * x)
//! })
//! .unwrap();
//!
//! // Get the best result
//! let best = study.best_trial().unwrap();
//! println!("Best value: {} at x={:?}", best.value, best.params);
//! ```
//!
//! # Creating a Study
//!
//! A [`Study`] manages optimization trials. Create one with an optimization direction:
//!
//! ```
//! use optimize::{Direction, RandomSampler, Study, TpeSampler};
//!
//! // Minimize with default random sampler
//! let study: Study<f64> = Study::new(Direction::Minimize);
//!
//! // Maximize with TPE sampler
//! let study: Study<f64> = Study::with_sampler(Direction::Maximize, TpeSampler::new());
//!
//! // With seeded sampler for reproducibility
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
//! ```
//!
//! # Suggesting Parameters
//!
//! Within the objective function, use [`Trial`] to suggest parameter values:
//!
//! ```
//! use optimize::{Direction, Study};
//!
//! let study: Study<f64> = Study::new(Direction::Minimize);
//!
//! study
//! .optimize(10, |trial| {
//! // Float parameters
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
//! let lr = trial.suggest_float_log("learning_rate", 1e-5, 1e-1)?;
//! let step = trial.suggest_float_step("step", 0.0, 1.0, 0.1)?;
//!
//! // Integer parameters
//! let n = trial.suggest_int("n_layers", 1, 10)?;
//! let batch = trial.suggest_int_log("batch_size", 16, 256)?;
//! let units = trial.suggest_int_step("units", 32, 512, 32)?;
//!
//! // Categorical parameters
//! let optimizer = trial.suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])?;
//!
//! // Return objective value
//! Ok::<_, optimize::TpeError>(x * n as f64)
//! })
//! .unwrap();
//! ```
//!
//! # Configuring TPE
//!
//! The [`TpeSampler`] can be configured using the builder pattern:
//!
//! ```
//! use optimize::TpeSampler;
//!
//! let sampler = TpeSampler::builder()
//! .gamma(0.15) // Quantile for good/bad split
//! .n_startup_trials(20) // Random trials before TPE
//! .n_ei_candidates(32) // Candidates to evaluate
//! .seed(42) // Reproducibility
//! .build();
//! ```
//!
//! # Async and Parallel Optimization
//!
//! With the `async` feature enabled, you can run trials asynchronously:
//!
//! ```ignore
//! use optimize::{Study, Direction};
//!
//! // Sequential async
//! study.optimize_async(10, |mut trial| async move {
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
//! Ok((trial, x * x))
//! }).await?;
//!
//! // Parallel with bounded concurrency
//! study.optimize_parallel(10, 4, |mut trial| async move {
//! let x = trial.suggest_float("x", 0.0, 1.0)?;
//! Ok((trial, x * x))
//! }).await?;
//! ```
//!
//! # Serialization
//!
//! With the `serde` feature enabled, studies can be serialized:
//!
//! ```ignore
//! use optimize::{Study, Direction, TpeSampler};
//!
//! // Save study state
//! let study: Study<f64> = Study::new(Direction::Minimize);
//! let json = serde_json::to_string(&study)?;
//!
//! // Load and continue
//! let mut study: Study<f64> = serde_json::from_str(&json)?;
//! study.set_sampler(TpeSampler::new()); // Restore sampler
//! study.optimize_with_sampler(10, |trial| { /* ... */ }).unwrap();
//! ```
//!
//! # Feature Flags
//!
//! - `serde`: Enable serialization/deserialization of studies and trials
//! - `async`: Enable async optimization methods (requires tokio)
mod distribution;
mod error;
mod kde;
mod param;
mod sampler;
mod study;
mod trial;
mod types;
pub use error::{Result, TpeError};
pub use sampler::{CompletedTrial, RandomSampler, Sampler, TpeSampler, TpeSamplerBuilder};
pub use study::Study;
pub use trial::Trial;
pub use types::{Direction, TrialState};
+20
View File
@@ -0,0 +1,20 @@
//! Parameter value storage types.
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
/// Represents a sampled parameter value.
///
/// This enum stores different parameter value types uniformly.
/// For categorical parameters, the `Categorical` variant stores
/// the index into the choices array.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum ParamValue {
/// A floating-point parameter value.
Float(f64),
/// An integer parameter value.
Int(i64),
/// A categorical parameter value, stored as an index into the choices array.
Categorical(usize),
}
+96
View File
@@ -0,0 +1,96 @@
//! Sampler trait and implementations for parameter sampling.
mod random;
pub mod tpe;
use std::collections::HashMap;
pub use random::RandomSampler;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
pub use tpe::{TpeSampler, TpeSamplerBuilder};
use crate::distribution::Distribution;
use crate::param::ParamValue;
/// A completed trial with its parameters, distributions, and objective value.
///
/// This struct stores the results of a completed trial, including all sampled
/// parameter values, their distributions, and the objective value returned
/// by the objective function.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct CompletedTrial<V = f64> {
/// The unique identifier for this trial.
pub id: u64,
/// The sampled parameter values, keyed by parameter name.
pub params: HashMap<String, ParamValue>,
/// The parameter distributions used, keyed by parameter name.
pub distributions: HashMap<String, Distribution>,
/// The objective value returned by the objective function.
pub value: V,
}
impl<V> CompletedTrial<V> {
/// Creates a new completed trial.
pub fn new(
id: u64,
params: HashMap<String, ParamValue>,
distributions: HashMap<String, Distribution>,
value: V,
) -> Self {
Self {
id,
params,
distributions,
value,
}
}
}
/// Trait for pluggable parameter sampling strategies.
///
/// Samplers are responsible for generating parameter values based on
/// the distribution and historical trial data. The trait requires
/// `Send + Sync` to support concurrent and async optimization.
///
/// # Examples
///
/// Implementing a custom sampler:
///
/// ```ignore
/// use optimize::{Sampler, ParamValue, Distribution, CompletedTrial};
///
/// struct MySampler;
///
/// impl Sampler for MySampler {
/// fn sample(
/// &self,
/// distribution: &Distribution,
/// trial_id: u64,
/// history: &[CompletedTrial],
/// ) -> ParamValue {
/// // Custom sampling logic here
/// todo!()
/// }
/// }
/// ```
pub trait Sampler: Send + Sync {
/// Samples a parameter value from the given distribution.
///
/// # Arguments
///
/// * `distribution` - The parameter distribution to sample from.
/// * `trial_id` - The unique ID of the trial being sampled for.
/// * `history` - Historical completed trials for informed sampling.
///
/// # Returns
///
/// A `ParamValue` sampled from the distribution.
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue;
}
+275
View File
@@ -0,0 +1,275 @@
//! Random sampler implementation.
use parking_lot::Mutex;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::sampler::{CompletedTrial, Sampler};
/// A simple random sampler that samples uniformly from distributions.
///
/// This sampler ignores the trial history and samples uniformly at random,
/// respecting log scale and step size constraints. It serves as a baseline
/// sampler and is used during the startup phase of more sophisticated samplers.
///
/// # Examples
///
/// ```
/// use optimize::RandomSampler;
///
/// // Create with default RNG
/// let sampler = RandomSampler::new();
///
/// // Create with a fixed seed for reproducibility
/// let sampler = RandomSampler::with_seed(42);
/// ```
pub struct RandomSampler {
rng: Mutex<StdRng>,
}
impl RandomSampler {
/// Creates a new random sampler with a default random seed.
pub fn new() -> Self {
Self {
rng: Mutex::new(StdRng::from_os_rng()),
}
}
/// Creates a new random sampler with a fixed seed for reproducibility.
///
/// Using the same seed will produce the same sequence of sampled values.
pub fn with_seed(seed: u64) -> Self {
Self {
rng: Mutex::new(StdRng::seed_from_u64(seed)),
}
}
}
impl Default for RandomSampler {
fn default() -> Self {
Self::new()
}
}
impl Sampler for RandomSampler {
fn sample(
&self,
distribution: &Distribution,
_trial_id: u64,
_history: &[CompletedTrial],
) -> ParamValue {
let mut rng = self.rng.lock();
match distribution {
Distribution::Float(d) => {
let value = if d.log_scale {
// Sample uniformly in log space
let log_low = d.low.ln();
let log_high = d.high.ln();
let log_value = rng.random_range(log_low..=log_high);
log_value.exp()
} else if let Some(step) = d.step {
// Sample from step grid
let n_steps = ((d.high - d.low) / step).floor() as i64;
let k = rng.random_range(0..=n_steps);
d.low + (k as f64) * step
} else {
// Uniform sampling
rng.random_range(d.low..=d.high)
};
ParamValue::Float(value)
}
Distribution::Int(d) => {
let value = if d.log_scale {
// Sample uniformly in log space, then round
let log_low = (d.low as f64).ln();
let log_high = (d.high as f64).ln();
let log_value = rng.random_range(log_low..=log_high);
let raw = log_value.exp().round() as i64;
// Clamp to bounds since rounding might push outside
raw.clamp(d.low, d.high)
} else if let Some(step) = d.step {
// Sample from step grid
let n_steps = (d.high - d.low) / step;
let k = rng.random_range(0..=n_steps);
d.low + k * step
} else {
// Uniform sampling
rng.random_range(d.low..=d.high)
};
ParamValue::Int(value)
}
Distribution::Categorical(d) => {
let index = rng.random_range(0..d.n_choices);
ParamValue::Categorical(index)
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
#[test]
fn test_random_sampler_float() {
let sampler = RandomSampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v));
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn test_random_sampler_float_log() {
let sampler = RandomSampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: 1e-5,
high: 1.0,
log_scale: true,
step: None,
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
if let ParamValue::Float(v) = value {
assert!((1e-5..=1.0).contains(&v));
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn test_random_sampler_float_step() {
let sampler = RandomSampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: Some(0.25),
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v));
// Check it's on the step grid
let k = ((v - 0.0) / 0.25).round() as i64;
let expected = 0.0 + k as f64 * 0.25;
assert!((v - expected).abs() < 1e-10);
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn test_random_sampler_int() {
let sampler = RandomSampler::with_seed(42);
let dist = Distribution::Int(IntDistribution {
low: 0,
high: 10,
log_scale: false,
step: None,
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
if let ParamValue::Int(v) = value {
assert!((0..=10).contains(&v));
} else {
panic!("Expected Int value");
}
}
}
#[test]
fn test_random_sampler_int_log() {
let sampler = RandomSampler::with_seed(42);
let dist = Distribution::Int(IntDistribution {
low: 1,
high: 1000,
log_scale: true,
step: None,
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
if let ParamValue::Int(v) = value {
assert!((1..=1000).contains(&v));
} else {
panic!("Expected Int value");
}
}
}
#[test]
fn test_random_sampler_int_step() {
let sampler = RandomSampler::with_seed(42);
let dist = Distribution::Int(IntDistribution {
low: 0,
high: 10,
log_scale: false,
step: Some(2),
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
if let ParamValue::Int(v) = value {
assert!((0..=10).contains(&v));
// Check it's on the step grid: 0, 2, 4, 6, 8, 10
assert!(v % 2 == 0);
} else {
panic!("Expected Int value");
}
}
}
#[test]
fn test_random_sampler_categorical() {
let sampler = RandomSampler::with_seed(42);
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 });
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
if let ParamValue::Categorical(idx) = value {
assert!(idx < 5);
} else {
panic!("Expected Categorical value");
}
}
}
#[test]
fn test_random_sampler_reproducibility() {
let sampler1 = RandomSampler::with_seed(42);
let sampler2 = RandomSampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
for _ in 0..10 {
let v1 = sampler1.sample(&dist, 0, &[]);
let v2 = sampler2.sample(&dist, 0, &[]);
assert_eq!(v1, v2);
}
}
}
+1019
View File
File diff suppressed because it is too large Load Diff
+1340
View File
File diff suppressed because it is too large Load Diff
+793
View File
@@ -0,0 +1,793 @@
//! Trial implementation for tracking sampled parameters and trial state.
use std::collections::HashMap;
use std::sync::Arc;
use parking_lot::RwLock;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use crate::distribution::{
CategoricalDistribution, Distribution, FloatDistribution, IntDistribution,
};
use crate::error::{Result, TpeError};
use crate::param::ParamValue;
use crate::sampler::{CompletedTrial, Sampler};
use crate::types::TrialState;
/// A trial represents a single evaluation of the objective function.
///
/// Each trial has a unique ID and stores the sampled parameters along with
/// their distributions. The trial progresses through states: Running -> Complete/Failed.
///
/// Trials use a sampler to generate parameter values. When created through
/// `Study::create_trial()`, the trial receives the study's sampler and access
/// to the history of completed trials for informed sampling.
#[derive(Clone)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Trial {
/// Unique identifier for this trial.
id: u64,
/// Current state of the trial.
state: TrialState,
/// Sampled parameter values, keyed by parameter name.
params: HashMap<String, ParamValue>,
/// Parameter distributions, keyed by parameter name.
distributions: HashMap<String, Distribution>,
/// The sampler to use for generating parameter values.
#[cfg_attr(feature = "serde", serde(skip))]
sampler: Option<Arc<dyn Sampler>>,
/// Access to the history of completed trials (shared with Study).
#[cfg_attr(feature = "serde", serde(skip))]
history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>,
}
impl std::fmt::Debug for Trial {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Trial")
.field("id", &self.id)
.field("state", &self.state)
.field("params", &self.params)
.field("distributions", &self.distributions)
.field("has_sampler", &self.sampler.is_some())
.field("has_history", &self.history.is_some())
.finish()
}
}
impl Trial {
/// Creates a new trial with the given ID.
///
/// The trial starts in the `Running` state with no parameters sampled.
/// This constructor creates a trial without a sampler, which will use
/// local random sampling for suggest_* methods.
///
/// For trials that use the study's sampler, use `Trial::with_sampler` instead.
///
/// # Arguments
///
/// * `id` - A unique identifier for this trial.
///
/// # Examples
///
/// ```
/// use optimize::Trial;
///
/// let trial = Trial::new(0);
/// assert_eq!(trial.id(), 0);
/// ```
pub fn new(id: u64) -> Self {
Self {
id,
state: TrialState::Running,
params: HashMap::new(),
distributions: HashMap::new(),
sampler: None,
history: None,
}
}
/// Creates a new trial with a sampler and access to trial history.
///
/// This constructor is used by `Study::create_trial()` to create trials
/// that use the study's sampler for informed parameter suggestions.
///
/// # Arguments
///
/// * `id` - A unique identifier for this trial.
/// * `sampler` - The sampler to use for generating parameter values.
/// * `history` - Shared access to the history of completed trials.
pub(crate) fn with_sampler(
id: u64,
sampler: Arc<dyn Sampler>,
history: Arc<RwLock<Vec<CompletedTrial<f64>>>>,
) -> Self {
Self {
id,
state: TrialState::Running,
params: HashMap::new(),
distributions: HashMap::new(),
sampler: Some(sampler),
history: Some(history),
}
}
/// Samples a value from the given distribution using the sampler.
///
/// If the trial has a sampler, it delegates to the sampler's sample method
/// with the history of completed trials. Otherwise, it uses the RandomSampler
/// as a fallback.
fn sample_value(&self, distribution: &Distribution) -> ParamValue {
if let (Some(sampler), Some(history)) = (&self.sampler, &self.history) {
let history_guard = history.read();
sampler.sample(distribution, self.id, &history_guard)
} else {
// Fallback to RandomSampler when no sampler is configured
use crate::sampler::RandomSampler;
let fallback = RandomSampler::new();
fallback.sample(distribution, self.id, &[])
}
}
/// Returns the unique ID of this trial.
pub fn id(&self) -> u64 {
self.id
}
/// Returns the current state of this trial.
pub fn state(&self) -> TrialState {
self.state
}
/// Returns a reference to the sampled parameters.
pub fn params(&self) -> &HashMap<String, ParamValue> {
&self.params
}
/// Returns a reference to the parameter distributions.
pub fn distributions(&self) -> &HashMap<String, Distribution> {
&self.distributions
}
/// Sets the trial state to Complete.
pub(crate) fn set_complete(&mut self) {
self.state = TrialState::Complete;
}
/// Sets the trial state to Failed.
pub(crate) fn set_failed(&mut self) {
self.state = TrialState::Failed;
}
/// Suggests a float parameter with the given bounds.
///
/// If the parameter has already been sampled with the same bounds, the cached value is returned.
/// If the parameter was sampled with different bounds, a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive).
/// * `high` - The upper bound (inclusive).
///
/// # Errors
///
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different bounds.
///
/// # Examples
///
/// ```
/// use optimize::Trial;
///
/// let mut trial = Trial::new(0);
/// let x = trial.suggest_float("x", 0.0, 1.0).unwrap();
/// assert!(x >= 0.0 && x <= 1.0);
///
/// // Calling again with same bounds returns cached value
/// let x2 = trial.suggest_float("x", 0.0, 1.0).unwrap();
/// assert_eq!(x, x2);
/// ```
pub fn suggest_float(&mut self, name: impl Into<String>, low: f64, high: f64) -> Result<f64> {
if low > high {
return Err(TpeError::InvalidBounds { low, high });
}
let name = name.into();
let distribution = FloatDistribution {
low,
high,
log_scale: false,
step: None,
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Float(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& !existing.log_scale
&& existing.step.is_none()
{
// Same distribution, return cached value
if let Some(ParamValue::Float(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(TpeError::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler
let dist = Distribution::Float(distribution);
let value = match self.sample_value(&dist) {
ParamValue::Float(v) => v,
_ => unreachable!("Float distribution should return Float value"),
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Float(value));
Ok(value)
}
/// Suggests a float parameter sampled on a logarithmic scale.
///
/// The value is sampled uniformly in log space, which is useful for parameters
/// that span multiple orders of magnitude (e.g., learning rates).
///
/// If the parameter has already been sampled with the same bounds and log_scale=true,
/// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive, must be positive).
/// * `high` - The upper bound (inclusive).
///
/// # Errors
///
/// Returns `InvalidLogBounds` if `low <= 0`.
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
///
/// # Examples
///
/// ```
/// use optimize::Trial;
///
/// let mut trial = Trial::new(0);
/// let lr = trial
/// .suggest_float_log("learning_rate", 1e-5, 1e-1)
/// .unwrap();
/// assert!(lr >= 1e-5 && lr <= 1e-1);
///
/// // Calling again with same bounds returns cached value
/// let lr2 = trial
/// .suggest_float_log("learning_rate", 1e-5, 1e-1)
/// .unwrap();
/// assert_eq!(lr, lr2);
/// ```
pub fn suggest_float_log(
&mut self,
name: impl Into<String>,
low: f64,
high: f64,
) -> Result<f64> {
if low <= 0.0 {
return Err(TpeError::InvalidLogBounds);
}
if low > high {
return Err(TpeError::InvalidBounds { low, high });
}
let name = name.into();
let distribution = FloatDistribution {
low,
high,
log_scale: true,
step: None,
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Float(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& existing.log_scale
&& existing.step.is_none()
{
// Same distribution, return cached value
if let Some(ParamValue::Float(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(TpeError::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler (sampler handles log-scale transformation)
let dist = Distribution::Float(distribution);
let value = match self.sample_value(&dist) {
ParamValue::Float(v) => v,
_ => unreachable!("Float distribution should return Float value"),
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Float(value));
Ok(value)
}
/// Suggests a float parameter that snaps to a step grid.
///
/// The value is sampled from the discrete set {low, low + step, low + 2*step, ...}
/// where each value is <= high.
///
/// If the parameter has already been sampled with the same configuration,
/// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive).
/// * `high` - The upper bound (inclusive).
/// * `step` - The step size (must be positive).
///
/// # Errors
///
/// Returns `InvalidStep` if `step <= 0`.
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
///
/// # Examples
///
/// ```
/// use optimize::Trial;
///
/// let mut trial = Trial::new(0);
/// let x = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap();
/// // x will be one of: 0.0, 0.25, 0.5, 0.75, 1.0
/// assert!(x >= 0.0 && x <= 1.0);
/// assert!((x / 0.25).fract().abs() < 1e-10 || (x / 0.25).fract().abs() > 1.0 - 1e-10);
///
/// // Calling again with same bounds returns cached value
/// let x2 = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap();
/// assert_eq!(x, x2);
/// ```
pub fn suggest_float_step(
&mut self,
name: impl Into<String>,
low: f64,
high: f64,
step: f64,
) -> Result<f64> {
if step <= 0.0 {
return Err(TpeError::InvalidStep);
}
if low > high {
return Err(TpeError::InvalidBounds { low, high });
}
let name = name.into();
let distribution = FloatDistribution {
low,
high,
log_scale: false,
step: Some(step),
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Float(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& !existing.log_scale
&& existing.step == Some(step)
{
// Same distribution, return cached value
if let Some(ParamValue::Float(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(TpeError::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler (sampler handles step-grid)
let dist = Distribution::Float(distribution);
let value = match self.sample_value(&dist) {
ParamValue::Float(v) => v,
_ => unreachable!("Float distribution should return Float value"),
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Float(value));
Ok(value)
}
/// Suggests an integer parameter with the given bounds.
///
/// The value is sampled uniformly from the range [low, high] inclusive.
///
/// If the parameter has already been sampled with the same bounds, the cached value is returned.
/// If the parameter was sampled with different bounds, a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive).
/// * `high` - The upper bound (inclusive).
///
/// # Errors
///
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different bounds.
///
/// # Examples
///
/// ```
/// use optimize::Trial;
///
/// let mut trial = Trial::new(0);
/// let n = trial.suggest_int("n_layers", 1, 10).unwrap();
/// assert!(n >= 1 && n <= 10);
///
/// // Calling again with same bounds returns cached value
/// let n2 = trial.suggest_int("n_layers", 1, 10).unwrap();
/// assert_eq!(n, n2);
/// ```
pub fn suggest_int(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
if low > high {
return Err(TpeError::InvalidBounds {
low: low as f64,
high: high as f64,
});
}
let name = name.into();
let distribution = IntDistribution {
low,
high,
log_scale: false,
step: None,
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Int(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& !existing.log_scale
&& existing.step.is_none()
{
// Same distribution, return cached value
if let Some(ParamValue::Int(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(TpeError::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler
let dist = Distribution::Int(distribution);
let value = match self.sample_value(&dist) {
ParamValue::Int(v) => v,
_ => unreachable!("Int distribution should return Int value"),
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Int(value));
Ok(value)
}
/// Suggests an integer parameter sampled on a logarithmic scale.
///
/// The value is sampled uniformly in log space, which is useful for parameters
/// that span multiple orders of magnitude (e.g., batch sizes).
///
/// If the parameter has already been sampled with the same bounds and log_scale=true,
/// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive, must be >= 1).
/// * `high` - The upper bound (inclusive).
///
/// # Errors
///
/// Returns `InvalidLogBounds` if `low < 1`.
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
///
/// # Examples
///
/// ```
/// use optimize::Trial;
///
/// let mut trial = Trial::new(0);
/// let batch_size = trial.suggest_int_log("batch_size", 1, 1024).unwrap();
/// assert!(batch_size >= 1 && batch_size <= 1024);
///
/// // Calling again with same bounds returns cached value
/// let batch_size2 = trial.suggest_int_log("batch_size", 1, 1024).unwrap();
/// assert_eq!(batch_size, batch_size2);
/// ```
pub fn suggest_int_log(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
if low < 1 {
return Err(TpeError::InvalidLogBounds);
}
if low > high {
return Err(TpeError::InvalidBounds {
low: low as f64,
high: high as f64,
});
}
let name = name.into();
let distribution = IntDistribution {
low,
high,
log_scale: true,
step: None,
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Int(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& existing.log_scale
&& existing.step.is_none()
{
// Same distribution, return cached value
if let Some(ParamValue::Int(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(TpeError::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler (sampler handles log-scale transformation)
let dist = Distribution::Int(distribution);
let value = match self.sample_value(&dist) {
ParamValue::Int(v) => v,
_ => unreachable!("Int distribution should return Int value"),
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Int(value));
Ok(value)
}
/// Suggests an integer parameter that snaps to a step grid.
///
/// The value is sampled from the discrete set {low, low + step, low + 2*step, ...}
/// where each value is <= high.
///
/// If the parameter has already been sampled with the same configuration,
/// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `low` - The lower bound (inclusive).
/// * `high` - The upper bound (inclusive).
/// * `step` - The step size (must be positive).
///
/// # Errors
///
/// Returns `InvalidStep` if `step <= 0`.
/// Returns `InvalidBounds` if `low > high`.
/// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
///
/// # Examples
///
/// ```
/// use optimize::Trial;
///
/// let mut trial = Trial::new(0);
/// let n = trial
/// .suggest_int_step("n_estimators", 100, 500, 50)
/// .unwrap();
/// // n will be one of: 100, 150, 200, 250, 300, 350, 400, 450, 500
/// assert!(n >= 100 && n <= 500);
/// assert!((n - 100) % 50 == 0);
///
/// // Calling again with same bounds returns cached value
/// let n2 = trial
/// .suggest_int_step("n_estimators", 100, 500, 50)
/// .unwrap();
/// assert_eq!(n, n2);
/// ```
pub fn suggest_int_step(
&mut self,
name: impl Into<String>,
low: i64,
high: i64,
step: i64,
) -> Result<i64> {
if step <= 0 {
return Err(TpeError::InvalidStep);
}
if low > high {
return Err(TpeError::InvalidBounds {
low: low as f64,
high: high as f64,
});
}
let name = name.into();
let distribution = IntDistribution {
low,
high,
log_scale: false,
step: Some(step),
};
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Int(existing) = existing_dist
&& existing.low == low
&& existing.high == high
&& !existing.log_scale
&& existing.step == Some(step)
{
// Same distribution, return cached value
if let Some(ParamValue::Int(value)) = self.params.get(&name) {
return Ok(*value);
}
}
// Distribution exists but doesn't match
return Err(TpeError::ParameterConflict {
name,
reason: "parameter was previously sampled with different bounds or type"
.to_string(),
});
}
// Sample using the sampler (sampler handles step-grid)
let dist = Distribution::Int(distribution);
let value = match self.sample_value(&dist) {
ParamValue::Int(v) => v,
_ => unreachable!("Int distribution should return Int value"),
};
// Store distribution and value
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Int(value));
Ok(value)
}
/// Suggests a categorical parameter from the given choices.
///
/// The value is selected uniformly at random from the provided choices.
/// Internally, the index of the selected choice is stored.
///
/// If the parameter has already been sampled with the same number of choices,
/// the cached value (same index) is returned. If the parameter was sampled with
/// a different number of choices, a `ParameterConflict` error is returned.
///
/// # Arguments
///
/// * `name` - The name of the parameter.
/// * `choices` - A slice of choices to select from.
///
/// # Type Parameters
///
/// * `T` - The type of the choices. Only requires `Clone`.
///
/// # Errors
///
/// Returns `EmptyChoices` if `choices` is empty.
/// Returns `ParameterConflict` if the parameter was previously sampled with a different number of choices.
///
/// # Examples
///
/// ```
/// use optimize::Trial;
///
/// let mut trial = Trial::new(0);
/// let optimizer = trial
/// .suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])
/// .unwrap();
/// assert!(["sgd", "adam", "rmsprop"].contains(&optimizer));
///
/// // Calling again with same choices returns cached value
/// let optimizer2 = trial
/// .suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])
/// .unwrap();
/// assert_eq!(optimizer, optimizer2);
/// ```
pub fn suggest_categorical<T: Clone>(
&mut self,
name: impl Into<String>,
choices: &[T],
) -> Result<T> {
if choices.is_empty() {
return Err(TpeError::EmptyChoices);
}
let name = name.into();
let n_choices = choices.len();
let distribution = CategoricalDistribution { n_choices };
// Check if parameter already exists
if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches
if let Distribution::Categorical(existing) = existing_dist
&& existing.n_choices == n_choices
{
// Same distribution, return cached value
if let Some(ParamValue::Categorical(index)) = self.params.get(&name) {
return Ok(choices[*index].clone());
}
}
// Distribution exists but doesn't match
return Err(TpeError::ParameterConflict {
name,
reason: "parameter was previously sampled with different number of choices or type"
.to_string(),
});
}
// Sample using the sampler
let dist = Distribution::Categorical(distribution);
let index = match self.sample_value(&dist) {
ParamValue::Categorical(idx) => idx,
_ => unreachable!("Categorical distribution should return Categorical value"),
};
// Store distribution and value (store the index)
self.distributions.insert(name.clone(), dist);
self.params.insert(name, ParamValue::Categorical(index));
Ok(choices[index].clone())
}
}
+26
View File
@@ -0,0 +1,26 @@
//! Core types for the optimize library.
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
/// The direction of optimization.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum Direction {
/// Minimize the objective value.
Minimize,
/// Maximize the objective value.
Maximize,
}
/// The state of a trial in its lifecycle.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum TrialState {
/// The trial is currently running.
Running,
/// The trial completed successfully.
Complete,
/// The trial failed with an error.
Failed,
}