Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 7c364ecfad | |||
| a488a6303e |
@@ -42,7 +42,9 @@ jobs:
|
||||
matrix:
|
||||
features:
|
||||
- ""
|
||||
- "serde"
|
||||
- "async"
|
||||
- "serde,async"
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Install Rust
|
||||
@@ -129,19 +131,6 @@ jobs:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: EmbarkStudios/cargo-deny-action@v2
|
||||
|
||||
machete:
|
||||
name: Unused Dependencies
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Install Rust
|
||||
run: |
|
||||
rustup override set stable
|
||||
rustup update stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
- name: Machete
|
||||
uses: bnjbvr/cargo-machete@7959c845782fed02ee69303126d4a12d64f1db18
|
||||
|
||||
semver:
|
||||
name: Semver Check
|
||||
runs-on: ubuntu-latest
|
||||
@@ -195,27 +184,6 @@ jobs:
|
||||
- name: Run tests
|
||||
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:
|
||||
name: Typos
|
||||
@@ -231,51 +199,3 @@ jobs:
|
||||
tool: typos-cli
|
||||
- name: Typos Check
|
||||
run: typos src/
|
||||
|
||||
publish:
|
||||
name: Publish to crates.io
|
||||
runs-on: ubuntu-latest
|
||||
if: github.event_name == 'push'
|
||||
needs:
|
||||
- fmt
|
||||
- clippy
|
||||
- test
|
||||
- docs
|
||||
- feature-check
|
||||
- coverage
|
||||
- msrv
|
||||
- deny
|
||||
- machete
|
||||
- semver
|
||||
- minimal-versions
|
||||
- cross-platform
|
||||
- cross-targets
|
||||
- typos
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Install Rust
|
||||
run: |
|
||||
rustup override set stable
|
||||
rustup update stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
|
||||
- name: Check if version already published
|
||||
id: check
|
||||
run: |
|
||||
CARGO_VERSION=$(cargo metadata --no-deps --format-version 1 | jq -r '.packages[0].version')
|
||||
echo "version=$CARGO_VERSION" >> "$GITHUB_OUTPUT"
|
||||
|
||||
HTTP_STATUS=$(curl -s -o /dev/null -w "%{http_code}" "https://crates.io/api/v1/crates/optimizer/$CARGO_VERSION")
|
||||
if [ "$HTTP_STATUS" = "200" ]; then
|
||||
echo "Version $CARGO_VERSION already exists on crates.io, skipping publish"
|
||||
echo "skip=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "Version $CARGO_VERSION not found on crates.io, will publish"
|
||||
echo "skip=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Publish
|
||||
if: steps.check.outputs.skip == 'false' && github.ref == 'refs/heads/master'
|
||||
run: cargo publish
|
||||
env:
|
||||
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
name: Publish
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- "v*"
|
||||
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
name: Publish to crates.io
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- name: Install Rust
|
||||
run: |
|
||||
rustup override set stable
|
||||
rustup update stable
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
|
||||
- name: Verify version matches tag
|
||||
run: |
|
||||
CARGO_VERSION=$(cargo metadata --no-deps --format-version 1 | jq -r '.packages[0].version')
|
||||
TAG_VERSION="${GITHUB_REF_NAME#v}"
|
||||
if [ "$CARGO_VERSION" != "$TAG_VERSION" ]; then
|
||||
echo "Error: Cargo.toml version ($CARGO_VERSION) doesn't match tag ($TAG_VERSION)"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Publish
|
||||
run: cargo publish
|
||||
env:
|
||||
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
|
||||
+5
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "optimizer"
|
||||
version = "0.3.0"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.88"
|
||||
license = "MIT"
|
||||
@@ -13,11 +13,15 @@ repository = "https://github.com/raimannma/rust-optimizer"
|
||||
rand = "0.9"
|
||||
thiserror = "2"
|
||||
parking_lot = "0.12"
|
||||
ordered-float = "5"
|
||||
serde = { version = "1", features = ["derive"], optional = true }
|
||||
tokio = { version = "1", features = ["sync", "rt-multi-thread"], optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
serde = ["dep:serde", "ordered-float/serde"]
|
||||
async = ["dep:tokio"]
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json = "1"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
|
||||
|
||||
@@ -12,6 +12,7 @@ A Rust library for black-box optimization using Tree-Parzen Estimator (TPE).
|
||||
- Float, integer, and categorical parameter types
|
||||
- Log-scale and stepped parameter sampling
|
||||
- Sync and async optimization with parallel trial evaluation
|
||||
- Serialization support for saving/loading study state
|
||||
|
||||
## Quick Start
|
||||
|
||||
@@ -34,6 +35,7 @@ println!("Best value: {} at x={:?}", best.value, best.params);
|
||||
|
||||
## Feature Flags
|
||||
|
||||
- `serde` - Enable serialization/deserialization of studies and trials
|
||||
- `async` - Enable async optimization methods (requires tokio)
|
||||
|
||||
## Documentation
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
//! Parameter distribution types.
|
||||
|
||||
#[cfg(feature = "serde")]
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Distribution for floating-point parameters.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub struct FloatDistribution {
|
||||
/// Lower bound (inclusive).
|
||||
pub low: f64,
|
||||
@@ -15,6 +19,7 @@ pub struct FloatDistribution {
|
||||
|
||||
/// Distribution for integer parameters.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub struct IntDistribution {
|
||||
/// Lower bound (inclusive).
|
||||
pub low: i64,
|
||||
@@ -28,6 +33,7 @@ pub struct IntDistribution {
|
||||
|
||||
/// Distribution for categorical parameters.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub struct CategoricalDistribution {
|
||||
/// Number of choices available.
|
||||
pub n_choices: usize,
|
||||
@@ -35,6 +41,7 @@ pub struct CategoricalDistribution {
|
||||
|
||||
/// Enum wrapping all parameter distribution types.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub enum Distribution {
|
||||
/// A floating-point distribution.
|
||||
Float(FloatDistribution),
|
||||
|
||||
+1
-22
@@ -38,28 +38,7 @@ pub enum TpeError {
|
||||
/// Returned when requesting the best trial but no trials have completed.
|
||||
#[error("no completed trials available")]
|
||||
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.
|
||||
pub type Result<T> = core::result::Result<T, TpeError>;
|
||||
pub type Result<T> = std::result::Result<T, TpeError>;
|
||||
|
||||
+34
-46
@@ -5,8 +5,6 @@
|
||||
|
||||
use rand::Rng;
|
||||
|
||||
use crate::error::{Result, TpeError};
|
||||
|
||||
/// A Gaussian kernel density estimator for continuous distributions.
|
||||
///
|
||||
/// KDE estimates a probability density function from a set of samples by
|
||||
@@ -30,7 +28,7 @@ use crate::error::{Result, TpeError};
|
||||
/// let sample = kde.sample(&mut rng);
|
||||
/// ```
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) struct KernelDensityEstimator {
|
||||
pub struct KernelDensityEstimator {
|
||||
/// The sample points used to construct the KDE.
|
||||
samples: Vec<f64>,
|
||||
/// The bandwidth (standard deviation) of the Gaussian kernels.
|
||||
@@ -40,45 +38,37 @@ pub(crate) struct KernelDensityEstimator {
|
||||
impl KernelDensityEstimator {
|
||||
/// Creates a new KDE with automatic bandwidth selection using Scott's rule.
|
||||
///
|
||||
/// Scott's rule sets bandwidth = n^(-1/5) * `std_dev`, which works well
|
||||
/// Scott's rule sets bandwidth = n^(-1/5) * std_dev, which works well
|
||||
/// for unimodal distributions close to normal.
|
||||
///
|
||||
/// # Errors
|
||||
/// # Panics
|
||||
///
|
||||
/// Returns `TpeError::EmptySamples` if `samples` is empty.
|
||||
pub(crate) fn new(samples: Vec<f64>) -> Result<Self> {
|
||||
if samples.is_empty() {
|
||||
return Err(TpeError::EmptySamples);
|
||||
}
|
||||
/// Panics if `samples` is empty.
|
||||
pub fn new(samples: Vec<f64>) -> Self {
|
||||
assert!(!samples.is_empty(), "KDE requires at least one sample");
|
||||
|
||||
let bandwidth = Self::scotts_rule(&samples);
|
||||
Ok(Self { samples, bandwidth })
|
||||
Self { samples, bandwidth }
|
||||
}
|
||||
|
||||
/// Creates a new KDE with a specified bandwidth.
|
||||
///
|
||||
/// Use this when you want explicit control over the smoothing parameter.
|
||||
///
|
||||
/// # Errors
|
||||
/// # Panics
|
||||
///
|
||||
/// Returns `TpeError::EmptySamples` if `samples` is empty.
|
||||
/// Returns `TpeError::InvalidBandwidth` if `bandwidth` is not positive.
|
||||
pub(crate) fn with_bandwidth(samples: Vec<f64>, bandwidth: f64) -> Result<Self> {
|
||||
if samples.is_empty() {
|
||||
return Err(TpeError::EmptySamples);
|
||||
}
|
||||
if bandwidth <= 0.0 {
|
||||
return Err(TpeError::InvalidBandwidth(bandwidth));
|
||||
}
|
||||
/// Panics if `samples` is empty or `bandwidth` is not positive.
|
||||
pub(crate) fn with_bandwidth(samples: Vec<f64>, bandwidth: f64) -> Self {
|
||||
assert!(!samples.is_empty(), "KDE requires at least one sample");
|
||||
assert!(bandwidth > 0.0, "Bandwidth must be positive");
|
||||
|
||||
Ok(Self { samples, bandwidth })
|
||||
Self { samples, bandwidth }
|
||||
}
|
||||
|
||||
/// Computes bandwidth using Scott's rule.
|
||||
///
|
||||
/// Scott's rule: h = n^(-1/5) * sigma
|
||||
/// where sigma is the sample standard deviation.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn scotts_rule(samples: &[f64]) -> f64 {
|
||||
let n = samples.len() as f64;
|
||||
let std_dev = Self::sample_std_dev(samples);
|
||||
@@ -93,7 +83,6 @@ impl KernelDensityEstimator {
|
||||
}
|
||||
|
||||
/// Computes the sample standard deviation.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
fn sample_std_dev(samples: &[f64]) -> f64 {
|
||||
let n = samples.len() as f64;
|
||||
let mean = samples.iter().sum::<f64>() / n;
|
||||
@@ -106,14 +95,13 @@ impl KernelDensityEstimator {
|
||||
/// The density is computed as the average of Gaussian kernels centered
|
||||
/// at each sample point:
|
||||
///
|
||||
/// f(x) = (1/n) * `sum_i` K((x - `x_i`) / h)
|
||||
/// f(x) = (1/n) * sum_i K((x - x_i) / h)
|
||||
///
|
||||
/// where K is the standard Gaussian kernel and h is the bandwidth.
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub(crate) fn pdf(&self, x: f64) -> f64 {
|
||||
pub fn pdf(&self, x: f64) -> f64 {
|
||||
let n = self.samples.len() as f64;
|
||||
let inv_bandwidth = 1.0 / self.bandwidth;
|
||||
let normalization = inv_bandwidth / (2.0 * core::f64::consts::PI).sqrt();
|
||||
let normalization = inv_bandwidth / (2.0 * std::f64::consts::PI).sqrt();
|
||||
|
||||
let density: f64 = self
|
||||
.samples
|
||||
@@ -132,7 +120,7 @@ impl KernelDensityEstimator {
|
||||
/// Sampling works by:
|
||||
/// 1. Uniformly selecting one of the kernel centers (samples)
|
||||
/// 2. Adding Gaussian noise with the bandwidth as standard deviation
|
||||
pub(crate) fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
|
||||
pub fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
|
||||
// Select a random sample to center the kernel on
|
||||
let idx = rng.random_range(0..self.samples.len());
|
||||
let center = self.samples[idx];
|
||||
@@ -142,7 +130,7 @@ impl KernelDensityEstimator {
|
||||
let u1: f64 = rng.random();
|
||||
let u2: f64 = rng.random();
|
||||
|
||||
let z = (-2.0 * u1.ln()).sqrt() * (2.0 * core::f64::consts::PI * u2).cos();
|
||||
let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
|
||||
center + z * self.bandwidth
|
||||
}
|
||||
|
||||
@@ -160,7 +148,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_kde_pdf_basic() {
|
||||
let samples = vec![0.0, 1.0, 2.0];
|
||||
let kde = KernelDensityEstimator::new(samples).unwrap();
|
||||
let kde = KernelDensityEstimator::new(samples);
|
||||
|
||||
// Density should be positive everywhere
|
||||
assert!(kde.pdf(0.0) > 0.0);
|
||||
@@ -176,17 +164,17 @@ mod tests {
|
||||
#[test]
|
||||
fn test_kde_pdf_integrates_to_one() {
|
||||
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
|
||||
let kde = KernelDensityEstimator::new(samples).unwrap();
|
||||
let kde = KernelDensityEstimator::new(samples);
|
||||
|
||||
// Numerical integration over a wide range
|
||||
let n_points = 10000;
|
||||
let low = -10.0;
|
||||
let high = 15.0;
|
||||
let dx = (high - low) / f64::from(n_points);
|
||||
let dx = (high - low) / n_points as f64;
|
||||
|
||||
let integral: f64 = (0..n_points)
|
||||
.map(|i| {
|
||||
let x = low + (f64::from(i) + 0.5) * dx;
|
||||
let x = low + (i as f64 + 0.5) * dx;
|
||||
kde.pdf(x) * dx
|
||||
})
|
||||
.sum();
|
||||
@@ -201,16 +189,16 @@ mod tests {
|
||||
#[test]
|
||||
fn test_kde_with_bandwidth() {
|
||||
let samples = vec![0.0, 1.0, 2.0];
|
||||
let kde = KernelDensityEstimator::with_bandwidth(samples, 0.5).unwrap();
|
||||
let kde = KernelDensityEstimator::with_bandwidth(samples, 0.5);
|
||||
|
||||
assert!((kde.bandwidth() - 0.5).abs() < f64::EPSILON);
|
||||
assert_eq!(kde.bandwidth(), 0.5);
|
||||
assert!(kde.pdf(1.0) > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kde_sample_in_reasonable_range() {
|
||||
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
|
||||
let kde = KernelDensityEstimator::new(samples).unwrap();
|
||||
let kde = KernelDensityEstimator::new(samples);
|
||||
let mut rng = rand::rng();
|
||||
|
||||
// Samples should generally be in a reasonable range around the data
|
||||
@@ -225,7 +213,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_kde_single_sample() {
|
||||
let samples = vec![5.0];
|
||||
let kde = KernelDensityEstimator::new(samples).unwrap();
|
||||
let kde = KernelDensityEstimator::new(samples);
|
||||
|
||||
// Should have positive density near the sample
|
||||
assert!(kde.pdf(5.0) > 0.0);
|
||||
@@ -235,7 +223,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_kde_identical_samples() {
|
||||
let samples = vec![3.0, 3.0, 3.0, 3.0];
|
||||
let kde = KernelDensityEstimator::new(samples).unwrap();
|
||||
let kde = KernelDensityEstimator::new(samples);
|
||||
|
||||
// Should handle degenerate case with identical samples
|
||||
assert!(kde.bandwidth() > 0.0);
|
||||
@@ -245,7 +233,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_scotts_rule_bandwidth() {
|
||||
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
|
||||
let kde = KernelDensityEstimator::new(samples).unwrap();
|
||||
let kde = KernelDensityEstimator::new(samples);
|
||||
|
||||
// n = 10, n^(-1/5) ≈ 0.631
|
||||
// std_dev ≈ 2.87
|
||||
@@ -258,23 +246,23 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "KDE requires at least one sample")]
|
||||
fn test_kde_empty_samples() {
|
||||
let samples: Vec<f64> = vec![];
|
||||
let result = KernelDensityEstimator::new(samples);
|
||||
assert!(matches!(result, Err(TpeError::EmptySamples)));
|
||||
KernelDensityEstimator::new(samples);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "Bandwidth must be positive")]
|
||||
fn test_kde_zero_bandwidth() {
|
||||
let samples = vec![1.0, 2.0, 3.0];
|
||||
let result = KernelDensityEstimator::with_bandwidth(samples, 0.0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
|
||||
KernelDensityEstimator::with_bandwidth(samples, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "Bandwidth must be positive")]
|
||||
fn test_kde_negative_bandwidth() {
|
||||
let samples = vec![1.0, 2.0, 3.0];
|
||||
let result = KernelDensityEstimator::with_bandwidth(samples, -1.0);
|
||||
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
|
||||
KernelDensityEstimator::with_bandwidth(samples, -1.0);
|
||||
}
|
||||
}
|
||||
|
||||
+27
-23
@@ -1,14 +1,3 @@
|
||||
#![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.
|
||||
//!
|
||||
//! This library provides an Optuna-like API for hyperparameter optimization
|
||||
@@ -18,15 +7,15 @@
|
||||
//! - Log-scale and stepped parameter sampling
|
||||
//! - Synchronous and async optimization
|
||||
//! - Parallel trial evaluation with bounded concurrency
|
||||
//! - Serialization for saving/loading study state
|
||||
//!
|
||||
//! # Quick Start
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::sampler::tpe::TpeSampler;
|
||||
//! use optimizer::{Direction, Study};
|
||||
//! use optimizer::{Direction, Study, TpeSampler};
|
||||
//!
|
||||
//! // Create a study with TPE sampler
|
||||
//! let sampler = TpeSampler::builder().seed(42).build().unwrap();
|
||||
//! let sampler = TpeSampler::builder().seed(42).build();
|
||||
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
//!
|
||||
//! // Optimize x^2 for 20 trials
|
||||
@@ -47,9 +36,7 @@
|
||||
//! A [`Study`] manages optimization trials. Create one with an optimization direction:
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::sampler::random::RandomSampler;
|
||||
//! use optimizer::sampler::tpe::TpeSampler;
|
||||
//! use optimizer::{Direction, Study};
|
||||
//! use optimizer::{Direction, RandomSampler, Study, TpeSampler};
|
||||
//!
|
||||
//! // Minimize with default random sampler
|
||||
//! let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
@@ -93,18 +80,17 @@
|
||||
//!
|
||||
//! # Configuring TPE
|
||||
//!
|
||||
//! The [`sampler::tpe::TpeSampler`] can be configured using the builder pattern:
|
||||
//! The [`TpeSampler`] can be configured using the builder pattern:
|
||||
//!
|
||||
//! ```
|
||||
//! use optimizer::sampler::tpe::TpeSampler;
|
||||
//! use optimizer::TpeSampler;
|
||||
//!
|
||||
//! let sampler = TpeSampler::builder()
|
||||
//! .gamma(0.15) // Quantile for good/bad split
|
||||
//! .n_startup_trials(20) // Random trials before TPE
|
||||
//! .n_ei_candidates(32) // Candidates to evaluate
|
||||
//! .seed(42) // Reproducibility
|
||||
//! .build()
|
||||
//! .unwrap();
|
||||
//! .build();
|
||||
//! ```
|
||||
//!
|
||||
//! # Async and Parallel Optimization
|
||||
@@ -127,21 +113,39 @@
|
||||
//! }).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
|
||||
//!
|
||||
//! - `serde`: Enable serialization/deserialization of studies and trials
|
||||
//! - `async`: Enable async optimization methods (requires tokio)
|
||||
|
||||
mod distribution;
|
||||
mod error;
|
||||
mod kde;
|
||||
mod param;
|
||||
pub mod sampler;
|
||||
mod sampler;
|
||||
mod study;
|
||||
mod trial;
|
||||
mod types;
|
||||
|
||||
pub use error::{Result, TpeError};
|
||||
pub use param::ParamValue;
|
||||
pub use sampler::{CompletedTrial, RandomSampler, Sampler, TpeSampler, TpeSamplerBuilder};
|
||||
pub use study::Study;
|
||||
pub use trial::Trial;
|
||||
pub use types::{Direction, TrialState};
|
||||
|
||||
@@ -1,11 +1,15 @@
|
||||
//! Parameter value storage types.
|
||||
|
||||
#[cfg(feature = "serde")]
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Represents a sampled parameter value.
|
||||
///
|
||||
/// This enum stores different parameter value types uniformly.
|
||||
/// For categorical parameters, the `Categorical` variant stores
|
||||
/// the index into the choices array.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub enum ParamValue {
|
||||
/// A floating-point parameter value.
|
||||
Float(f64),
|
||||
|
||||
+7
-1
@@ -1,10 +1,15 @@
|
||||
//! Sampler trait and implementations for parameter sampling.
|
||||
|
||||
pub mod random;
|
||||
mod random;
|
||||
pub mod tpe;
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
pub use random::RandomSampler;
|
||||
#[cfg(feature = "serde")]
|
||||
use serde::{Deserialize, Serialize};
|
||||
pub use tpe::{TpeSampler, TpeSamplerBuilder};
|
||||
|
||||
use crate::distribution::Distribution;
|
||||
use crate::param::ParamValue;
|
||||
|
||||
@@ -14,6 +19,7 @@ use crate::param::ParamValue;
|
||||
/// parameter values, their distributions, and the objective value returned
|
||||
/// by the objective function.
|
||||
#[derive(Clone, Debug)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub struct CompletedTrial<V = f64> {
|
||||
/// The unique identifier for this trial.
|
||||
pub id: u64,
|
||||
|
||||
@@ -17,7 +17,7 @@ use crate::sampler::{CompletedTrial, Sampler};
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::RandomSampler;
|
||||
///
|
||||
/// // Create with default RNG
|
||||
/// let sampler = RandomSampler::new();
|
||||
@@ -31,7 +31,6 @@ pub struct RandomSampler {
|
||||
|
||||
impl RandomSampler {
|
||||
/// Creates a new random sampler with a default random seed.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
rng: Mutex::new(StdRng::from_os_rng()),
|
||||
@@ -41,7 +40,6 @@ impl RandomSampler {
|
||||
/// Creates a new random sampler with a fixed seed for reproducibility.
|
||||
///
|
||||
/// Using the same seed will produce the same sequence of sampled values.
|
||||
#[must_use]
|
||||
pub fn with_seed(seed: u64) -> Self {
|
||||
Self {
|
||||
rng: Mutex::new(StdRng::seed_from_u64(seed)),
|
||||
@@ -56,7 +54,6 @@ impl Default for RandomSampler {
|
||||
}
|
||||
|
||||
impl Sampler for RandomSampler {
|
||||
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
|
||||
fn sample(
|
||||
&self,
|
||||
distribution: &Distribution,
|
||||
@@ -113,7 +110,6 @@ impl Sampler for RandomSampler {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
|
||||
|
||||
+99
-148
@@ -9,7 +9,6 @@ use rand::rngs::StdRng;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
use crate::distribution::Distribution;
|
||||
use crate::error::{Result, TpeError};
|
||||
use crate::kde::KernelDensityEstimator;
|
||||
use crate::param::ParamValue;
|
||||
use crate::sampler::{CompletedTrial, Sampler};
|
||||
@@ -28,7 +27,7 @@ use crate::sampler::{CompletedTrial, Sampler};
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSampler;
|
||||
/// use optimizer::TpeSampler;
|
||||
///
|
||||
/// // Create with default settings
|
||||
/// let sampler = TpeSampler::new();
|
||||
@@ -39,8 +38,7 @@ use crate::sampler::{CompletedTrial, Sampler};
|
||||
/// .n_startup_trials(20)
|
||||
/// .n_ei_candidates(32)
|
||||
/// .seed(42)
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
pub struct TpeSampler {
|
||||
/// Fraction of trials to consider as "good" (gamma quantile).
|
||||
@@ -60,10 +58,9 @@ impl TpeSampler {
|
||||
///
|
||||
/// Default settings:
|
||||
/// - gamma: 0.25 (top 25% of trials are considered "good")
|
||||
/// - `n_startup_trials`: 10 (random sampling for first 10 trials)
|
||||
/// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample)
|
||||
/// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth)
|
||||
#[must_use]
|
||||
/// - n_startup_trials: 10 (random sampling for first 10 trials)
|
||||
/// - n_ei_candidates: 24 (evaluate 24 candidates per sample)
|
||||
/// - kde_bandwidth: None (uses Scott's rule for automatic bandwidth)
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
gamma: 0.25,
|
||||
@@ -79,17 +76,15 @@ impl TpeSampler {
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSampler;
|
||||
/// use optimizer::TpeSampler;
|
||||
///
|
||||
/// let sampler = TpeSampler::builder()
|
||||
/// .gamma(0.15)
|
||||
/// .n_startup_trials(20)
|
||||
/// .n_ei_candidates(32)
|
||||
/// .seed(42)
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn builder() -> TpeSamplerBuilder {
|
||||
TpeSamplerBuilder::new()
|
||||
}
|
||||
@@ -104,24 +99,22 @@ impl TpeSampler {
|
||||
/// * `kde_bandwidth` - Optional fixed bandwidth for KDE. If None, uses Scott's rule.
|
||||
/// * `seed` - Optional seed for reproducibility.
|
||||
///
|
||||
/// # Errors
|
||||
/// # Panics
|
||||
///
|
||||
/// Returns `TpeError::InvalidGamma` if gamma is not in (0.0, 1.0).
|
||||
/// Returns `TpeError::InvalidBandwidth` if `kde_bandwidth` is Some but not positive.
|
||||
/// Panics if gamma is not in (0.0, 1.0) or if kde_bandwidth is Some but not positive.
|
||||
pub fn with_config(
|
||||
gamma: f64,
|
||||
n_startup_trials: usize,
|
||||
n_ei_candidates: usize,
|
||||
kde_bandwidth: Option<f64>,
|
||||
seed: Option<u64>,
|
||||
) -> Result<Self> {
|
||||
if gamma <= 0.0 || gamma >= 1.0 {
|
||||
return Err(TpeError::InvalidGamma(gamma));
|
||||
}
|
||||
if let Some(bw) = kde_bandwidth
|
||||
&& bw <= 0.0
|
||||
{
|
||||
return Err(TpeError::InvalidBandwidth(bw));
|
||||
) -> Self {
|
||||
assert!(
|
||||
gamma > 0.0 && gamma < 1.0,
|
||||
"gamma must be in (0.0, 1.0), got {gamma}"
|
||||
);
|
||||
if let Some(bw) = kde_bandwidth {
|
||||
assert!(bw > 0.0, "kde_bandwidth must be positive, got {bw}");
|
||||
}
|
||||
|
||||
let rng = match seed {
|
||||
@@ -129,24 +122,19 @@ impl TpeSampler {
|
||||
None => StdRng::from_os_rng(),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
Self {
|
||||
gamma,
|
||||
n_startup_trials,
|
||||
n_ei_candidates,
|
||||
kde_bandwidth,
|
||||
rng: Mutex::new(rng),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 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).
|
||||
#[allow(
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss
|
||||
)]
|
||||
fn split_trials<'a>(
|
||||
&self,
|
||||
history: &'a [CompletedTrial],
|
||||
@@ -161,7 +149,7 @@ impl TpeSampler {
|
||||
history[a]
|
||||
.value
|
||||
.partial_cmp(&history[b].value)
|
||||
.unwrap_or(core::cmp::Ordering::Equal)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
});
|
||||
|
||||
// Calculate the split point (gamma quantile)
|
||||
@@ -183,11 +171,6 @@ impl TpeSampler {
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
match distribution {
|
||||
Distribution::Float(d) => {
|
||||
@@ -258,11 +241,6 @@ impl TpeSampler {
|
||||
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
|
||||
let mut best_candidate = internal_low;
|
||||
let mut best_ratio = f64::NEG_INFINITY;
|
||||
@@ -311,19 +289,15 @@ impl TpeSampler {
|
||||
}
|
||||
|
||||
/// Samples using TPE for integer distributions.
|
||||
#[allow(
|
||||
clippy::too_many_arguments,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_truncation
|
||||
)]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn sample_tpe_int(
|
||||
&self,
|
||||
low: i64,
|
||||
high: i64,
|
||||
log_scale: bool,
|
||||
step: Option<i64>,
|
||||
good_values: &[i64],
|
||||
bad_values: &[i64],
|
||||
good_values: Vec<i64>,
|
||||
bad_values: Vec<i64>,
|
||||
rng: &mut StdRng,
|
||||
) -> i64 {
|
||||
// Convert to floats for KDE
|
||||
@@ -357,24 +331,23 @@ impl TpeSampler {
|
||||
}
|
||||
|
||||
/// Samples using TPE for categorical distributions.
|
||||
#[allow(clippy::cast_precision_loss, clippy::unused_self)]
|
||||
fn sample_tpe_categorical(
|
||||
&self,
|
||||
n_choices: usize,
|
||||
good_indices: &[usize],
|
||||
bad_indices: &[usize],
|
||||
good_indices: Vec<usize>,
|
||||
bad_indices: Vec<usize>,
|
||||
rng: &mut StdRng,
|
||||
) -> usize {
|
||||
// Count occurrences in good and bad groups
|
||||
let mut good_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 {
|
||||
good_counts[idx] += 1;
|
||||
}
|
||||
}
|
||||
for &idx in bad_indices {
|
||||
for &idx in &bad_indices {
|
||||
if idx < n_choices {
|
||||
bad_counts[idx] += 1;
|
||||
}
|
||||
@@ -422,15 +395,14 @@ impl Default for TpeSampler {
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
/// use optimizer::TpeSamplerBuilder;
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .gamma(0.15)
|
||||
/// .n_startup_trials(20)
|
||||
/// .n_ei_candidates(32)
|
||||
/// .seed(42)
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TpeSamplerBuilder {
|
||||
@@ -446,11 +418,10 @@ impl TpeSamplerBuilder {
|
||||
///
|
||||
/// Default settings:
|
||||
/// - gamma: 0.25 (top 25% of trials are considered "good")
|
||||
/// - `n_startup_trials`: 10 (random sampling for first 10 trials)
|
||||
/// - `n_ei_candidates`: 24 (evaluate 24 candidates per sample)
|
||||
/// - `kde_bandwidth`: None (uses Scott's rule for automatic bandwidth)
|
||||
/// - n_startup_trials: 10 (random sampling for first 10 trials)
|
||||
/// - n_ei_candidates: 24 (evaluate 24 candidates per sample)
|
||||
/// - kde_bandwidth: None (uses Scott's rule for automatic bandwidth)
|
||||
/// - seed: None (use OS-provided entropy)
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
gamma: 0.25,
|
||||
@@ -470,23 +441,24 @@ impl TpeSamplerBuilder {
|
||||
///
|
||||
/// * `gamma` - Quantile value, must be in (0.0, 1.0).
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if gamma is not in (0.0, 1.0).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
/// use optimizer::TpeSamplerBuilder;
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .gamma(0.10) // Use top 10% as "good" trials
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
///
|
||||
/// # 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 {
|
||||
assert!(
|
||||
gamma > 0.0 && gamma < 1.0,
|
||||
"gamma must be in (0.0, 1.0), got {gamma}"
|
||||
);
|
||||
self.gamma = gamma;
|
||||
self
|
||||
}
|
||||
@@ -504,14 +476,12 @@ impl TpeSamplerBuilder {
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
/// use optimizer::TpeSamplerBuilder;
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .n_startup_trials(20) // Random sample first 20 trials
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn n_startup_trials(mut self, n: usize) -> Self {
|
||||
self.n_startup_trials = n;
|
||||
self
|
||||
@@ -530,14 +500,12 @@ impl TpeSamplerBuilder {
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
/// use optimizer::TpeSamplerBuilder;
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .n_ei_candidates(48) // Evaluate more candidates
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn n_ei_candidates(mut self, n: usize) -> Self {
|
||||
self.n_ei_candidates = n;
|
||||
self
|
||||
@@ -555,23 +523,24 @@ impl TpeSamplerBuilder {
|
||||
///
|
||||
/// * `bandwidth` - The fixed bandwidth (standard deviation) for Gaussian kernels.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// Panics if bandwidth is not positive.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
/// use optimizer::TpeSamplerBuilder;
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .kde_bandwidth(0.5) // Fixed bandwidth of 0.5
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
///
|
||||
/// # 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 {
|
||||
assert!(
|
||||
bandwidth > 0.0,
|
||||
"kde_bandwidth must be positive, got {bandwidth}"
|
||||
);
|
||||
self.kde_bandwidth = Some(bandwidth);
|
||||
self
|
||||
}
|
||||
@@ -585,14 +554,12 @@ impl TpeSamplerBuilder {
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
/// use optimizer::TpeSamplerBuilder;
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .seed(42) // Reproducible results
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn seed(mut self, seed: u64) -> Self {
|
||||
self.seed = Some(seed);
|
||||
self
|
||||
@@ -600,25 +567,19 @@ impl TpeSamplerBuilder {
|
||||
|
||||
/// 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
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
/// use optimizer::TpeSamplerBuilder;
|
||||
///
|
||||
/// let sampler = TpeSamplerBuilder::new()
|
||||
/// .gamma(0.15)
|
||||
/// .n_startup_trials(20)
|
||||
/// .n_ei_candidates(32)
|
||||
/// .seed(42)
|
||||
/// .build()
|
||||
/// .unwrap();
|
||||
/// .build();
|
||||
/// ```
|
||||
pub fn build(self) -> Result<TpeSampler> {
|
||||
pub fn build(self) -> TpeSampler {
|
||||
TpeSampler::with_config(
|
||||
self.gamma,
|
||||
self.n_startup_trials,
|
||||
@@ -636,7 +597,6 @@ impl Default for TpeSamplerBuilder {
|
||||
}
|
||||
|
||||
impl Sampler for TpeSampler {
|
||||
#[allow(clippy::too_many_lines)]
|
||||
fn sample(
|
||||
&self,
|
||||
distribution: &Distribution,
|
||||
@@ -733,8 +693,8 @@ impl Sampler for TpeSampler {
|
||||
d.high,
|
||||
d.log_scale,
|
||||
d.step,
|
||||
&good_values,
|
||||
&bad_values,
|
||||
good_values,
|
||||
bad_values,
|
||||
&mut rng,
|
||||
);
|
||||
ParamValue::Int(value)
|
||||
@@ -765,7 +725,7 @@ impl Sampler for TpeSampler {
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -773,11 +733,6 @@ impl Sampler for TpeSampler {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[allow(
|
||||
clippy::similar_names,
|
||||
clippy::cast_sign_loss,
|
||||
clippy::cast_precision_loss
|
||||
)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
@@ -801,34 +756,34 @@ mod tests {
|
||||
#[test]
|
||||
fn test_tpe_sampler_new() {
|
||||
let sampler = TpeSampler::new();
|
||||
assert!((sampler.gamma - 0.25).abs() < f64::EPSILON);
|
||||
assert_eq!(sampler.gamma, 0.25);
|
||||
assert_eq!(sampler.n_startup_trials, 10);
|
||||
assert_eq!(sampler.n_ei_candidates, 24);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_sampler_with_config() {
|
||||
let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42)).unwrap();
|
||||
assert!((sampler.gamma - 0.15).abs() < f64::EPSILON);
|
||||
let sampler = TpeSampler::with_config(0.15, 20, 32, None, Some(42));
|
||||
assert_eq!(sampler.gamma, 0.15);
|
||||
assert_eq!(sampler.n_startup_trials, 20);
|
||||
assert_eq!(sampler.n_ei_candidates, 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "gamma must be in (0.0, 1.0)")]
|
||||
fn test_tpe_sampler_invalid_gamma_zero() {
|
||||
let result = TpeSampler::with_config(0.0, 10, 24, None, None);
|
||||
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
|
||||
TpeSampler::with_config(0.0, 10, 24, None, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "gamma must be in (0.0, 1.0)")]
|
||||
fn test_tpe_sampler_invalid_gamma_one() {
|
||||
let result = TpeSampler::with_config(1.0, 10, 24, None, None);
|
||||
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
|
||||
TpeSampler::with_config(1.0, 10, 24, None, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tpe_startup_random_sampling() {
|
||||
let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)).unwrap();
|
||||
let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42));
|
||||
let dist = Distribution::Float(FloatDistribution {
|
||||
low: 0.0,
|
||||
high: 1.0,
|
||||
@@ -851,7 +806,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_split_trials() {
|
||||
let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42)).unwrap();
|
||||
let sampler = TpeSampler::with_config(0.25, 10, 24, None, Some(42));
|
||||
|
||||
let dist = Distribution::Float(FloatDistribution {
|
||||
low: 0.0,
|
||||
@@ -865,8 +820,8 @@ mod tests {
|
||||
.map(|i| {
|
||||
create_trial(
|
||||
i as u64,
|
||||
f64::from(i),
|
||||
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
i as f64,
|
||||
vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())],
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
@@ -885,7 +840,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_samples_float_with_history() {
|
||||
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();
|
||||
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42));
|
||||
|
||||
let dist = Distribution::Float(FloatDistribution {
|
||||
low: 0.0,
|
||||
@@ -897,7 +852,7 @@ mod tests {
|
||||
// Create history where low values (near 0.2) are "good"
|
||||
let history: Vec<CompletedTrial> = (0..20)
|
||||
.map(|i| {
|
||||
let x = f64::from(i) / 20.0;
|
||||
let x = i as f64 / 20.0;
|
||||
// Objective is (x - 0.2)^2, minimized at x=0.2
|
||||
let value = (x - 0.2).powi(2);
|
||||
create_trial(
|
||||
@@ -927,7 +882,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_categorical_sampling() {
|
||||
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();
|
||||
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42));
|
||||
|
||||
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 });
|
||||
|
||||
@@ -967,7 +922,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_int_sampling() {
|
||||
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42)).unwrap();
|
||||
let sampler = TpeSampler::with_config(0.25, 5, 24, None, Some(42));
|
||||
|
||||
let dist = Distribution::Int(IntDistribution {
|
||||
low: 0,
|
||||
@@ -1013,14 +968,14 @@ mod tests {
|
||||
.map(|i| {
|
||||
create_trial(
|
||||
i as u64,
|
||||
f64::from(i),
|
||||
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
i as f64,
|
||||
vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())],
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
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)).unwrap();
|
||||
let sampler1 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345));
|
||||
let sampler2 = TpeSampler::with_config(0.25, 5, 24, None, Some(12345));
|
||||
|
||||
for i in 0..10 {
|
||||
let v1 = sampler1.sample(&dist, i, &history);
|
||||
@@ -1032,8 +987,8 @@ mod tests {
|
||||
#[test]
|
||||
fn test_tpe_sampler_builder_default() {
|
||||
let builder = TpeSamplerBuilder::new();
|
||||
let sampler = builder.build().unwrap();
|
||||
assert!((sampler.gamma - 0.25).abs() < f64::EPSILON);
|
||||
let sampler = builder.build();
|
||||
assert_eq!(sampler.gamma, 0.25);
|
||||
assert_eq!(sampler.n_startup_trials, 10);
|
||||
assert_eq!(sampler.n_ei_candidates, 24);
|
||||
}
|
||||
@@ -1045,9 +1000,8 @@ mod tests {
|
||||
.n_startup_trials(20)
|
||||
.n_ei_candidates(32)
|
||||
.seed(42)
|
||||
.build()
|
||||
.unwrap();
|
||||
assert!((sampler.gamma - 0.15).abs() < f64::EPSILON);
|
||||
.build();
|
||||
assert_eq!(sampler.gamma, 0.15);
|
||||
assert_eq!(sampler.n_startup_trials, 20);
|
||||
assert_eq!(sampler.n_ei_candidates, 32);
|
||||
}
|
||||
@@ -1058,9 +1012,8 @@ mod tests {
|
||||
.gamma(0.10)
|
||||
.n_startup_trials(15)
|
||||
.n_ei_candidates(48)
|
||||
.build()
|
||||
.unwrap();
|
||||
assert!((sampler.gamma - 0.10).abs() < f64::EPSILON);
|
||||
.build();
|
||||
assert_eq!(sampler.gamma, 0.10);
|
||||
assert_eq!(sampler.n_startup_trials, 15);
|
||||
assert_eq!(sampler.n_ei_candidates, 48);
|
||||
}
|
||||
@@ -1068,16 +1021,16 @@ mod tests {
|
||||
#[test]
|
||||
fn test_tpe_sampler_builder_partial() {
|
||||
// Test setting only some options
|
||||
let sampler = TpeSamplerBuilder::new().gamma(0.20).build().unwrap();
|
||||
assert!((sampler.gamma - 0.20).abs() < f64::EPSILON);
|
||||
let sampler = TpeSamplerBuilder::new().gamma(0.20).build();
|
||||
assert_eq!(sampler.gamma, 0.20);
|
||||
assert_eq!(sampler.n_startup_trials, 10); // default
|
||||
assert_eq!(sampler.n_ei_candidates, 24); // default
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "gamma must be in (0.0, 1.0)")]
|
||||
fn test_tpe_sampler_builder_invalid_gamma() {
|
||||
let result = TpeSamplerBuilder::new().gamma(1.5).build();
|
||||
assert!(matches!(result, Err(TpeError::InvalidGamma(_))));
|
||||
TpeSamplerBuilder::new().gamma(1.5).build();
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1089,12 +1042,12 @@ mod tests {
|
||||
step: None,
|
||||
});
|
||||
|
||||
let history: Vec<CompletedTrial> = (0..20u32)
|
||||
let history: Vec<CompletedTrial> = (0..20)
|
||||
.map(|i| {
|
||||
create_trial(
|
||||
u64::from(i),
|
||||
f64::from(i),
|
||||
vec![("x", ParamValue::Float(f64::from(i) / 20.0), dist.clone())],
|
||||
i as u64,
|
||||
i as f64,
|
||||
vec![("x", ParamValue::Float(i as f64 / 20.0), dist.clone())],
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
@@ -1102,13 +1055,11 @@ mod tests {
|
||||
let sampler1 = TpeSampler::builder()
|
||||
.seed(99999)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
.build();
|
||||
let sampler2 = TpeSampler::builder()
|
||||
.seed(99999)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
.build();
|
||||
|
||||
for i in 0..10 {
|
||||
let v1 = sampler1.sample(&dist, i, &history);
|
||||
|
||||
+287
-72
@@ -1,18 +1,25 @@
|
||||
//! Study implementation for managing optimization trials.
|
||||
|
||||
#[cfg(feature = "async")]
|
||||
use core::future::Future;
|
||||
use core::ops::ControlFlow;
|
||||
use core::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::future::Future;
|
||||
use std::ops::ControlFlow;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use parking_lot::RwLock;
|
||||
#[cfg(feature = "serde")]
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::sampler::random::RandomSampler;
|
||||
use crate::sampler::{CompletedTrial, Sampler};
|
||||
use crate::sampler::{CompletedTrial, RandomSampler, Sampler};
|
||||
use crate::trial::Trial;
|
||||
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.
|
||||
///
|
||||
/// The study is parameterized by the objective value type `V`, which defaults to `f64`.
|
||||
@@ -22,6 +29,13 @@ use crate::types::Direction;
|
||||
/// When `V = f64`, the study passes trial history to the sampler for informed
|
||||
/// 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
|
||||
///
|
||||
/// ```
|
||||
@@ -65,7 +79,6 @@ where
|
||||
/// let study: Study<f64> = Study::new(Direction::Minimize);
|
||||
/// assert_eq!(study.direction(), Direction::Minimize);
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn new(direction: Direction) -> Self {
|
||||
Self::with_sampler(direction, RandomSampler::new())
|
||||
}
|
||||
@@ -80,8 +93,7 @@ where
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
/// let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
|
||||
@@ -103,6 +115,9 @@ where
|
||||
|
||||
/// 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
|
||||
///
|
||||
/// * `sampler` - The sampler to use for parameter sampling.
|
||||
@@ -110,9 +125,9 @@ where
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::tpe::TpeSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, Study, TpeSampler};
|
||||
///
|
||||
/// // After deserializing a study, restore the TPE sampler
|
||||
/// let mut study: Study<f64> = Study::new(Direction::Minimize);
|
||||
/// study.set_sampler(TpeSampler::new());
|
||||
/// ```
|
||||
@@ -317,15 +332,17 @@ where
|
||||
match self.direction {
|
||||
Direction::Minimize => {
|
||||
// Reverse ordering: smaller values are "greater" for max_by
|
||||
ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse)
|
||||
ordering
|
||||
.map(|o| o.reverse())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
}
|
||||
Direction::Maximize => {
|
||||
// Normal ordering: larger values are "greater" for max_by
|
||||
ordering.unwrap_or(core::cmp::Ordering::Equal)
|
||||
ordering.unwrap_or(std::cmp::Ordering::Equal)
|
||||
}
|
||||
}
|
||||
})
|
||||
.ok_or(crate::TpeError::NoCompletedTrials)?;
|
||||
.expect("trials is not empty");
|
||||
|
||||
Ok(best.clone())
|
||||
}
|
||||
@@ -392,8 +409,7 @@ where
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// // Minimize x^2
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
@@ -413,7 +429,7 @@ where
|
||||
/// ```
|
||||
pub fn optimize<F, E>(&self, n_trials: usize, mut objective: F) -> crate::Result<()>
|
||||
where
|
||||
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
|
||||
F: FnMut(&mut Trial) -> std::result::Result<V, E>,
|
||||
E: ToString,
|
||||
{
|
||||
for _ in 0..n_trials {
|
||||
@@ -460,8 +476,7 @@ where
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// # #[cfg(feature = "async")]
|
||||
/// # async fn example() -> optimizer::Result<()> {
|
||||
@@ -491,7 +506,7 @@ where
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
F: Fn(Trial) -> Fut,
|
||||
Fut: Future<Output = core::result::Result<(Trial, V), E>>,
|
||||
Fut: Future<Output = std::result::Result<(Trial, V), E>>,
|
||||
E: ToString,
|
||||
{
|
||||
for _ in 0..n_trials {
|
||||
@@ -537,13 +552,11 @@ where
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `TpeError::TaskError` if the semaphore is closed or a spawned task panics.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// # #[cfg(feature = "async")]
|
||||
/// # async fn example() -> optimizer::Result<()> {
|
||||
@@ -574,7 +587,7 @@ where
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
F: Fn(Trial) -> Fut + Send + Sync + 'static,
|
||||
Fut: Future<Output = core::result::Result<(Trial, V), E>> + Send,
|
||||
Fut: Future<Output = std::result::Result<(Trial, V), E>> + Send,
|
||||
E: ToString + Send + 'static,
|
||||
V: Send + 'static,
|
||||
{
|
||||
@@ -586,11 +599,7 @@ where
|
||||
let mut handles = Vec::with_capacity(n_trials);
|
||||
|
||||
for _ in 0..n_trials {
|
||||
let permit = semaphore
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?;
|
||||
let permit = semaphore.clone().acquire_owned().await.unwrap();
|
||||
let trial = self.create_trial();
|
||||
let objective = Arc::clone(&objective);
|
||||
|
||||
@@ -605,10 +614,7 @@ where
|
||||
|
||||
// Wait for all tasks and record results
|
||||
for handle in handles {
|
||||
match handle
|
||||
.await
|
||||
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?
|
||||
{
|
||||
match handle.await.unwrap() {
|
||||
Ok((trial, value)) => {
|
||||
self.complete_trial(trial, value);
|
||||
}
|
||||
@@ -645,15 +651,13 @@ where
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if no trials completed successfully
|
||||
/// 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
|
||||
///
|
||||
/// ```
|
||||
/// use std::ops::ControlFlow;
|
||||
///
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// // Stop early when we find a good enough value
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
@@ -688,7 +692,7 @@ where
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
V: Clone,
|
||||
F: FnMut(&mut Trial) -> core::result::Result<V, E>,
|
||||
F: FnMut(&mut Trial) -> std::result::Result<V, E>,
|
||||
C: FnMut(&Study<V>, &CompletedTrial<V>) -> ControlFlow<()>,
|
||||
E: ToString,
|
||||
{
|
||||
@@ -701,11 +705,7 @@ where
|
||||
|
||||
// Get the just-completed trial for the callback
|
||||
let trials = self.completed_trials.read();
|
||||
let Some(completed) = trials.last() else {
|
||||
return Err(crate::TpeError::Internal(
|
||||
"completed trial not found after adding",
|
||||
));
|
||||
};
|
||||
let completed = trials.last().expect("just added a trial");
|
||||
|
||||
// Call the callback and check if we should stop
|
||||
// Note: We need to drop the read lock before calling callback
|
||||
@@ -746,8 +746,7 @@ impl Study<f64> {
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// // With a seeded sampler for reproducibility
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
@@ -788,8 +787,7 @@ impl Study<f64> {
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// // Minimize x^2 with sampler integration
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
@@ -811,7 +809,7 @@ impl Study<f64> {
|
||||
mut objective: F,
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
F: FnMut(&mut Trial) -> core::result::Result<f64, E>,
|
||||
F: FnMut(&mut Trial) -> std::result::Result<f64, E>,
|
||||
E: ToString,
|
||||
{
|
||||
for _ in 0..n_trials {
|
||||
@@ -852,15 +850,13 @@ impl Study<f64> {
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if no trials completed successfully.
|
||||
/// Returns `TpeError::Internal` if a completed trial is not found after adding (internal invariant violation).
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use std::ops::ControlFlow;
|
||||
///
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// // Optimize with sampler integration and early stopping
|
||||
/// let sampler = RandomSampler::with_seed(42);
|
||||
@@ -893,7 +889,7 @@ impl Study<f64> {
|
||||
mut callback: C,
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
F: FnMut(&mut Trial) -> core::result::Result<f64, E>,
|
||||
F: FnMut(&mut Trial) -> std::result::Result<f64, E>,
|
||||
C: FnMut(&Study<f64>, &CompletedTrial<f64>) -> ControlFlow<()>,
|
||||
E: ToString,
|
||||
{
|
||||
@@ -906,11 +902,7 @@ impl Study<f64> {
|
||||
|
||||
// Get the just-completed trial for the callback
|
||||
let trials = self.completed_trials.read();
|
||||
let Some(completed) = trials.last() else {
|
||||
return Err(crate::TpeError::Internal(
|
||||
"completed trial not found after adding",
|
||||
));
|
||||
};
|
||||
let completed = trials.last().expect("just added a trial");
|
||||
|
||||
// Call the callback and check if we should stop
|
||||
// Note: We need to drop the read lock before calling callback
|
||||
@@ -958,8 +950,7 @@ impl Study<f64> {
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// # #[cfg(feature = "async")]
|
||||
/// # async fn example() -> optimizer::Result<()> {
|
||||
@@ -989,7 +980,7 @@ impl Study<f64> {
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
F: Fn(Trial) -> Fut,
|
||||
Fut: Future<Output = core::result::Result<(Trial, f64), E>>,
|
||||
Fut: Future<Output = std::result::Result<(Trial, f64), E>>,
|
||||
E: ToString,
|
||||
{
|
||||
for _ in 0..n_trials {
|
||||
@@ -1035,13 +1026,11 @@ impl Study<f64> {
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `TpeError::NoCompletedTrials` if all trials failed (no successful trials).
|
||||
/// Returns `TpeError::TaskError` if the semaphore is closed or a spawned task panics.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::sampler::random::RandomSampler;
|
||||
/// use optimizer::{Direction, Study};
|
||||
/// use optimizer::{Direction, RandomSampler, Study};
|
||||
///
|
||||
/// # #[cfg(feature = "async")]
|
||||
/// # async fn example() -> optimizer::Result<()> {
|
||||
@@ -1054,7 +1043,7 @@ impl Study<f64> {
|
||||
/// let x = trial.suggest_float("x", -10.0, 10.0)?;
|
||||
/// // Async objective function (e.g., network request)
|
||||
/// let value = x * x;
|
||||
/// Ok::<_, optimizer::TpeError>((trial, value))
|
||||
/// Ok::<_, optimize::TpeError>((trial, value))
|
||||
/// })
|
||||
/// .await?;
|
||||
///
|
||||
@@ -1072,7 +1061,7 @@ impl Study<f64> {
|
||||
) -> crate::Result<()>
|
||||
where
|
||||
F: Fn(Trial) -> Fut + Send + Sync + 'static,
|
||||
Fut: Future<Output = core::result::Result<(Trial, f64), E>> + Send,
|
||||
Fut: Future<Output = std::result::Result<(Trial, f64), E>> + Send,
|
||||
E: ToString + Send + 'static,
|
||||
{
|
||||
use tokio::sync::Semaphore;
|
||||
@@ -1083,11 +1072,7 @@ impl Study<f64> {
|
||||
let mut handles = Vec::with_capacity(n_trials);
|
||||
|
||||
for _ in 0..n_trials {
|
||||
let permit = semaphore
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?;
|
||||
let permit = semaphore.clone().acquire_owned().await.unwrap();
|
||||
let trial = self.create_trial_with_sampler();
|
||||
let objective = Arc::clone(&objective);
|
||||
|
||||
@@ -1102,10 +1087,7 @@ impl Study<f64> {
|
||||
|
||||
// Wait for all tasks and record results
|
||||
for handle in handles {
|
||||
match handle
|
||||
.await
|
||||
.map_err(|e| crate::TpeError::TaskError(e.to_string()))?
|
||||
{
|
||||
match handle.await.unwrap() {
|
||||
Ok((trial, value)) => {
|
||||
self.complete_trial(trial, value);
|
||||
}
|
||||
@@ -1123,3 +1105,236 @@ impl Study<f64> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
+38
-48
@@ -4,6 +4,8 @@ use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use parking_lot::RwLock;
|
||||
#[cfg(feature = "serde")]
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::distribution::{
|
||||
CategoricalDistribution, Distribution, FloatDistribution, IntDistribution,
|
||||
@@ -22,6 +24,7 @@ use crate::types::TrialState;
|
||||
/// `Study::create_trial()`, the trial receives the study's sampler and access
|
||||
/// to the history of completed trials for informed sampling.
|
||||
#[derive(Clone)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub struct Trial {
|
||||
/// Unique identifier for this trial.
|
||||
id: u64,
|
||||
@@ -32,13 +35,15 @@ pub struct Trial {
|
||||
/// Parameter distributions, keyed by parameter name.
|
||||
distributions: HashMap<String, Distribution>,
|
||||
/// The sampler to use for generating parameter values.
|
||||
#[cfg_attr(feature = "serde", serde(skip))]
|
||||
sampler: Option<Arc<dyn Sampler>>,
|
||||
/// Access to the history of completed trials (shared with Study).
|
||||
#[cfg_attr(feature = "serde", serde(skip))]
|
||||
history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>,
|
||||
}
|
||||
|
||||
impl core::fmt::Debug for Trial {
|
||||
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
||||
impl std::fmt::Debug for Trial {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("Trial")
|
||||
.field("id", &self.id)
|
||||
.field("state", &self.state)
|
||||
@@ -71,7 +76,6 @@ impl Trial {
|
||||
/// let trial = Trial::new(0);
|
||||
/// assert_eq!(trial.id(), 0);
|
||||
/// ```
|
||||
#[must_use]
|
||||
pub fn new(id: u64) -> Self {
|
||||
Self {
|
||||
id,
|
||||
@@ -111,7 +115,7 @@ impl Trial {
|
||||
/// Samples a value from the given distribution using the sampler.
|
||||
///
|
||||
/// If the trial has a sampler, it delegates to the sampler's sample method
|
||||
/// with the history of completed trials. Otherwise, it uses the `RandomSampler`
|
||||
/// with the history of completed trials. Otherwise, it uses the RandomSampler
|
||||
/// as a fallback.
|
||||
fn sample_value(&self, distribution: &Distribution) -> ParamValue {
|
||||
if let (Some(sampler), Some(history)) = (&self.sampler, &self.history) {
|
||||
@@ -119,32 +123,28 @@ impl Trial {
|
||||
sampler.sample(distribution, self.id, &history_guard)
|
||||
} else {
|
||||
// Fallback to RandomSampler when no sampler is configured
|
||||
use crate::sampler::random::RandomSampler;
|
||||
use crate::sampler::RandomSampler;
|
||||
let fallback = RandomSampler::new();
|
||||
fallback.sample(distribution, self.id, &[])
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the unique ID of this trial.
|
||||
#[must_use]
|
||||
pub fn id(&self) -> u64 {
|
||||
self.id
|
||||
}
|
||||
|
||||
/// Returns the current state of this trial.
|
||||
#[must_use]
|
||||
pub fn state(&self) -> TrialState {
|
||||
self.state
|
||||
}
|
||||
|
||||
/// Returns a reference to the sampled parameters.
|
||||
#[must_use]
|
||||
pub fn params(&self) -> &HashMap<String, ParamValue> {
|
||||
&self.params
|
||||
}
|
||||
|
||||
/// Returns a reference to the parameter distributions.
|
||||
#[must_use]
|
||||
pub fn distributions(&self) -> &HashMap<String, Distribution> {
|
||||
&self.distributions
|
||||
}
|
||||
@@ -205,8 +205,8 @@ impl Trial {
|
||||
if let Some(existing_dist) = self.distributions.get(&name) {
|
||||
// Verify the distribution matches
|
||||
if let Distribution::Float(existing) = existing_dist
|
||||
&& (existing.low - low).abs() < f64::EPSILON
|
||||
&& (existing.high - high).abs() < f64::EPSILON
|
||||
&& existing.low == low
|
||||
&& existing.high == high
|
||||
&& !existing.log_scale
|
||||
&& existing.step.is_none()
|
||||
{
|
||||
@@ -225,10 +225,9 @@ impl Trial {
|
||||
|
||||
// Sample using the sampler
|
||||
let dist = Distribution::Float(distribution);
|
||||
let ParamValue::Float(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Float distribution should return Float value",
|
||||
));
|
||||
let value = match self.sample_value(&dist) {
|
||||
ParamValue::Float(v) => v,
|
||||
_ => unreachable!("Float distribution should return Float value"),
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -243,7 +242,7 @@ impl Trial {
|
||||
/// The value is sampled uniformly in log space, which is useful for parameters
|
||||
/// that span multiple orders of magnitude (e.g., learning rates).
|
||||
///
|
||||
/// If the parameter has already been sampled with the same bounds and `log_scale=true`,
|
||||
/// If the parameter has already been sampled with the same bounds and log_scale=true,
|
||||
/// the cached value is returned. If the parameter was sampled with different configuration,
|
||||
/// a `ParameterConflict` error is returned.
|
||||
///
|
||||
@@ -302,8 +301,8 @@ impl Trial {
|
||||
if let Some(existing_dist) = self.distributions.get(&name) {
|
||||
// Verify the distribution matches
|
||||
if let Distribution::Float(existing) = existing_dist
|
||||
&& (existing.low - low).abs() < f64::EPSILON
|
||||
&& (existing.high - high).abs() < f64::EPSILON
|
||||
&& existing.low == low
|
||||
&& existing.high == high
|
||||
&& existing.log_scale
|
||||
&& existing.step.is_none()
|
||||
{
|
||||
@@ -322,10 +321,9 @@ impl Trial {
|
||||
|
||||
// Sample using the sampler (sampler handles log-scale transformation)
|
||||
let dist = Distribution::Float(distribution);
|
||||
let ParamValue::Float(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Float distribution should return Float value",
|
||||
));
|
||||
let value = match self.sample_value(&dist) {
|
||||
ParamValue::Float(v) => v,
|
||||
_ => unreachable!("Float distribution should return Float value"),
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -399,8 +397,8 @@ impl Trial {
|
||||
if let Some(existing_dist) = self.distributions.get(&name) {
|
||||
// Verify the distribution matches
|
||||
if let Distribution::Float(existing) = existing_dist
|
||||
&& (existing.low - low).abs() < f64::EPSILON
|
||||
&& (existing.high - high).abs() < f64::EPSILON
|
||||
&& existing.low == low
|
||||
&& existing.high == high
|
||||
&& !existing.log_scale
|
||||
&& existing.step == Some(step)
|
||||
{
|
||||
@@ -419,10 +417,9 @@ impl Trial {
|
||||
|
||||
// Sample using the sampler (sampler handles step-grid)
|
||||
let dist = Distribution::Float(distribution);
|
||||
let ParamValue::Float(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Float distribution should return Float value",
|
||||
));
|
||||
let value = match self.sample_value(&dist) {
|
||||
ParamValue::Float(v) => v,
|
||||
_ => unreachable!("Float distribution should return Float value"),
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -463,7 +460,6 @@ impl Trial {
|
||||
/// let n2 = trial.suggest_int("n_layers", 1, 10).unwrap();
|
||||
/// assert_eq!(n, n2);
|
||||
/// ```
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub fn suggest_int(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
|
||||
if low > high {
|
||||
return Err(TpeError::InvalidBounds {
|
||||
@@ -504,10 +500,9 @@ impl Trial {
|
||||
|
||||
// Sample using the sampler
|
||||
let dist = Distribution::Int(distribution);
|
||||
let ParamValue::Int(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Int distribution should return Int value",
|
||||
));
|
||||
let value = match self.sample_value(&dist) {
|
||||
ParamValue::Int(v) => v,
|
||||
_ => unreachable!("Int distribution should return Int value"),
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -522,7 +517,7 @@ impl Trial {
|
||||
/// The value is sampled uniformly in log space, which is useful for parameters
|
||||
/// that span multiple orders of magnitude (e.g., batch sizes).
|
||||
///
|
||||
/// If the parameter has already been sampled with the same bounds and `log_scale=true`,
|
||||
/// If the parameter has already been sampled with the same bounds and log_scale=true,
|
||||
/// the cached value is returned. If the parameter was sampled with different configuration,
|
||||
/// a `ParameterConflict` error is returned.
|
||||
///
|
||||
@@ -551,7 +546,6 @@ impl Trial {
|
||||
/// let batch_size2 = trial.suggest_int_log("batch_size", 1, 1024).unwrap();
|
||||
/// assert_eq!(batch_size, batch_size2);
|
||||
/// ```
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub fn suggest_int_log(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
|
||||
if low < 1 {
|
||||
return Err(TpeError::InvalidLogBounds);
|
||||
@@ -596,10 +590,9 @@ impl Trial {
|
||||
|
||||
// Sample using the sampler (sampler handles log-scale transformation)
|
||||
let dist = Distribution::Int(distribution);
|
||||
let ParamValue::Int(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Int distribution should return Int value",
|
||||
));
|
||||
let value = match self.sample_value(&dist) {
|
||||
ParamValue::Int(v) => v,
|
||||
_ => unreachable!("Int distribution should return Int value"),
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -650,7 +643,6 @@ impl Trial {
|
||||
/// .unwrap();
|
||||
/// assert_eq!(n, n2);
|
||||
/// ```
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
pub fn suggest_int_step(
|
||||
&mut self,
|
||||
name: impl Into<String>,
|
||||
@@ -701,10 +693,9 @@ impl Trial {
|
||||
|
||||
// Sample using the sampler (sampler handles step-grid)
|
||||
let dist = Distribution::Int(distribution);
|
||||
let ParamValue::Int(value) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Int distribution should return Int value",
|
||||
));
|
||||
let value = match self.sample_value(&dist) {
|
||||
ParamValue::Int(v) => v,
|
||||
_ => unreachable!("Int distribution should return Int value"),
|
||||
};
|
||||
|
||||
// Store distribution and value
|
||||
@@ -788,10 +779,9 @@ impl Trial {
|
||||
|
||||
// Sample using the sampler
|
||||
let dist = Distribution::Categorical(distribution);
|
||||
let ParamValue::Categorical(index) = self.sample_value(&dist) else {
|
||||
return Err(TpeError::Internal(
|
||||
"Categorical distribution should return Categorical value",
|
||||
));
|
||||
let index = match self.sample_value(&dist) {
|
||||
ParamValue::Categorical(idx) => idx,
|
||||
_ => unreachable!("Categorical distribution should return Categorical value"),
|
||||
};
|
||||
|
||||
// Store distribution and value (store the index)
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
//! Core types for the optimizer library.
|
||||
|
||||
#[cfg(feature = "serde")]
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// The direction of optimization.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub enum Direction {
|
||||
/// Minimize the objective value.
|
||||
Minimize,
|
||||
@@ -11,6 +15,7 @@ pub enum Direction {
|
||||
|
||||
/// The state of a trial in its lifecycle.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
|
||||
pub enum TrialState {
|
||||
/// The trial is currently running.
|
||||
Running,
|
||||
|
||||
+3
-13
@@ -4,9 +4,7 @@
|
||||
|
||||
#![cfg(feature = "async")]
|
||||
|
||||
use optimizer::sampler::random::RandomSampler;
|
||||
use optimizer::sampler::tpe::TpeSampler;
|
||||
use optimizer::{Direction, Study, TpeError};
|
||||
use optimizer::{Direction, RandomSampler, Study, TpeError, TpeSampler};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optimize_async_basic() {
|
||||
@@ -28,11 +26,7 @@ async fn test_optimize_async_basic() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optimize_async_with_sampler() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -67,11 +61,7 @@ async fn test_optimize_parallel() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optimize_parallel_with_sampler() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
|
||||
+157
-56
@@ -1,14 +1,6 @@
|
||||
//! Integration tests for the optimizer library.
|
||||
|
||||
#![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};
|
||||
use optimizer::{Direction, RandomSampler, Study, TpeError, TpeSampler, Trial};
|
||||
|
||||
// =============================================================================
|
||||
// Test: optimize simple quadratic function with TPE, finds near-optimal
|
||||
@@ -22,8 +14,7 @@ fn test_tpe_optimizes_quadratic_function() {
|
||||
.seed(42)
|
||||
.n_startup_trials(5) // Quick startup for test
|
||||
.n_ei_candidates(24)
|
||||
.build()
|
||||
.unwrap();
|
||||
.build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -49,11 +40,7 @@ fn test_tpe_optimizes_quadratic_function() {
|
||||
fn test_tpe_optimizes_multivariate_function() {
|
||||
// Minimize f(x, y) = x^2 + y^2 where x, y ∈ [-5, 5]
|
||||
// Optimal: (0, 0), f(0, 0) = 0
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(123)
|
||||
.n_startup_trials(10)
|
||||
.build()
|
||||
.unwrap();
|
||||
let sampler = TpeSampler::builder().seed(123).n_startup_trials(10).build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -79,11 +66,7 @@ fn test_tpe_optimizes_multivariate_function() {
|
||||
fn test_tpe_maximization() {
|
||||
// Maximize f(x) = -(x - 2)^2 + 10 where x ∈ [-10, 10]
|
||||
// Optimal: x = 2, f(2) = 10
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(456)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let sampler = TpeSampler::builder().seed(456).n_startup_trials(5).build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
|
||||
|
||||
@@ -567,11 +550,7 @@ fn test_invalid_step_errors() {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_with_categorical_parameter() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
|
||||
|
||||
@@ -603,11 +582,7 @@ fn test_tpe_with_categorical_parameter() {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_with_integer_parameters() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(789)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let sampler = TpeSampler::builder().seed(789).n_startup_trials(5).build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -781,11 +756,7 @@ fn test_study_set_sampler() {
|
||||
let mut study: Study<f64> = Study::new(Direction::Minimize);
|
||||
|
||||
// Initially uses RandomSampler, now switch to TPE
|
||||
let tpe = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let tpe = TpeSampler::builder().seed(42).n_startup_trials(5).build();
|
||||
study.set_sampler(tpe);
|
||||
|
||||
// Should work with the new sampler
|
||||
@@ -892,10 +863,10 @@ fn test_trial_debug_format() {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_sampler_builder_default_trait() {
|
||||
use optimizer::sampler::tpe::TpeSamplerBuilder;
|
||||
use optimizer::TpeSamplerBuilder;
|
||||
|
||||
let builder = TpeSamplerBuilder::default();
|
||||
let sampler = builder.build().unwrap();
|
||||
let sampler = builder.build();
|
||||
|
||||
// Should have default values
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
@@ -930,8 +901,7 @@ fn test_tpe_with_fixed_kde_bandwidth() {
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.kde_bandwidth(0.5)
|
||||
.build()
|
||||
.unwrap();
|
||||
.build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -947,9 +917,9 @@ fn test_tpe_with_fixed_kde_bandwidth() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic(expected = "kde_bandwidth must be positive")]
|
||||
fn test_tpe_sampler_invalid_kde_bandwidth() {
|
||||
let result = TpeSampler::with_config(0.25, 10, 24, Some(-1.0), None);
|
||||
assert!(matches!(result, Err(TpeError::InvalidBandwidth(_))));
|
||||
TpeSampler::with_config(0.25, 10, 24, Some(-1.0), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -958,8 +928,7 @@ fn test_tpe_split_trials_with_two_trials() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(2) // TPE kicks in after 2 trials
|
||||
.build()
|
||||
.unwrap();
|
||||
.build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -975,11 +944,7 @@ fn test_tpe_split_trials_with_two_trials() {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_with_log_scale_int() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -997,11 +962,7 @@ fn test_tpe_with_log_scale_int() {
|
||||
|
||||
#[test]
|
||||
fn test_tpe_with_step_distributions() {
|
||||
let sampler = TpeSampler::builder()
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.build()
|
||||
.unwrap();
|
||||
let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -1079,8 +1040,7 @@ fn test_tpe_empty_good_or_bad_values_fallback() {
|
||||
.seed(42)
|
||||
.n_startup_trials(5)
|
||||
.gamma(0.1) // Very small gamma means few "good" trials
|
||||
.build()
|
||||
.unwrap();
|
||||
.build();
|
||||
|
||||
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
|
||||
|
||||
@@ -1182,3 +1142,144 @@ fn test_best_trial_with_nan_values() {
|
||||
let best = study.best_trial();
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,2 +0,0 @@
|
||||
[default.extend-words]
|
||||
Tpe = "Tpe"
|
||||
Reference in New Issue
Block a user