Manuel Raimann
804 lines
27 KiB
Rust
804 lines
27 KiB
Rust
//! Trial implementation for tracking sampled parameters and trial state.
|
|
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
use parking_lot::RwLock;
|
|
|
|
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)]
|
|
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.
|
|
sampler: Option<Arc<dyn Sampler>>,
|
|
/// Access to the history of completed trials (shared with Study).
|
|
history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>,
|
|
}
|
|
|
|
impl core::fmt::Debug for Trial {
|
|
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::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 optimizer::Trial;
|
|
///
|
|
/// let trial = Trial::new(0);
|
|
/// assert_eq!(trial.id(), 0);
|
|
/// ```
|
|
#[must_use]
|
|
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::random::RandomSampler;
|
|
let fallback = RandomSampler::new();
|
|
fallback.sample(distribution, self.id, &[])
|
|
}
|
|
}
|
|
|
|
/// Returns the unique ID of this trial.
|
|
#[must_use]
|
|
pub fn id(&self) -> u64 {
|
|
self.id
|
|
}
|
|
|
|
/// Returns the current state of this trial.
|
|
#[must_use]
|
|
pub fn state(&self) -> TrialState {
|
|
self.state
|
|
}
|
|
|
|
/// Returns a reference to the sampled parameters.
|
|
#[must_use]
|
|
pub fn params(&self) -> &HashMap<String, ParamValue> {
|
|
&self.params
|
|
}
|
|
|
|
/// Returns a reference to the parameter distributions.
|
|
#[must_use]
|
|
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 optimizer::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).abs() < f64::EPSILON
|
|
&& (existing.high - high).abs() < f64::EPSILON
|
|
&& !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 ParamValue::Float(value) = self.sample_value(&dist) else {
|
|
return Err(TpeError::Internal(
|
|
"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 optimizer::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).abs() < f64::EPSILON
|
|
&& (existing.high - high).abs() < f64::EPSILON
|
|
&& 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 ParamValue::Float(value) = self.sample_value(&dist) else {
|
|
return Err(TpeError::Internal(
|
|
"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 optimizer::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).abs() < f64::EPSILON
|
|
&& (existing.high - high).abs() < f64::EPSILON
|
|
&& !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 ParamValue::Float(value) = self.sample_value(&dist) else {
|
|
return Err(TpeError::Internal(
|
|
"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 optimizer::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);
|
|
/// ```
|
|
#[allow(clippy::cast_precision_loss)]
|
|
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 ParamValue::Int(value) = self.sample_value(&dist) else {
|
|
return Err(TpeError::Internal(
|
|
"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 optimizer::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);
|
|
/// ```
|
|
#[allow(clippy::cast_precision_loss)]
|
|
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 ParamValue::Int(value) = self.sample_value(&dist) else {
|
|
return Err(TpeError::Internal(
|
|
"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 optimizer::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);
|
|
/// ```
|
|
#[allow(clippy::cast_precision_loss)]
|
|
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 ParamValue::Int(value) = self.sample_value(&dist) else {
|
|
return Err(TpeError::Internal(
|
|
"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 optimizer::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 ParamValue::Categorical(index) = self.sample_value(&dist) else {
|
|
return Err(TpeError::Internal(
|
|
"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())
|
|
}
|
|
}
|