Initial Implementation

Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "optimize"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+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"

rustfmt.toml

@@ -0,0 +1,10 @@
+# Nightly Features
+imports_granularity = "Module"
+group_imports = "StdExternalCrate"
+#format_strings = true
+format_code_in_doc_comments = true
+
+# Stable Features
+merge_derives = true
+use_field_init_shorthand = true
+use_try_shorthand = true

src/distribution.rs

@@ -0,0 +1,52 @@
+//! Parameter distribution types.
+
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
+
+/// Distribution for floating-point parameters.
+#[derive(Clone, Debug, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct FloatDistribution {
+    /// Lower bound (inclusive).
+    pub low: f64,
+    /// Upper bound (inclusive).
+    pub high: f64,
+    /// Whether to sample in log space.
+    pub log_scale: bool,
+    /// Optional step size for discretization.
+    pub step: Option<f64>,
+}
+
+/// Distribution for integer parameters.
+#[derive(Clone, Debug, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct IntDistribution {
+    /// Lower bound (inclusive).
+    pub low: i64,
+    /// Upper bound (inclusive).
+    pub high: i64,
+    /// Whether to sample in log space.
+    pub log_scale: bool,
+    /// Optional step size for discretization.
+    pub step: Option<i64>,
+}
+
+/// Distribution for categorical parameters.
+#[derive(Clone, Debug, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct CategoricalDistribution {
+    /// Number of choices available.
+    pub n_choices: usize,
+}
+
+/// Enum wrapping all parameter distribution types.
+#[derive(Clone, Debug, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum Distribution {
+    /// A floating-point distribution.
+    Float(FloatDistribution),
+    /// An integer distribution.
+    Int(IntDistribution),
+    /// A categorical distribution.
+    Categorical(CategoricalDistribution),
+}

src/error.rs

@@ -0,0 +1,44 @@
+//! Error types for the optimize library.
+
+use thiserror::Error;
+
+/// The error type for TPE operations.
+#[derive(Debug, Error)]
+pub enum TpeError {
+    /// Returned when the lower bound is greater than the upper bound.
+    #[error("invalid bounds: low ({low}) must be less than or equal to high ({high})")]
+    InvalidBounds {
+        /// The lower bound value.
+        low: f64,
+        /// The upper bound value.
+        high: f64,
+    },
+
+    /// Returned when log scale is used with non-positive bounds.
+    #[error("invalid log bounds: low must be positive for log scale")]
+    InvalidLogBounds,
+
+    /// Returned when step size is not positive.
+    #[error("invalid step: step must be positive")]
+    InvalidStep,
+
+    /// Returned when categorical choices are empty.
+    #[error("categorical choices cannot be empty")]
+    EmptyChoices,
+
+    /// Returned when a parameter is suggested with a different configuration.
+    #[error("parameter conflict for '{name}': {reason}")]
+    ParameterConflict {
+        /// The name of the conflicting parameter.
+        name: String,
+        /// The reason for the conflict.
+        reason: String,
+    },
+
+    /// Returned when requesting the best trial but no trials have completed.
+    #[error("no completed trials available")]
+    NoCompletedTrials,
+}
+
+/// A specialized Result type for TPE operations.
+pub type Result<T> = std::result::Result<T, TpeError>;

src/kde.rs

@@ -0,0 +1,267 @@
+//! Kernel Density Estimation for continuous parameters.
+//!
+//! This module provides a Gaussian kernel density estimator used by the TPE
+//! sampler to model probability distributions over good and bad trial regions.
+
+use rand::Rng;
+
+/// A Gaussian kernel density estimator for continuous distributions.
+///
+/// KDE estimates a probability density function from a set of samples by
+/// placing Gaussian kernels centered at each sample point. This is used
+/// in TPE to model the distributions l(x) (good trials) and g(x) (bad trials).
+///
+/// # Examples
+///
+/// ```ignore
+/// use crate::kde::KernelDensityEstimator;
+///
+/// let samples = vec![1.0, 2.0, 3.0, 4.0, 5.0];
+/// let kde = KernelDensityEstimator::new(samples);
+///
+/// // Get probability density at a point
+/// let density = kde.pdf(2.5);
+/// assert!(density > 0.0);
+///
+/// // Sample from the estimated distribution
+/// let mut rng = rand::rng();
+/// let sample = kde.sample(&mut rng);
+/// ```
+#[derive(Clone, Debug)]
+pub struct KernelDensityEstimator {
+    /// The sample points used to construct the KDE.
+    samples: Vec<f64>,
+    /// The bandwidth (standard deviation) of the Gaussian kernels.
+    bandwidth: f64,
+}
+
+impl KernelDensityEstimator {
+    /// Creates a new KDE with automatic bandwidth selection using Scott's rule.
+    ///
+    /// Scott's rule sets bandwidth = n^(-1/5) * std_dev, which works well
+    /// for unimodal distributions close to normal.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `samples` is empty.
+    pub fn new(samples: Vec<f64>) -> Self {
+        assert!(!samples.is_empty(), "KDE requires at least one sample");
+
+        let bandwidth = Self::scotts_rule(&samples);
+        Self { samples, bandwidth }
+    }
+
+    /// Creates a new KDE with a specified bandwidth.
+    ///
+    /// Use this when you want explicit control over the smoothing parameter.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `samples` is empty or `bandwidth` is not positive.
+    pub fn with_bandwidth(samples: Vec<f64>, bandwidth: f64) -> Self {
+        assert!(!samples.is_empty(), "KDE requires at least one sample");
+        assert!(bandwidth > 0.0, "Bandwidth must be positive");
+
+        Self { samples, bandwidth }
+    }
+
+    /// Computes bandwidth using Scott's rule.
+    ///
+    /// Scott's rule: h = n^(-1/5) * sigma
+    /// where sigma is the sample standard deviation.
+    fn scotts_rule(samples: &[f64]) -> f64 {
+        let n = samples.len() as f64;
+        let std_dev = Self::sample_std_dev(samples);
+
+        // For degenerate case where all samples are identical,
+        // use a small positive bandwidth
+        if std_dev < f64::EPSILON {
+            return 1.0;
+        }
+
+        n.powf(-0.2) * std_dev
+    }
+
+    /// Computes the sample standard deviation.
+    fn sample_std_dev(samples: &[f64]) -> f64 {
+        let n = samples.len() as f64;
+        let mean = samples.iter().sum::<f64>() / n;
+        let variance = samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n;
+        variance.sqrt()
+    }
+
+    /// Returns the probability density at point `x`.
+    ///
+    /// The density is computed as the average of Gaussian kernels centered
+    /// at each sample point:
+    ///
+    /// f(x) = (1/n) * sum_i K((x - x_i) / h)
+    ///
+    /// where K is the standard Gaussian kernel and h is the bandwidth.
+    pub fn pdf(&self, x: f64) -> f64 {
+        let n = self.samples.len() as f64;
+        let inv_bandwidth = 1.0 / self.bandwidth;
+        let normalization = inv_bandwidth / (2.0 * std::f64::consts::PI).sqrt();
+
+        let density: f64 = self
+            .samples
+            .iter()
+            .map(|&xi| {
+                let z = (x - xi) * inv_bandwidth;
+                normalization * (-0.5 * z * z).exp()
+            })
+            .sum();
+
+        density / n
+    }
+
+    /// Samples a value from the estimated density distribution.
+    ///
+    /// Sampling works by:
+    /// 1. Uniformly selecting one of the kernel centers (samples)
+    /// 2. Adding Gaussian noise with the bandwidth as standard deviation
+    pub fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
+        // Select a random sample to center the kernel on
+        let idx = rng.random_range(0..self.samples.len());
+        let center = self.samples[idx];
+
+        // Add Gaussian noise with bandwidth as standard deviation
+        // Using Box-Muller transform for Gaussian sampling
+        let u1: f64 = rng.random();
+        let u2: f64 = rng.random();
+
+        let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
+        center + z * self.bandwidth
+    }
+
+    /// Returns the bandwidth of this KDE.
+    pub fn bandwidth(&self) -> f64 {
+        self.bandwidth
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_kde_pdf_basic() {
+        let samples = vec![0.0, 1.0, 2.0];
+        let kde = KernelDensityEstimator::new(samples);
+
+        // Density should be positive everywhere
+        assert!(kde.pdf(0.0) > 0.0);
+        assert!(kde.pdf(1.0) > 0.0);
+        assert!(kde.pdf(2.0) > 0.0);
+
+        // Density should be higher near sample points
+        let mid_density = kde.pdf(1.0);
+        let far_density = kde.pdf(10.0);
+        assert!(mid_density > far_density);
+    }
+
+    #[test]
+    fn test_kde_pdf_integrates_to_one() {
+        let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
+        let kde = KernelDensityEstimator::new(samples);
+
+        // Numerical integration over a wide range
+        let n_points = 10000;
+        let low = -10.0;
+        let high = 15.0;
+        let dx = (high - low) / n_points as f64;
+
+        let integral: f64 = (0..n_points)
+            .map(|i| {
+                let x = low + (i as f64 + 0.5) * dx;
+                kde.pdf(x) * dx
+            })
+            .sum();
+
+        // Should be approximately 1.0 (within numerical error)
+        assert!(
+            (integral - 1.0).abs() < 0.01,
+            "Integral = {integral}, expected ~1.0"
+        );
+    }
+
+    #[test]
+    fn test_kde_with_bandwidth() {
+        let samples = vec![0.0, 1.0, 2.0];
+        let kde = KernelDensityEstimator::with_bandwidth(samples, 0.5);
+
+        assert_eq!(kde.bandwidth(), 0.5);
+        assert!(kde.pdf(1.0) > 0.0);
+    }
+
+    #[test]
+    fn test_kde_sample_in_reasonable_range() {
+        let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
+        let kde = KernelDensityEstimator::new(samples);
+        let mut rng = rand::rng();
+
+        // Samples should generally be in a reasonable range around the data
+        for _ in 0..100 {
+            let s = kde.sample(&mut rng);
+            // With high probability, samples should be within a few bandwidths
+            // of the data range. Use a generous range to avoid flaky tests.
+            assert!(s > -10.0 && s < 15.0, "Sample {s} outside expected range");
+        }
+    }
+
+    #[test]
+    fn test_kde_single_sample() {
+        let samples = vec![5.0];
+        let kde = KernelDensityEstimator::new(samples);
+
+        // Should have positive density near the sample
+        assert!(kde.pdf(5.0) > 0.0);
+        assert!(kde.pdf(4.5) > 0.0);
+    }
+
+    #[test]
+    fn test_kde_identical_samples() {
+        let samples = vec![3.0, 3.0, 3.0, 3.0];
+        let kde = KernelDensityEstimator::new(samples);
+
+        // Should handle degenerate case with identical samples
+        assert!(kde.bandwidth() > 0.0);
+        assert!(kde.pdf(3.0) > 0.0);
+    }
+
+    #[test]
+    fn test_scotts_rule_bandwidth() {
+        let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
+        let kde = KernelDensityEstimator::new(samples);
+
+        // n = 10, n^(-1/5) ≈ 0.631
+        // std_dev ≈ 2.87
+        // bandwidth ≈ 0.631 * 2.87 ≈ 1.81
+        let bandwidth = kde.bandwidth();
+        assert!(
+            bandwidth > 1.0 && bandwidth < 3.0,
+            "Bandwidth {bandwidth} outside expected range"
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "KDE requires at least one sample")]
+    fn test_kde_empty_samples() {
+        let samples: Vec<f64> = vec![];
+        KernelDensityEstimator::new(samples);
+    }
+
+    #[test]
+    #[should_panic(expected = "Bandwidth must be positive")]
+    fn test_kde_zero_bandwidth() {
+        let samples = vec![1.0, 2.0, 3.0];
+        KernelDensityEstimator::with_bandwidth(samples, 0.0);
+    }
+
+    #[test]
+    #[should_panic(expected = "Bandwidth must be positive")]
+    fn test_kde_negative_bandwidth() {
+        let samples = vec![1.0, 2.0, 3.0];
+        KernelDensityEstimator::with_bandwidth(samples, -1.0);
+    }
+}

