gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs(sampler,pruner): add implementation guides for custom samplers and pruners

Browse Source 
- 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)
This commit is contained in:
Manuel Raimann
2026-02-12 16:31:25 +01:00
parent 8d06d213db
commit cfa8426312
2 changed files with 322 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pruner/mod.rs
+103
View File
@@ -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<HashMap<u64, u64>>` 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<dyn Pruner>` 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<dyn Pruner>`, 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<HashMap<u64, (f64, u64)>>,
/// }
///
/// 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.
///
src/sampler/mod.rs
+219
View File
@@ -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<State>` 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(&param_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<dyn Sampler>`, 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<f64> = Study::with_sampler(Direction::Minimize, MySampler);
//! ```
//!
//! The sampler is wrapped in `Arc<dyn Sampler>` 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.
///