6 Commits

Author SHA1 Message Date
Manuel Raimann e44ab3669a chore: release v0.3.0 2026-01-30 19:21:42 +01:00
Manuel Raimann 7a3f89360e Refactor 2026-01-30 19:21:35 +01:00
Manuel Raimann 6912cc83d9 feat: add cross-target compilation checks in CI 2026-01-30 18:46:31 +01:00
Manuel Raimann fae57e48f3 chore: release v0.2.0 2026-01-30 18:43:45 +01:00
Manuel Raimann daab3ea202 Remove Serde 2026-01-30 18:43:45 +01:00
Manuel Raimann aaf880e1c7 fix: remove unused dependency 'ordered-float' from Cargo.toml 2026-01-30 18:43:45 +01:00
16 changed files with 456 additions and 678 deletions
+22 -2
View File
@@ -42,9 +42,7 @@ jobs:
matrix: matrix:
features: features:
- "" - ""
- "serde"
- "async" - "async"
- "serde,async"
steps: steps:
- uses: actions/checkout@v6 - uses: actions/checkout@v6
- name: Install Rust - name: Install Rust
@@ -197,6 +195,27 @@ jobs:
- name: Run tests - name: Run tests
run: cargo test --all-features run: cargo test --all-features
cross-targets:
name: ${{ matrix.target }}
runs-on: ubuntu-latest
strategy:
fail-fast: false
matrix:
target:
- aarch64-unknown-linux-gnu
- i686-unknown-linux-gnu
- powerpc64le-unknown-linux-gnu
- s390x-unknown-linux-gnu
steps:
- uses: actions/checkout@v6
- name: Install Rust
run: |
rustup override set stable
rustup update stable
rustup target add ${{ matrix.target }}
- uses: Swatinem/rust-cache@v2
- name: Check compilation
run: cargo check --all-features --target ${{ matrix.target }}
typos: typos:
name: Typos name: Typos
@@ -230,6 +249,7 @@ jobs:
- semver - semver
- minimal-versions - minimal-versions
- cross-platform - cross-platform
- cross-targets
- typos - typos
steps: steps:
- uses: actions/checkout@v6 - uses: actions/checkout@v6
+1 -5
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "optimizer" name = "optimizer"
version = "0.1.1" version = "0.3.0"
edition = "2024" edition = "2024"
rust-version = "1.88" rust-version = "1.88"
license = "MIT" license = "MIT"
@@ -13,15 +13,11 @@ repository = "https://github.com/raimannma/rust-optimizer"
rand = "0.9" rand = "0.9"
thiserror = "2" thiserror = "2"
parking_lot = "0.12" parking_lot = "0.12"
ordered-float = "5"
serde = { version = "1", features = ["derive"], optional = true }
tokio = { version = "1", features = ["sync", "rt-multi-thread"], optional = true } tokio = { version = "1", features = ["sync", "rt-multi-thread"], optional = true }
[features] [features]
default = [] default = []
serde = ["dep:serde", "ordered-float/serde"]
async = ["dep:tokio"] async = ["dep:tokio"]
[dev-dependencies] [dev-dependencies]
serde_json = "1"
tokio = { version = "1", features = ["rt-multi-thread", "macros"] } tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
-2
View File
@@ -12,7 +12,6 @@ A Rust library for black-box optimization using Tree-Parzen Estimator (TPE).
- Float, integer, and categorical parameter types - Float, integer, and categorical parameter types
- Log-scale and stepped parameter sampling - Log-scale and stepped parameter sampling
- Sync and async optimization with parallel trial evaluation - Sync and async optimization with parallel trial evaluation
- Serialization support for saving/loading study state
## Quick Start ## Quick Start
@@ -35,7 +34,6 @@ println!("Best value: {} at x={:?}", best.value, best.params);
## Feature Flags ## Feature Flags
- `serde` - Enable serialization/deserialization of studies and trials
- `async` - Enable async optimization methods (requires tokio) - `async` - Enable async optimization methods (requires tokio)
## Documentation ## Documentation
-7
View File
@@ -1,11 +1,7 @@
//! Parameter distribution types. //! Parameter distribution types.
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
/// Distribution for floating-point parameters. /// Distribution for floating-point parameters.
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct FloatDistribution { pub struct FloatDistribution {
/// Lower bound (inclusive). /// Lower bound (inclusive).
pub low: f64, pub low: f64,
@@ -19,7 +15,6 @@ pub struct FloatDistribution {
/// Distribution for integer parameters. /// Distribution for integer parameters.
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct IntDistribution { pub struct IntDistribution {
/// Lower bound (inclusive). /// Lower bound (inclusive).
pub low: i64, pub low: i64,
@@ -33,7 +28,6 @@ pub struct IntDistribution {
/// Distribution for categorical parameters. /// Distribution for categorical parameters.
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct CategoricalDistribution { pub struct CategoricalDistribution {
/// Number of choices available. /// Number of choices available.
pub n_choices: usize, pub n_choices: usize,
@@ -41,7 +35,6 @@ pub struct CategoricalDistribution {
/// Enum wrapping all parameter distribution types. /// Enum wrapping all parameter distribution types.
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum Distribution { pub enum Distribution {
/// A floating-point distribution. /// A floating-point distribution.
Float(FloatDistribution), Float(FloatDistribution),
+22 -1
View File
@@ -38,7 +38,28 @@ pub enum TpeError {
/// Returned when requesting the best trial but no trials have completed. /// Returned when requesting the best trial but no trials have completed.
#[error("no completed trials available")] #[error("no completed trials available")]
NoCompletedTrials, NoCompletedTrials,
/// Returned when gamma is not in the valid range (0.0, 1.0).
#[error("invalid gamma: {0} must be in (0.0, 1.0)")]
InvalidGamma(f64),
/// Returned when bandwidth is not positive.
#[error("invalid bandwidth: {0} must be positive")]
InvalidBandwidth(f64),
/// Returned when KDE is created with empty samples.
#[error("KDE requires at least one sample")]
EmptySamples,
/// Returned when an internal invariant is violated.
#[error("internal error: {0}")]
Internal(&'static str),
/// Returned when an async task fails.
#[cfg(feature = "async")]
#[error("async task error: {0}")]
TaskError(String),
} }
/// A specialized Result type for TPE operations. /// A specialized Result type for TPE operations.
pub type Result<T> = std::result::Result<T, TpeError>; pub type Result<T> = core::result::Result<T, TpeError>;
+46 -34
View File
@@ -5,6 +5,8 @@
use rand::Rng; use rand::Rng;
use crate::error::{Result, TpeError};
/// A Gaussian kernel density estimator for continuous distributions. /// A Gaussian kernel density estimator for continuous distributions.
/// ///
/// KDE estimates a probability density function from a set of samples by /// KDE estimates a probability density function from a set of samples by
@@ -28,7 +30,7 @@ use rand::Rng;
/// let sample = kde.sample(&mut rng); /// let sample = kde.sample(&mut rng);
/// ``` /// ```
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct KernelDensityEstimator { pub(crate) struct KernelDensityEstimator {
/// The sample points used to construct the KDE. /// The sample points used to construct the KDE.
samples: Vec<f64>, samples: Vec<f64>,
/// The bandwidth (standard deviation) of the Gaussian kernels. /// The bandwidth (standard deviation) of the Gaussian kernels.
@@ -38,37 +40,45 @@ pub struct KernelDensityEstimator {
impl KernelDensityEstimator { impl KernelDensityEstimator {
/// Creates a new KDE with automatic bandwidth selection using Scott's rule. /// 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 /// Scott's rule sets bandwidth = n^(-1/5) * `std_dev`, which works well
/// for unimodal distributions close to normal. /// for unimodal distributions close to normal.
/// ///
/// # Panics /// # Errors
/// ///
/// Panics if `samples` is empty. /// Returns `TpeError::EmptySamples` if `samples` is empty.
pub fn new(samples: Vec<f64>) -> Self { pub(crate) fn new(samples: Vec<f64>) -> Result<Self> {
assert!(!samples.is_empty(), "KDE requires at least one sample"); if samples.is_empty() {
return Err(TpeError::EmptySamples);
}
let bandwidth = Self::scotts_rule(&samples); let bandwidth = Self::scotts_rule(&samples);
Self { samples, bandwidth } Ok(Self { samples, bandwidth })
} }
/// Creates a new KDE with a specified bandwidth. /// Creates a new KDE with a specified bandwidth.
/// ///
/// Use this when you want explicit control over the smoothing parameter. /// Use this when you want explicit control over the smoothing parameter.
/// ///
/// # Panics /// # Errors
/// ///
/// Panics if `samples` is empty or `bandwidth` is not positive. /// Returns `TpeError::EmptySamples` if `samples` is empty.
pub(crate) fn with_bandwidth(samples: Vec<f64>, bandwidth: f64) -> Self { /// Returns `TpeError::InvalidBandwidth` if `bandwidth` is not positive.
assert!(!samples.is_empty(), "KDE requires at least one sample"); pub(crate) fn with_bandwidth(samples: Vec<f64>, bandwidth: f64) -> Result<Self> {
assert!(bandwidth > 0.0, "Bandwidth must be positive"); if samples.is_empty() {
return Err(TpeError::EmptySamples);
}
if bandwidth <= 0.0 {
return Err(TpeError::InvalidBandwidth(bandwidth));
}
Self { samples, bandwidth } Ok(Self { samples, bandwidth })
} }
/// Computes bandwidth using Scott's rule. /// Computes bandwidth using Scott's rule.
/// ///
/// Scott's rule: h = n^(-1/5) * sigma /// Scott's rule: h = n^(-1/5) * sigma
/// where sigma is the sample standard deviation. /// where sigma is the sample standard deviation.
#[allow(clippy::cast_precision_loss)]
fn scotts_rule(samples: &[f64]) -> f64 { fn scotts_rule(samples: &[f64]) -> f64 {
let n = samples.len() as f64; let n = samples.len() as f64;
let std_dev = Self::sample_std_dev(samples); let std_dev = Self::sample_std_dev(samples);
@@ -83,6 +93,7 @@ impl KernelDensityEstimator {
} }
/// Computes the sample standard deviation. /// Computes the sample standard deviation.
#[allow(clippy::cast_precision_loss)]
fn sample_std_dev(samples: &[f64]) -> f64 { fn sample_std_dev(samples: &[f64]) -> f64 {
let n = samples.len() as f64; let n = samples.len() as f64;
let mean = samples.iter().sum::<f64>() / n; let mean = samples.iter().sum::<f64>() / n;
@@ -95,13 +106,14 @@ impl KernelDensityEstimator {
/// The density is computed as the average of Gaussian kernels centered /// The density is computed as the average of Gaussian kernels centered
/// at each sample point: /// at each sample point:
/// ///
/// f(x) = (1/n) * sum_i K((x - x_i) / h) /// f(x) = (1/n) * `sum_i` K((x - `x_i`) / h)
/// ///
/// where K is the standard Gaussian kernel and h is the bandwidth. /// where K is the standard Gaussian kernel and h is the bandwidth.
pub fn pdf(&self, x: f64) -> f64 { #[allow(clippy::cast_precision_loss)]
pub(crate) fn pdf(&self, x: f64) -> f64 {
let n = self.samples.len() as f64; let n = self.samples.len() as f64;
let inv_bandwidth = 1.0 / self.bandwidth; let inv_bandwidth = 1.0 / self.bandwidth;
let normalization = inv_bandwidth / (2.0 * std::f64::consts::PI).sqrt(); let normalization = inv_bandwidth / (2.0 * core::f64::consts::PI).sqrt();
let density: f64 = self let density: f64 = self
.samples .samples
@@ -120,7 +132,7 @@ impl KernelDensityEstimator {
/// Sampling works by: /// Sampling works by:
/// 1. Uniformly selecting one of the kernel centers (samples) /// 1. Uniformly selecting one of the kernel centers (samples)
/// 2. Adding Gaussian noise with the bandwidth as standard deviation /// 2. Adding Gaussian noise with the bandwidth as standard deviation
pub fn sample<R: Rng>(&self, rng: &mut R) -> f64 { pub(crate) fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
// Select a random sample to center the kernel on // Select a random sample to center the kernel on
let idx = rng.random_range(0..self.samples.len()); let idx = rng.random_range(0..self.samples.len());
let center = self.samples[idx]; let center = self.samples[idx];
@@ -130,7 +142,7 @@ impl KernelDensityEstimator {
let u1: f64 = rng.random(); let u1: f64 = rng.random();
let u2: f64 = rng.random(); let u2: f64 = rng.random();
let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos(); let z = (-2.0 * u1.ln()).sqrt() * (2.0 * core::f64::consts::PI * u2).cos();
center + z * self.bandwidth center + z * self.bandwidth
} }
@@ -148,7 +160,7 @@ mod tests {
#[test] #[test]
fn test_kde_pdf_basic() { fn test_kde_pdf_basic() {
let samples = vec![0.0, 1.0, 2.0]; let samples = vec![0.0, 1.0, 2.0];
let kde = KernelDensityEstimator::new(samples); let kde = KernelDensityEstimator::new(samples).unwrap();
// Density should be positive everywhere // Density should be positive everywhere
assert!(kde.pdf(0.0) > 0.0); assert!(kde.pdf(0.0) > 0.0);
@@ -164,17 +176,17 @@ mod tests {
#[test] #[test]
fn test_kde_pdf_integrates_to_one() { fn test_kde_pdf_integrates_to_one() {
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0]; let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
let kde = KernelDensityEstimator::new(samples); let kde = KernelDensityEstimator::new(samples).unwrap();
// Numerical integration over a wide range // Numerical integration over a wide range
let n_points = 10000; let n_points = 10000;
let low = -10.0; let low = -10.0;
let high = 15.0; let high = 15.0;
let dx = (high - low) / n_points as f64; let dx = (high - low) / f64::from(n_points);
let integral: f64 = (0..n_points) let integral: f64 = (0..n_points)
.map(|i| { .map(|i| {
let x = low + (i as f64 + 0.5) * dx; let x = low + (f64::from(i) + 0.5) * dx;
kde.pdf(x) * dx kde.pdf(x) * dx
}) })
.sum(); .sum();
@@ -189,16 +201,16 @@ mod tests {
#[test] #[test]
fn test_kde_with_bandwidth() { fn test_kde_with_bandwidth() {
let samples = vec![0.0, 1.0, 2.0]; let samples = vec![0.0, 1.0, 2.0];
let kde = KernelDensityEstimator::with_bandwidth(samples, 0.5); let kde = KernelDensityEstimator::with_bandwidth(samples, 0.5).unwrap();
assert_eq!(kde.bandwidth(), 0.5); assert!((kde.bandwidth() - 0.5).abs() < f64::EPSILON);
assert!(kde.pdf(1.0) > 0.0); assert!(kde.pdf(1.0) > 0.0);
} }
#[test] #[test]
fn test_kde_sample_in_reasonable_range() { fn test_kde_sample_in_reasonable_range() {
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0]; let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
let kde = KernelDensityEstimator::new(samples); let kde = KernelDensityEstimator::new(samples).unwrap();
let mut rng = rand::rng(); let mut rng = rand::rng();
// Samples should generally be in a reasonable range around the data // Samples should generally be in a reasonable range around the data
@@ -213,7 +225,7 @@ mod tests {
#[test] #[test]
fn test_kde_single_sample() { fn test_kde_single_sample() {
let samples = vec![5.0]; let samples = vec![5.0];
let kde = KernelDensityEstimator::new(samples); let kde = KernelDensityEstimator::new(samples).unwrap();
// Should have positive density near the sample // Should have positive density near the sample
assert!(kde.pdf(5.0) > 0.0); assert!(kde.pdf(5.0) > 0.0);
@@ -223,7 +235,7 @@ mod tests {
#[test] #[test]
fn test_kde_identical_samples() { fn test_kde_identical_samples() {
let samples = vec![3.0, 3.0, 3.0, 3.0]; let samples = vec![3.0, 3.0, 3.0, 3.0];
let kde = KernelDensityEstimator::new(samples); let kde = KernelDensityEstimator::new(samples).unwrap();
// Should handle degenerate case with identical samples // Should handle degenerate case with identical samples
assert!(kde.bandwidth() > 0.0); assert!(kde.bandwidth() > 0.0);
@@ -233,7 +245,7 @@ mod tests {
#[test] #[test]
fn test_scotts_rule_bandwidth() { 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 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); let kde = KernelDensityEstimator::new(samples).unwrap();
// n = 10, n^(-1/5) ≈ 0.631 // n = 10, n^(-1/5) ≈ 0.631
// std_dev ≈ 2.87 // std_dev ≈ 2.87
@@ -246,23 +258,23 @@ mod tests {
} }
#[test] #[test]
#[should_panic(expected = "KDE requires at least one sample")]
fn test_kde_empty_samples() { fn test_kde_empty_samples() {
let samples: Vec<f64> = vec![]; let samples: Vec<f64> = vec![];
KernelDensityEstimator::new(samples); let result = KernelDensityEstimator::new(samples);
assert!(matches!(result, Err(TpeError::EmptySamples)));
} }
#[test] #[test]
#[should_panic(expected = "Bandwidth must be positive")]
fn test_kde_zero_bandwidth() { fn test_kde_zero_bandwidth() {
let samples = vec![1.0, 2.0, 3.0]; let samples = vec![1.0, 2.0, 3.0];
KernelDensityEstimator::with_bandwidth(samples, 0.0); let result = KernelDensityEstimator::with_bandwidth(samples, 0.0);
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
} }
#[test] #[test]
#[should_panic(expected = "Bandwidth must be positive")]
fn test_kde_negative_bandwidth() { fn test_kde_negative_bandwidth() {
let samples = vec![1.0, 2.0, 3.0]; let samples = vec![1.0, 2.0, 3.0];
KernelDensityEstimator::with_bandwidth(samples, -1.0); let result = KernelDensityEstimator::with_bandwidth(samples, -1.0);
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
} }
} }
+23 -27
View File
@@ -1,3 +1,14 @@
#![forbid(unsafe_code)]
#![deny(clippy::all)]
#![deny(unreachable_pub)]
#![deny(clippy::correctness)]
#![deny(clippy::suspicious)]
#![deny(clippy::style)]
#![deny(clippy::complexity)]
#![deny(clippy::perf)]
#![deny(clippy::pedantic)]
#![deny(clippy::std_instead_of_core)]
//! A Tree-Parzen Estimator (TPE) library for black-box optimization. //! A Tree-Parzen Estimator (TPE) library for black-box optimization.
//! //!
//! This library provides an Optuna-like API for hyperparameter optimization //! This library provides an Optuna-like API for hyperparameter optimization
@@ -7,15 +18,15 @@
//! - Log-scale and stepped parameter sampling //! - Log-scale and stepped parameter sampling
//! - Synchronous and async optimization //! - Synchronous and async optimization
//! - Parallel trial evaluation with bounded concurrency //! - Parallel trial evaluation with bounded concurrency
//! - Serialization for saving/loading study state
//! //!
//! # Quick Start //! # Quick Start
//! //!
//! ``` //! ```
//! use optimizer::{Direction, Study, TpeSampler}; //! use optimizer::sampler::tpe::TpeSampler;
//! use optimizer::{Direction, Study};
//! //!
//! // Create a study with TPE sampler //! // Create a study with TPE sampler
//! let sampler = TpeSampler::builder().seed(42).build(); //! let sampler = TpeSampler::builder().seed(42).build().unwrap();
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); //! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! //!
//! // Optimize x^2 for 20 trials //! // Optimize x^2 for 20 trials
@@ -36,7 +47,9 @@
//! A [`Study`] manages optimization trials. Create one with an optimization direction: //! A [`Study`] manages optimization trials. Create one with an optimization direction:
//! //!
//! ``` //! ```
//! use optimizer::{Direction, RandomSampler, Study, TpeSampler}; //! use optimizer::sampler::random::RandomSampler;
//! use optimizer::sampler::tpe::TpeSampler;
//! use optimizer::{Direction, Study};
//! //!
//! // Minimize with default random sampler //! // Minimize with default random sampler
//! let study: Study<f64> = Study::new(Direction::Minimize); //! let study: Study<f64> = Study::new(Direction::Minimize);
@@ -80,17 +93,18 @@
//! //!
//! # Configuring TPE //! # Configuring TPE
//! //!
//! The [`TpeSampler`] can be configured using the builder pattern: //! The [`sampler::tpe::TpeSampler`] can be configured using the builder pattern:
//! //!
//! ``` //! ```
//! use optimizer::TpeSampler; //! use optimizer::sampler::tpe::TpeSampler;
//! //!
//! let sampler = TpeSampler::builder() //! let sampler = TpeSampler::builder()
//! .gamma(0.15) // Quantile for good/bad split //! .gamma(0.15) // Quantile for good/bad split
//! .n_startup_trials(20) // Random trials before TPE //! .n_startup_trials(20) // Random trials before TPE
//! .n_ei_candidates(32) // Candidates to evaluate //! .n_ei_candidates(32) // Candidates to evaluate
//! .seed(42) // Reproducibility //! .seed(42) // Reproducibility
//! .build(); //! .build()
//! .unwrap();
//! ``` //! ```
//! //!
//! # Async and Parallel Optimization //! # Async and Parallel Optimization
@@ -113,39 +127,21 @@
//! }).await?; //! }).await?;
//! ``` //! ```
//! //!
//! # Serialization
//!
//! With the `serde` feature enabled, studies can be serialized:
//!
//! ```ignore
//! use optimizer::{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 //! # Feature Flags
//! //!
//! - `serde`: Enable serialization/deserialization of studies and trials
//! - `async`: Enable async optimization methods (requires tokio) //! - `async`: Enable async optimization methods (requires tokio)
mod distribution; mod distribution;
mod error; mod error;
mod kde; mod kde;
mod param; mod param;
mod sampler; pub mod sampler;
mod study; mod study;
mod trial; mod trial;
mod types; mod types;
pub use error::{Result, TpeError}; pub use error::{Result, TpeError};
pub use sampler::{CompletedTrial, RandomSampler, Sampler, TpeSampler, TpeSamplerBuilder}; pub use param::ParamValue;
pub use study::Study; pub use study::Study;
pub use trial::Trial; pub use trial::Trial;
pub use types::{Direction, TrialState}; pub use types::{Direction, TrialState};
-4
View File
@@ -1,15 +1,11 @@
//! Parameter value storage types. //! Parameter value storage types.
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
/// Represents a sampled parameter value. /// Represents a sampled parameter value.
/// ///
/// This enum stores different parameter value types uniformly. /// This enum stores different parameter value types uniformly.
/// For categorical parameters, the `Categorical` variant stores /// For categorical parameters, the `Categorical` variant stores
/// the index into the choices array. /// the index into the choices array.
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum ParamValue { pub enum ParamValue {
/// A floating-point parameter value. /// A floating-point parameter value.
Float(f64), Float(f64),
+1 -7
View File
@@ -1,15 +1,10 @@
//! Sampler trait and implementations for parameter sampling. //! Sampler trait and implementations for parameter sampling.
mod random; pub mod random;
pub mod tpe; pub mod tpe;
use std::collections::HashMap; 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::distribution::Distribution;
use crate::param::ParamValue; use crate::param::ParamValue;
@@ -19,7 +14,6 @@ use crate::param::ParamValue;
/// parameter values, their distributions, and the objective value returned /// parameter values, their distributions, and the objective value returned
/// by the objective function. /// by the objective function.
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct CompletedTrial<V = f64> { pub struct CompletedTrial<V = f64> {
/// The unique identifier for this trial. /// The unique identifier for this trial.
pub id: u64, pub id: u64,
+5 -1
View File
@@ -17,7 +17,7 @@ use crate::sampler::{CompletedTrial, Sampler};
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::RandomSampler; /// use optimizer::sampler::random::RandomSampler;
/// ///
/// // Create with default RNG /// // Create with default RNG
/// let sampler = RandomSampler::new(); /// let sampler = RandomSampler::new();
@@ -31,6 +31,7 @@ pub struct RandomSampler {
impl RandomSampler { impl RandomSampler {
/// Creates a new random sampler with a default random seed. /// Creates a new random sampler with a default random seed.
#[must_use]
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
rng: Mutex::new(StdRng::from_os_rng()), rng: Mutex::new(StdRng::from_os_rng()),
@@ -40,6 +41,7 @@ impl RandomSampler {
/// Creates a new random sampler with a fixed seed for reproducibility. /// Creates a new random sampler with a fixed seed for reproducibility.
/// ///
/// Using the same seed will produce the same sequence of sampled values. /// Using the same seed will produce the same sequence of sampled values.
#[must_use]
pub fn with_seed(seed: u64) -> Self { pub fn with_seed(seed: u64) -> Self {
Self { Self {
rng: Mutex::new(StdRng::seed_from_u64(seed)), rng: Mutex::new(StdRng::seed_from_u64(seed)),
@@ -54,6 +56,7 @@ impl Default for RandomSampler {
} }
impl Sampler for RandomSampler { impl Sampler for RandomSampler {
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
fn sample( fn sample(
&self, &self,
distribution: &Distribution, distribution: &Distribution,
@@ -110,6 +113,7 @@ impl Sampler for RandomSampler {
} }
#[cfg(test)] #[cfg(test)]
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
mod tests { mod tests {
use super::*; use super::*;
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution}; use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
+148 -99
View File
@@ -9,6 +9,7 @@ use rand::rngs::StdRng;
use rand::{Rng, SeedableRng}; use rand::{Rng, SeedableRng};
use crate::distribution::Distribution; use crate::distribution::Distribution;
use crate::error::{Result, TpeError};
use crate::kde::KernelDensityEstimator; use crate::kde::KernelDensityEstimator;
use crate::param::ParamValue; use crate::param::ParamValue;
use crate::sampler::{CompletedTrial, Sampler}; use crate::sampler::{CompletedTrial, Sampler};
@@ -27,7 +28,7 @@ use crate::sampler::{CompletedTrial, Sampler};
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSampler; /// use optimizer::sampler::tpe::TpeSampler;
/// ///
/// // Create with default settings /// // Create with default settings
/// let sampler = TpeSampler::new(); /// let sampler = TpeSampler::new();
@@ -38,7 +39,8 @@ use crate::sampler::{CompletedTrial, Sampler};
/// .n_startup_trials(20) /// .n_startup_trials(20)
/// .n_ei_candidates(32) /// .n_ei_candidates(32)
/// .seed(42) /// .seed(42)
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
pub struct TpeSampler { pub struct TpeSampler {
/// Fraction of trials to consider as "good" (gamma quantile). /// Fraction of trials to consider as "good" (gamma quantile).
@@ -58,9 +60,10 @@ impl TpeSampler {
/// ///
/// Default settings: /// Default settings:
/// - gamma: 0.25 (top 25% of trials are considered "good") /// - gamma: 0.25 (top 25% of trials are considered "good")
/// - n_startup_trials: 10 (random sampling for first 10 trials) /// - `n_startup_trials`: 10 (random sampling for first 10 trials)
/// - n_ei_candidates: 24 (evaluate 24 candidates per sample) /// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample)
/// - kde_bandwidth: None (uses Scott's rule for automatic bandwidth) /// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth)
#[must_use]
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
gamma: 0.25, gamma: 0.25,
@@ -76,15 +79,17 @@ impl TpeSampler {
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSampler; /// use optimizer::sampler::tpe::TpeSampler;
/// ///
/// let sampler = TpeSampler::builder() /// let sampler = TpeSampler::builder()
/// .gamma(0.15) /// .gamma(0.15)
/// .n_startup_trials(20) /// .n_startup_trials(20)
/// .n_ei_candidates(32) /// .n_ei_candidates(32)
/// .seed(42) /// .seed(42)
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
#[must_use]
pub fn builder() -> TpeSamplerBuilder { pub fn builder() -> TpeSamplerBuilder {
TpeSamplerBuilder::new() TpeSamplerBuilder::new()
} }
@@ -99,22 +104,24 @@ impl TpeSampler {
/// * `kde_bandwidth` - Optional fixed bandwidth for KDE. If None, uses Scott's rule. /// * `kde_bandwidth` - Optional fixed bandwidth for KDE. If None, uses Scott's rule.
/// * `seed` - Optional seed for reproducibility. /// * `seed` - Optional seed for reproducibility.
/// ///
/// # Panics /// # Errors
/// ///
/// Panics if gamma is not in (0.0, 1.0) or if kde_bandwidth is Some but not positive. /// Returns `TpeError::InvalidGamma` if gamma is not in (0.0, 1.0).
/// Returns `TpeError::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
pub fn with_config( pub fn with_config(
gamma: f64, gamma: f64,
n_startup_trials: usize, n_startup_trials: usize,
n_ei_candidates: usize, n_ei_candidates: usize,
kde_bandwidth: Option<f64>, kde_bandwidth: Option<f64>,
seed: Option<u64>, seed: Option<u64>,
) -> Self { ) -> Result<Self> {
assert!( if gamma <= 0.0 || gamma >= 1.0 {
gamma > 0.0 && gamma < 1.0, return Err(TpeError::InvalidGamma(gamma));
"gamma must be in (0.0, 1.0), got {gamma}" }
); if let Some(bw) = kde_bandwidth
if let Some(bw) = kde_bandwidth { && bw <= 0.0
assert!(bw > 0.0, "kde_bandwidth must be positive, got {bw}"); {
return Err(TpeError::InvalidBandwidth(bw));
} }
let rng = match seed { let rng = match seed {
@@ -122,19 +129,24 @@ impl TpeSampler {
None => StdRng::from_os_rng(), None => StdRng::from_os_rng(),
}; };
Self { Ok(Self {
gamma, gamma,
n_startup_trials, n_startup_trials,
n_ei_candidates, n_ei_candidates,
kde_bandwidth, kde_bandwidth,
rng: Mutex::new(rng), rng: Mutex::new(rng),
} })
} }
/// Splits trials into good and bad groups based on the gamma quantile. /// Splits trials into good and bad groups based on the gamma quantile.
/// ///
/// Returns (good_trials, bad_trials) where good_trials contains trials /// Returns (`good_trials`, `bad_trials`) where `good_trials` contains trials
/// with values below the gamma quantile (for minimization). /// with values below the gamma quantile (for minimization).
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
fn split_trials<'a>( fn split_trials<'a>(
&self, &self,
history: &'a [CompletedTrial], history: &'a [CompletedTrial],
@@ -149,7 +161,7 @@ impl TpeSampler {
history[a] history[a]
.value .value
.partial_cmp(&history[b].value) .partial_cmp(&history[b].value)
.unwrap_or(std::cmp::Ordering::Equal) .unwrap_or(core::cmp::Ordering::Equal)
}); });
// Calculate the split point (gamma quantile) // Calculate the split point (gamma quantile)
@@ -171,6 +183,11 @@ impl TpeSampler {
} }
/// Samples uniformly from a distribution (used during startup phase). /// Samples uniformly from a distribution (used during startup phase).
#[allow(
clippy::cast_possible_truncation,
clippy::cast_precision_loss,
clippy::unused_self
)]
fn sample_uniform(&self, distribution: &Distribution, rng: &mut StdRng) -> ParamValue { fn sample_uniform(&self, distribution: &Distribution, rng: &mut StdRng) -> ParamValue {
match distribution { match distribution {
Distribution::Float(d) => { Distribution::Float(d) => {
@@ -241,6 +258,11 @@ impl TpeSampler {
None => KernelDensityEstimator::new(bad_internal), None => KernelDensityEstimator::new(bad_internal),
}; };
// If KDE construction fails, fall back to uniform sampling
let (Ok(l_kde), Ok(g_kde)) = (l_kde, g_kde) else {
return rng.random_range(low..=high);
};
// Generate candidates from l(x) and select the one with best l(x)/g(x) ratio // Generate candidates from l(x) and select the one with best l(x)/g(x) ratio
let mut best_candidate = internal_low; let mut best_candidate = internal_low;
let mut best_ratio = f64::NEG_INFINITY; let mut best_ratio = f64::NEG_INFINITY;
@@ -289,15 +311,19 @@ impl TpeSampler {
} }
/// Samples using TPE for integer distributions. /// Samples using TPE for integer distributions.
#[allow(clippy::too_many_arguments)] #[allow(
clippy::too_many_arguments,
clippy::cast_precision_loss,
clippy::cast_possible_truncation
)]
fn sample_tpe_int( fn sample_tpe_int(
&self, &self,
low: i64, low: i64,
high: i64, high: i64,
log_scale: bool, log_scale: bool,
step: Option<i64>, step: Option<i64>,
good_values: Vec<i64>, good_values: &[i64],
bad_values: Vec<i64>, bad_values: &[i64],
rng: &mut StdRng, rng: &mut StdRng,
) -> i64 { ) -> i64 {
// Convert to floats for KDE // Convert to floats for KDE
@@ -331,23 +357,24 @@ impl TpeSampler {
} }
/// Samples using TPE for categorical distributions. /// Samples using TPE for categorical distributions.
#[allow(clippy::cast_precision_loss, clippy::unused_self)]
fn sample_tpe_categorical( fn sample_tpe_categorical(
&self, &self,
n_choices: usize, n_choices: usize,
good_indices: Vec<usize>, good_indices: &[usize],
bad_indices: Vec<usize>, bad_indices: &[usize],
rng: &mut StdRng, rng: &mut StdRng,
) -> usize { ) -> usize {
// Count occurrences in good and bad groups // Count occurrences in good and bad groups
let mut good_counts = vec![0usize; n_choices]; let mut good_counts = vec![0usize; n_choices];
let mut bad_counts = vec![0usize; n_choices]; let mut bad_counts = vec![0usize; n_choices];
for &idx in &good_indices { for &idx in good_indices {
if idx < n_choices { if idx < n_choices {
good_counts[idx] += 1; good_counts[idx] += 1;
} }
} }
for &idx in &bad_indices { for &idx in bad_indices {
if idx < n_choices { if idx < n_choices {
bad_counts[idx] += 1; bad_counts[idx] += 1;
} }
@@ -395,14 +422,15 @@ impl Default for TpeSampler {
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSamplerBuilder; /// use optimizer::sampler::tpe::TpeSamplerBuilder;
/// ///
/// let sampler = TpeSamplerBuilder::new() /// let sampler = TpeSamplerBuilder::new()
/// .gamma(0.15) /// .gamma(0.15)
/// .n_startup_trials(20) /// .n_startup_trials(20)
/// .n_ei_candidates(32) /// .n_ei_candidates(32)
/// .seed(42) /// .seed(42)
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct TpeSamplerBuilder { pub struct TpeSamplerBuilder {
@@ -418,10 +446,11 @@ impl TpeSamplerBuilder {
/// ///
/// Default settings: /// Default settings:
/// - gamma: 0.25 (top 25% of trials are considered "good") /// - gamma: 0.25 (top 25% of trials are considered "good")
/// - n_startup_trials: 10 (random sampling for first 10 trials) /// - `n_startup_trials`: 10 (random sampling for first 10 trials)
/// - n_ei_candidates: 24 (evaluate 24 candidates per sample) /// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample)
/// - kde_bandwidth: None (uses Scott's rule for automatic bandwidth) /// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth)
/// - seed: None (use OS-provided entropy) /// - seed: None (use OS-provided entropy)
#[must_use]
pub fn new() -> Self { pub fn new() -> Self {
Self { Self {
gamma: 0.25, gamma: 0.25,
@@ -441,24 +470,23 @@ impl TpeSamplerBuilder {
/// ///
/// * `gamma` - Quantile value, must be in (0.0, 1.0). /// * `gamma` - Quantile value, must be in (0.0, 1.0).
/// ///
/// # Panics
///
/// Panics if gamma is not in (0.0, 1.0).
///
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSamplerBuilder; /// use optimizer::sampler::tpe::TpeSamplerBuilder;
/// ///
/// let sampler = TpeSamplerBuilder::new() /// let sampler = TpeSamplerBuilder::new()
/// .gamma(0.10) // Use top 10% as "good" trials /// .gamma(0.10) // Use top 10% as "good" trials
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
///
/// # Note
///
/// Validation happens at `build()` time. If gamma is not in (0.0, 1.0),
/// `build()` will return `Err(TpeError::InvalidGamma)`.
#[must_use]
pub fn gamma(mut self, gamma: f64) -> Self { pub fn gamma(mut self, gamma: f64) -> Self {
assert!(
gamma > 0.0 && gamma < 1.0,
"gamma must be in (0.0, 1.0), got {gamma}"
);
self.gamma = gamma; self.gamma = gamma;
self self
} }
@@ -476,12 +504,14 @@ impl TpeSamplerBuilder {
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSamplerBuilder; /// use optimizer::sampler::tpe::TpeSamplerBuilder;
/// ///
/// let sampler = TpeSamplerBuilder::new() /// let sampler = TpeSamplerBuilder::new()
/// .n_startup_trials(20) // Random sample first 20 trials /// .n_startup_trials(20) // Random sample first 20 trials
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
#[must_use]
pub fn n_startup_trials(mut self, n: usize) -> Self { pub fn n_startup_trials(mut self, n: usize) -> Self {
self.n_startup_trials = n; self.n_startup_trials = n;
self self
@@ -500,12 +530,14 @@ impl TpeSamplerBuilder {
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSamplerBuilder; /// use optimizer::sampler::tpe::TpeSamplerBuilder;
/// ///
/// let sampler = TpeSamplerBuilder::new() /// let sampler = TpeSamplerBuilder::new()
/// .n_ei_candidates(48) // Evaluate more candidates /// .n_ei_candidates(48) // Evaluate more candidates
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
#[must_use]
pub fn n_ei_candidates(mut self, n: usize) -> Self { pub fn n_ei_candidates(mut self, n: usize) -> Self {
self.n_ei_candidates = n; self.n_ei_candidates = n;
self self
@@ -523,24 +555,23 @@ impl TpeSamplerBuilder {
/// ///
/// * `bandwidth` - The fixed bandwidth (standard deviation) for Gaussian kernels. /// * `bandwidth` - The fixed bandwidth (standard deviation) for Gaussian kernels.
/// ///
/// # Panics
///
/// Panics if bandwidth is not positive.
///
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSamplerBuilder; /// use optimizer::sampler::tpe::TpeSamplerBuilder;
/// ///
/// let sampler = TpeSamplerBuilder::new() /// let sampler = TpeSamplerBuilder::new()
/// .kde_bandwidth(0.5) // Fixed bandwidth of 0.5 /// .kde_bandwidth(0.5) // Fixed bandwidth of 0.5
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
///
/// # Note
///
/// Validation happens at `build()` time. If bandwidth is not positive,
/// `build()` will return `Err(TpeError::InvalidBandwidth)`.
#[must_use]
pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self { pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self {
assert!(
bandwidth > 0.0,
"kde_bandwidth must be positive, got {bandwidth}"
);
self.kde_bandwidth = Some(bandwidth); self.kde_bandwidth = Some(bandwidth);
self self
} }
@@ -554,12 +585,14 @@ impl TpeSamplerBuilder {
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSamplerBuilder; /// use optimizer::sampler::tpe::TpeSamplerBuilder;
/// ///
/// let sampler = TpeSamplerBuilder::new() /// let sampler = TpeSamplerBuilder::new()
/// .seed(42) // Reproducible results /// .seed(42) // Reproducible results
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
#[must_use]
pub fn seed(mut self, seed: u64) -> Self { pub fn seed(mut self, seed: u64) -> Self {
self.seed = Some(seed); self.seed = Some(seed);
self self
@@ -567,19 +600,25 @@ impl TpeSamplerBuilder {
/// Builds the configured [`TpeSampler`]. /// Builds the configured [`TpeSampler`].
/// ///
/// # Errors
///
/// Returns `TpeError::InvalidGamma` if gamma is not in (0.0, 1.0).
/// Returns `TpeError::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
///
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::TpeSamplerBuilder; /// use optimizer::sampler::tpe::TpeSamplerBuilder;
/// ///
/// let sampler = TpeSamplerBuilder::new() /// let sampler = TpeSamplerBuilder::new()
/// .gamma(0.15) /// .gamma(0.15)
/// .n_startup_trials(20) /// .n_startup_trials(20)
/// .n_ei_candidates(32) /// .n_ei_candidates(32)
/// .seed(42) /// .seed(42)
/// .build(); /// .build()
/// .unwrap();
/// ``` /// ```
pub fn build(self) -> TpeSampler { pub fn build(self) -> Result<TpeSampler> {
TpeSampler::with_config( TpeSampler::with_config(
self.gamma, self.gamma,
self.n_startup_trials, self.n_startup_trials,
@@ -597,6 +636,7 @@ impl Default for TpeSamplerBuilder {
} }
impl Sampler for TpeSampler { impl Sampler for TpeSampler {
#[allow(clippy::too_many_lines)]
fn sample( fn sample(
&self, &self,
distribution: &Distribution, distribution: &Distribution,
@@ -693,8 +733,8 @@ impl Sampler for TpeSampler {
d.high, d.high,
d.log_scale, d.log_scale,
d.step, d.step,
good_values, &good_values,
bad_values, &bad_values,
&mut rng, &mut rng,
); );
ParamValue::Int(value) ParamValue::Int(value)
@@ -725,7 +765,7 @@ impl Sampler for TpeSampler {
} }
let index = let index =
self.sample_tpe_categorical(d.n_choices, good_indices, bad_indices, &mut rng); self.sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, &mut rng);
ParamValue::Categorical(index) ParamValue::Categorical(index)
} }
} }
@@ -733,6 +773,11 @@ impl Sampler for TpeSampler {
} }
#[cfg(test)] #[cfg(test)]
#[allow(
clippy::similar_names,
clippy::cast_sign_loss,
clippy::cast_precision_loss
)]
mod tests { mod tests {
use std::collections::HashMap; use std::collections::HashMap;
@@ -756,34 +801,34 @@ mod tests {
#[test] #[test]
fn test_tpe_sampler_new() { fn test_tpe_sampler_new() {
let sampler = TpeSampler::new(); let sampler = TpeSampler::new();
assert_eq!(sampler.gamma, 0.25); assert!((sampler.gamma - 0.25).abs() < f64::EPSILON);
assert_eq!(sampler.n_startup_trials, 10); assert_eq!(sampler.n_startup_trials, 10);
assert_eq!(sampler.n_ei_candidates, 24); assert_eq!(sampler.n_ei_candidates, 24);
} }
#[test] #[test]
fn test_tpe_sampler_with_config() { fn test_tpe_sampler_with_config() {
let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42)); let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42)).unwrap();
assert_eq!(sampler.gamma, 0.15); assert!((sampler.gamma - 0.15).abs() < f64::EPSILON);
assert_eq!(sampler.n_startup_trials, 20); assert_eq!(sampler.n_startup_trials, 20);
assert_eq!(sampler.n_ei_candidates, 32); assert_eq!(sampler.n_ei_candidates, 32);
} }
#[test] #[test]
#[should_panic(expected = "gamma must be in (0.0, 1.0)")]
fn test_tpe_sampler_invalid_gamma_zero() { fn test_tpe_sampler_invalid_gamma_zero() {
TpeSampler::with_config(0.0, 10, 24, None, None); let result = TpeSampler::with_config(0.0, 10, 24, None, None);
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
} }
#[test] #[test]
#[should_panic(expected = "gamma must be in (0.0, 1.0)")]
fn test_tpe_sampler_invalid_gamma_one() { fn test_tpe_sampler_invalid_gamma_one() {
TpeSampler::with_config(1.0, 10, 24, None, None); let result = TpeSampler::with_config(1.0, 10, 24, None, None);
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
} }
#[test] #[test]
fn test_tpe_startup_random_sampling() { fn test_tpe_startup_random_sampling() {
let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)); let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)).unwrap();
let dist = Distribution::Float(FloatDistribution { let dist = Distribution::Float(FloatDistribution {
low: 0.0, low: 0.0,
high: 1.0, high: 1.0,
@@ -806,7 +851,7 @@ mod tests {
#[test] #[test]
fn test_tpe_split_trials() { fn test_tpe_split_trials() {
let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)); let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)).unwrap();
let dist = Distribution::Float(FloatDistribution { let dist = Distribution::Float(FloatDistribution {
low: 0.0, low: 0.0,
@@ -820,8 +865,8 @@ mod tests {
.map(|i| { .map(|i| {
create_trial( create_trial(
i as u64, i as u64,
i as f64, f64::from(i),
vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())], vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
) )
}) })
.collect(); .collect();
@@ -840,7 +885,7 @@ mod tests {
#[test] #[test]
fn test_tpe_samples_float_with_history() { fn test_tpe_samples_float_with_history() {
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)); let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();
let dist = Distribution::Float(FloatDistribution { let dist = Distribution::Float(FloatDistribution {
low: 0.0, low: 0.0,
@@ -852,7 +897,7 @@ mod tests {
// Create history where low values (near 0.2) are "good" // Create history where low values (near 0.2) are "good"
let history: Vec<CompletedTrial> = (0..20) let history: Vec<CompletedTrial> = (0..20)
.map(|i| { .map(|i| {
let x = i as f64 / 20.0; let x = f64::from(i) / 20.0;
// Objective is (x - 0.2)^2, minimized at x=0.2 // Objective is (x - 0.2)^2, minimized at x=0.2
let value = (x - 0.2).powi(2); let value = (x - 0.2).powi(2);
create_trial( create_trial(
@@ -882,7 +927,7 @@ mod tests {
#[test] #[test]
fn test_tpe_categorical_sampling() { fn test_tpe_categorical_sampling() {
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)); let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 }); let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 });
@@ -922,7 +967,7 @@ mod tests {
#[test] #[test]
fn test_tpe_int_sampling() { fn test_tpe_int_sampling() {
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)); let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();
let dist = Distribution::Int(IntDistribution { let dist = Distribution::Int(IntDistribution {
low: 0, low: 0,
@@ -968,14 +1013,14 @@ mod tests {
.map(|i| { .map(|i| {
create_trial( create_trial(
i as u64, i as u64,
i as f64, f64::from(i),
vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())], vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
) )
}) })
.collect(); .collect();
let sampler1 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345)); let sampler1 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345)).unwrap();
let sampler2 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345)); let sampler2 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345)).unwrap();
for i in 0..10 { for i in 0..10 {
let v1 = sampler1.sample(&dist, i, &history); let v1 = sampler1.sample(&dist, i, &history);
@@ -987,8 +1032,8 @@ mod tests {
#[test] #[test]
fn test_tpe_sampler_builder_default() { fn test_tpe_sampler_builder_default() {
let builder = TpeSamplerBuilder::new(); let builder = TpeSamplerBuilder::new();
let sampler = builder.build(); let sampler = builder.build().unwrap();
assert_eq!(sampler.gamma, 0.25); assert!((sampler.gamma - 0.25).abs() < f64::EPSILON);
assert_eq!(sampler.n_startup_trials, 10); assert_eq!(sampler.n_startup_trials, 10);
assert_eq!(sampler.n_ei_candidates, 24); assert_eq!(sampler.n_ei_candidates, 24);
} }
@@ -1000,8 +1045,9 @@ mod tests {
.n_startup_trials(20) .n_startup_trials(20)
.n_ei_candidates(32) .n_ei_candidates(32)
.seed(42) .seed(42)
.build(); .build()
assert_eq!(sampler.gamma, 0.15); .unwrap();
assert!((sampler.gamma - 0.15).abs() < f64::EPSILON);
assert_eq!(sampler.n_startup_trials, 20); assert_eq!(sampler.n_startup_trials, 20);
assert_eq!(sampler.n_ei_candidates, 32); assert_eq!(sampler.n_ei_candidates, 32);
} }
@@ -1012,8 +1058,9 @@ mod tests {
.gamma(0.10) .gamma(0.10)
.n_startup_trials(15) .n_startup_trials(15)
.n_ei_candidates(48) .n_ei_candidates(48)
.build(); .build()
assert_eq!(sampler.gamma, 0.10); .unwrap();
assert!((sampler.gamma - 0.10).abs() < f64::EPSILON);
assert_eq!(sampler.n_startup_trials, 15); assert_eq!(sampler.n_startup_trials, 15);
assert_eq!(sampler.n_ei_candidates, 48); assert_eq!(sampler.n_ei_candidates, 48);
} }
@@ -1021,16 +1068,16 @@ mod tests {
#[test] #[test]
fn test_tpe_sampler_builder_partial() { fn test_tpe_sampler_builder_partial() {
// Test setting only some options // Test setting only some options
let sampler = TpeSamplerBuilder::new().gamma(0.20).build(); let sampler = TpeSamplerBuilder::new().gamma(0.20).build().unwrap();
assert_eq!(sampler.gamma, 0.20); assert!((sampler.gamma - 0.20).abs() < f64::EPSILON);
assert_eq!(sampler.n_startup_trials, 10); // default assert_eq!(sampler.n_startup_trials, 10); // default
assert_eq!(sampler.n_ei_candidates, 24); // default assert_eq!(sampler.n_ei_candidates, 24); // default
} }
#[test] #[test]
#[should_panic(expected = "gamma must be in (0.0, 1.0)")]
fn test_tpe_sampler_builder_invalid_gamma() { fn test_tpe_sampler_builder_invalid_gamma() {
TpeSamplerBuilder::new().gamma(1.5).build(); let result = TpeSamplerBuilder::new().gamma(1.5).build();
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
} }
#[test] #[test]
@@ -1042,12 +1089,12 @@ mod tests {
step: None, step: None,
}); });
let history: Vec<CompletedTrial> = (0..20) let history: Vec<CompletedTrial> = (0..20u32)
.map(|i| { .map(|i| {
create_trial( create_trial(
i as u64, u64::from(i),
i as f64, f64::from(i),
vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())], vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
) )
}) })
.collect(); .collect();
@@ -1055,11 +1102,13 @@ mod tests {
let sampler1 = TpeSampler::builder() let sampler1 = TpeSampler::builder()
.seed(99999) .seed(99999)
.n_startup_trials(5) .n_startup_trials(5)
.build(); .build()
.unwrap();
let sampler2 = TpeSampler::builder() let sampler2 = TpeSampler::builder()
.seed(99999) .seed(99999)
.n_startup_trials(5) .n_startup_trials(5)
.build(); .build()
.unwrap();
for i in 0..10 { for i in 0..10 {
let v1 = sampler1.sample(&dist, i, &history); let v1 = sampler1.sample(&dist, i, &history);
+71 -286
View File
@@ -1,25 +1,18 @@
//! Study implementation for managing optimization trials. //! Study implementation for managing optimization trials.
#[cfg(feature = "async")] #[cfg(feature = "async")]
use std::future::Future; use core::future::Future;
use std::ops::ControlFlow; use core::ops::ControlFlow;
use core::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc; use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use parking_lot::RwLock; use parking_lot::RwLock;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use crate::sampler::{CompletedTrial, RandomSampler, Sampler}; use crate::sampler::random::RandomSampler;
use crate::sampler::{CompletedTrial, Sampler};
use crate::trial::Trial; use crate::trial::Trial;
use crate::types::Direction; use crate::types::Direction;
/// Helper function to create default sampler for serde deserialization.
#[cfg(feature = "serde")]
fn default_sampler() -> Arc<dyn Sampler> {
Arc::new(RandomSampler::new())
}
/// A study manages the optimization process, tracking trials and their results. /// A study manages the optimization process, tracking trials and their results.
/// ///
/// The study is parameterized by the objective value type `V`, which defaults to `f64`. /// The study is parameterized by the objective value type `V`, which defaults to `f64`.
@@ -29,13 +22,6 @@ fn default_sampler() -> Arc<dyn Sampler> {
/// When `V = f64`, the study passes trial history to the sampler for informed /// When `V = f64`, the study passes trial history to the sampler for informed
/// parameter suggestions (e.g., TPE sampler uses history to guide sampling). /// parameter suggestions (e.g., TPE sampler uses history to guide sampling).
/// ///
/// # Serialization
///
/// When the `serde` feature is enabled, the study can be serialized and deserialized.
/// The completed trials and trial ID counter are preserved, allowing optimization to
/// continue after deserialization. The sampler is not serialized; upon deserialization,
/// a default `RandomSampler` is used. Use `Study::set_sampler()` to restore a custom sampler.
///
/// # Examples /// # Examples
/// ///
/// ``` /// ```
@@ -79,6 +65,7 @@ where
/// let study: Study<f64> = Study::new(Direction::Minimize); /// let study: Study<f64> = Study::new(Direction::Minimize);
/// assert_eq!(study.direction(), Direction::Minimize); /// assert_eq!(study.direction(), Direction::Minimize);
/// ``` /// ```
#[must_use]
pub fn new(direction: Direction) -> Self { pub fn new(direction: Direction) -> Self {
Self::with_sampler(direction, RandomSampler::new()) Self::with_sampler(direction, RandomSampler::new())
} }
@@ -93,7 +80,8 @@ where
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// let sampler = RandomSampler::with_seed(42); /// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler); /// let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
@@ -115,9 +103,6 @@ where
/// Sets a new sampler for the study. /// Sets a new sampler for the study.
/// ///
/// This method is useful after deserializing a study when you want to use
/// a custom sampler (e.g., TPE) instead of the default `RandomSampler`.
///
/// # Arguments /// # Arguments
/// ///
/// * `sampler` - The sampler to use for parameter sampling. /// * `sampler` - The sampler to use for parameter sampling.
@@ -125,9 +110,9 @@ where
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, Study, TpeSampler}; /// use optimizer::sampler::tpe::TpeSampler;
/// use optimizer::{Direction, Study};
/// ///
/// // After deserializing a study, restore the TPE sampler
/// let mut study: Study<f64> = Study::new(Direction::Minimize); /// let mut study: Study<f64> = Study::new(Direction::Minimize);
/// study.set_sampler(TpeSampler::new()); /// study.set_sampler(TpeSampler::new());
/// ``` /// ```
@@ -332,17 +317,15 @@ where
match self.direction { match self.direction {
Direction::Minimize => { Direction::Minimize => {
// Reverse ordering: smaller values are "greater" for max_by // Reverse ordering: smaller values are "greater" for max_by
ordering ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse)
.map(|o| o.reverse())
.unwrap_or(std::cmp::Ordering::Equal)
} }
Direction::Maximize => { Direction::Maximize => {
// Normal ordering: larger values are "greater" for max_by // Normal ordering: larger values are "greater" for max_by
ordering.unwrap_or(std::cmp::Ordering::Equal) ordering.unwrap_or(core::cmp::Ordering::Equal)
} }
} }
}) })
.expect("trials is not empty"); .ok_or(crate::TpeError::NoCompletedTrials)?;
Ok(best.clone()) Ok(best.clone())
} }
@@ -409,7 +392,8 @@ where
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// // Minimize x^2 /// // Minimize x^2
/// let sampler = RandomSampler::with_seed(42); /// let sampler = RandomSampler::with_seed(42);
@@ -429,7 +413,7 @@ where
/// ``` /// ```
pub fn optimize<F, E>(&self, n_trials: usize, mut objective: F) -> crate::Result<()> pub fn optimize<F, E>(&self, n_trials: usize, mut objective: F) -> crate::Result<()>
where where
F: FnMut(&mut Trial) -> std::result::Result<V, E>, F: FnMut(&mut Trial) -> core::result::Result<V, E>,
E: ToString, E: ToString,
{ {
for _ in 0..n_trials { for _ in 0..n_trials {
@@ -476,7 +460,8 @@ where
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// # #[cfg(feature = "async")] /// # #[cfg(feature = "async")]
/// # async fn example() -> optimizer::Result<()> { /// # async fn example() -> optimizer::Result<()> {
@@ -506,7 +491,7 @@ where
) -> crate::Result<()> ) -> crate::Result<()>
where where
F: Fn(Trial) -> Fut, F: Fn(Trial) -> Fut,
Fut: Future<Output = std::result::Result<(Trial, V), E>>, Fut: Future<Output = core::result::Result<(Trial, V), E>>,
E: ToString, E: ToString,
{ {
for _ in 0..n_trials { for _ in 0..n_trials {
@@ -552,11 +537,13 @@ where
/// # Errors /// # Errors
/// ///
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials). /// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
/// Returns `TpeError::TaskError` if the semaphore is closed or a spawned task panics.
/// ///
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// # #[cfg(feature = "async")] /// # #[cfg(feature = "async")]
/// # async fn example() -> optimizer::Result<()> { /// # async fn example() -> optimizer::Result<()> {
@@ -587,7 +574,7 @@ where
) -> crate::Result<()> ) -> crate::Result<()>
where where
F: Fn(Trial) -> Fut + Send + Sync + 'static, F: Fn(Trial) -> Fut + Send + Sync + 'static,
Fut: Future<Output = std::result::Result<(Trial, V), E>> + Send, Fut: Future<Output = core::result::Result<(Trial, V), E>> + Send,
E: ToString + Send + 'static, E: ToString + Send + 'static,
V: Send + 'static, V: Send + 'static,
{ {
@@ -599,7 +586,11 @@ where
let mut handles = Vec::with_capacity(n_trials); let mut handles = Vec::with_capacity(n_trials);
for _ in 0..n_trials { for _ in 0..n_trials {
let permit = semaphore.clone().acquire_owned().await.unwrap(); let permit = semaphore
.clone()
.acquire_owned()
.await
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?;
let trial = self.create_trial(); let trial = self.create_trial();
let objective = Arc::clone(&objective); let objective = Arc::clone(&objective);
@@ -614,7 +605,10 @@ where
// Wait for all tasks and record results // Wait for all tasks and record results
for handle in handles { for handle in handles {
match handle.await.unwrap() { match handle
.await
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?
{
Ok((trial, value)) => { Ok((trial, value)) => {
self.complete_trial(trial, value); self.complete_trial(trial, value);
} }
@@ -651,13 +645,15 @@ where
/// ///
/// Returns `TpeError::NoCompletedTrials` if no trials completed successfully /// Returns `TpeError::NoCompletedTrials` if no trials completed successfully
/// before optimization stopped (either by completing all trials or early stopping). /// before optimization stopped (either by completing all trials or early stopping).
/// Returns `TpeError::Internal` if a completed trial is not found after adding (internal invariant violation).
/// ///
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use std::ops::ControlFlow; /// use std::ops::ControlFlow;
/// ///
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// // Stop early when we find a good enough value /// // Stop early when we find a good enough value
/// let sampler = RandomSampler::with_seed(42); /// let sampler = RandomSampler::with_seed(42);
@@ -692,7 +688,7 @@ where
) -> crate::Result<()> ) -> crate::Result<()>
where where
V: Clone, V: Clone,
F: FnMut(&mut Trial) -> std::result::Result<V, E>, F: FnMut(&mut Trial) -> core::result::Result<V, E>,
C: FnMut(&Study<V>, &CompletedTrial<V>) -> ControlFlow<()>, C: FnMut(&Study<V>, &CompletedTrial<V>) -> ControlFlow<()>,
E: ToString, E: ToString,
{ {
@@ -705,7 +701,11 @@ where
// Get the just-completed trial for the callback // Get the just-completed trial for the callback
let trials = self.completed_trials.read(); let trials = self.completed_trials.read();
let completed = trials.last().expect("just added a trial"); let Some(completed) = trials.last() else {
return Err(crate::TpeError::Internal(
"completed trial not found after adding",
));
};
// Call the callback and check if we should stop // Call the callback and check if we should stop
// Note: We need to drop the read lock before calling callback // Note: We need to drop the read lock before calling callback
@@ -746,7 +746,8 @@ impl Study<f64> {
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// // With a seeded sampler for reproducibility /// // With a seeded sampler for reproducibility
/// let sampler = RandomSampler::with_seed(42); /// let sampler = RandomSampler::with_seed(42);
@@ -787,7 +788,8 @@ impl Study<f64> {
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// // Minimize x^2 with sampler integration /// // Minimize x^2 with sampler integration
/// let sampler = RandomSampler::with_seed(42); /// let sampler = RandomSampler::with_seed(42);
@@ -809,7 +811,7 @@ impl Study<f64> {
mut objective: F, mut objective: F,
) -> crate::Result<()> ) -> crate::Result<()>
where where
F: FnMut(&mut Trial) -> std::result::Result<f64, E>, F: FnMut(&mut Trial) -> core::result::Result<f64, E>,
E: ToString, E: ToString,
{ {
for _ in 0..n_trials { for _ in 0..n_trials {
@@ -850,13 +852,15 @@ impl Study<f64> {
/// # Errors /// # Errors
/// ///
/// Returns `TpeError::NoCompletedTrials` if no trials completed successfully. /// Returns `TpeError::NoCompletedTrials` if no trials completed successfully.
/// Returns `TpeError::Internal` if a completed trial is not found after adding (internal invariant violation).
/// ///
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use std::ops::ControlFlow; /// use std::ops::ControlFlow;
/// ///
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// // Optimize with sampler integration and early stopping /// // Optimize with sampler integration and early stopping
/// let sampler = RandomSampler::with_seed(42); /// let sampler = RandomSampler::with_seed(42);
@@ -889,7 +893,7 @@ impl Study<f64> {
mut callback: C, mut callback: C,
) -> crate::Result<()> ) -> crate::Result<()>
where where
F: FnMut(&mut Trial) -> std::result::Result<f64, E>, F: FnMut(&mut Trial) -> core::result::Result<f64, E>,
C: FnMut(&Study<f64>, &CompletedTrial<f64>) -> ControlFlow<()>, C: FnMut(&Study<f64>, &CompletedTrial<f64>) -> ControlFlow<()>,
E: ToString, E: ToString,
{ {
@@ -902,7 +906,11 @@ impl Study<f64> {
// Get the just-completed trial for the callback // Get the just-completed trial for the callback
let trials = self.completed_trials.read(); let trials = self.completed_trials.read();
let completed = trials.last().expect("just added a trial"); let Some(completed) = trials.last() else {
return Err(crate::TpeError::Internal(
"completed trial not found after adding",
));
};
// Call the callback and check if we should stop // Call the callback and check if we should stop
// Note: We need to drop the read lock before calling callback // Note: We need to drop the read lock before calling callback
@@ -950,7 +958,8 @@ impl Study<f64> {
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// # #[cfg(feature = "async")] /// # #[cfg(feature = "async")]
/// # async fn example() -> optimizer::Result<()> { /// # async fn example() -> optimizer::Result<()> {
@@ -980,7 +989,7 @@ impl Study<f64> {
) -> crate::Result<()> ) -> crate::Result<()>
where where
F: Fn(Trial) -> Fut, F: Fn(Trial) -> Fut,
Fut: Future<Output = std::result::Result<(Trial, f64), E>>, Fut: Future<Output = core::result::Result<(Trial, f64), E>>,
E: ToString, E: ToString,
{ {
for _ in 0..n_trials { for _ in 0..n_trials {
@@ -1026,11 +1035,13 @@ impl Study<f64> {
/// # Errors /// # Errors
/// ///
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials). /// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
/// Returns `TpeError::TaskError` if the semaphore is closed or a spawned task panics.
/// ///
/// # Examples /// # Examples
/// ///
/// ``` /// ```
/// use optimizer::{Direction, RandomSampler, Study}; /// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
/// ///
/// # #[cfg(feature = "async")] /// # #[cfg(feature = "async")]
/// # async fn example() -> optimizer::Result<()> { /// # async fn example() -> optimizer::Result<()> {
@@ -1061,7 +1072,7 @@ impl Study<f64> {
) -> crate::Result<()> ) -> crate::Result<()>
where where
F: Fn(Trial) -> Fut + Send + Sync + 'static, F: Fn(Trial) -> Fut + Send + Sync + 'static,
Fut: Future<Output = std::result::Result<(Trial, f64), E>> + Send, Fut: Future<Output = core::result::Result<(Trial, f64), E>> + Send,
E: ToString + Send + 'static, E: ToString + Send + 'static,
{ {
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
@@ -1072,7 +1083,11 @@ impl Study<f64> {
let mut handles = Vec::with_capacity(n_trials); let mut handles = Vec::with_capacity(n_trials);
for _ in 0..n_trials { for _ in 0..n_trials {
let permit = semaphore.clone().acquire_owned().await.unwrap(); let permit = semaphore
.clone()
.acquire_owned()
.await
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?;
let trial = self.create_trial_with_sampler(); let trial = self.create_trial_with_sampler();
let objective = Arc::clone(&objective); let objective = Arc::clone(&objective);
@@ -1087,7 +1102,10 @@ impl Study<f64> {
// Wait for all tasks and record results // Wait for all tasks and record results
for handle in handles { for handle in handles {
match handle.await.unwrap() { match handle
.await
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?
{
Ok((trial, value)) => { Ok((trial, value)) => {
self.complete_trial(trial, value); self.complete_trial(trial, value);
} }
@@ -1105,236 +1123,3 @@ impl Study<f64> {
Ok(()) Ok(())
} }
} }
// Manual Serialize implementation for Study<V> when serde feature is enabled.
#[cfg(feature = "serde")]
impl<V> Serialize for Study<V>
where
V: PartialOrd + Serialize,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
use serde::ser::SerializeStruct;
let mut state = serializer.serialize_struct("Study", 3)?;
state.serialize_field("direction", &self.direction)?;
// Serialize the Vec inside the Arc<RwLock<>>
let trials = self.completed_trials.read();
state.serialize_field("completed_trials", &*trials)?;
state.serialize_field("next_trial_id", &self.next_trial_id.load(Ordering::SeqCst))?;
state.end()
}
}
// Manual Deserialize implementation for Study<V> when serde feature is enabled.
#[cfg(feature = "serde")]
impl<'de, V> Deserialize<'de> for Study<V>
where
V: PartialOrd + Deserialize<'de>,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
use std::fmt;
use std::marker::PhantomData;
use serde::de::{self, MapAccess, Visitor};
#[derive(serde::Deserialize)]
#[serde(field_identifier, rename_all = "snake_case")]
enum Field {
Direction,
CompletedTrials,
NextTrialId,
}
struct StudyVisitor<V>(PhantomData<V>);
impl<'de, V> Visitor<'de> for StudyVisitor<V>
where
V: PartialOrd + Deserialize<'de>,
{
type Value = Study<V>;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("struct Study")
}
fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
where
A: MapAccess<'de>,
{
let mut direction = None;
let mut completed_trials: Option<Vec<CompletedTrial<V>>> = None;
let mut next_trial_id = None;
while let Some(key) = map.next_key()? {
match key {
Field::Direction => {
if direction.is_some() {
return Err(de::Error::duplicate_field("direction"));
}
direction = Some(map.next_value()?);
}
Field::CompletedTrials => {
if completed_trials.is_some() {
return Err(de::Error::duplicate_field("completed_trials"));
}
completed_trials = Some(map.next_value()?);
}
Field::NextTrialId => {
if next_trial_id.is_some() {
return Err(de::Error::duplicate_field("next_trial_id"));
}
next_trial_id = Some(map.next_value()?);
}
}
}
let direction = direction.ok_or_else(|| de::Error::missing_field("direction"))?;
let completed_trials =
completed_trials.ok_or_else(|| de::Error::missing_field("completed_trials"))?;
let next_trial_id: u64 =
next_trial_id.ok_or_else(|| de::Error::missing_field("next_trial_id"))?;
Ok(Study {
direction,
sampler: default_sampler(),
completed_trials: Arc::new(RwLock::new(completed_trials)),
next_trial_id: AtomicU64::new(next_trial_id),
})
}
}
const FIELDS: &[&str] = &["direction", "completed_trials", "next_trial_id"];
deserializer.deserialize_struct("Study", FIELDS, StudyVisitor(PhantomData))
}
}
#[cfg(all(test, feature = "serde"))]
mod serde_tests {
use super::*;
#[test]
fn test_study_serde_round_trip() {
// Create a study and add some trials
let study: Study<f64> = Study::new(Direction::Minimize);
// Run some optimization
study
.optimize(5, |trial| {
let x = trial.suggest_float("x", 0.0, 10.0)?;
let y = trial.suggest_int("y", 1, 5)?;
Ok::<_, crate::TpeError>(x + y as f64)
})
.unwrap();
// Serialize to JSON
let serialized = serde_json::to_string(&study).unwrap();
// Deserialize from JSON
let deserialized: Study<f64> = serde_json::from_str(&serialized).unwrap();
// Verify the data is preserved
assert_eq!(deserialized.direction(), study.direction());
assert_eq!(deserialized.n_trials(), study.n_trials());
// Verify the best trial is the same
let original_best = study.best_trial().unwrap();
let deserialized_best = deserialized.best_trial().unwrap();
assert_eq!(original_best.id, deserialized_best.id);
// Use approximate comparison for floats due to JSON serialization precision
assert!((original_best.value - deserialized_best.value).abs() < 1e-10);
// Check that all param keys match
assert_eq!(original_best.params.len(), deserialized_best.params.len());
for (key, original_val) in &original_best.params {
let deserialized_val = deserialized_best.params.get(key).unwrap();
match (original_val, deserialized_val) {
(crate::param::ParamValue::Float(a), crate::param::ParamValue::Float(b)) => {
assert!((a - b).abs() < 1e-10, "Float param {key} differs");
}
_ => assert_eq!(original_val, deserialized_val),
}
}
// Verify we can continue optimization on the deserialized study
let initial_count = deserialized.n_trials();
deserialized
.optimize(3, |trial| {
let x = trial.suggest_float("x", 0.0, 10.0)?;
let y = trial.suggest_int("y", 1, 5)?;
Ok::<_, crate::TpeError>(x + y as f64)
})
.unwrap();
// Verify new trials were added
assert_eq!(deserialized.n_trials(), initial_count + 3);
}
#[test]
fn test_study_serde_preserves_trial_ids() {
let study: Study<f64> = Study::new(Direction::Maximize);
// Add 5 trials
study
.optimize(5, |trial| {
let x = trial.suggest_float("x", -1.0, 1.0)?;
Ok::<_, crate::TpeError>(x * x)
})
.unwrap();
// Serialize and deserialize
let serialized = serde_json::to_string(&study).unwrap();
let deserialized: Study<f64> = serde_json::from_str(&serialized).unwrap();
// Create a new trial - its ID should continue from where we left off
let new_trial = deserialized.create_trial();
assert_eq!(new_trial.id(), 5); // Next trial should be ID 5
}
#[test]
fn test_completed_trial_serde() {
use std::collections::HashMap;
use crate::distribution::{Distribution, FloatDistribution, IntDistribution};
use crate::param::ParamValue;
let mut params = HashMap::new();
params.insert("x".to_string(), ParamValue::Float(0.5));
params.insert("n".to_string(), ParamValue::Int(42));
let mut distributions = HashMap::new();
distributions.insert(
"x".to_string(),
Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
}),
);
distributions.insert(
"n".to_string(),
Distribution::Int(IntDistribution {
low: 1,
high: 100,
log_scale: false,
step: None,
}),
);
let completed = CompletedTrial::new(42, params.clone(), distributions.clone(), 0.75);
// Serialize and deserialize
let serialized = serde_json::to_string(&completed).unwrap();
let deserialized: CompletedTrial<f64> = serde_json::from_str(&serialized).unwrap();
assert_eq!(deserialized.id, 42);
assert_eq!(deserialized.value, 0.75);
assert_eq!(deserialized.params, params);
assert_eq!(deserialized.distributions, distributions);
}
}
+48 -38
View File
@@ -4,8 +4,6 @@ use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use parking_lot::RwLock; use parking_lot::RwLock;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use crate::distribution::{ use crate::distribution::{
CategoricalDistribution, Distribution, FloatDistribution, IntDistribution, CategoricalDistribution, Distribution, FloatDistribution, IntDistribution,
@@ -24,7 +22,6 @@ use crate::types::TrialState;
/// `Study::create_trial()`, the trial receives the study's sampler and access /// `Study::create_trial()`, the trial receives the study's sampler and access
/// to the history of completed trials for informed sampling. /// to the history of completed trials for informed sampling.
#[derive(Clone)] #[derive(Clone)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Trial { pub struct Trial {
/// Unique identifier for this trial. /// Unique identifier for this trial.
id: u64, id: u64,
@@ -35,15 +32,13 @@ pub struct Trial {
/// Parameter distributions, keyed by parameter name. /// Parameter distributions, keyed by parameter name.
distributions: HashMap<String, Distribution>, distributions: HashMap<String, Distribution>,
/// The sampler to use for generating parameter values. /// The sampler to use for generating parameter values.
#[cfg_attr(feature = "serde", serde(skip))]
sampler: Option<Arc<dyn Sampler>>, sampler: Option<Arc<dyn Sampler>>,
/// Access to the history of completed trials (shared with Study). /// Access to the history of completed trials (shared with Study).
#[cfg_attr(feature = "serde", serde(skip))]
history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>, history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>,
} }
impl std::fmt::Debug for Trial { impl core::fmt::Debug for Trial {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Trial") f.debug_struct("Trial")
.field("id", &self.id) .field("id", &self.id)
.field("state", &self.state) .field("state", &self.state)
@@ -76,6 +71,7 @@ impl Trial {
/// let trial = Trial::new(0); /// let trial = Trial::new(0);
/// assert_eq!(trial.id(), 0); /// assert_eq!(trial.id(), 0);
/// ``` /// ```
#[must_use]
pub fn new(id: u64) -> Self { pub fn new(id: u64) -> Self {
Self { Self {
id, id,
@@ -115,7 +111,7 @@ impl Trial {
/// Samples a value from the given distribution using the sampler. /// Samples a value from the given distribution using the sampler.
/// ///
/// If the trial has a sampler, it delegates to the sampler's sample method /// 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 /// with the history of completed trials. Otherwise, it uses the `RandomSampler`
/// as a fallback. /// as a fallback.
fn sample_value(&self, distribution: &Distribution) -> ParamValue { fn sample_value(&self, distribution: &Distribution) -> ParamValue {
if let (Some(sampler), Some(history)) = (&self.sampler, &self.history) { if let (Some(sampler), Some(history)) = (&self.sampler, &self.history) {
@@ -123,28 +119,32 @@ impl Trial {
sampler.sample(distribution, self.id, &history_guard) sampler.sample(distribution, self.id, &history_guard)
} else { } else {
// Fallback to RandomSampler when no sampler is configured // Fallback to RandomSampler when no sampler is configured
use crate::sampler::RandomSampler; use crate::sampler::random::RandomSampler;
let fallback = RandomSampler::new(); let fallback = RandomSampler::new();
fallback.sample(distribution, self.id, &[]) fallback.sample(distribution, self.id, &[])
} }
} }
/// Returns the unique ID of this trial. /// Returns the unique ID of this trial.
#[must_use]
pub fn id(&self) -> u64 { pub fn id(&self) -> u64 {
self.id self.id
} }
/// Returns the current state of this trial. /// Returns the current state of this trial.
#[must_use]
pub fn state(&self) -> TrialState { pub fn state(&self) -> TrialState {
self.state self.state
} }
/// Returns a reference to the sampled parameters. /// Returns a reference to the sampled parameters.
#[must_use]
pub fn params(&self) -> &HashMap<String, ParamValue> { pub fn params(&self) -> &HashMap<String, ParamValue> {
&self.params &self.params
} }
/// Returns a reference to the parameter distributions. /// Returns a reference to the parameter distributions.
#[must_use]
pub fn distributions(&self) -> &HashMap<String, Distribution> { pub fn distributions(&self) -> &HashMap<String, Distribution> {
&self.distributions &self.distributions
} }
@@ -205,8 +205,8 @@ impl Trial {
if let Some(existing_dist) = self.distributions.get(&name) { if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches // Verify the distribution matches
if let Distribution::Float(existing) = existing_dist if let Distribution::Float(existing) = existing_dist
&& existing.low == low && (existing.low - low).abs() < f64::EPSILON
&& existing.high == high && (existing.high - high).abs() < f64::EPSILON
&& !existing.log_scale && !existing.log_scale
&& existing.step.is_none() && existing.step.is_none()
{ {
@@ -225,9 +225,10 @@ impl Trial {
// Sample using the sampler // Sample using the sampler
let dist = Distribution::Float(distribution); let dist = Distribution::Float(distribution);
let value = match self.sample_value(&dist) { let ParamValue::Float(value) = self.sample_value(&dist) else {
ParamValue::Float(v) => v, return Err(TpeError::Internal(
_ => unreachable!("Float distribution should return Float value"), "Float distribution should return Float value",
));
}; };
// Store distribution and value // Store distribution and value
@@ -242,7 +243,7 @@ impl Trial {
/// The value is sampled uniformly in log space, which is useful for parameters /// The value is sampled uniformly in log space, which is useful for parameters
/// that span multiple orders of magnitude (e.g., learning rates). /// 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, /// 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, /// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned. /// a `ParameterConflict` error is returned.
/// ///
@@ -301,8 +302,8 @@ impl Trial {
if let Some(existing_dist) = self.distributions.get(&name) { if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches // Verify the distribution matches
if let Distribution::Float(existing) = existing_dist if let Distribution::Float(existing) = existing_dist
&& existing.low == low && (existing.low - low).abs() < f64::EPSILON
&& existing.high == high && (existing.high - high).abs() < f64::EPSILON
&& existing.log_scale && existing.log_scale
&& existing.step.is_none() && existing.step.is_none()
{ {
@@ -321,9 +322,10 @@ impl Trial {
// Sample using the sampler (sampler handles log-scale transformation) // Sample using the sampler (sampler handles log-scale transformation)
let dist = Distribution::Float(distribution); let dist = Distribution::Float(distribution);
let value = match self.sample_value(&dist) { let ParamValue::Float(value) = self.sample_value(&dist) else {
ParamValue::Float(v) => v, return Err(TpeError::Internal(
_ => unreachable!("Float distribution should return Float value"), "Float distribution should return Float value",
));
}; };
// Store distribution and value // Store distribution and value
@@ -397,8 +399,8 @@ impl Trial {
if let Some(existing_dist) = self.distributions.get(&name) { if let Some(existing_dist) = self.distributions.get(&name) {
// Verify the distribution matches // Verify the distribution matches
if let Distribution::Float(existing) = existing_dist if let Distribution::Float(existing) = existing_dist
&& existing.low == low && (existing.low - low).abs() < f64::EPSILON
&& existing.high == high && (existing.high - high).abs() < f64::EPSILON
&& !existing.log_scale && !existing.log_scale
&& existing.step == Some(step) && existing.step == Some(step)
{ {
@@ -417,9 +419,10 @@ impl Trial {
// Sample using the sampler (sampler handles step-grid) // Sample using the sampler (sampler handles step-grid)
let dist = Distribution::Float(distribution); let dist = Distribution::Float(distribution);
let value = match self.sample_value(&dist) { let ParamValue::Float(value) = self.sample_value(&dist) else {
ParamValue::Float(v) => v, return Err(TpeError::Internal(
_ => unreachable!("Float distribution should return Float value"), "Float distribution should return Float value",
));
}; };
// Store distribution and value // Store distribution and value
@@ -460,6 +463,7 @@ impl Trial {
/// let n2 = trial.suggest_int("n_layers", 1, 10).unwrap(); /// let n2 = trial.suggest_int("n_layers", 1, 10).unwrap();
/// assert_eq!(n, n2); /// 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> { pub fn suggest_int(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
if low > high { if low > high {
return Err(TpeError::InvalidBounds { return Err(TpeError::InvalidBounds {
@@ -500,9 +504,10 @@ impl Trial {
// Sample using the sampler // Sample using the sampler
let dist = Distribution::Int(distribution); let dist = Distribution::Int(distribution);
let value = match self.sample_value(&dist) { let ParamValue::Int(value) = self.sample_value(&dist) else {
ParamValue::Int(v) => v, return Err(TpeError::Internal(
_ => unreachable!("Int distribution should return Int value"), "Int distribution should return Int value",
));
}; };
// Store distribution and value // Store distribution and value
@@ -517,7 +522,7 @@ impl Trial {
/// The value is sampled uniformly in log space, which is useful for parameters /// The value is sampled uniformly in log space, which is useful for parameters
/// that span multiple orders of magnitude (e.g., batch sizes). /// 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, /// 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, /// the cached value is returned. If the parameter was sampled with different configuration,
/// a `ParameterConflict` error is returned. /// a `ParameterConflict` error is returned.
/// ///
@@ -546,6 +551,7 @@ impl Trial {
/// let batch_size2 = trial.suggest_int_log("batch_size", 1, 1024).unwrap(); /// let batch_size2 = trial.suggest_int_log("batch_size", 1, 1024).unwrap();
/// assert_eq!(batch_size, batch_size2); /// 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> { pub fn suggest_int_log(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
if low < 1 { if low < 1 {
return Err(TpeError::InvalidLogBounds); return Err(TpeError::InvalidLogBounds);
@@ -590,9 +596,10 @@ impl Trial {
// Sample using the sampler (sampler handles log-scale transformation) // Sample using the sampler (sampler handles log-scale transformation)
let dist = Distribution::Int(distribution); let dist = Distribution::Int(distribution);
let value = match self.sample_value(&dist) { let ParamValue::Int(value) = self.sample_value(&dist) else {
ParamValue::Int(v) => v, return Err(TpeError::Internal(
_ => unreachable!("Int distribution should return Int value"), "Int distribution should return Int value",
));
}; };
// Store distribution and value // Store distribution and value
@@ -643,6 +650,7 @@ impl Trial {
/// .unwrap(); /// .unwrap();
/// assert_eq!(n, n2); /// assert_eq!(n, n2);
/// ``` /// ```
#[allow(clippy::cast_precision_loss)]
pub fn suggest_int_step( pub fn suggest_int_step(
&mut self, &mut self,
name: impl Into<String>, name: impl Into<String>,
@@ -693,9 +701,10 @@ impl Trial {
// Sample using the sampler (sampler handles step-grid) // Sample using the sampler (sampler handles step-grid)
let dist = Distribution::Int(distribution); let dist = Distribution::Int(distribution);
let value = match self.sample_value(&dist) { let ParamValue::Int(value) = self.sample_value(&dist) else {
ParamValue::Int(v) => v, return Err(TpeError::Internal(
_ => unreachable!("Int distribution should return Int value"), "Int distribution should return Int value",
));
}; };
// Store distribution and value // Store distribution and value
@@ -779,9 +788,10 @@ impl Trial {
// Sample using the sampler // Sample using the sampler
let dist = Distribution::Categorical(distribution); let dist = Distribution::Categorical(distribution);
let index = match self.sample_value(&dist) { let ParamValue::Categorical(index) = self.sample_value(&dist) else {
ParamValue::Categorical(idx) => idx, return Err(TpeError::Internal(
_ => unreachable!("Categorical distribution should return Categorical value"), "Categorical distribution should return Categorical value",
));
}; };
// Store distribution and value (store the index) // Store distribution and value (store the index)
-5
View File
@@ -1,11 +1,7 @@
//! Core types for the optimizer library. //! Core types for the optimizer library.
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
/// The direction of optimization. /// The direction of optimization.
#[derive(Clone, Copy, Debug, PartialEq, Eq)] #[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum Direction { pub enum Direction {
/// Minimize the objective value. /// Minimize the objective value.
Minimize, Minimize,
@@ -15,7 +11,6 @@ pub enum Direction {
/// The state of a trial in its lifecycle. /// The state of a trial in its lifecycle.
#[derive(Clone, Copy, Debug, PartialEq, Eq)] #[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum TrialState { pub enum TrialState {
/// The trial is currently running. /// The trial is currently running.
Running, Running,
+13 -3
View File
@@ -4,7 +4,9 @@
#![cfg(feature = "async")] #![cfg(feature = "async")]
use optimizer::{Direction, RandomSampler, Study, TpeError, TpeSampler}; use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, Study, TpeError};
#[tokio::test] #[tokio::test]
async fn test_optimize_async_basic() { async fn test_optimize_async_basic() {
@@ -26,7 +28,11 @@ async fn test_optimize_async_basic() {
#[tokio::test] #[tokio::test]
async fn test_optimize_async_with_sampler() { async fn test_optimize_async_with_sampler() {
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build(); let sampler = TpeSampler::builder()
.seed(42)
.n_startup_trials(5)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -61,7 +67,11 @@ async fn test_optimize_parallel() {
#[tokio::test] #[tokio::test]
async fn test_optimize_parallel_with_sampler() { async fn test_optimize_parallel_with_sampler() {
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build(); let sampler = TpeSampler::builder()
.seed(42)
.n_startup_trials(5)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
+56 -157
View File
@@ -1,6 +1,14 @@
//! Integration tests for the optimizer library. //! Integration tests for the optimizer library.
use optimizer::{Direction, RandomSampler, Study, TpeError, TpeSampler, Trial}; #![allow(
clippy::cast_sign_loss,
clippy::cast_precision_loss,
clippy::cast_possible_truncation
)]
use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, Study, TpeError, Trial};
// ============================================================================= // =============================================================================
// Test: optimize simple quadratic function with TPE, finds near-optimal // Test: optimize simple quadratic function with TPE, finds near-optimal
@@ -14,7 +22,8 @@ fn test_tpe_optimizes_quadratic_function() {
.seed(42) .seed(42)
.n_startup_trials(5) // Quick startup for test .n_startup_trials(5) // Quick startup for test
.n_ei_candidates(24) .n_ei_candidates(24)
.build(); .build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -40,7 +49,11 @@ fn test_tpe_optimizes_quadratic_function() {
fn test_tpe_optimizes_multivariate_function() { fn test_tpe_optimizes_multivariate_function() {
// Minimize f(x, y) = x^2 + y^2 where x, y ∈ [-5, 5] // Minimize f(x, y) = x^2 + y^2 where x, y ∈ [-5, 5]
// Optimal: (0, 0), f(0, 0) = 0 // Optimal: (0, 0), f(0, 0) = 0
let sampler = TpeSampler::builder().seed(123).n_startup_trials(10).build(); let sampler = TpeSampler::builder()
.seed(123)
.n_startup_trials(10)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -66,7 +79,11 @@ fn test_tpe_optimizes_multivariate_function() {
fn test_tpe_maximization() { fn test_tpe_maximization() {
// Maximize f(x) = -(x - 2)^2 + 10 where x ∈ [-10, 10] // Maximize f(x) = -(x - 2)^2 + 10 where x ∈ [-10, 10]
// Optimal: x = 2, f(2) = 10 // Optimal: x = 2, f(2) = 10
let sampler = TpeSampler::builder().seed(456).n_startup_trials(5).build(); let sampler = TpeSampler::builder()
.seed(456)
.n_startup_trials(5)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
@@ -550,7 +567,11 @@ fn test_invalid_step_errors() {
#[test] #[test]
fn test_tpe_with_categorical_parameter() { fn test_tpe_with_categorical_parameter() {
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build(); let sampler = TpeSampler::builder()
.seed(42)
.n_startup_trials(5)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
@@ -582,7 +603,11 @@ fn test_tpe_with_categorical_parameter() {
#[test] #[test]
fn test_tpe_with_integer_parameters() { fn test_tpe_with_integer_parameters() {
let sampler = TpeSampler::builder().seed(789).n_startup_trials(5).build(); let sampler = TpeSampler::builder()
.seed(789)
.n_startup_trials(5)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -756,7 +781,11 @@ fn test_study_set_sampler() {
let mut study: Study<f64> = Study::new(Direction::Minimize); let mut study: Study<f64> = Study::new(Direction::Minimize);
// Initially uses RandomSampler, now switch to TPE // Initially uses RandomSampler, now switch to TPE
let tpe = TpeSampler::builder().seed(42).n_startup_trials(5).build(); let tpe = TpeSampler::builder()
.seed(42)
.n_startup_trials(5)
.build()
.unwrap();
study.set_sampler(tpe); study.set_sampler(tpe);
// Should work with the new sampler // Should work with the new sampler
@@ -863,10 +892,10 @@ fn test_trial_debug_format() {
#[test] #[test]
fn test_tpe_sampler_builder_default_trait() { fn test_tpe_sampler_builder_default_trait() {
use optimizer::TpeSamplerBuilder; use optimizer::sampler::tpe::TpeSamplerBuilder;
let builder = TpeSamplerBuilder::default(); let builder = TpeSamplerBuilder::default();
let sampler = builder.build(); let sampler = builder.build().unwrap();
// Should have default values // Should have default values
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -901,7 +930,8 @@ fn test_tpe_with_fixed_kde_bandwidth() {
.seed(42) .seed(42)
.n_startup_trials(5) .n_startup_trials(5)
.kde_bandwidth(0.5) .kde_bandwidth(0.5)
.build(); .build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -917,9 +947,9 @@ fn test_tpe_with_fixed_kde_bandwidth() {
} }
#[test] #[test]
#[should_panic(expected = "kde_bandwidth must be positive")]
fn test_tpe_sampler_invalid_kde_bandwidth() { fn test_tpe_sampler_invalid_kde_bandwidth() {
TpeSampler::with_config(0.25, 10, 24, Some(-1.0), None); let result = TpeSampler::with_config(0.25, 10, 24, Some(-1.0), None);
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
} }
#[test] #[test]
@@ -928,7 +958,8 @@ fn test_tpe_split_trials_with_two_trials() {
let sampler = TpeSampler::builder() let sampler = TpeSampler::builder()
.seed(42) .seed(42)
.n_startup_trials(2) // TPE kicks in after 2 trials .n_startup_trials(2) // TPE kicks in after 2 trials
.build(); .build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -944,7 +975,11 @@ fn test_tpe_split_trials_with_two_trials() {
#[test] #[test]
fn test_tpe_with_log_scale_int() { fn test_tpe_with_log_scale_int() {
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build(); let sampler = TpeSampler::builder()
.seed(42)
.n_startup_trials(5)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -962,7 +997,11 @@ fn test_tpe_with_log_scale_int() {
#[test] #[test]
fn test_tpe_with_step_distributions() { fn test_tpe_with_step_distributions() {
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build(); let sampler = TpeSampler::builder()
.seed(42)
.n_startup_trials(5)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -1040,7 +1079,8 @@ fn test_tpe_empty_good_or_bad_values_fallback() {
.seed(42) .seed(42)
.n_startup_trials(5) .n_startup_trials(5)
.gamma(0.1) // Very small gamma means few "good" trials .gamma(0.1) // Very small gamma means few "good" trials
.build(); .build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler); let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
@@ -1142,144 +1182,3 @@ fn test_best_trial_with_nan_values() {
let best = study.best_trial(); let best = study.best_trial();
assert!(best.is_ok()); assert!(best.is_ok());
} }
// =============================================================================
// Serde tests (only run when serde feature is enabled)
// =============================================================================
#[cfg(feature = "serde")]
mod serde_tests {
use super::*;
#[test]
fn test_direction_serde() {
// Test Direction serialization
let min = Direction::Minimize;
let max = Direction::Maximize;
let min_json = serde_json::to_string(&min).unwrap();
let max_json = serde_json::to_string(&max).unwrap();
let min_deser: Direction = serde_json::from_str(&min_json).unwrap();
let max_deser: Direction = serde_json::from_str(&max_json).unwrap();
assert_eq!(min, min_deser);
assert_eq!(max, max_deser);
}
#[test]
fn test_study_serde_with_categorical() {
let study: Study<f64> = Study::new(Direction::Minimize);
study
.optimize(5, |trial| {
let x = trial.suggest_float("x", 0.0, 10.0)?;
let opt = trial.suggest_categorical("opt", &["a", "b", "c"])?;
let _ = opt;
Ok::<_, TpeError>(x)
})
.unwrap();
// Serialize
let json = serde_json::to_string(&study).unwrap();
// Deserialize
let loaded: Study<f64> = serde_json::from_str(&json).unwrap();
assert_eq!(loaded.n_trials(), 5);
assert_eq!(loaded.direction(), Direction::Minimize);
}
#[test]
fn test_study_serde_with_all_param_types() {
let study: Study<f64> = Study::new(Direction::Maximize);
study
.optimize(3, |trial| {
let x = trial.suggest_float("x", 0.0, 10.0)?;
let y = trial.suggest_float_log("y", 0.001, 1.0)?;
let z = trial.suggest_float_step("z", 0.0, 1.0, 0.1)?;
let a = trial.suggest_int("a", 1, 10)?;
let b = trial.suggest_int_log("b", 1, 100)?;
let c = trial.suggest_int_step("c", 0, 100, 10)?;
let d = trial.suggest_categorical("d", &["p", "q"])?;
let _ = (y, z, b, c, d);
Ok::<_, TpeError>(x + a as f64)
})
.unwrap();
let json = serde_json::to_string(&study).unwrap();
let loaded: Study<f64> = serde_json::from_str(&json).unwrap();
assert_eq!(loaded.n_trials(), 3);
assert_eq!(loaded.direction(), Direction::Maximize);
// Verify we can continue optimization
loaded
.optimize(2, |trial| {
let x = trial.suggest_float("x", 0.0, 10.0)?;
Ok::<_, TpeError>(x)
})
.unwrap();
assert_eq!(loaded.n_trials(), 5);
}
#[test]
fn test_study_serde_empty() {
// Test serializing a study with no trials
let study: Study<f64> = Study::new(Direction::Minimize);
let json = serde_json::to_string(&study).unwrap();
let loaded: Study<f64> = serde_json::from_str(&json).unwrap();
assert_eq!(loaded.n_trials(), 0);
assert_eq!(loaded.direction(), Direction::Minimize);
assert!(loaded.best_trial().is_err());
}
#[test]
fn test_study_serde_with_custom_value_type() {
// Test Study with i32 value type
let study: Study<i32> = Study::new(Direction::Minimize);
study
.optimize(5, |trial| {
let n = trial.suggest_int("n", 1, 100)?;
Ok::<_, TpeError>(n as i32)
})
.unwrap();
let json = serde_json::to_string(&study).unwrap();
let loaded: Study<i32> = serde_json::from_str(&json).unwrap();
assert_eq!(loaded.n_trials(), 5);
let best = loaded.best_trial().unwrap();
assert!(best.value >= 1 && best.value <= 100);
}
#[test]
fn test_completed_trial_access_after_serde() {
let study: Study<f64> = Study::new(Direction::Minimize);
study
.optimize(3, |trial| {
let x = trial.suggest_float("x", 0.0, 10.0)?;
Ok::<_, TpeError>(x * x)
})
.unwrap();
let json = serde_json::to_string(&study).unwrap();
let loaded: Study<f64> = serde_json::from_str(&json).unwrap();
// Access all trials
let trials = loaded.trials();
assert_eq!(trials.len(), 3);
for trial in &trials {
assert!(trial.params.contains_key("x"));
assert!(trial.distributions.contains_key("x"));
assert!(trial.value >= 0.0); // x^2 is non-negative
}
}
}