src/lib.rs

@@ -0,0 +1,151 @@
+//! A Tree-Parzen Estimator (TPE) library for black-box optimization.
+//!
+//! This library provides an Optuna-like API for hyperparameter optimization
+//! using the Tree-Parzen Estimator algorithm. It supports:
+//!
+//! - Float, integer, and categorical parameter types
+//! - Log-scale and stepped parameter sampling
+//! - Synchronous and async optimization
+//! - Parallel trial evaluation with bounded concurrency
+//! - Serialization for saving/loading study state
+//!
+//! # Quick Start
+//!
+//! ```
+//! use optimize::{Direction, Study, TpeSampler};
+//!
+//! // Create a study with TPE sampler
+//! let sampler = TpeSampler::builder().seed(42).build();
+//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
+//!
+//! // Optimize x^2 for 20 trials
+//! study
+//!     .optimize_with_sampler(20, |trial| {
+//!         let x = trial.suggest_float("x", -10.0, 10.0)?;
+//!         Ok::<_, optimize::TpeError>(x * x)
+//!     })
+//!     .unwrap();
+//!
+//! // Get the best result
+//! let best = study.best_trial().unwrap();
+//! println!("Best value: {} at x={:?}", best.value, best.params);
+//! ```
+//!
+//! # Creating a Study
+//!
+//! A [`Study`] manages optimization trials. Create one with an optimization direction:
+//!
+//! ```
+//! use optimize::{Direction, RandomSampler, Study, TpeSampler};
+//!
+//! // Minimize with default random sampler
+//! let study: Study<f64> = Study::new(Direction::Minimize);
+//!
+//! // Maximize with TPE sampler
+//! let study: Study<f64> = Study::with_sampler(Direction::Maximize, TpeSampler::new());
+//!
+//! // With seeded sampler for reproducibility
+//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
+//! ```
+//!
+//! # Suggesting Parameters
+//!
+//! Within the objective function, use [`Trial`] to suggest parameter values:
+//!
+//! ```
+//! use optimize::{Direction, Study};
+//!
+//! let study: Study<f64> = Study::new(Direction::Minimize);
+//!
+//! study
+//!     .optimize(10, |trial| {
+//!         // Float parameters
+//!         let x = trial.suggest_float("x", 0.0, 1.0)?;
+//!         let lr = trial.suggest_float_log("learning_rate", 1e-5, 1e-1)?;
+//!         let step = trial.suggest_float_step("step", 0.0, 1.0, 0.1)?;
+//!
+//!         // Integer parameters
+//!         let n = trial.suggest_int("n_layers", 1, 10)?;
+//!         let batch = trial.suggest_int_log("batch_size", 16, 256)?;
+//!         let units = trial.suggest_int_step("units", 32, 512, 32)?;
+//!
+//!         // Categorical parameters
+//!         let optimizer = trial.suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])?;
+//!
+//!         // Return objective value
+//!         Ok::<_, optimize::TpeError>(x * n as f64)
+//!     })
+//!     .unwrap();
+//! ```
+//!
+//! # Configuring TPE
+//!
+//! The [`TpeSampler`] can be configured using the builder pattern:
+//!
+//! ```
+//! use optimize::TpeSampler;
+//!
+//! let sampler = TpeSampler::builder()
+//!     .gamma(0.15)           // Quantile for good/bad split
+//!     .n_startup_trials(20)  // Random trials before TPE
+//!     .n_ei_candidates(32)   // Candidates to evaluate
+//!     .seed(42)              // Reproducibility
+//!     .build();
+//! ```
+//!
+//! # Async and Parallel Optimization
+//!
+//! With the `async` feature enabled, you can run trials asynchronously:
+//!
+//! ```ignore
+//! use optimize::{Study, Direction};
+//!
+//! // Sequential async
+//! study.optimize_async(10, |mut trial| async move {
+//!     let x = trial.suggest_float("x", 0.0, 1.0)?;
+//!     Ok((trial, x * x))
+//! }).await?;
+//!
+//! // Parallel with bounded concurrency
+//! study.optimize_parallel(10, 4, |mut trial| async move {
+//!     let x = trial.suggest_float("x", 0.0, 1.0)?;
+//!     Ok((trial, x * x))
+//! }).await?;
+//! ```
+//!
+//! # Serialization
+//!
+//! With the `serde` feature enabled, studies can be serialized:
+//!
+//! ```ignore
+//! use optimize::{Study, Direction, TpeSampler};
+//!
+//! // Save study state
+//! let study: Study<f64> = Study::new(Direction::Minimize);
+//! let json = serde_json::to_string(&study)?;
+//!
+//! // Load and continue
+//! let mut study: Study<f64> = serde_json::from_str(&json)?;
+//! study.set_sampler(TpeSampler::new());  // Restore sampler
+//! study.optimize_with_sampler(10, |trial| { /* ... */ }).unwrap();
+//! ```
+//!
+//! # Feature Flags
+//!
+//! - `serde`: Enable serialization/deserialization of studies and trials
+//! - `async`: Enable async optimization methods (requires tokio)
+
+mod distribution;
+mod error;
+mod kde;
+mod param;
+mod sampler;
+mod study;
+mod trial;
+mod types;
+
+pub use error::{Result, TpeError};
+pub use sampler::{CompletedTrial, RandomSampler, Sampler, TpeSampler, TpeSamplerBuilder};
+pub use study::Study;
+pub use trial::Trial;
+pub use types::{Direction, TrialState};

src/param.rs

@@ -0,0 +1,20 @@
+//! Parameter value storage types.
+
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
+
+/// Represents a sampled parameter value.
+///
+/// This enum stores different parameter value types uniformly.
+/// For categorical parameters, the `Categorical` variant stores
+/// the index into the choices array.
+#[derive(Clone, Debug, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum ParamValue {
+    /// A floating-point parameter value.
+    Float(f64),
+    /// An integer parameter value.
+    Int(i64),
+    /// A categorical parameter value, stored as an index into the choices array.
+    Categorical(usize),
+}

src/sampler/mod.rs

