From cfa8426312eefac5e6c152df19e93a8530481819 Mon Sep 17 00:00:00 2001 From: Manuel Raimann Date: Thu, 12 Feb 2026 16:31:25 +0100 Subject: [PATCH] docs(sampler,pruner): add implementation guides for custom samplers and pruners - Expand sampler module docs with available samplers tables, custom sampler walkthrough, stateless/stateful patterns, cold start handling, history reading, thread safety, and testing guidance - Expand pruner module docs with stateful/stateless classification, warmup parameters, decorator composition, thread safety, and testing - Add code examples to both trait docs (NoisySampler, StalePruner) --- src/pruner/mod.rs | 103 +++++++++++++++++++++ src/sampler/mod.rs | 219 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 322 insertions(+) diff --git a/src/pruner/mod.rs b/src/pruner/mod.rs index 03d7f01..02fb055 100644 --- a/src/pruner/mod.rs +++ b/src/pruner/mod.rs @@ -41,6 +41,61 @@ //! Start with [`MedianPruner`] for most use cases. Switch to [`WilcoxonPruner`] //! if your intermediate values are noisy, or to [`HyperbandPruner`] if you want //! automatic budget scheduling. +//! +//! # Stateful vs stateless pruners +//! +//! **Stateless** pruners make their decision purely from the arguments passed +//! to [`Pruner::should_prune`] — they hold no mutable per-trial state. +//! [`MedianPruner`], [`PercentilePruner`], [`ThresholdPruner`], +//! [`WilcoxonPruner`], and [`NopPruner`] are all stateless. +//! +//! **Stateful** pruners track information across calls. [`PatientPruner`] +//! uses `Mutex>` to count consecutive prune signals per +//! trial. [`HyperbandPruner`] uses `Mutex` and `AtomicU64` for bracket +//! assignment state. When writing a stateful pruner, wrap mutable state in a +//! `Mutex` and key it by `trial_id` to keep trials independent. +//! +//! # Cold start and warmup +//! +//! Two builder parameters control when pruning begins: +//! +//! - **`n_warmup_steps`** — skip pruning before step N *within a trial*, +//! giving the objective time to stabilize. +//! - **`n_min_trials`** — require N completed trials before pruning any trial, +//! ensuring a meaningful comparison baseline. +//! +//! See [`MedianPruner`] for the canonical implementation of both parameters. +//! Custom pruners should expose similar knobs when applicable. +//! +//! # Composing pruners +//! +//! [`PatientPruner`] demonstrates the decorator pattern: it wraps any +//! `Box` and adds patience logic on top. Custom pruners can use +//! the same pattern to layer multiple pruning conditions — for example, +//! combining a statistical test with a hard threshold. +//! +//! # Thread safety +//! +//! The [`Pruner`] trait requires `Send + Sync`. +//! [`Study`](crate::Study) stores the pruner as `Arc`, so +//! multiple threads may call [`Pruner::should_prune`] concurrently. +//! +//! - **Stateless pruners** satisfy `Send + Sync` automatically. +//! - **Stateful pruners** should use `std::sync::Mutex` or +//! `parking_lot::Mutex` to protect mutable state, keyed by `trial_id`. +//! +//! # Testing custom pruners +//! +//! Recommended test categories: +//! +//! 1. **Never-prune baseline** — empty history and early steps should not +//! prune. +//! 2. **Known-prune scenario** — a clearly worse trial should be pruned. +//! 3. **Known-keep scenario** — a well-performing trial should survive. +//! 4. **Warmup respected** — pruning must be suppressed during warmup steps +//! and while the minimum trial count has not been reached. +//! 5. **Per-trial independence** — stateful pruners must not leak state +//! between different `trial_id` values. mod hyperband; mod median; @@ -93,6 +148,54 @@ use crate::sampler::CompletedTrial; /// } /// } /// ``` +/// +/// A stateful pruner that tracks per-trial state with a `Mutex`: +/// +/// ``` +/// use std::collections::HashMap; +/// use std::sync::Mutex; +/// use optimizer::pruner::Pruner; +/// use optimizer::sampler::CompletedTrial; +/// +/// /// Prune after the value worsens for `max_stale` consecutive steps. +/// struct StalePruner { +/// max_stale: u64, +/// // Per-trial: (previous_value, consecutive_stale_count) +/// state: Mutex>, +/// } +/// +/// impl StalePruner { +/// fn new(max_stale: u64) -> Self { +/// Self { max_stale, state: Mutex::new(HashMap::new()) } +/// } +/// } +/// +/// impl Pruner for StalePruner { +/// fn should_prune( +/// &self, +/// trial_id: u64, +/// _step: u64, +/// intermediate_values: &[(u64, f64)], +/// _completed_trials: &[CompletedTrial], +/// ) -> bool { +/// let Some(&(_, current)) = intermediate_values.last() else { +/// return false; +/// }; +/// let mut state = self.state.lock().unwrap(); +/// let entry = state.entry(trial_id).or_insert((current, 0)); +/// if current >= entry.0 { +/// entry.1 += 1; +/// } else { +/// entry.1 = 0; +/// } +/// entry.0 = current; +/// entry.1 >= self.max_stale +/// } +/// } +/// ``` +/// +/// See the [module-level documentation](self) for a comprehensive guide +/// covering warmup, composition, thread safety, and testing. pub trait Pruner: Send + Sync { /// Decide whether to prune a trial at the given step. /// diff --git a/src/sampler/mod.rs b/src/sampler/mod.rs index ff83bc6..6f5ba17 100644 --- a/src/sampler/mod.rs +++ b/src/sampler/mod.rs @@ -1,4 +1,179 @@ //! Sampler trait and implementations for parameter sampling. +//! +//! A sampler generates parameter values for each trial. It receives a +//! [`Distribution`] describing the parameter space, a monotonically increasing +//! `trial_id`, and the list of all [`CompletedTrial`]s so far, and returns a +//! [`ParamValue`] that matches the distribution variant. +//! +//! # Available samplers +//! +//! ## Single-objective +//! +//! | Sampler | Algorithm | Best for | +//! |---------|-----------|----------| +//! | [`RandomSampler`] | Uniform independent sampling | Baselines, startup phases | +//! | [`TpeSampler`] | Tree-Parzen Estimator | General-purpose Bayesian optimization | +//! | [`TpeSampler`] (multivariate) | Multivariate TPE with tree-structured Parzen | Correlated parameters | +//! | [`GridSampler`] | Exhaustive grid evaluation | Small discrete spaces | +//! | [`SobolSampler`]\* | Quasi-random Sobol sequences | Uniform coverage without model | +//! | [`CmaEsSampler`]\* | Covariance Matrix Adaptation | Continuous, non-separable problems | +//! | [`GpSampler`]\* | Gaussian Process with EI | Expensive, low-dimensional functions | +//! | [`DESampler`] | Differential Evolution | Population-based, multi-modal landscapes | +//! | [`BohbSampler`] | Bayesian Optimization + `HyperBand` | Combined sampling and pruning | +//! +//! \*Requires a feature flag (`sobol`, `cma-es`, or `gp`). +//! +//! ## Multi-objective +//! +//! | Sampler | Algorithm | Best for | +//! |---------|-----------|----------| +//! | [`Nsga2Sampler`] | NSGA-II | General multi-objective with 2-3 objectives | +//! | [`Nsga3Sampler`] | NSGA-III | Many-objective (4+ objectives) | +//! | [`MoeadSampler`] | MOEA/D with decomposition | Structured Pareto front exploration | +//! | [`MotpeSampler`] | Multi-objective TPE | Bayesian multi-objective | +//! +//! # Implementing a custom sampler +//! +//! Implement the [`Sampler`] trait with its single method: +//! +//! ```rust +//! use optimizer::sampler::{Sampler, CompletedTrial}; +//! use optimizer::distribution::Distribution; +//! use optimizer::param::ParamValue; +//! +//! /// A sampler that always picks the midpoint of each distribution. +//! struct MidpointSampler; +//! +//! impl Sampler for MidpointSampler { +//! fn sample( +//! &self, +//! distribution: &Distribution, +//! _trial_id: u64, +//! _history: &[CompletedTrial], +//! ) -> ParamValue { +//! match distribution { +//! Distribution::Float(fd) => { +//! ParamValue::Float((fd.low + fd.high) / 2.0) +//! } +//! Distribution::Int(id) => { +//! ParamValue::Int((id.low + id.high) / 2) +//! } +//! Distribution::Categorical(cd) => { +//! ParamValue::Categorical(cd.n_choices / 2) +//! } +//! } +//! } +//! } +//! ``` +//! +//! The arguments to [`Sampler::sample`]: +//! +//! - **`distribution`** — a [`Distribution::Float`], [`Distribution::Int`], or +//! [`Distribution::Categorical`] that describes the parameter bounds, +//! log-scale flag, and optional step size. +//! - **`trial_id`** — a monotonically increasing identifier. Useful for +//! deterministic RNG seeding (see [Stateless vs stateful samplers]). +//! - **`history`** — all completed trials so far. May be empty on the first +//! trial. Model-based samplers use this to guide future sampling. +//! - **Return value** — the [`ParamValue`] variant *must* match the +//! distribution variant (`Float` → `ParamValue::Float`, etc.). +//! +//! [Stateless vs stateful samplers]: #stateless-vs-stateful-samplers +//! +//! # Stateless vs stateful samplers +//! +//! **Stateless** samplers derive all randomness from a deterministic function +//! of `seed + trial_id + distribution`. They use an [`AtomicU64`] call-sequence +//! counter to disambiguate multiple calls within the same trial, but need no +//! `Mutex`. See [`RandomSampler`] and [`TpeSampler`] for this pattern. +//! +//! **Stateful** samplers maintain mutable state (e.g. a population pool) +//! across calls. Wrap mutable state in `parking_lot::Mutex` and lock +//! for the duration of [`Sampler::sample`]. See [`DESampler`] and +//! [`GridSampler`] for this pattern. +//! +//! [`AtomicU64`]: core::sync::atomic::AtomicU64 +//! +//! # Cold start handling +//! +//! Model-based samplers need completed trials before their surrogate model is +//! useful. The standard pattern is to check `history.len() < n_startup_trials` +//! and fall back to random sampling during the startup phase. Expose this as a +//! builder parameter so users can tune the trade-off between exploration and +//! exploitation. See [`TpeSampler`] for a reference implementation. +//! +//! # Reading trial history +//! +//! The `history` slice contains only completed trials (never pending ones). +//! Common operations: +//! +//! - **Extract a parameter value:** +//! `trial.params.get(¶m_id)` returns `Option<&ParamValue>`. +//! - **Find the best trial:** +//! `history.iter().min_by(|a, b| a.value.partial_cmp(&b.value).unwrap())`. +//! - **Filter by state:** +//! `history.iter().filter(|t| t.state == TrialState::Complete)`. +//! - **Check feasibility:** +//! `trial.is_feasible()` returns `true` when all constraints are ≤ 0. +//! +//! # Thread safety +//! +//! The [`Sampler`] trait requires `Send + Sync`. [`Study`](crate::Study) stores +//! the sampler as `Arc`, so multiple threads may call +//! [`Sampler::sample`] concurrently. +//! +//! - **Stateless:** `AtomicU64` counters satisfy `Send + Sync` without locking. +//! - **Stateful:** use `parking_lot::Mutex` (the crate convention) or +//! `std::sync::Mutex` to protect mutable state. +//! +//! # Testing custom samplers +//! +//! Recommended test categories: +//! +//! 1. **Bounds compliance** — sample many values and assert they fall within +//! the distribution range. +//! 2. **Step / log-scale correctness** — verify that discretized and +//! log-scaled distributions produce valid values. +//! 3. **Reproducibility** — the same seed must produce the same output. +//! 4. **History sensitivity** — model-based samplers should produce different +//! (better) samples as history grows. +//! 5. **Empty history** — `sample()` must not panic when `history` is empty. +//! +//! # Using a custom sampler with Study +//! +//! ```rust +//! use optimizer::{Direction, Study}; +//! use optimizer::sampler::{Sampler, CompletedTrial}; +//! use optimizer::distribution::Distribution; +//! use optimizer::param::ParamValue; +//! +//! struct MySampler; +//! impl Sampler for MySampler { +//! fn sample( +//! &self, +//! distribution: &Distribution, +//! _trial_id: u64, +//! _history: &[CompletedTrial], +//! ) -> ParamValue { +//! match distribution { +//! Distribution::Float(fd) => ParamValue::Float(fd.low), +//! Distribution::Int(id) => ParamValue::Int(id.low), +//! Distribution::Categorical(_) => ParamValue::Categorical(0), +//! } +//! } +//! } +//! +//! let study: Study = Study::with_sampler(Direction::Minimize, MySampler); +//! ``` +//! +//! The sampler is wrapped in `Arc` internally. +//! +//! # Reference implementations +//! +//! - [`RandomSampler`] — simplest sampler; stateless, ignores history. +//! - [`TpeSampler`] — model-based with cold start fallback. +//! - [`DESampler`] — stateful, population-based. +//! - [`GridSampler`] — deterministic, exhaustive search. pub mod bohb; #[cfg(feature = "cma-es")] @@ -280,6 +455,50 @@ impl PendingTrial { /// 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. +/// +/// # Implementing a custom sampler +/// +/// ``` +/// use optimizer::sampler::{Sampler, CompletedTrial}; +/// use optimizer::distribution::Distribution; +/// use optimizer::param::ParamValue; +/// +/// struct NoisySampler { +/// noise_scale: f64, +/// seed: u64, +/// } +/// +/// impl Sampler for NoisySampler { +/// fn sample( +/// &self, +/// distribution: &Distribution, +/// trial_id: u64, +/// history: &[CompletedTrial], +/// ) -> ParamValue { +/// // Find the best value seen so far, or fall back to the midpoint +/// match distribution { +/// Distribution::Float(fd) => { +/// let center = if history.is_empty() { +/// (fd.low + fd.high) / 2.0 +/// } else { +/// history.iter() +/// .filter_map(|t| t.params.values().next()) +/// .filter_map(|v| if let ParamValue::Float(f) = v { Some(*f) } else { None }) +/// .next() +/// .unwrap_or((fd.low + fd.high) / 2.0) +/// }; +/// let noise = (trial_id as f64 * 0.1).sin() * self.noise_scale; +/// ParamValue::Float(center + noise) +/// } +/// Distribution::Int(id) => ParamValue::Int((id.low + id.high) / 2), +/// Distribution::Categorical(cd) => ParamValue::Categorical(trial_id as usize % cd.n_choices), +/// } +/// } +/// } +/// ``` +/// +/// See the [module-level documentation](self) for a comprehensive guide +/// covering cold start handling, thread safety patterns, and testing. pub trait Sampler: Send + Sync { /// Samples a parameter value from the given distribution. ///