@@ -0,0 +1,96 @@
+//! Sampler trait and implementations for parameter sampling.
+
+mod random;
+pub mod tpe;
+
+use std::collections::HashMap;
+
+pub use random::RandomSampler;
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
+pub use tpe::{TpeSampler, TpeSamplerBuilder};
+
+use crate::distribution::Distribution;
+use crate::param::ParamValue;
+
+/// A completed trial with its parameters, distributions, and objective value.
+///
+/// This struct stores the results of a completed trial, including all sampled
+/// parameter values, their distributions, and the objective value returned
+/// by the objective function.
+#[derive(Clone, Debug)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct CompletedTrial<V = f64> {
+    /// The unique identifier for this trial.
+    pub id: u64,
+    /// The sampled parameter values, keyed by parameter name.
+    pub params: HashMap<String, ParamValue>,
+    /// The parameter distributions used, keyed by parameter name.
+    pub distributions: HashMap<String, Distribution>,
+    /// The objective value returned by the objective function.
+    pub value: V,
+}
+
+impl<V> CompletedTrial<V> {
+    /// Creates a new completed trial.
+    pub fn new(
+        id: u64,
+        params: HashMap<String, ParamValue>,
+        distributions: HashMap<String, Distribution>,
+        value: V,
+    ) -> Self {
+        Self {
+            id,
+            params,
+            distributions,
+            value,
+        }
+    }
+}
+
+/// Trait for pluggable parameter sampling strategies.
+///
+/// Samplers are responsible for generating parameter values based on
+/// the distribution and historical trial data. The trait requires
+/// `Send + Sync` to support concurrent and async optimization.
+///
+/// # Examples
+///
+/// Implementing a custom sampler:
+///
+/// ```ignore
+/// use optimize::{Sampler, ParamValue, Distribution, CompletedTrial};
+///
+/// struct MySampler;
+///
+/// impl Sampler for MySampler {
+///     fn sample(
+///         &self,
+///         distribution: &Distribution,
+///         trial_id: u64,
+///         history: &[CompletedTrial],
+///     ) -> ParamValue {
+///         // Custom sampling logic here
+///         todo!()
+///     }
+/// }
+/// ```
+pub trait Sampler: Send + Sync {
+    /// Samples a parameter value from the given distribution.
+    ///
+    /// # Arguments
+    ///
+    /// * `distribution` - The parameter distribution to sample from.
+    /// * `trial_id` - The unique ID of the trial being sampled for.
+    /// * `history` - Historical completed trials for informed sampling.
+    ///
+    /// # Returns
+    ///
+    /// A `ParamValue` sampled from the distribution.
+    fn sample(
+        &self,
+        distribution: &Distribution,
+        trial_id: u64,
+        history: &[CompletedTrial],
+    ) -> ParamValue;
+}

src/sampler/random.rs

@@ -0,0 +1,275 @@
+//! Random sampler implementation.
+
+use parking_lot::Mutex;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+
+use crate::distribution::Distribution;
+use crate::param::ParamValue;
+use crate::sampler::{CompletedTrial, Sampler};
+
+/// A simple random sampler that samples uniformly from distributions.
+///
+/// This sampler ignores the trial history and samples uniformly at random,
+/// respecting log scale and step size constraints. It serves as a baseline
+/// sampler and is used during the startup phase of more sophisticated samplers.
+///
+/// # Examples
+///
+/// ```
+/// use optimize::RandomSampler;
+///
+/// // Create with default RNG
+/// let sampler = RandomSampler::new();
+///
+/// // Create with a fixed seed for reproducibility
+/// let sampler = RandomSampler::with_seed(42);
+/// ```
+pub struct RandomSampler {
+    rng: Mutex<StdRng>,
+}
+
+impl RandomSampler {
+    /// Creates a new random sampler with a default random seed.
+    pub fn new() -> Self {
+        Self {
+            rng: Mutex::new(StdRng::from_os_rng()),
+        }
+    }
+
+    /// Creates a new random sampler with a fixed seed for reproducibility.
+    ///
+    /// Using the same seed will produce the same sequence of sampled values.
+    pub fn with_seed(seed: u64) -> Self {
+        Self {
+            rng: Mutex::new(StdRng::seed_from_u64(seed)),
+        }
+    }
+}
+
+impl Default for RandomSampler {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Sampler for RandomSampler {
+    fn sample(
+        &self,
+        distribution: &Distribution,
+        _trial_id: u64,
+        _history: &[CompletedTrial],
+    ) -> ParamValue {
+        let mut rng = self.rng.lock();
+
+        match distribution {
+            Distribution::Float(d) => {
+                let value = if d.log_scale {
+                    // Sample uniformly in log space
+                    let log_low = d.low.ln();
+                    let log_high = d.high.ln();
+                    let log_value = rng.random_range(log_low..=log_high);
+                    log_value.exp()
+                } else if let Some(step) = d.step {
+                    // Sample from step grid
+                    let n_steps = ((d.high - d.low) / step).floor() as i64;
+                    let k = rng.random_range(0..=n_steps);
+                    d.low + (k as f64) * step
+                } else {
+                    // Uniform sampling
+                    rng.random_range(d.low..=d.high)
+                };
+                ParamValue::Float(value)
+            }
+            Distribution::Int(d) => {
+                let value = if d.log_scale {
+                    // Sample uniformly in log space, then round
+                    let log_low = (d.low as f64).ln();
+                    let log_high = (d.high as f64).ln();
+                    let log_value = rng.random_range(log_low..=log_high);
+                    let raw = log_value.exp().round() as i64;
+                    // Clamp to bounds since rounding might push outside
+                    raw.clamp(d.low, d.high)
+                } else if let Some(step) = d.step {
+                    // Sample from step grid
+                    let n_steps = (d.high - d.low) / step;
+                    let k = rng.random_range(0..=n_steps);
+                    d.low + k * step
+                } else {
+                    // Uniform sampling
+                    rng.random_range(d.low..=d.high)
+                };
+                ParamValue::Int(value)
+            }
+            Distribution::Categorical(d) => {
+                let index = rng.random_range(0..d.n_choices);
+                ParamValue::Categorical(index)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
+
+    #[test]
+    fn test_random_sampler_float() {
+        let sampler = RandomSampler::with_seed(42);
+        let dist = Distribution::Float(FloatDistribution {
+            low: 0.0,
+            high: 1.0,
+            log_scale: false,
+            step: None,
+        });
+
+        for _ in 0..100 {
+            let value = sampler.sample(&dist, 0, &[]);
+            if let ParamValue::Float(v) = value {
+                assert!((0.0..=1.0).contains(&v));
+            } else {
+                panic!("Expected Float value");
+            }
+        }
+    }
+
+    #[test]
+    fn test_random_sampler_float_log() {
+        let sampler = RandomSampler::with_seed(42);
+        let dist = Distribution::Float(FloatDistribution {
+            low: 1e-5,
+            high: 1.0,
+            log_scale: true,
+            step: None,
+        });
+
+        for _ in 0..100 {
+            let value = sampler.sample(&dist, 0, &[]);
+            if let ParamValue::Float(v) = value {
+                assert!((1e-5..=1.0).contains(&v));
+            } else {
+                panic!("Expected Float value");
+            }
+        }
+    }
+
+    #[test]
+    fn test_random_sampler_float_step() {
+        let sampler = RandomSampler::with_seed(42);
+        let dist = Distribution::Float(FloatDistribution {
+            low: 0.0,
+            high: 1.0,
+            log_scale: false,
+            step: Some(0.25),
+        });
+
+        for _ in 0..100 {
+            let value = sampler.sample(&dist, 0, &[]);
+            if let ParamValue::Float(v) = value {
+                assert!((0.0..=1.0).contains(&v));
+                // Check it's on the step grid
+                let k = ((v - 0.0) / 0.25).round() as i64;
+                let expected = 0.0 + k as f64 * 0.25;
+                assert!((v - expected).abs() < 1e-10);
+            } else {
+                panic!("Expected Float value");
+            }
+        }
+    }
+
+    #[test]
+    fn test_random_sampler_int() {
+        let sampler = RandomSampler::with_seed(42);
+        let dist = Distribution::Int(IntDistribution {
+            low: 0,
+            high: 10,
+            log_scale: false,
+            step: None,
+        });
+
+        for _ in 0..100 {
+            let value = sampler.sample(&dist, 0, &[]);
+            if let ParamValue::Int(v) = value {
+                assert!((0..=10).contains(&v));
+            } else {
+                panic!("Expected Int value");
+            }
+        }
+    }
+
+    #[test]
+    fn test_random_sampler_int_log() {
+        let sampler = RandomSampler::with_seed(42);
+        let dist = Distribution::Int(IntDistribution {
+            low: 1,
+            high: 1000,
+            log_scale: true,
+            step: None,
+        });
+
+        for _ in 0..100 {
+            let value = sampler.sample(&dist, 0, &[]);
+            if let ParamValue::Int(v) = value {
+                assert!((1..=1000).contains(&v));
+            } else {
+                panic!("Expected Int value");
+            }
+        }
+    }
+
+    #[test]
+    fn test_random_sampler_int_step() {
+        let sampler = RandomSampler::with_seed(42);
+        let dist = Distribution::Int(IntDistribution {
+            low: 0,
+            high: 10,
+            log_scale: false,
+            step: Some(2),
+        });
+
+        for _ in 0..100 {
+            let value = sampler.sample(&dist, 0, &[]);
+            if let ParamValue::Int(v) = value {
+                assert!((0..=10).contains(&v));
+                // Check it's on the step grid: 0, 2, 4, 6, 8, 10
+                assert!(v % 2 == 0);
+            } else {
+                panic!("Expected Int value");
+            }
+        }
+    }
+
+    #[test]
+    fn test_random_sampler_categorical() {
+        let sampler = RandomSampler::with_seed(42);
+        let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 });
+
+        for _ in 0..100 {
+            let value = sampler.sample(&dist, 0, &[]);
+            if let ParamValue::Categorical(idx) = value {
+                assert!(idx < 5);
+            } else {
+                panic!("Expected Categorical value");
+            }
+        }
+    }
+
+    #[test]
+    fn test_random_sampler_reproducibility() {
+        let sampler1 = RandomSampler::with_seed(42);
+        let sampler2 = RandomSampler::with_seed(42);
+        let dist = Distribution::Float(FloatDistribution {
+            low: 0.0,
+            high: 1.0,
+            log_scale: false,
+            step: None,
+        });
+
+        for _ in 0..10 {
+            let v1 = sampler1.sample(&dist, 0, &[]);
+            let v2 = sampler2.sample(&dist, 0, &[]);
+            assert_eq!(v1, v2);
+        }
+    }
+}

src/trial.rs

@@ -0,0 +1,793 @@
+//! Trial implementation for tracking sampled parameters and trial state.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use parking_lot::RwLock;
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
+
+use crate::distribution::{
+    CategoricalDistribution, Distribution, FloatDistribution, IntDistribution,
+};
+use crate::error::{Result, TpeError};
+use crate::param::ParamValue;
+use crate::sampler::{CompletedTrial, Sampler};
+use crate::types::TrialState;
+
+/// A trial represents a single evaluation of the objective function.
+///
+/// Each trial has a unique ID and stores the sampled parameters along with
+/// their distributions. The trial progresses through states: Running -> Complete/Failed.
+///
+/// Trials use a sampler to generate parameter values. When created through
+/// `Study::create_trial()`, the trial receives the study's sampler and access
+/// to the history of completed trials for informed sampling.
+#[derive(Clone)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct Trial {
+    /// Unique identifier for this trial.
+    id: u64,
+    /// Current state of the trial.
+    state: TrialState,
+    /// Sampled parameter values, keyed by parameter name.
+    params: HashMap<String, ParamValue>,
+    /// Parameter distributions, keyed by parameter name.
+    distributions: HashMap<String, Distribution>,
+    /// The sampler to use for generating parameter values.
+    #[cfg_attr(feature = "serde", serde(skip))]
+    sampler: Option<Arc<dyn Sampler>>,
+    /// Access to the history of completed trials (shared with Study).
+    #[cfg_attr(feature = "serde", serde(skip))]
+    history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>,
+}
+
+impl std::fmt::Debug for Trial {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Trial")
+            .field("id", &self.id)
+            .field("state", &self.state)
+            .field("params", &self.params)
+            .field("distributions", &self.distributions)
+            .field("has_sampler", &self.sampler.is_some())
+            .field("has_history", &self.history.is_some())
+            .finish()
+    }
+}
+
+impl Trial {
+    /// Creates a new trial with the given ID.
+    ///
+    /// The trial starts in the `Running` state with no parameters sampled.
+    /// This constructor creates a trial without a sampler, which will use
+    /// local random sampling for suggest_* methods.
+    ///
+    /// For trials that use the study's sampler, use `Trial::with_sampler` instead.
+    ///
+    /// # Arguments
+    ///
+    /// * `id` - A unique identifier for this trial.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use optimize::Trial;
+    ///
+    /// let trial = Trial::new(0);
+    /// assert_eq!(trial.id(), 0);
+    /// ```
+    pub fn new(id: u64) -> Self {
+        Self {
+            id,
+            state: TrialState::Running,
+            params: HashMap::new(),
+            distributions: HashMap::new(),
+            sampler: None,
+            history: None,
+        }
+    }
+
+    /// Creates a new trial with a sampler and access to trial history.
+    ///
+    /// This constructor is used by `Study::create_trial()` to create trials
+    /// that use the study's sampler for informed parameter suggestions.
+    ///
+    /// # Arguments
+    ///
+    /// * `id` - A unique identifier for this trial.
+    /// * `sampler` - The sampler to use for generating parameter values.
+    /// * `history` - Shared access to the history of completed trials.
+    pub(crate) fn with_sampler(
+        id: u64,
+        sampler: Arc<dyn Sampler>,
+        history: Arc<RwLock<Vec<CompletedTrial<f64>>>>,
+    ) -> Self {
+        Self {
+            id,
+            state: TrialState::Running,
+            params: HashMap::new(),
+            distributions: HashMap::new(),
+            sampler: Some(sampler),
+            history: Some(history),
+        }
+    }
+
+    /// Samples a value from the given distribution using the sampler.
+    ///
+    /// If the trial has a sampler, it delegates to the sampler's sample method
+    /// with the history of completed trials. Otherwise, it uses the RandomSampler
+    /// as a fallback.
+    fn sample_value(&self, distribution: &Distribution) -> ParamValue {
+        if let (Some(sampler), Some(history)) = (&self.sampler, &self.history) {
+            let history_guard = history.read();
+            sampler.sample(distribution, self.id, &history_guard)
+        } else {
+            // Fallback to RandomSampler when no sampler is configured
+            use crate::sampler::RandomSampler;
+            let fallback = RandomSampler::new();
+            fallback.sample(distribution, self.id, &[])
+        }
+    }
+
+    /// Returns the unique ID of this trial.
+    pub fn id(&self) -> u64 {
+        self.id
+    }
+
+    /// Returns the current state of this trial.
+    pub fn state(&self) -> TrialState {
+        self.state
+    }
+
+    /// Returns a reference to the sampled parameters.
+    pub fn params(&self) -> &HashMap<String, ParamValue> {
+        &self.params
+    }
+
+    /// Returns a reference to the parameter distributions.
+    pub fn distributions(&self) -> &HashMap<String, Distribution> {
+        &self.distributions
+    }
+
+    /// Sets the trial state to Complete.
+    pub(crate) fn set_complete(&mut self) {
+        self.state = TrialState::Complete;
+    }
+
+    /// Sets the trial state to Failed.
+    pub(crate) fn set_failed(&mut self) {
+        self.state = TrialState::Failed;
+    }
+
+    /// Suggests a float parameter with the given bounds.
+    ///
+    /// If the parameter has already been sampled with the same bounds, the cached value is returned.
+    /// If the parameter was sampled with different bounds, a `ParameterConflict` error is returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The name of the parameter.
+    /// * `low` - The lower bound (inclusive).
+    /// * `high` - The upper bound (inclusive).
+    ///
+    /// # Errors
+    ///
+    /// Returns `InvalidBounds` if `low > high`.
+    /// Returns `ParameterConflict` if the parameter was previously sampled with different bounds.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use optimize::Trial;
+    ///
+    /// let mut trial = Trial::new(0);
+    /// let x = trial.suggest_float("x", 0.0, 1.0).unwrap();
+    /// assert!(x >= 0.0 && x <= 1.0);
+    ///
+    /// // Calling again with same bounds returns cached value
+    /// let x2 = trial.suggest_float("x", 0.0, 1.0).unwrap();
+    /// assert_eq!(x, x2);
+    /// ```
+    pub fn suggest_float(&mut self, name: impl Into<String>, low: f64, high: f64) -> Result<f64> {
+        if low > high {
+            return Err(TpeError::InvalidBounds { low, high });
+        }
+
+        let name = name.into();
+        let distribution = FloatDistribution {
+            low,
+            high,
+            log_scale: false,
+            step: None,
+        };
+
+        // Check if parameter already exists
+        if let Some(existing_dist) = self.distributions.get(&name) {
+            // Verify the distribution matches
+            if let Distribution::Float(existing) = existing_dist
+                && existing.low == low
+                && existing.high == high
+                && !existing.log_scale
+                && existing.step.is_none()
+            {
+                // Same distribution, return cached value
+                if let Some(ParamValue::Float(value)) = self.params.get(&name) {
+                    return Ok(*value);
+                }
+            }
+            // Distribution exists but doesn't match
+            return Err(TpeError::ParameterConflict {
+                name,
+                reason: "parameter was previously sampled with different bounds or type"
+                    .to_string(),
+            });
+        }
+
+        // Sample using the sampler
+        let dist = Distribution::Float(distribution);
+        let value = match self.sample_value(&dist) {
+            ParamValue::Float(v) => v,
+            _ => unreachable!("Float distribution should return Float value"),
+        };
+
+        // Store distribution and value
+        self.distributions.insert(name.clone(), dist);
+        self.params.insert(name, ParamValue::Float(value));
+
+        Ok(value)
+    }
+
+    /// Suggests a float parameter sampled on a logarithmic scale.
+    ///
+    /// The value is sampled uniformly in log space, which is useful for parameters
+    /// that span multiple orders of magnitude (e.g., learning rates).
+    ///
+    /// If the parameter has already been sampled with the same bounds and log_scale=true,
+    /// the cached value is returned. If the parameter was sampled with different configuration,
+    /// a `ParameterConflict` error is returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The name of the parameter.
+    /// * `low` - The lower bound (inclusive, must be positive).
+    /// * `high` - The upper bound (inclusive).
+    ///
+    /// # Errors
+    ///
+    /// Returns `InvalidLogBounds` if `low <= 0`.
+    /// Returns `InvalidBounds` if `low > high`.
+    /// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use optimize::Trial;
+    ///
+    /// let mut trial = Trial::new(0);
+    /// let lr = trial
+    ///     .suggest_float_log("learning_rate", 1e-5, 1e-1)
+    ///     .unwrap();
+    /// assert!(lr >= 1e-5 && lr <= 1e-1);
+    ///
+    /// // Calling again with same bounds returns cached value
+    /// let lr2 = trial
+    ///     .suggest_float_log("learning_rate", 1e-5, 1e-1)
+    ///     .unwrap();
+    /// assert_eq!(lr, lr2);
+    /// ```
+    pub fn suggest_float_log(
+        &mut self,
+        name: impl Into<String>,
+        low: f64,
+        high: f64,
+    ) -> Result<f64> {
+        if low <= 0.0 {
+            return Err(TpeError::InvalidLogBounds);
+        }
+
+        if low > high {
+            return Err(TpeError::InvalidBounds { low, high });
+        }
+
+        let name = name.into();
+        let distribution = FloatDistribution {
+            low,
+            high,
+            log_scale: true,
+            step: None,
+        };
+
+        // Check if parameter already exists
+        if let Some(existing_dist) = self.distributions.get(&name) {
+            // Verify the distribution matches
+            if let Distribution::Float(existing) = existing_dist
+                && existing.low == low
+                && existing.high == high
+                && existing.log_scale
+                && existing.step.is_none()
+            {
+                // Same distribution, return cached value
+                if let Some(ParamValue::Float(value)) = self.params.get(&name) {
+                    return Ok(*value);
+                }
+            }
+            // Distribution exists but doesn't match
+            return Err(TpeError::ParameterConflict {
+                name,
+                reason: "parameter was previously sampled with different bounds or type"
+                    .to_string(),
+            });
+        }
+
+        // Sample using the sampler (sampler handles log-scale transformation)
+        let dist = Distribution::Float(distribution);
+        let value = match self.sample_value(&dist) {
+            ParamValue::Float(v) => v,
+            _ => unreachable!("Float distribution should return Float value"),
+        };
+
+        // Store distribution and value
+        self.distributions.insert(name.clone(), dist);
+        self.params.insert(name, ParamValue::Float(value));
+
+        Ok(value)
+    }
+
+    /// Suggests a float parameter that snaps to a step grid.
+    ///
+    /// The value is sampled from the discrete set {low, low + step, low + 2*step, ...}
+    /// where each value is <= high.
+    ///
+    /// If the parameter has already been sampled with the same configuration,
+    /// the cached value is returned. If the parameter was sampled with different configuration,
+    /// a `ParameterConflict` error is returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The name of the parameter.
+    /// * `low` - The lower bound (inclusive).
+    /// * `high` - The upper bound (inclusive).
+    /// * `step` - The step size (must be positive).
+    ///
+    /// # Errors
+    ///
+    /// Returns `InvalidStep` if `step <= 0`.
+    /// Returns `InvalidBounds` if `low > high`.
+    /// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use optimize::Trial;
+    ///
+    /// let mut trial = Trial::new(0);
+    /// let x = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap();
+    /// // x will be one of: 0.0, 0.25, 0.5, 0.75, 1.0
+    /// assert!(x >= 0.0 && x <= 1.0);
+    /// assert!((x / 0.25).fract().abs() < 1e-10 || (x / 0.25).fract().abs() > 1.0 - 1e-10);
+    ///
+    /// // Calling again with same bounds returns cached value
+    /// let x2 = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap();
+    /// assert_eq!(x, x2);
+    /// ```
+    pub fn suggest_float_step(
+        &mut self,
+        name: impl Into<String>,
+        low: f64,
+        high: f64,
+        step: f64,
+    ) -> Result<f64> {
+        if step <= 0.0 {
+            return Err(TpeError::InvalidStep);
+        }
+
+        if low > high {
+            return Err(TpeError::InvalidBounds { low, high });
+        }
+
+        let name = name.into();
+        let distribution = FloatDistribution {
+            low,
+            high,
+            log_scale: false,
+            step: Some(step),
+        };
+
+        // Check if parameter already exists
+        if let Some(existing_dist) = self.distributions.get(&name) {
+            // Verify the distribution matches
+            if let Distribution::Float(existing) = existing_dist
+                && existing.low == low
+                && existing.high == high
+                && !existing.log_scale
+                && existing.step == Some(step)
+            {
+                // Same distribution, return cached value
+                if let Some(ParamValue::Float(value)) = self.params.get(&name) {
+                    return Ok(*value);
+                }
+            }
+            // Distribution exists but doesn't match
+            return Err(TpeError::ParameterConflict {
+                name,
+                reason: "parameter was previously sampled with different bounds or type"
+                    .to_string(),
+            });
+        }
+
+        // Sample using the sampler (sampler handles step-grid)
+        let dist = Distribution::Float(distribution);
+        let value = match self.sample_value(&dist) {
+            ParamValue::Float(v) => v,
+            _ => unreachable!("Float distribution should return Float value"),
+        };
+
+        // Store distribution and value
+        self.distributions.insert(name.clone(), dist);
+        self.params.insert(name, ParamValue::Float(value));
+
+        Ok(value)
+    }
+
+    /// Suggests an integer parameter with the given bounds.
+    ///
+    /// The value is sampled uniformly from the range [low, high] inclusive.
+    ///
+    /// If the parameter has already been sampled with the same bounds, the cached value is returned.
+    /// If the parameter was sampled with different bounds, a `ParameterConflict` error is returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The name of the parameter.
+    /// * `low` - The lower bound (inclusive).
+    /// * `high` - The upper bound (inclusive).
+    ///
+    /// # Errors
+    ///
+    /// Returns `InvalidBounds` if `low > high`.
+    /// Returns `ParameterConflict` if the parameter was previously sampled with different bounds.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use optimize::Trial;
+    ///
+    /// let mut trial = Trial::new(0);
+    /// let n = trial.suggest_int("n_layers", 1, 10).unwrap();
+    /// assert!(n >= 1 && n <= 10);
+    ///
+    /// // Calling again with same bounds returns cached value
+    /// let n2 = trial.suggest_int("n_layers", 1, 10).unwrap();
+    /// assert_eq!(n, n2);
+    /// ```
+    pub fn suggest_int(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
+        if low > high {
+            return Err(TpeError::InvalidBounds {
+                low: low as f64,
+                high: high as f64,
+            });
+        }
+
+        let name = name.into();
+        let distribution = IntDistribution {
+            low,
+            high,
+            log_scale: false,
+            step: None,
+        };
+
+        // Check if parameter already exists
+        if let Some(existing_dist) = self.distributions.get(&name) {
+            // Verify the distribution matches
+            if let Distribution::Int(existing) = existing_dist
+                && existing.low == low
+                && existing.high == high
+                && !existing.log_scale
+                && existing.step.is_none()
+            {
+                // Same distribution, return cached value
+                if let Some(ParamValue::Int(value)) = self.params.get(&name) {
+                    return Ok(*value);
+                }
+            }
+            // Distribution exists but doesn't match
+            return Err(TpeError::ParameterConflict {
+                name,
+                reason: "parameter was previously sampled with different bounds or type"
+                    .to_string(),
+            });
+        }
+
+        // Sample using the sampler
+        let dist = Distribution::Int(distribution);
+        let value = match self.sample_value(&dist) {
+            ParamValue::Int(v) => v,
+            _ => unreachable!("Int distribution should return Int value"),
+        };
+
+        // Store distribution and value
+        self.distributions.insert(name.clone(), dist);
+        self.params.insert(name, ParamValue::Int(value));
+
+        Ok(value)
+    }
+
+    /// Suggests an integer parameter sampled on a logarithmic scale.
+    ///
+    /// The value is sampled uniformly in log space, which is useful for parameters
+    /// that span multiple orders of magnitude (e.g., batch sizes).
+    ///
+    /// If the parameter has already been sampled with the same bounds and log_scale=true,
+    /// the cached value is returned. If the parameter was sampled with different configuration,
+    /// a `ParameterConflict` error is returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The name of the parameter.
+    /// * `low` - The lower bound (inclusive, must be >= 1).
+    /// * `high` - The upper bound (inclusive).
+    ///
+    /// # Errors
+    ///
+    /// Returns `InvalidLogBounds` if `low < 1`.
+    /// Returns `InvalidBounds` if `low > high`.
+    /// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use optimize::Trial;
+    ///
+    /// let mut trial = Trial::new(0);
+    /// let batch_size = trial.suggest_int_log("batch_size", 1, 1024).unwrap();
+    /// assert!(batch_size >= 1 && batch_size <= 1024);
+    ///
+    /// // Calling again with same bounds returns cached value
+    /// let batch_size2 = trial.suggest_int_log("batch_size", 1, 1024).unwrap();
+    /// assert_eq!(batch_size, batch_size2);
+    /// ```
+    pub fn suggest_int_log(&mut self, name: impl Into<String>, low: i64, high: i64) -> Result<i64> {
+        if low < 1 {
+            return Err(TpeError::InvalidLogBounds);
+        }
+
+        if low > high {
+            return Err(TpeError::InvalidBounds {
+                low: low as f64,
+                high: high as f64,
+            });
+        }
+
+        let name = name.into();
+        let distribution = IntDistribution {
+            low,
+            high,
+            log_scale: true,
+            step: None,
+        };
+
+        // Check if parameter already exists
+        if let Some(existing_dist) = self.distributions.get(&name) {
+            // Verify the distribution matches
+            if let Distribution::Int(existing) = existing_dist
+                && existing.low == low
+                && existing.high == high
+                && existing.log_scale
+                && existing.step.is_none()
+            {
+                // Same distribution, return cached value
+                if let Some(ParamValue::Int(value)) = self.params.get(&name) {
+                    return Ok(*value);
+                }
+            }
+            // Distribution exists but doesn't match
+            return Err(TpeError::ParameterConflict {
+                name,
+                reason: "parameter was previously sampled with different bounds or type"
+                    .to_string(),
+            });
+        }
+
+        // Sample using the sampler (sampler handles log-scale transformation)
+        let dist = Distribution::Int(distribution);
+        let value = match self.sample_value(&dist) {
+            ParamValue::Int(v) => v,
+            _ => unreachable!("Int distribution should return Int value"),
+        };
+
+        // Store distribution and value
+        self.distributions.insert(name.clone(), dist);
+        self.params.insert(name, ParamValue::Int(value));
+
+        Ok(value)
+    }
+
+    /// Suggests an integer parameter that snaps to a step grid.
+    ///
+    /// The value is sampled from the discrete set {low, low + step, low + 2*step, ...}
+    /// where each value is <= high.
+    ///
+    /// If the parameter has already been sampled with the same configuration,
+    /// the cached value is returned. If the parameter was sampled with different configuration,
+    /// a `ParameterConflict` error is returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The name of the parameter.
+    /// * `low` - The lower bound (inclusive).
+    /// * `high` - The upper bound (inclusive).
+    /// * `step` - The step size (must be positive).
+    ///
+    /// # Errors
+    ///
+    /// Returns `InvalidStep` if `step <= 0`.
+    /// Returns `InvalidBounds` if `low > high`.
+    /// Returns `ParameterConflict` if the parameter was previously sampled with different configuration.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use optimize::Trial;
+    ///
+    /// let mut trial = Trial::new(0);
+    /// let n = trial
+    ///     .suggest_int_step("n_estimators", 100, 500, 50)
+    ///     .unwrap();
+    /// // n will be one of: 100, 150, 200, 250, 300, 350, 400, 450, 500
+    /// assert!(n >= 100 && n <= 500);
+    /// assert!((n - 100) % 50 == 0);
+    ///
+    /// // Calling again with same bounds returns cached value
+    /// let n2 = trial
+    ///     .suggest_int_step("n_estimators", 100, 500, 50)
+    ///     .unwrap();
+    /// assert_eq!(n, n2);
+    /// ```
+    pub fn suggest_int_step(
+        &mut self,
+        name: impl Into<String>,
+        low: i64,
+        high: i64,
+        step: i64,
+    ) -> Result<i64> {
+        if step <= 0 {
+            return Err(TpeError::InvalidStep);
+        }
+
+        if low > high {
+            return Err(TpeError::InvalidBounds {
+                low: low as f64,
+                high: high as f64,
+            });
+        }
+
+        let name = name.into();
+        let distribution = IntDistribution {
+            low,
+            high,
+            log_scale: false,
+            step: Some(step),
+        };
+
+        // Check if parameter already exists
+        if let Some(existing_dist) = self.distributions.get(&name) {
+            // Verify the distribution matches
+            if let Distribution::Int(existing) = existing_dist
+                && existing.low == low
+                && existing.high == high
+                && !existing.log_scale
+                && existing.step == Some(step)
+            {
+                // Same distribution, return cached value
+                if let Some(ParamValue::Int(value)) = self.params.get(&name) {
+                    return Ok(*value);
+                }
+            }
+            // Distribution exists but doesn't match
+            return Err(TpeError::ParameterConflict {
+                name,
+                reason: "parameter was previously sampled with different bounds or type"
+                    .to_string(),
+            });
+        }
+
+        // Sample using the sampler (sampler handles step-grid)
+        let dist = Distribution::Int(distribution);
+        let value = match self.sample_value(&dist) {
+            ParamValue::Int(v) => v,
+            _ => unreachable!("Int distribution should return Int value"),
+        };
+
+        // Store distribution and value
+        self.distributions.insert(name.clone(), dist);
+        self.params.insert(name, ParamValue::Int(value));
+
+        Ok(value)
+    }
+
+    /// Suggests a categorical parameter from the given choices.
+    ///
+    /// The value is selected uniformly at random from the provided choices.
+    /// Internally, the index of the selected choice is stored.
+    ///
+    /// If the parameter has already been sampled with the same number of choices,
+    /// the cached value (same index) is returned. If the parameter was sampled with
+    /// a different number of choices, a `ParameterConflict` error is returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The name of the parameter.
+    /// * `choices` - A slice of choices to select from.
+    ///
+    /// # Type Parameters
+    ///
+    /// * `T` - The type of the choices. Only requires `Clone`.
+    ///
+    /// # Errors
+    ///
+    /// Returns `EmptyChoices` if `choices` is empty.
+    /// Returns `ParameterConflict` if the parameter was previously sampled with a different number of choices.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use optimize::Trial;
+    ///
+    /// let mut trial = Trial::new(0);
+    /// let optimizer = trial
+    ///     .suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])
+    ///     .unwrap();
+    /// assert!(["sgd", "adam", "rmsprop"].contains(&optimizer));
+    ///
+    /// // Calling again with same choices returns cached value
+    /// let optimizer2 = trial
+    ///     .suggest_categorical("optimizer", &["sgd", "adam", "rmsprop"])
+    ///     .unwrap();
+    /// assert_eq!(optimizer, optimizer2);
+    /// ```
+    pub fn suggest_categorical<T: Clone>(
+        &mut self,
+        name: impl Into<String>,
+        choices: &[T],
+    ) -> Result<T> {
+        if choices.is_empty() {
+            return Err(TpeError::EmptyChoices);
+        }
+
+        let name = name.into();
+        let n_choices = choices.len();
+        let distribution = CategoricalDistribution { n_choices };
+
+        // Check if parameter already exists
+        if let Some(existing_dist) = self.distributions.get(&name) {
+            // Verify the distribution matches
+            if let Distribution::Categorical(existing) = existing_dist
+                && existing.n_choices == n_choices
+            {
+                // Same distribution, return cached value
+                if let Some(ParamValue::Categorical(index)) = self.params.get(&name) {
+                    return Ok(choices[*index].clone());
+                }
+            }
+            // Distribution exists but doesn't match
+            return Err(TpeError::ParameterConflict {
+                name,
+                reason: "parameter was previously sampled with different number of choices or type"
+                    .to_string(),
+            });
+        }
+
+        // Sample using the sampler
+        let dist = Distribution::Categorical(distribution);
+        let index = match self.sample_value(&dist) {
+            ParamValue::Categorical(idx) => idx,
+            _ => unreachable!("Categorical distribution should return Categorical value"),
+        };
+
+        // Store distribution and value (store the index)
+        self.distributions.insert(name.clone(), dist);
+        self.params.insert(name, ParamValue::Categorical(index));
+
+        Ok(choices[index].clone())
+    }
+}

src/types.rs

@@ -0,0 +1,26 @@
+//! Core types for the optimize library.
+
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
+
+/// The direction of optimization.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum Direction {
+    /// Minimize the objective value.
+    Minimize,
+    /// Maximize the objective value.
+    Maximize,
+}
+
+/// The state of a trial in its lifecycle.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum TrialState {
+    /// The trial is currently running.
+    Running,
+    /// The trial completed successfully.
+    Complete,
+    /// The trial failed with an error.
+    Failed,
+}

tests/integration.rs

@@ -0,0 +1,747 @@
+//! Integration tests for the optimize library.
+
+use optimize::{Direction, RandomSampler, Study, TpeError, TpeSampler, Trial};
+
+// =============================================================================
+// Test: optimize simple quadratic function with TPE, finds near-optimal
+// =============================================================================
+
+#[test]
+fn test_tpe_optimizes_quadratic_function() {
+    // Minimize f(x) = (x - 3)^2 where x ∈ [-10, 10]
+    // Optimal: x = 3, f(3) = 0
+    let sampler = TpeSampler::builder()
+        .seed(42)
+        .n_startup_trials(5) // Quick startup for test
+        .n_ei_candidates(24)
+        .build();
+
+    let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
+
+    study
+        .optimize_with_sampler(50, |trial| {
+            let x = trial.suggest_float("x", -10.0, 10.0)?;
+            Ok::<_, TpeError>((x - 3.0).powi(2))
+        })
+        .expect("optimization should succeed");
+
+    let best = study.best_trial().expect("should have at least one trial");
+
+    // TPE should find a value close to optimal (x ≈ 3)
+    // We expect the best value to be small (close to 0)
+    assert!(
+        best.value < 1.0,
+        "TPE should find near-optimal: best value {} should be < 1.0",
+        best.value
+    );
+}
+
+#[test]
+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();
+
+    let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
+
+    study
+        .optimize_with_sampler(100, |trial| {
+            let x = trial.suggest_float("x", -5.0, 5.0)?;
+            let y = trial.suggest_float("y", -5.0, 5.0)?;
+            Ok::<_, TpeError>(x * x + y * y)
+        })
+        .expect("optimization should succeed");
+
+    let best = study.best_trial().expect("should have at least one trial");
+
+    // TPE should find a reasonably good solution
+    assert!(
+        best.value < 5.0,
+        "TPE should find near-optimal: best value {} should be < 5.0",
+        best.value
+    );
+}
+
+#[test]
+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();
+
+    let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
+
+    study
+        .optimize_with_sampler(50, |trial| {
+            let x = trial.suggest_float("x", -10.0, 10.0)?;
+            Ok::<_, TpeError>(-(x - 2.0).powi(2) + 10.0)
+        })
+        .expect("optimization should succeed");
+
+    let best = study.best_trial().expect("should have at least one trial");
+
+    // For maximization, best value should be better than a random baseline
+    // The function ranges from -90 (at x=-10 or x=10, when far from x=2) to 10 (at x=2)
+    // A random approach would average around 0, so finding >5 is a reasonable check
+    assert!(
+        best.value > 5.0,
+        "TPE should find reasonably good solution: best value {} should be > 5.0",
+        best.value
+    );
+}
+
+// =============================================================================
+// Test: RandomSampler samples uniformly across range
+// =============================================================================
+
+#[test]
+fn test_random_sampler_uniform_float_distribution() {
+    // Test that RandomSampler samples uniformly by running multiple trials
+    // and checking the distribution of sampled values
+    let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
+
+    let n_samples = 1000;
+    let mut samples = Vec::with_capacity(n_samples);
+
+    study
+        .optimize(n_samples, |trial| {
+            let x = trial.suggest_float("x", 0.0, 1.0)?;
+            samples.push(x);
+            Ok::<_, TpeError>(x)
+        })
+        .unwrap();
+
+    // All samples should be in range
+    for &s in &samples {
+        assert!((0.0..=1.0).contains(&s), "sample {s} out of range [0, 1]");
+    }
+
+    // Check distribution is roughly uniform by looking at quartiles
+    samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
+
+    let q1 = samples[n_samples / 4];
+    let q2 = samples[n_samples / 2];
+    let q3 = samples[3 * n_samples / 4];
+
+    // For uniform distribution, quartiles should be approximately 0.25, 0.5, 0.75
+    assert!((q1 - 0.25).abs() < 0.1, "Q1 {q1} should be close to 0.25");
+    assert!(
+        (q2 - 0.5).abs() < 0.1,
+        "Q2 (median) {q2} should be close to 0.5"
+    );
+    assert!((q3 - 0.75).abs() < 0.1, "Q3 {q3} should be close to 0.75");
+}
+
+#[test]
+fn test_random_sampler_uniform_int_distribution() {
+    let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(123));
+
+    let n_samples = 1000;
+    let mut counts = [0u32; 10]; // counts for values 1-10
+
+    study
+        .optimize(n_samples, |trial| {
+            let n = trial.suggest_int("n", 1, 10)?;
+            assert!((1..=10).contains(&n), "sample {n} out of range [1, 10]");
+            counts[(n - 1) as usize] += 1;
+            Ok::<_, TpeError>(n as f64)
+        })
+        .unwrap();
+
+    // Each value should appear roughly n_samples / 10 times
+    let expected = n_samples as f64 / 10.0;
+    for (i, &count) in counts.iter().enumerate() {
+        let diff = (count as f64 - expected).abs() / expected;
+        assert!(
+            diff < 0.3,
+            "value {} appeared {} times, expected ~{}, diff = {:.1}%",
+            i + 1,
+            count,
+            expected,
+            diff * 100.0
+        );
+    }
+}
+
+#[test]
+fn test_random_sampler_uniform_categorical_distribution() {
+    let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(456));
+
+    let n_samples = 1000;
+    let mut counts = [0u32; 4];
+    let choices = ["a", "b", "c", "d"];
+
+    study
+        .optimize(n_samples, |trial| {
+            let choice = trial.suggest_categorical("cat", &choices)?;
+            let idx = choices.iter().position(|&c| c == choice).unwrap();
+            counts[idx] += 1;
+            Ok::<_, TpeError>(idx as f64)
+        })
+        .unwrap();
+
+    // Each category should appear roughly n_samples / 4 times
+    let expected = n_samples as f64 / 4.0;
+    for (i, &count) in counts.iter().enumerate() {
+        let diff = (count as f64 - expected).abs() / expected;
+        assert!(
+            diff < 0.25,
+            "category {} appeared {} times, expected ~{}, diff = {:.1}%",
+            i,
+            count,
+            expected,
+            diff * 100.0
+        );
+    }
+}
+
+#[test]
+fn test_random_sampler_reproducibility() {
+    // Two studies with the same seed should produce the same sequence
+    // NOTE: We must use optimize_with_sampler() for Study<f64> to get sampler integration
+    let study1: Study<f64> =
+        Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999));
+    let study2: Study<f64> =
+        Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(999));
+
+    let mut values1 = Vec::new();
+    let mut values2 = Vec::new();
+
+    study1
+        .optimize_with_sampler(100, |trial| {
+            let x = trial.suggest_float("x", 0.0, 100.0)?;
+            values1.push(x);
+            Ok::<_, TpeError>(x)
+        })
+        .unwrap();
+
+    study2
+        .optimize_with_sampler(100, |trial| {
+            let x = trial.suggest_float("x", 0.0, 100.0)?;
+            values2.push(x);
+            Ok::<_, TpeError>(x)
+        })
+        .unwrap();
+
+    for (i, (v1, v2)) in values1.iter().zip(values2.iter()).enumerate() {
+        assert_eq!(
+            v1, v2,
+            "values at trial {i} should be identical with same seed: {v1} vs {v2}"
+        );
+    }
+}
+
+// =============================================================================
+// Test: suggest_* methods return cached values on repeated calls
+// =============================================================================
+
+#[test]
+fn test_suggest_float_caching() {
+    let mut trial = Trial::new(0);
+
+    let x1 = trial.suggest_float("x", 0.0, 10.0).unwrap();
+    let x2 = trial.suggest_float("x", 0.0, 10.0).unwrap();
+    let x3 = trial.suggest_float("x", 0.0, 10.0).unwrap();
+
+    assert_eq!(x1, x2, "repeated suggest_float should return cached value");
+    assert_eq!(x2, x3, "repeated suggest_float should return cached value");
+}
+
+#[test]
+fn test_suggest_float_log_caching() {
+    let mut trial = Trial::new(0);
+
+    let x1 = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap();
+    let x2 = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap();
+
+    assert_eq!(
+        x1, x2,
+        "repeated suggest_float_log should return cached value"
+    );
+}
+
+#[test]
+fn test_suggest_float_step_caching() {
+    let mut trial = Trial::new(0);
+
+    let x1 = trial.suggest_float_step("step", 0.0, 1.0, 0.1).unwrap();
+    let x2 = trial.suggest_float_step("step", 0.0, 1.0, 0.1).unwrap();
+
+    assert_eq!(
+        x1, x2,
+        "repeated suggest_float_step should return cached value"
+    );
+}
+
+#[test]
+fn test_suggest_int_caching() {
+    let mut trial = Trial::new(0);
+
+    let n1 = trial.suggest_int("n", 1, 100).unwrap();
+    let n2 = trial.suggest_int("n", 1, 100).unwrap();
+
+    assert_eq!(n1, n2, "repeated suggest_int should return cached value");
+}
+
+#[test]
+fn test_suggest_int_log_caching() {
+    let mut trial = Trial::new(0);
+
+    let n1 = trial.suggest_int_log("batch", 1, 1024).unwrap();
+    let n2 = trial.suggest_int_log("batch", 1, 1024).unwrap();
+
+    assert_eq!(
+        n1, n2,
+        "repeated suggest_int_log should return cached value"
+    );
+}
+
+#[test]
+fn test_suggest_int_step_caching() {
+    let mut trial = Trial::new(0);
+
+    let n1 = trial.suggest_int_step("units", 32, 512, 32).unwrap();
+    let n2 = trial.suggest_int_step("units", 32, 512, 32).unwrap();
+
+    assert_eq!(
+        n1, n2,
+        "repeated suggest_int_step should return cached value"
+    );
+}
+
+#[test]
+fn test_suggest_categorical_caching() {
+    let mut trial = Trial::new(0);
+
+    let choices = ["sgd", "adam", "rmsprop"];
+    let c1 = trial.suggest_categorical("optimizer", &choices).unwrap();
+    let c2 = trial.suggest_categorical("optimizer", &choices).unwrap();
+
+    assert_eq!(
+        c1, c2,
+        "repeated suggest_categorical should return cached value"
+    );
+}
+
+#[test]
+fn test_multiple_parameters_independent_caching() {
+    let mut trial = Trial::new(0);
+
+    // Suggest multiple parameters
+    let x = trial.suggest_float("x", 0.0, 1.0).unwrap();
+    let y = trial.suggest_float("y", 0.0, 1.0).unwrap();
+    let n = trial.suggest_int("n", 1, 10).unwrap();
+    let opt = trial.suggest_categorical("opt", &["a", "b"]).unwrap();
+
+    // All should be cached independently
+    assert_eq!(x, trial.suggest_float("x", 0.0, 1.0).unwrap());
+    assert_eq!(y, trial.suggest_float("y", 0.0, 1.0).unwrap());
+    assert_eq!(n, trial.suggest_int("n", 1, 10).unwrap());
+    assert_eq!(opt, trial.suggest_categorical("opt", &["a", "b"]).unwrap());
+}
+
+// =============================================================================
+// Test: parameter conflict returns error
+// =============================================================================
+
+#[test]
+fn test_parameter_conflict_float_different_bounds() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_float("x", 0.0, 1.0).unwrap();
+    let result = trial.suggest_float("x", 0.0, 2.0); // Different upper bound
+
+    assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
+}
+
+#[test]
+fn test_parameter_conflict_float_vs_log() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_float("x", 0.1, 1.0).unwrap();
+    let result = trial.suggest_float_log("x", 0.1, 1.0); // Same bounds but log scale
+
+    assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
+}
+
+#[test]
+fn test_parameter_conflict_float_vs_step() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_float("x", 0.0, 1.0).unwrap();
+    let result = trial.suggest_float_step("x", 0.0, 1.0, 0.1); // Same bounds but with step
+
+    assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
+}
+
+#[test]
+fn test_parameter_conflict_int_different_bounds() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_int("n", 1, 10).unwrap();
+    let result = trial.suggest_int("n", 1, 20); // Different upper bound
+
+    assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
+}
+
+#[test]
+fn test_parameter_conflict_int_vs_log() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_int("n", 1, 100).unwrap();
+    let result = trial.suggest_int_log("n", 1, 100); // Same bounds but log scale
+
+    assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
+}
+
+#[test]
+fn test_parameter_conflict_categorical_different_n_choices() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_categorical("opt", &["a", "b", "c"]).unwrap();
+    let result = trial.suggest_categorical("opt", &["x", "y"]); // Different number of choices
+
+    assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
+}
+
+#[test]
+fn test_parameter_conflict_float_vs_int() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_float("x", 0.0, 10.0).unwrap();
+    let result = trial.suggest_int("x", 0, 10); // Different type
+
+    assert!(matches!(result, Err(TpeError::ParameterConflict { .. })));
+}
+
+#[test]
+fn test_parameter_conflict_returns_name() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_float("my_param", 0.0, 1.0).unwrap();
+    let result = trial.suggest_float("my_param", 0.0, 2.0);
+
+    match result {
+        Err(TpeError::ParameterConflict { name, .. }) => {
+            assert_eq!(name, "my_param");
+        }
+        _ => panic!("expected ParameterConflict error"),
+    }
+}
+
+// =============================================================================
+// Test: empty categorical returns error
+// =============================================================================
+
+#[test]
+fn test_empty_categorical_returns_error() {
+    let mut trial = Trial::new(0);
+    let empty: &[&str] = &[];
+
+    let result = trial.suggest_categorical("opt", empty);
+
+    assert!(matches!(result, Err(TpeError::EmptyChoices)));
+}
+
+#[test]
+fn test_empty_categorical_vec_returns_error() {
+    let mut trial = Trial::new(0);
+    let empty: Vec<i32> = vec![];
+
+    let result = trial.suggest_categorical("numbers", &empty);
+
+    assert!(matches!(result, Err(TpeError::EmptyChoices)));
+}
+
+// =============================================================================
+// Additional integration tests
+// =============================================================================
+
+#[test]
+fn test_study_basic_workflow() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+
+    study
+        .optimize(10, |trial| {
+            let x = trial.suggest_float("x", -5.0, 5.0)?;
+            Ok::<_, TpeError>(x * x)
+        })
+        .expect("optimization should succeed");
+
+    assert_eq!(study.n_trials(), 10);
+    let best = study.best_trial().expect("should have best trial");
+    assert!(best.value >= 0.0, "x^2 should be non-negative");
+}
+
+#[test]
+fn test_study_with_failures() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+
+    // Every other trial fails
+    let mut counter = 0;
+    study
+        .optimize(10, |trial| {
+            counter += 1;
+            if counter % 2 == 0 {
+                return Err::<f64, &str>("intentional failure");
+            }
+            let x = trial
+                .suggest_float("x", -5.0, 5.0)
+                .map_err(|_| "param error")?;
+            Ok(x * x)
+        })
+        .expect("optimization should succeed with some failures");
+
+    // Only half the trials should have succeeded
+    assert_eq!(study.n_trials(), 5, "only 5 trials should have completed");
+}
+
+#[test]
+fn test_no_completed_trials_error() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+
+    let result = study.best_trial();
+    assert!(matches!(result, Err(TpeError::NoCompletedTrials)));
+}
+
+#[test]
+fn test_invalid_bounds_errors() {
+    let mut trial = Trial::new(0);
+
+    // low > high for float
+    let result = trial.suggest_float("x", 10.0, 5.0);
+    assert!(matches!(result, Err(TpeError::InvalidBounds { .. })));
+
+    // low > high for int
+    let result = trial.suggest_int("n", 100, 50);
+    assert!(matches!(result, Err(TpeError::InvalidBounds { .. })));
+}
+
+#[test]
+fn test_invalid_log_bounds_errors() {
+    let mut trial = Trial::new(0);
+
+    // low <= 0 for log float
+    let result = trial.suggest_float_log("x", 0.0, 1.0);
+    assert!(matches!(result, Err(TpeError::InvalidLogBounds)));
+
+    let result = trial.suggest_float_log("y", -1.0, 1.0);
+    assert!(matches!(result, Err(TpeError::InvalidLogBounds)));
+
+    // low < 1 for log int
+    let result = trial.suggest_int_log("n", 0, 100);
+    assert!(matches!(result, Err(TpeError::InvalidLogBounds)));
+}
+
+#[test]
+fn test_invalid_step_errors() {
+    let mut trial = Trial::new(0);
+
+    // step <= 0 for float
+    let result = trial.suggest_float_step("x", 0.0, 1.0, 0.0);
+    assert!(matches!(result, Err(TpeError::InvalidStep)));
+
+    let result = trial.suggest_float_step("y", 0.0, 1.0, -0.1);
+    assert!(matches!(result, Err(TpeError::InvalidStep)));
+
+    // step <= 0 for int
+    let result = trial.suggest_int_step("n", 0, 100, 0);
+    assert!(matches!(result, Err(TpeError::InvalidStep)));
+}
+
+#[test]
+fn test_tpe_with_categorical_parameter() {
+    let sampler = TpeSampler::builder().seed(42).n_startup_trials(5).build();
+
+    let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
+
+    // Optimization where the best choice depends on the categorical
+    study
+        .optimize_with_sampler(30, |trial| {
+            let choice = trial.suggest_categorical("model", &["linear", "quadratic", "cubic"])?;
+            let x = trial.suggest_float("x", 0.0, 2.0)?;
+
+            // cubic model is best at x=1
+            let value = match choice {
+                "linear" => x,
+                "quadratic" => x * x,
+                "cubic" => -((x - 1.0).powi(2)) + 10.0, // peak at x=1, max value 10
+                _ => unreachable!(),
+            };
+            Ok::<_, TpeError>(value)
+        })
+        .expect("optimization should succeed");
+
+    let best = study.best_trial().expect("should have best trial");
+    // The optimizer should find that "cubic" with x≈1 is best
+    assert!(
+        best.value > 5.0,
+        "should find good solution, got {}",
+        best.value
+    );
+}
+
+#[test]
+fn test_tpe_with_integer_parameters() {
+    let sampler = TpeSampler::builder().seed(789).n_startup_trials(5).build();
+
+    let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
+
+    // Minimize (n - 7)^2 where n ∈ [1, 10]
+    study
+        .optimize_with_sampler(30, |trial| {
+            let n = trial.suggest_int("n", 1, 10)?;
+            Ok::<_, TpeError>(((n - 7) as f64).powi(2))
+        })
+        .expect("optimization should succeed");
+
+    let best = study.best_trial().expect("should have best trial");
+
+    // Best value should be small (n close to 7)
+    assert!(
+        best.value < 5.0,
+        "should find n close to 7, best value = {}",
+        best.value
+    );
+}
+
+#[test]
+fn test_callback_early_stopping() {
+    use std::cell::Cell;
+    use std::ops::ControlFlow;
+
+    let study: Study<f64> = Study::new(Direction::Minimize);
+    let trials_run = Cell::new(0);
+
+    study
+        .optimize_with_callback(
+            100,
+            |trial| {
+                trials_run.set(trials_run.get() + 1);
+                let x = trial.suggest_float("x", 0.0, 10.0)?;
+                Ok::<_, TpeError>(x)
+            },
+            |_study, _trial| {
+                // Stop after 5 trials
+                if trials_run.get() >= 5 {
+                    ControlFlow::Break(())
+                } else {
+                    ControlFlow::Continue(())
+                }
+            },
+        )
+        .expect("optimization should succeed");
+
+    assert_eq!(study.n_trials(), 5, "should have stopped after 5 trials");
+}
+
+#[test]
+fn test_study_trials_iteration() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+
+    study
+        .optimize(5, |trial| {
+            let x = trial.suggest_float("x", 0.0, 1.0)?;
+            Ok::<_, TpeError>(x)
+        })
+        .unwrap();
+
+    let trials = study.trials();
+    assert_eq!(trials.len(), 5);
+
+    for trial in &trials {
+        assert!(
+            !trial.params.is_empty(),
+            "each trial should have parameters"
+        );
+        assert!(
+            trial.params.contains_key("x"),
+            "each trial should have parameter 'x'"
+        );
+    }
+}
+
+#[test]
+fn test_study_direction() {
+    let study_min: Study<f64> = Study::new(Direction::Minimize);
+    assert_eq!(study_min.direction(), Direction::Minimize);
+
+    let study_max: Study<f64> = Study::new(Direction::Maximize);
+    assert_eq!(study_max.direction(), Direction::Maximize);
+}
+
+#[test]
+fn test_trial_state() {
+    use optimize::TrialState;
+
+    let trial = Trial::new(0);
+    assert_eq!(trial.state(), TrialState::Running);
+}
+
+#[test]
+fn test_trial_params_access() {
+    let mut trial = Trial::new(0);
+
+    trial.suggest_float("x", 0.0, 1.0).unwrap();
+    trial.suggest_int("n", 1, 10).unwrap();
+
+    let params = trial.params();
+    assert_eq!(params.len(), 2);
+    assert!(params.contains_key("x"));
+    assert!(params.contains_key("n"));
+}
+
+#[test]
+fn test_log_scale_float_range() {
+    let mut trial = Trial::new(0);
+
+    let lr = trial.suggest_float_log("lr", 1e-5, 1e-1).unwrap();
+    assert!(
+        (1e-5..=1e-1).contains(&lr),
+        "log-scale value {lr} out of range"
+    );
+}
+
+#[test]
+fn test_step_float_snaps_to_grid() {
+    let mut trial = Trial::new(0);
+
+    let x = trial.suggest_float_step("x", 0.0, 1.0, 0.25).unwrap();
+
+    // x should be one of: 0.0, 0.25, 0.5, 0.75, 1.0
+    let valid_values = [0.0, 0.25, 0.5, 0.75, 1.0];
+    let is_valid = valid_values.iter().any(|&v| (x - v).abs() < 1e-10);
+    assert!(is_valid, "stepped float {x} should snap to grid");
+}
+
+#[test]
+fn test_step_int_snaps_to_grid() {
+    let mut trial = Trial::new(0);
+
+    let n = trial.suggest_int_step("n", 0, 100, 25).unwrap();
+
+    // n should be one of: 0, 25, 50, 75, 100
+    assert!(
+        n % 25 == 0 && (0..=100).contains(&n),
+        "stepped int {n} should snap to grid"
+    );
+}
+
+#[test]
+fn test_best_value() {
+    let study: Study<f64> = Study::new(Direction::Minimize);
+
+    study
+        .optimize(10, |trial| {
+            let x = trial.suggest_float("x", 0.0, 10.0)?;
+            Ok::<_, TpeError>(x)
+        })
+        .unwrap();
+
+    let best_value = study.best_value().expect("should have best value");
+    let best_trial = study.best_trial().expect("should have best trial");
+
+    assert_eq!(
+        best_value, best_trial.value,
+        "best_value should match best_trial.value"
+    );
+}