140 Commits

Author SHA1 Message Date
Manuel Raimann db192eb1cb chore: release v1.0.0 2026-02-13 10:36:16 +01:00
Manuel Raimann 80a7c27551 ci: publish optimizer-derive before optimizer to fix crates.io dependency 2026-02-13 10:36:16 +01:00
Manuel Raimann 86db6361d6 fix(sampler): clamp multivariate TPE candidates to parameter bounds
- Clamp KDE candidates to parameter bounds before evaluating l(x)/g(x),
  matching the univariate TPE behavior; without this, candidates scored
  well at out-of-bounds locations but became suboptimal when clamped
- Sort HashMap iterations by ParamId before consuming the seeded RNG to
  eliminate non-deterministic sampling caused by global ParamId counter
- Replace wall-clock timing assertion in async concurrency test with
  atomic max-active counter to avoid CI flakiness
- Gate unused HashSet import behind cfg(feature = "async")
2026-02-13 10:21:39 +01:00
Manuel 11a8534b38 fix: address 18 bugs found during codebase audit (#8)
* fix: address 18 bugs found during codebase audit

High severity:
- TPE/MOTPE: match parameters by exact distribution equality instead of
  flat-mapping over all param values, preventing cross-parameter mixing
- MultivariateTpeSampler: find_matching_param now uses search space
  distributions for exact matching instead of type+range heuristic
- JournalStorage: write_to_file no longer advances file_offset (left to
  refresh), both operations serialized under single io_lock mutex,
  refresh uses fetch_max and deduplicates by trial ID

Medium severity:
- NSGA-III: use actual Pareto front ranks for tournament selection
  instead of artificial cyclic indices
- sample_random: apply step quantization after log-scale sampling
- internal_bounds: return None for non-positive log-scale bounds
- SobolSampler: use per-trial dimension HashMap for concurrent safety
- JournalStorage refresh: protect with io_lock mutex, use fetch_max
- n_trials(): filter by TrialState::Complete as documented
- FloatParam: reject NaN/Infinity in validate()
- Pruners: assert n_min_trials >= 1, guard compute_percentile on empty
- Visualization: escape_js for importance chart parameter names

Low severity:
- save(): use peek_next_trial_id() from Storage trait
- csv_escape: handle carriage return per RFC 4180
- from_internal: use saturating arithmetic for stepped Int distributions
- BoolParam: bounds-check categorical index < 2
- min_max: skip NaN values with safe fallback

* ci: trigger CI on pull requests targeting any branch
2026-02-13 10:03:37 +01:00
Manuel Raimann 1cd18c16f9 fix: address review findings in grid sampler, parallel optimization, and CSV export
- Rename stale "GridSearchSampler" panic message to "GridSampler"
- Assert concurrency > 0 in optimize_parallel to prevent deadlock
- Fix inaccurate comment in CSV export (empty for all non-complete trials, not just pruned)
2026-02-12 17:57:29 +01:00
Manuel Raimann 9f2ce2a482 fix(docs): remove private intra-doc link from JournalStorage doc comment 2026-02-12 17:05:12 +01:00
Manuel Raimann 8c025b4525 fix(journal): canonicalize file paths in JournalStorage constructors
Resolve symlinks and `../` traversals best-effort in `new()` and
`open()`, falling back to the original path if canonicalization fails
(e.g. file doesn't exist yet).
2026-02-12 16:34:27 +01:00
Manuel Raimann cfa8426312 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)
2026-02-12 16:31:25 +01:00
Manuel Raimann 8d06d213db docs(journal): document file-size DoS potential in JournalStorage 2026-02-12 16:26:41 +01:00
Manuel Raimann 479554ffb2 fix(docs): convert JournalStorage intra-doc link to plain backtick text 2026-02-12 16:24:26 +01:00
Manuel Raimann fd5c5453e5 test(async): add concurrency verification tests for optimize_parallel
- Timing-based proof that trials run in parallel, not sequentially
- Atomic counter proof that multiple workers overlap simultaneously
- Panic propagation returns TaskError with panic message
- Partial failures with parallel path store only successful trials
2026-02-12 16:22:38 +01:00
Manuel Raimann 7c0cce160a test(stress): add ignored stress tests for large-scale scenarios
- 10k sequential trials with RandomSampler
- 128 parameters with TPE sampler
- 128 concurrent workers with optimize_parallel
- 5k trials with TPE and 32 parallel workers
2026-02-12 16:18:55 +01:00
Manuel Raimann 8394a747a5 test(journal): add corrupted and malicious JSONL file tests 2026-02-12 16:07:31 +01:00
Manuel Raimann 0e0e9ff38f test(sobol): add integration tests for SobolSampler via Study API 2026-02-12 16:02:46 +01:00
Manuel Raimann 12efbaabab fix(sampler): return None instead of panicking on parameter type mismatch in CompletedTrial::get() 2026-02-12 16:00:05 +01:00
Manuel Raimann 3d3dcb4c26 feat(error): mark Error enum as #[non_exhaustive] 2026-02-12 15:56:16 +01:00
Manuel Raimann cc50e1ad4b refactor: reduce #[allow] attributes and extract per-distribution helpers
- Remove stale #[allow(dead_code)] and #[allow(too_many_lines)] attributes
- Replace #[allow(dead_code)] with #[cfg(test)] for test-only methods
- Delete unused fill_remaining_independent() and ParamMeta.dist field
- Extract sample_float/int/categorical helpers from TpeSampler, MotpeSampler,
  and SamplingEngine to shorten match-heavy sample() methods
- Extract try_fit_kdes() helper to consolidate repeated fallback blocks
- Refactor sample_tpe_float/int to take &FloatDistribution/&IntDistribution
  instead of destructured args, removing #[allow(too_many_arguments)]
2026-02-12 15:54:29 +01:00
Manuel Raimann 38a20bbc42 fix(docs): resolve all broken rustdoc intra-doc links
Fix direct path references for in-scope items (GammaStrategy,
CompletedTrial) and convert feature-gated items to plain backtick
formatting so docs build cleanly without --all-features.
2026-02-12 15:47:53 +01:00
Manuel Raimann f0798ca58a refactor(study): move export_json() from persistence to export module 2026-02-12 15:42:34 +01:00
Manuel Raimann a502d7e1b8 refactor(tpe): split multivariate sampler into focused submodules
- Extract engine.rs (core sampling logic) and trials.rs (trial processing)
- Deduplicate TPE sampling functions by delegating to tpe::common
- Gate persistence.rs imports with #[cfg(feature = "serde")]
2026-02-12 15:38:23 +01:00
Manuel Raimann d80d2bb1fb refactor(sampler): extract shared utilities to reduce duplication
- Add sampler/common.rs with distribution helpers (internal_bounds,
  from_internal, to_internal, sample_random) used by 8 samplers
- Add sampler/tpe/common.rs with TPE sampling functions
  (sample_tpe_float, sample_tpe_int, sample_tpe_categorical)
  shared by TpeSampler, MultivariateTpeSampler, and MotpeSampler
- Remove ~940 lines of near-identical code across sampler modules
2026-02-12 15:25:27 +01:00
Manuel Raimann 4781107ede refactor(study): split monolithic study.rs into focused submodules
- mod.rs: core struct (pub(crate) fields), constructors, trial management
- builder.rs: StudyBuilder fluent API
- optimize.rs: sync optimization loop
- async_impl.rs: optimize_async/optimize_parallel (feature-gated)
- analysis.rs: best_trial, top_trials, param_importance, fanova
- export.rs: CSV, summary, Display, export_html
- persistence.rs: StudySnapshot, save/load, with_journal
- iter.rs: iter(), IntoIterator
2026-02-12 15:03:02 +01:00
Manuel Raimann 9ab94002d9 perf(tpe): eliminate unnecessary Vec allocations in KDE/TPE sampling
- In-place log transformation in sample_tpe_float (saves 2 allocs per log-scale float sample)
- Pass owned Vec<i64> to sample_tpe_int to reduce peak memory during int-to-float conversion
- Stack-allocate categorical count arrays for <=32 choices in sample_tpe_categorical
2026-02-12 14:51:03 +01:00
Manuel Raimann 1b158e4d1b perf: pre-allocate vectors in analysis methods
- Use Vec::with_capacity() in param_importance() and fanova_with_config()
  where iteration count is known or bounded
- Fix flaky MultivariateTpeSampler doctest by increasing trials from 30
  to 50
2026-02-12 14:46:46 +01:00
Manuel Raimann ddbbe294bc perf(journal): use incremental refresh instead of re-reading entire file
Track a byte offset so refresh() only parses lines appended since the
last read, reducing total I/O from O(n²) to O(n) over an n-trial run.
2026-02-12 14:44:19 +01:00
Manuel Raimann 5da63d8976 perf: reduce CompletedTrial cloning overhead
- Add Trial::into_completed() and into_multi_objective_trial() to move
  fields instead of cloning 5 HashMaps/Vecs per trial completion
- Fire after_trial callback before pushing to storage, eliminating the
  clone-from-storage pattern at all 4 call sites
- Optimize top_trials(n) to sort indices and clone only N trials instead
  of cloning all completed trials
- Remove unused set_complete/set_pruned methods
2026-02-12 14:44:19 +01:00
Manuel Raimann 74cf0643eb perf: replace RNG mutex with per-call seed derivation in stateless samplers
- Replace Mutex<fastrand::Rng> with a stored seed + MurmurHash3 mixer
  in RandomSampler, TpeSampler, and MotpeSampler so parallel workers
  no longer serialize on a shared lock
- Add mix_seed() and distribution_fingerprint() to rng_util for
  deterministic per-call RNG derivation from (seed, trial_id, distribution)
- Include an AtomicU64 call counter to disambiguate parameters that
  share the same distribution within a trial
2026-02-12 14:44:19 +01:00
Manuel Raimann baebbae64c docs: add pre-commit guidelines for formatting and testing 2026-02-12 14:44:19 +01:00
Manuel Raimann 750bbad0d8 test: assert pruned count unconditionally in summary test
Remove the conditional guard on pruned trial assertions since the test
explicitly prunes 2 trials. Assert the expected count first, then check
summary content unconditionally to catch regressions.
2026-02-12 14:44:19 +01:00
Manuel Raimann 271748700f perf(tpe): use quickselect instead of full sort in split_trials 2026-02-12 14:44:19 +01:00
Manuel Raimann c561caaa34 fix: validate NaN/Inf in deserialized trials during journal loading
- Add `CompletedTrial::validate()` checking all f64 fields are finite
- Call validate after deserializing each trial in journal storage
- Make `distribution` and `param` modules public (types already in public fields)
2026-02-12 14:44:19 +01:00
Manuel 04d822479e Update Cargo.toml
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-02-12 14:12:52 +01:00
Manuel Raimann 8713a93636 refactor: adjust indentation for better readability in pruning.rs 2026-02-12 13:59:14 +01:00
Manuel Raimann 3723f91849 refactor: update rustfmt configuration by commenting out unused options 2026-02-12 13:56:09 +01:00
Manuel Raimann 968eb67d07 fix: return error instead of panicking on out-of-bounds categorical index 2026-02-12 13:55:00 +01:00
Manuel 2d88526d8d Update examples/pruning.rs
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-02-12 13:52:26 +01:00
Manuel Raimann cc3937d4f7 Add LICENSE.md 2026-02-12 13:49:49 +01:00
Manuel Raimann 6f1c070b7d test: actually use NaN values in test_best_trial_with_nan_values 2026-02-12 13:38:54 +01:00
Manuel 12f7f25baa Update README.md
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-02-12 13:29:44 +01:00
Manuel Raimann e68a3788be refactor: move RandomMultiObjectiveSampler into sampler::random 2026-02-12 13:21:37 +01:00
Manuel Raimann d20f09c66a refactor: shorten sampler type names
- GridSearchSampler → GridSampler
- DifferentialEvolutionSampler → DESampler
- DifferentialEvolutionStrategy → DEStrategy
- DifferentialEvolutionSamplerBuilder → DESamplerBuilder
2026-02-12 13:19:06 +01:00
Manuel Raimann 47b5f9cec8 feat: unify optimize and optimize_with via blanket Objective impl
- Add blanket `impl Objective<V> for Fn(&mut Trial) -> Result<V, E>`
  so closures work directly with `optimize`
- Rewrite optimize, optimize_async, optimize_parallel to accept
  `impl Objective<V>` with before_trial/after_trial hooks
- Remove optimize_with, optimize_with_async, optimize_with_parallel
- Remove max_retries and retry logic from Objective trait
- Add explicit closure type annotations for HRTB inference
- Convert FnMut test closures to Fn via RefCell/Cell
2026-02-12 13:09:14 +01:00
Manuel Raimann c20a53dfba fix: bump minimum tokio version to 1.30 for JoinSet::len() 2026-02-12 12:52:02 +01:00
Manuel Raimann d81d1de4ff refactor(tests): split integration.rs into focused subfolders
- Delete 20 duplicate tests already covered by parameter_tests.rs
- Move 11 pure Trial unit tests into src/trial.rs
- Split remaining 84 integration tests into tests/study/ (9 modules)
- Group sampler tests into tests/sampler/ (7 modules)
- Group pruner tests into tests/pruner/ (2 modules)
2026-02-12 12:49:24 +01:00
Manuel Raimann 964b4d5749 feat: add Objective trait and unify optimize API
- Add `Objective<V>` trait with lifecycle hooks (`before_trial`,
  `after_trial`, `max_retries`) in new `src/objective.rs`
- Replace 14+ optimize variants with 6 methods: `optimize`,
  `optimize_with`, and async/parallel counterparts
- `optimize*` methods accept closures directly (FnMut for sync, Fn for
  async); `optimize_with*` methods accept `impl Objective<V>` for
  struct-based objectives with hooks and retries
- Remove `optimize_until`, `optimize_with_callback`,
  `optimize_with_retries`, `optimize_with_checkpoint`, and all
  deprecated `_with_sampler` methods
2026-02-12 12:49:24 +01:00
Manuel Raimann ee59c9cdd0 refactor(examples): split multi-concept examples into focused single-topic files
- Split pruning_and_callbacks into pruning and early_stopping
- Split advanced_features into async_parallel, journal_storage, ask_and_tell, multi_objective
- Each example now requires only its own feature flag
- Trim sampler_comparison winner logic and verbose header
- Update CI workflow and README to match new example names
2026-02-12 12:49:24 +01:00
Manuel Raimann 6a4b27c46f refactor: move all type re-exports out of crate root into their modules
- Sampler, pruner, storage, parameter, and multi_objective types are no
  longer re-exported at the crate root; access via module paths instead
  (e.g. optimizer::sampler::TpeSampler, optimizer::parameter::FloatParam)
- Prelude continues to re-export everything for convenience
- Add module-level pub use re-exports in sampler/mod.rs and parameter.rs
- Update derive macro to reference optimizer::parameter::Categorical
- Rename sampler::differential_evolution to sampler::de
- Fix all downstream imports in tests, benches, and doctests
- Fix all rustdoc intra-doc links to use explicit module paths
2026-02-12 12:49:24 +01:00
Manuel Raimann cc0dce5daf ci: add doc tests and update examples to match current codebase 2026-02-12 12:49:24 +01:00
Manuel Raimann 1f1e9d1779 refactor(docs): Documentation Overhaul 2026-02-12 12:49:24 +01:00
Manuel Raimann 3581187cd3 chore: release v0.9.1 2026-02-12 08:17:28 +01:00
Manuel Raimann 349ffcbed3 fix: update journal file handling to support reading and writing 2026-02-12 08:13:48 +01:00
Manuel Raimann 88e2ac0ff0 chore: release v0.9.0 2026-02-12 00:01:44 +01:00
Manuel Raimann d122582bf7 Revert "feat: add SQLite storage backend for multi-process optimization"
This reverts commit ed80c59795.

# Conflicts:
#	src/lib.rs
#	src/storage/sqlite.rs
2026-02-12 00:01:27 +01:00
Manuel Raimann 705687a42e feat: add NSGA-III and MOEA/D samplers for many-objective optimization
Extract shared evolutionary algorithm infrastructure (genetic operators,
candidate management, Das-Dennis reference points) from NSGA-II into a
new genetic.rs module, then build two new multi-objective samplers on top:

- NSGA-III: reference-point-based niching for well-distributed fronts
  on 3+ objective problems (Das-Dennis structured points, normalization,
  perpendicular distance association, niching selection)

- MOEA/D: decomposition-based optimization with three scalarization
  methods (Tchebycheff, WeightedSum, PBI), weight-vector neighborhoods,
  and neighborhood-based mating selection

Both implement MultiObjectiveSampler with builder pattern, seeded RNG,
and SBX crossover / polynomial mutation via the shared genetic module.
2026-02-11 23:53:20 +01:00
Manuel Raimann 03deedd775 feat: add Zlib license to deny.toml 2026-02-11 23:40:22 +01:00
Manuel Raimann 047a788afd fix: remove unnecessary feature flag for visualization example in CI 2026-02-11 23:39:43 +01:00
Manuel Raimann d2e36ec304 refactor: move next_trial_id counter from Study into Storage trait 2026-02-11 23:32:15 +01:00
Manuel Raimann d14f4e0792 feat: add StudyBuilder for fluent study construction 2026-02-11 23:23:56 +01:00
Manuel Raimann ed80c59795 feat: add SQLite storage backend for multi-process optimization 2026-02-11 23:21:47 +01:00
Manuel Raimann ed1e9e31da fix: use atomic counter for temp paths in journal tests
On Windows, SystemTime::now().as_nanos() has ~15ms resolution,
causing parallel tests to generate identical file paths and
corrupt each other's data.
2026-02-11 23:17:09 +01:00
Manuel Raimann 36c6262dbc refactor: remove visualization and fanova feature flags
Always build visualization and fanova code unconditionally.
2026-02-11 23:06:28 +01:00
Manuel Raimann 0c295b4bc7 ci: simplify test command by removing feature matrix 2026-02-11 23:03:38 +01:00
Manuel Raimann 0a6f2345a8 feat: add Storage trait and JSONL journal backend
Replace the internal Vec<CompletedTrial<V>> with a pluggable Storage<V>
trait. MemoryStorage is the default (no behavior change for existing
users). Behind the `journal` feature flag, JournalStorage persists
trials to a JSONL file with fs2 file locking for multi-process safety.

- Storage<V> trait with push(), trials_arc(), refresh() methods
- MemoryStorage<V> wraps Arc<RwLock<Vec<CompletedTrial<V>>>>
- JournalStorage<V> appends JSON lines with exclusive file locks
- Study::with_sampler_and_storage() general constructor
- Study::with_journal() convenience constructor (journal feature)
- Refresh from storage on create_trial() for multi-process discovery
- MSRV bumped to 1.89
2026-02-11 22:58:38 +01:00
Manuel Raimann 24a0bdd473 feat: add Differential Evolution sampler
Population-based optimizer with mutation + crossover. Supports three
strategies (Rand1, Best1, CurrentToBest1), configurable F and CR,
and auto-sized populations. No extra dependencies needed.
2026-02-11 22:04:23 +01:00
Manuel Raimann 8239cc58a1 refactor: replace rand 0.10 with fastrand 2.3
fastrand is smaller, faster, and has no dependencies. Add rng_util
helper for f64 range generation since fastrand lacks a built-in
equivalent. Migrate all samplers, KDE modules, and fANOVA to use
fastrand's concrete Rng type instead of rand's trait-based generics.
2026-02-11 21:54:34 +01:00
Manuel Raimann 906e5296de fix: seed importance tests to eliminate flakiness
The no_effect_parameter test (and others) used unseeded random sampling,
causing spurious correlations between the noise parameter and the
objective in unlucky runs. Seed all importance tests with
RandomSampler::with_seed(42) for deterministic results.
2026-02-11 21:30:21 +01:00
Manuel Raimann 947701e83f feat: add Gaussian Process sampler with Expected Improvement
Implement a classical Bayesian optimization sampler using a GP surrogate
with Matérn 5/2 kernel and Expected Improvement acquisition function.
Feature-gated behind `gp = ["dep:nalgebra"]`.
2026-02-11 21:27:40 +01:00
Manuel Raimann b3bf4ccb71 chore: parallelize feature-check CI with 4 matrix partitions
Use cargo hack's --partition flag to split the feature powerset across
4 parallel jobs, reducing wall-clock time for the feature-check step.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 21:13:39 +01:00
Manuel Raimann 7bf64d6cf6 chore: add bench-check to CI workflow 2026-02-11 21:06:56 +01:00
Manuel Raimann 624283c766 chore: release v0.8.1 2026-02-11 21:05:12 +01:00
Manuel Raimann 64c7aa2448 fix: resolve broken rustdoc intra-doc links for Error::NoCompletedTrials 2026-02-11 21:04:40 +01:00
Manuel Raimann abb1b4c229 chore: release v0.8.0 2026-02-11 20:53:30 +01:00
Manuel Raimann dff58340a4 feat: add fANOVA parameter importance via random forest
Implement functional ANOVA decomposition behind the `fanova` feature
flag. A self-contained random forest is trained on trial data, then
marginal predictions are used to compute per-parameter main effects
and pairwise interaction effects, normalized to sum to 1.0.

Adds Study::fanova() / fanova_with_config(), FanovaResult, and
FanovaConfig. Includes unit and integration tests covering dominant
parameters, interaction detection, consistency with correlation-based
importance, and error handling.
2026-02-11 20:51:24 +01:00
Manuel Raimann e359392a00 feat: add HTML visualization reports with Plotly.js charts
Add a `visualization` feature flag that generates self-contained HTML
reports with interactive Plotly.js charts for offline visualization of
optimization results. Charts include optimization history, slice plots,
parallel coordinates, parameter importance, trial timeline, and
intermediate values (learning curves).
2026-02-11 20:12:35 +01:00
Manuel Raimann 0e54356345 fix: rename variable to pass typos check 2026-02-11 20:04:51 +01:00
Manuel Raimann d851aad548 feat: add CSV and JSON data export for visualization
Add to_csv(), export_csv(), and export_json() methods to Study for
exporting trial data to external visualization tools. CSV export works
without extra dependencies; JSON export requires the serde feature.
2026-02-11 20:02:35 +01:00
Manuel Raimann 95402dc9b6 feat: add Pareto front analysis utilities
Expose public pareto module with hypervolume indicator, non-dominated
sorting, pareto front filtering, and crowding distance functions.
2026-02-11 19:58:07 +01:00
Manuel Raimann f873722763 ci: add benchmark CI jobs and missing examples
Add bench-check job (push): smoke-tests benchmarks compile and run.
Add bench-compare job (PRs): runs benchmarks on base and head commits,
compares with critcmp, and posts a comparison comment on the PR.
Add missing examples (benchmark_convergence, parameter_api) to the
examples job.
2026-02-11 19:50:24 +01:00
Manuel Raimann ba31df69c7 feat: add Multi-Objective TPE (MOTPE) sampler
Extend TPE to handle multi-objective optimization using Pareto-based
splitting. MOTPE uses non-dominated sorting to define "good" (front 0)
vs "bad" (dominated) regions for the KDE models, replacing the
single-objective gamma-based split.
2026-02-11 19:43:01 +01:00
Manuel Raimann bcc4549e66 feat: add multi-objective optimization with NSGA-II
Add MultiObjectiveStudy for optimizing multiple objectives simultaneously,
backed by NSGA-II (Non-dominated Sorting Genetic Algorithm II) with SBX
crossover, polynomial mutation, and constraint-aware dominance.

New public API:
- MultiObjectiveStudy with optimize(), pareto_front(), ask()/tell()
- MultiObjectiveTrial with get(), is_feasible(), user attributes
- MultiObjectiveSampler trait for custom MO samplers
- Nsga2Sampler with builder for population size, crossover/mutation params
- ObjectiveDimensionMismatch error variant
2026-02-11 19:35:18 +01:00
Manuel Raimann 7d79111d81 chore: release v0.7.2 2026-02-11 19:23:45 +01:00
Manuel Raimann a9ea097377 chore: add advisory ignore for unmaintained transitive dependency 2026-02-11 19:23:31 +01:00
Manuel Raimann 7d70220da0 chore: update nalgebra dependency to version 0.34 2026-02-11 19:18:42 +01:00
Manuel Raimann b5c0d6353b chore: release v0.7.1 2026-02-11 19:18:14 +01:00
Manuel Raimann 08fde0c707 chore: update tracing dependency to version 0.1.29 2026-02-11 19:17:53 +01:00
Manuel Raimann 147a64708d chore: release v0.7.0 2026-02-11 19:05:52 +01:00
Manuel Raimann b36137bc96 feat: add benchmark suite with criterion and standard test functions
Add micro-benchmarks for TPE, Random, and Grid samplers with varying
history sizes, end-to-end optimization benchmarks on Sphere and
Rosenbrock, and a convergence tracking example that outputs CSV data.
Includes 6 standard test functions (Sphere, Rosenbrock, Rastrigin,
Ackley, Branin, Hartmann6) with correctness tests at known optima.
2026-02-11 19:05:33 +01:00
Manuel Raimann f3941a4b3e chore: remove unused proc-macro2 dependency from optimizer-derive 2026-02-11 18:57:30 +01:00
Manuel Raimann 0cae3b4227 feat: add parameter importance via Spearman rank correlation
Compute per-parameter importance scores by measuring the absolute
Spearman rank correlation between each parameter's values and the
objective across completed trials. Scores are normalized to sum to 1.0
and returned sorted by descending importance.
2026-02-11 18:56:19 +01:00
Manuel Raimann 09480a973b feat: add constraint handling with feasibility-aware trial ranking
Add constraint support so that optimization problems with constraints
(e.g., "model size < 100MB") prefer feasible solutions. Convention:
constraint value <= 0.0 means feasible.

- Add constraint_values field to Trial with set_constraints/getter
- Add constraints field to CompletedTrial with is_feasible() method
- Propagate constraints through complete_trial/prune_trial
- Make best_trial() and top_trials() constraint-aware: feasible trials
  rank above infeasible, infeasible ranked by total violation
2026-02-11 18:49:08 +01:00
Manuel Raimann aa9734587a feat: add BOHB sampler for budget-aware Bayesian optimization
BOHB combines TPE's model-guided sampling with Hyperband's budget-aware
evaluation by conditioning the TPE model on trials at specific budget
levels, producing better-calibrated parameter proposals.
2026-02-11 18:42:09 +01:00
Manuel Raimann a1dd6d393c feat: add CMA-ES sampler behind cma-es feature flag
Implement CMA-ES (Covariance Matrix Adaptation Evolution Strategy) as a
new sampler for continuous optimization. Uses nalgebra for matrix
operations and implements the full Hansen 2016 algorithm: mean update,
evolution paths, covariance matrix adaptation, and cumulative step-size
adaptation.

Key features:
- Builder API with configurable sigma0, population_size, and seed
- Three-phase state machine: discovery, active sampling, generation update
- Rejection sampling with bounds clipping fallback
- Log-scale and step parameter support via internal-space mapping
- Categorical parameters handled via random sampling fallback
- Numerical robustness: eigenvalue clamping, symmetry enforcement
2026-02-11 18:31:12 +01:00
Manuel Raimann a53ba4f472 feat: add Sobol quasi-random sampler behind sobol feature flag
Add SobolSampler using Owen-scrambled Sobol sequences (sobol_burley
crate) for better uniform coverage of the parameter space compared to
random sampling. Supports all distribution types including log-scale
and stepped parameters.
2026-02-11 17:59:48 +01:00
Manuel Raimann 9b4a8321b5 feat: implement IntoIterator for &Study and add iter() method
Enables idiomatic `for trial in &study` iteration over completed trials.
2026-02-11 17:54:09 +01:00
Manuel Raimann db0314a1d1 feat: add optimize_with_retries() for automatic retry of failed trials 2026-02-11 17:51:36 +01:00
Manuel Raimann df5fcedecf feat: add tracing integration behind tracing feature flag
Add optional structured logging via the tracing crate:
- info-level spans on all optimize* methods (direction, n_trials/duration)
- info events for trial completed, new best value, trial pruned
- debug events for trial failed and parameter sampled
- Zero overhead when the feature is disabled
2026-02-11 17:47:13 +01:00
Manuel Raimann 5061df9876 feat: add optimize_with_checkpoint() and atomic save writes
Adds a convenience method that combines optimize_with_callback + save
to automatically checkpoint every N trials. Also makes save() use
atomic writes (write to temp file, then rename) to prevent corrupt
files on crash.
2026-02-11 17:37:57 +01:00
Manuel Raimann 5fe0a75f78 feat: add serde serialization support behind serde feature flag
Add Serialize/Deserialize derives to all public data types (ParamValue,
Distribution, Direction, TrialState, ParamId, AttrValue, CompletedTrial)
gated behind a `serde` feature flag. Introduce StudySnapshot struct and
Study::save()/Study::load() for persisting and restoring study state as
human-readable JSON.
2026-02-11 17:33:48 +01:00
Manuel Raimann dcf3403261 feat: add Study::summary() and Display impl
Add a human-readable summary method and Display trait for Study<V>
where V: Display. The summary shows optimization direction, trial
counts (with complete/pruned breakdown), best value, and best
parameters with their labels.

Also derive PartialOrd + Ord on ParamId for deterministic parameter
ordering in summary output.
2026-02-11 17:27:43 +01:00
Manuel Raimann bb79d98027 chore: release v0.6.0 2026-02-11 17:24:10 +01:00
Manuel Raimann f4b0631178 feat: add enqueue_trial() for pre-specified parameter evaluation
Add Study::enqueue() to push specific parameter configurations onto a
FIFO queue. The next call to ask()/create_trial() or the next iteration
in optimize() pops the front entry and injects it into the trial so
that suggest_param() returns the pre-filled value instead of sampling.

Parameters missing from an enqueued map fall back to normal sampling,
and once the queue is drained regular sampler-driven trials resume.
2026-02-11 17:23:25 +01:00
Manuel Raimann 52f3c074dc feat: add trial user attributes for logging and analysis
Add AttrValue enum (Float, Int, String, Bool) with From impls and
set_user_attr/user_attr/user_attrs methods on both Trial and
CompletedTrial. Attributes set during optimization are propagated
through complete_trial and prune_trial. AttrValue is re-exported
at crate root and in the prelude.
2026-02-11 17:18:17 +01:00
Manuel Raimann f67188e3a7 feat: add ask() and tell() methods for ask-and-tell interface 2026-02-11 17:11:38 +01:00
Manuel Raimann 72da883add feat: add Study::top_trials(n) for retrieving best N trials
Returns the top N completed trials sorted by objective value,
respecting the study's optimization direction. Pruned and failed
trials are excluded.
2026-02-11 17:08:10 +01:00
Manuel Raimann bc278d83a3 feat: add timeout-based optimization methods (optimize_until)
Add duration-based variants of all optimization methods that run trials
until a wall-clock deadline rather than for a fixed trial count:

- optimize_until(duration, objective)
- optimize_until_with_callback(duration, objective, callback)
- optimize_until_async(duration, objective)
- optimize_until_parallel(duration, concurrency, objective)
2026-02-11 17:06:08 +01:00
Manuel Raimann c586650df4 feat: add From<RangeInclusive> for FloatParam and IntParam 2026-02-11 17:02:18 +01:00
Manuel Raimann cc78a92ed6 feat: add Study::minimize() and Study::maximize() constructor shortcuts 2026-02-11 17:00:35 +01:00
Manuel Raimann 12c35e7cb4 feat: add WilcoxonPruner for statistics-based trial pruning 2026-02-11 16:59:05 +01:00
Manuel Raimann a885d10436 feat: add HyperbandPruner for multi-bracket trial pruning 2026-02-11 16:52:59 +01:00
Manuel Raimann 41150057d6 feat: add SuccessiveHalvingPruner for SHA-based trial pruning 2026-02-11 16:47:57 +01:00
Manuel Raimann a522fb2f34 feat: add PatientPruner for patience-based trial pruning
Wraps any inner Pruner and requires it to recommend pruning for N
consecutive steps before actually pruning. Useful for noisy
intermediate values where a single bad step shouldn't trigger pruning.
2026-02-11 16:40:00 +01:00
Manuel Raimann ec04757740 feat: add PercentilePruner for configurable percentile-based trial pruning
Generalizes MedianPruner with a configurable percentile threshold.
MedianPruner now reuses the shared compute_percentile function at 50%.
2026-02-11 16:36:57 +01:00
Manuel Raimann 432c74b927 feat: add MedianPruner for statistics-based trial pruning
Prunes trials whose intermediate values are worse than the median
of completed trials at the same step. Supports configurable warmup
steps and minimum trial count before pruning activates.
2026-02-11 16:32:03 +01:00
Manuel Raimann 8199ba6b43 feat: add ThresholdPruner for fixed-bound trial pruning 2026-02-11 16:28:22 +01:00
Manuel Raimann abe15bdd7b feat: add TrialPruned error variant and Pruned trial state
- Add `Pruned` variant to `TrialState`
- Add `Error::TrialPruned` variant and standalone `TrialPruned` struct
  with `From<TrialPruned> for Error` for ergonomic `?` usage
- Add `state` field to `CompletedTrial` (defaults to `Complete`)
- Add `Study::prune_trial()` and `Study::n_pruned_trials()`
- `optimize()` and `optimize_with_callback()` detect `TrialPruned`
  errors via Any downcasting and record pruned trials instead of
  failing them
- `best_trial()` / `best_value()` now filter to only `Complete` trials
- Re-export `TrialPruned` from crate root and prelude
2026-02-11 16:24:57 +01:00
Manuel Raimann 4d8af3242b feat: add intermediate values and pruner integration to Trial
- Add report(step, value) and should_prune() methods to Trial
- Store intermediate_values in Trial, copied to CompletedTrial on completion
- Pass pruner through trial factory so trials can consult it
- Rebuild trial factory when pruner is changed via set_pruner()
2026-02-11 16:08:25 +01:00
Manuel Raimann 2651e61b2c feat: add Pruner trait and NopPruner default implementation
Introduce the pruning extension point for the trial pruning system.
The Pruner trait allows deciding whether to stop a trial early based
on intermediate values. NopPruner (never prunes) is the default.
2026-02-11 15:46:56 +01:00
Manuel Raimann d88280c152 chore: release v0.5.1 2026-02-10 09:30:58 +01:00
Manuel Raimann 0e8fc82201 refactor: update random number generation to use rand::make_rng() and upgrade rand to version 0.10 2026-02-10 09:30:39 +01:00
Manuel Raimann d7ad3f60db chore: release v0.5.0 2026-02-06 18:55:06 +01:00
Manuel Raimann 5dc81fa0ab Implement Parameters API
- Add `.name()` builder method on all 5 parameter types for custom labels
- Add `CompletedTrial::get(&param)` for typed parameter access
- Add `Display` impl on `ParamValue`
- Add prelude module at `optimizer::prelude::*`
- Shadow `_with_sampler` methods on `Study<f64>` so `optimize()` auto-uses
  the configured sampler; deprecate `_with_sampler` variants
- Use runtime `Any` downcasting with `trial_factory` to avoid E0592
- Update all examples and tests to use the new API

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-06 18:54:55 +01:00
Manuel Raimann 55e50e6afb Implement Parameters API 2026-02-06 17:15:47 +01:00
Manuel Raimann 2fef1ab3c3 docs: add missing derive feature flag to README
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-06 16:51:13 +01:00
Manuel Raimann aa1718fed7 docs: update quick start example to use typed Parameter API
The old example used the removed string-based trial.suggest_float()
API. Updated to use FloatParam::new() with param.suggest(trial).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-06 16:50:59 +01:00
Manuel Raimann f9227df96e docs: update features list with boolean, enum, and derive support
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-06 16:50:43 +01:00
Manuel Raimann 02ec9498c4 refactor: reorganize sampler module and update imports 2026-02-04 13:14:22 +01:00
Manuel Raimann 183a78c09e chore: remove log dependency from Cargo.toml 2026-02-02 17:33:03 +01:00
Manuel Raimann ce817c1f38 feat: add 'consts' to default extended words in typos configuration 2026-02-02 17:32:20 +01:00
Manuel Raimann b9f8844e60 chore: release v0.4.0 2026-02-02 17:31:35 +01:00
Manuel Raimann d1c8925fc3 refactor: remove warn_independent_sampling option from MultivariateTpeSampler 2026-02-02 17:31:25 +01:00
Manuel Raimann 4ce4311c68 Implement Multivariant TPE 2026-02-02 17:25:50 +01:00
Manuel Raimann 293b992e75 feat: implement gamma strategies for TPE sampler with examples 2026-02-02 13:40:33 +01:00
Manuel Raimann af8a9f7638 feat: add examples for async API parameter optimization and ML hyperparameter tuning 2026-01-31 11:56:50 +01:00
Manuel Raimann b9a90ebec3 Format 2026-01-31 11:56:37 +01:00
Manuel Raimann 24e921b7b9 style: format suggest method for improved readability 2026-01-31 11:45:58 +01:00
Manuel Raimann 6363540269 fix: handle end value in suggest method to avoid panic on underflow 2026-01-31 11:44:33 +01:00
Manuel Raimann 9f3c0f0d3b chore: release v0.3.1 2026-01-31 11:34:36 +01:00
Manuel Raimann 747e1928b9 feat: add tests for suggest_bool and suggest_range methods 2026-01-31 11:34:36 +01:00
Manuel Raimann 15a43d9485 feat: add SuggestableRange trait and suggest_range method for parameter suggestion from ranges 2026-01-31 11:34:36 +01:00
Manuel Raimann 1091429697 feat: add suggest_bool method for boolean parameter suggestion 2026-01-31 11:34:36 +01:00
124 changed files with 36023 additions and 3836 deletions
+101
View File
@@ -0,0 +1,101 @@
---
allowed-tools: Bash(git add:*), Bash(git status:*), Bash(git commit:*), Bash(git diff:*), Bash(git log:*), Bash(git config:*), Bash(git push:*), Bash(git branch:*), Bash(git rev-parse:*), Bash(gh pr:*)
description: Create a git commit, push, and open a PR
---
## Context
- Current git status: !`git status`
- Current git diff (staged and unstaged changes): !`git diff HEAD`
- Current branch: !`git branch --show-current`
- Remote tracking: !`git rev-parse --abbrev-ref --symbolic-full-name @{u} 2>/dev/null || echo "no upstream"`
- Recent commits (for style reference): !`git log --oneline -10`
- Author commits (for style reference): !`git log --author="$(git config user.email)" --oneline -10`
- All commits on this branch not on main: !`git log --oneline main..HEAD 2>/dev/null`
- User input (optional, can specify target branch): $ARGUMENTS
## Your task
Based on the above changes, create a git commit, push, and open a pull request.
1. If there are no staged changes, stage only the relevant changed files by name (never use `git add -A` or `git add .`).
2. Do NOT commit files that likely contain secrets (.env, credentials.json, etc).
3. Write a commit message following the guidelines below, then commit.
4. Push to the remote. Use `git push -u origin <branch>` if there is no upstream set.
5. Create a pull request using `gh pr create` targeting the base branch:
- If the user specified a target branch in their input, use that.
- Otherwise, default to `main`.
6. The PR title should match the commit title (or summarize all branch commits if multiple).
7. Write the PR body following the PR format below.
8. Return the PR URL when done.
## Commit message guidelines
- Use imperative mood (e.g., "add feature" not "added feature")
- First line: brief summary, max 72 characters
- Focus on the "why" and "what", not the "how"
- Be specific but concise
- Many commits do not need a body if the title is self-explanatory
- Litmus test: "Would a developer understand this commit from the title + diff?" If yes, skip the body.
- Do NOT include "Generated with ...", "Co-Authored-By ...", or any AI attribution
### Conventional commit prefixes
Match the prefix to the nature of the change. These must align with `cliff.toml` commit parsers so they appear correctly in the changelog.
**Changelog: "Added"**
- `feat:` / `feat(scope):` — new feature
**Changelog: "Fixed"**
- `fix:` / `fix(scope):` — bug fix
**Changelog: "Changed"**
- `refactor:` / `refactor(scope):` — code restructuring without behavior change
- `perf:` / `perf(scope):` — performance improvement
- `docs:` / `docs(scope):` — documentation changes (README, guides, etc.)
- `style:` / `style(scope):` — formatting only (no logic change)
- `chore:` / `chore(scope):` — tooling, deps, config, scripts, AI config (.claude/, CLAUDE.md)
**Changelog: "Removed"**
- `revert:` / `revert(scope):` — revert a previous commit
**Excluded from changelog**
- `test:` / `test(scope):` — adding or updating tests
- `ci:` / `ci(scope):` — CI/CD pipeline changes
### Body rules
- When a body is needed (multiple important things in one commit), use bullet points.
- Body is separated from the title by a blank line.
- One bullet point per concept/change.
- Don't explain obvious things like "added unit tests for X".
### Commit format
Always pass the commit message via a HEREDOC:
```
git commit -m "$(cat <<'EOF'
<title line>
<optional body>
EOF
)"
```
## PR format
Use a HEREDOC for the body. Look at ALL commits on the branch (not just the latest) to write the summary.
```
gh pr create --title "<title>" --base <target-branch> --body "$(cat <<'EOF'
## Summary
<1-3 bullet points covering all branch commits>
EOF
)"
```
## Before you commit
1. Run `cargo +nightly fmt --all && cargo +nightly clippy --all-features --all-targets --fix --allow-dirty` to auto-format and fix lint issues.
2. Run all tests and doctests to ensure nothing is broken.
+82
View File
@@ -0,0 +1,82 @@
---
allowed-tools: Bash(git add:*), Bash(git status:*), Bash(git commit:*), Bash(git diff:*), Bash(git log:*), Bash(git config:*), Bash(git push:*), Bash(git branch:*), Bash(git rev-parse:*)
description: Create a git commit and push to remote
---
## Context
- Current git status: !`git status`
- Current git diff (staged and unstaged changes): !`git diff HEAD`
- Current branch: !`git branch --show-current`
- Remote tracking: !`git rev-parse --abbrev-ref --symbolic-full-name @{u} 2>/dev/null || echo "no upstream"`
- Recent commits (for style reference): !`git log --oneline -10`
- Author commits (for style reference): !`git log --author="$(git config user.email)" --oneline -10`
- User input (optional): $ARGUMENTS
## Your task
Based on the above changes, create a single git commit and push it to the remote.
1. If there are no staged changes, stage only the relevant changed files by name (never use `git add -A` or `git add .`).
2. Do NOT commit files that likely contain secrets (.env, credentials.json, etc).
3. Write a commit message following the guidelines below, then commit.
4. Push to the remote. Use `git push -u origin <branch>` if there is no upstream set.
## Commit message guidelines
- Use imperative mood (e.g., "add feature" not "added feature")
- First line: brief summary, max 72 characters
- Focus on the "why" and "what", not the "how"
- Be specific but concise
- Many commits do not need a body if the title is self-explanatory
- Litmus test: "Would a developer understand this commit from the title + diff?" If yes, skip the body.
- Do NOT include "Generated with ...", "Co-Authored-By ...", or any AI attribution
### Conventional commit prefixes
Match the prefix to the nature of the change. These must align with `cliff.toml` commit parsers so they appear correctly in the changelog.
**Changelog: "Added"**
- `feat:` / `feat(scope):` — new feature
**Changelog: "Fixed"**
- `fix:` / `fix(scope):` — bug fix
**Changelog: "Changed"**
- `refactor:` / `refactor(scope):` — code restructuring without behavior change
- `perf:` / `perf(scope):` — performance improvement
- `docs:` / `docs(scope):` — documentation changes (README, guides, etc.)
- `style:` / `style(scope):` — formatting only (no logic change)
- `chore:` / `chore(scope):` — tooling, deps, config, scripts, AI config (.claude/, CLAUDE.md)
**Changelog: "Removed"**
- `revert:` / `revert(scope):` — revert a previous commit
**Excluded from changelog**
- `test:` / `test(scope):` — adding or updating tests
- `ci:` / `ci(scope):` — CI/CD pipeline changes
### Body rules
- When a body is needed (multiple important things in one commit), use bullet points.
- Body is separated from the title by a blank line.
- One bullet point per concept/change.
- Don't explain obvious things like "added unit tests for X".
### Commit format
Always pass the commit message via a HEREDOC:
```
git commit -m "$(cat <<'EOF'
<title line>
<optional body>
EOF
)"
```
## Before you commit
1. Run `cargo +nightly fmt --all && cargo +nightly clippy --all-features --all-targets --fix --allow-dirty` to auto-format and fix lint issues.
2. Run all tests and doctests to ensure nothing is broken.
+81
View File
@@ -0,0 +1,81 @@
---
allowed-tools: Bash(git add:*), Bash(git status:*), Bash(git commit:*), Bash(git diff:*), Bash(git log:*), Bash(git config:*), Bash(git branch:*)
description: Create a git commit
---
## Context
- Current git status: !`git status`
- Current git diff (staged and unstaged changes): !`git diff HEAD`
- Current branch: !`git branch --show-current`
- Recent commits (for style reference): !`git log --oneline -10`
- Author commits (for style reference): !`git log --author="$(git config user.email)" --oneline -10`
- User input (optional): $ARGUMENTS
## Your task
Based on the above changes, create a single git commit.
1. If there are no staged changes, stage only the relevant changed files by name (never use `git add -A` or `git add .`).
2. Do NOT commit files that likely contain secrets (.env, credentials.json, etc).
3. Write a commit message following the guidelines below, then commit.
4. Do NOT push to remote.
## Commit message guidelines
- Use imperative mood (e.g., "add feature" not "added feature")
- First line: brief summary, max 72 characters
- Focus on the "why" and "what", not the "how"
- Be specific but concise
- Many commits do not need a body if the title is self-explanatory
- Litmus test: "Would a developer understand this commit from the title + diff?" If yes, skip the body.
- Do NOT include "Generated with ...", "Co-Authored-By ...", or any AI attribution
### Conventional commit prefixes
Match the prefix to the nature of the change. These must align with `cliff.toml` commit parsers so they appear correctly in the changelog.
**Changelog: "Added"**
- `feat:` / `feat(scope):` — new feature
**Changelog: "Fixed"**
- `fix:` / `fix(scope):` — bug fix
**Changelog: "Changed"**
- `refactor:` / `refactor(scope):` — code restructuring without behavior change
- `perf:` / `perf(scope):` — performance improvement
- `docs:` / `docs(scope):` — documentation changes (README, guides, etc.)
- `style:` / `style(scope):` — formatting only (no logic change)
- `chore:` / `chore(scope):` — tooling, deps, config, scripts, AI config (.claude/, CLAUDE.md)
**Changelog: "Removed"**
- `revert:` / `revert(scope):` — revert a previous commit
**Excluded from changelog**
- `test:` / `test(scope):` — adding or updating tests
- `ci:` / `ci(scope):` — CI/CD pipeline changes
### Body rules
- When a body is needed (multiple important things in one commit), use bullet points.
- Body is separated from the title by a blank line.
- One bullet point per concept/change.
- Don't explain obvious things like "added unit tests for X".
### Commit format
Always pass the commit message via a HEREDOC:
```
git commit -m "$(cat <<'EOF'
<title line>
<optional body>
EOF
)"
```
## Before you commit
1. Run `cargo +nightly fmt --all && cargo +nightly clippy --all-features --all-targets --fix --allow-dirty` to auto-format and fix lint issues.
2. Run all tests and doctests to ensure nothing is broken.
+167 -19
View File
@@ -4,10 +4,10 @@ on:
push:
branches: [main, master]
pull_request:
branches: [main, master]
permissions:
contents: read
pull-requests: write
env:
CARGO_TERM_COLOR: always
@@ -41,11 +41,6 @@ jobs:
test:
name: Test
runs-on: ubuntu-latest
strategy:
matrix:
features:
- ""
- "async"
steps:
- uses: actions/checkout@v6
- name: Install Rust
@@ -54,12 +49,38 @@ jobs:
rustup update stable
- uses: Swatinem/rust-cache@v2
- name: Run tests
run: cargo test --verbose --all-features --all-targets
- name: Run doc tests
run: cargo test --verbose --all-features --doc
examples:
name: Examples
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- name: Install Rust
run: |
if [ -z "${{ matrix.features }}" ]; then
cargo test --verbose
else
cargo test --verbose --features "${{ matrix.features }}"
fi
rustup override set stable
rustup update stable
- uses: Swatinem/rust-cache@v2
- name: Run basic_optimization
run: cargo run --example basic_optimization
- name: Run parameter_types
run: cargo run --example parameter_types --features derive
- name: Run sampler_comparison
run: cargo run --example sampler_comparison
- name: Run pruning
run: cargo run --example pruning
- name: Run early_stopping
run: cargo run --example early_stopping
- name: Run async_parallel
run: cargo run --example async_parallel --features async
- name: Run journal_storage
run: cargo run --example journal_storage --features journal
- name: Run ask_and_tell
run: cargo run --example ask_and_tell
- name: Run multi_objective
run: cargo run --example multi_objective
docs:
name: Docs
@@ -76,8 +97,12 @@ jobs:
RUSTDOCFLAGS: -D warnings
feature-check:
name: Feature Combinations
name: Feature Combinations (${{ matrix.partition }}/4)
runs-on: ubuntu-latest
strategy:
fail-fast: false
matrix:
partition: [1, 2, 3, 4]
steps:
- uses: actions/checkout@v6
- name: Install Rust
@@ -88,7 +113,7 @@ jobs:
uses: taiki-e/install-action@cargo-hack
- uses: Swatinem/rust-cache@v2
- name: Check feature powerset
run: cargo hack check --feature-powerset --no-dev-deps
run: cargo hack check --feature-powerset --no-dev-deps --partition ${{ matrix.partition }}/4
coverage:
name: Coverage
@@ -113,14 +138,14 @@ jobs:
fail_ci_if_error: true
msrv:
name: MSRV (1.88)
name: MSRV (1.89)
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- name: Install Rust 1.88
- name: Install Rust 1.89
run: |
rustup override set 1.88
rustup update 1.88
rustup override set 1.89
rustup update 1.89
- uses: Swatinem/rust-cache@v2
- name: Check compilation
run: cargo check --all-features
@@ -235,6 +260,99 @@ jobs:
- name: Typos Check
run: typos src/
bench-check:
name: Benchmark Smoke Test
if: github.event_name == 'push'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- name: Install Rust
run: |
rustup override set stable
rustup update stable
- uses: Swatinem/rust-cache@v2
- name: Run benchmarks (smoke test)
run: cargo bench --all-features --bench samplers --bench optimization
bench-compare:
name: Benchmark Comparison
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
steps:
- name: Install Rust
run: |
rustup override set stable
rustup update stable
- name: Install critcmp
run: cargo install critcmp
# --- Run benchmarks on the BASE (target) branch ---
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.base.sha }}
clean: false
- uses: Swatinem/rust-cache@v2
with:
key: bench-base
- name: Bench baseline
run: cargo bench --all-features --bench samplers --bench optimization -- --save-baseline base
# --- Run benchmarks on the HEAD (PR) branch ---
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha }}
clean: false
- uses: Swatinem/rust-cache@v2
with:
key: bench-head
- name: Bench PR head
run: cargo bench --all-features --bench samplers --bench optimization -- --save-baseline head
# --- Compare and post comment ---
- name: Compare benchmarks
id: compare
run: |
EOF=$(dd if=/dev/urandom bs=15 count=1 status=none | base64)
echo "result<<$EOF" >> "$GITHUB_OUTPUT"
critcmp base head --color never >> "$GITHUB_OUTPUT"
echo "$EOF" >> "$GITHUB_OUTPUT"
- name: Find existing comment
uses: peter-evans/find-comment@v4
id: find-comment
with:
issue-number: ${{ github.event.pull_request.number }}
comment-author: github-actions[bot]
body-includes: "## Benchmark Comparison"
- name: Post or update PR comment
uses: peter-evans/create-or-update-comment@v5
with:
comment-id: ${{ steps.find-comment.outputs.comment-id }}
issue-number: ${{ github.event.pull_request.number }}
edit-mode: replace
body: |
## Benchmark Comparison
**Base:** `${{ github.event.pull_request.base.sha }}` | **Head:** `${{ github.event.pull_request.head.sha }}`
<details>
<summary>Click to expand full results</summary>
```
${{ steps.compare.outputs.result }}
```
</details>
> Benchmarks run on `ubuntu-latest` via [Criterion](https://github.com/bheisler/criterion.rs) + [critcmp](https://github.com/BurntSushi/critcmp).
> Results may vary due to shared CI runners. Look for consistent >5% changes.
publish:
name: Publish to crates.io
runs-on: ubuntu-latest
@@ -243,6 +361,7 @@ jobs:
- fmt
- clippy
- test
- examples
- docs
- feature-check
- coverage
@@ -254,6 +373,7 @@ jobs:
- cross-platform
- cross-targets
- typos
- bench-check
steps:
- uses: actions/checkout@v6
- name: Install Rust
@@ -265,8 +385,10 @@ jobs:
- name: Check if version already published
id: check
run: |
CARGO_VERSION=$(cargo metadata --no-deps --format-version 1 | jq -r '.packages[0].version')
CARGO_VERSION=$(cargo metadata --no-deps --format-version 1 | jq -r '.packages[] | select(.name == "optimizer") | .version')
DERIVE_VERSION=$(cargo metadata --no-deps --format-version 1 | jq -r '.packages[] | select(.name == "optimizer-derive") | .version')
echo "version=$CARGO_VERSION" >> "$GITHUB_OUTPUT"
echo "derive_version=$DERIVE_VERSION" >> "$GITHUB_OUTPUT"
HTTP_STATUS=$(curl -s -o /dev/null -w "%{http_code}" "https://crates.io/api/v1/crates/optimizer/$CARGO_VERSION")
if [ "$HTTP_STATUS" = "200" ]; then
@@ -277,7 +399,33 @@ jobs:
echo "skip=false" >> "$GITHUB_OUTPUT"
fi
- name: Publish
DERIVE_HTTP_STATUS=$(curl -s -o /dev/null -w "%{http_code}" "https://crates.io/api/v1/crates/optimizer-derive/$DERIVE_VERSION")
if [ "$DERIVE_HTTP_STATUS" = "200" ]; then
echo "optimizer-derive $DERIVE_VERSION already exists on crates.io"
echo "skip_derive=true" >> "$GITHUB_OUTPUT"
else
echo "optimizer-derive $DERIVE_VERSION not found on crates.io, will publish"
echo "skip_derive=false" >> "$GITHUB_OUTPUT"
fi
- name: Publish optimizer-derive
if: steps.check.outputs.skip == 'false' && steps.check.outputs.skip_derive == 'false'
run: |
cargo publish -p optimizer-derive --allow-dirty 2>&1 | tee publish_output.txt || {
if grep -q "already uploaded" publish_output.txt; then
echo "optimizer-derive already published, skipping"
exit 0
fi
exit 1
}
env:
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}
- name: Wait for crates.io index update
if: steps.check.outputs.skip == 'false' && steps.check.outputs.skip_derive == 'false'
run: sleep 30
- name: Publish optimizer
if: steps.check.outputs.skip == 'false'
run: |
cargo publish --allow-dirty 2>&1 | tee publish_output.txt || {
+230
View File
@@ -0,0 +1,230 @@
# Changelog
All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [0.9.1] - 2026-02-12
### Fixed
- Fix journal file handling to support both reading and writing modes
## [0.9.0] - 2026-02-12
### Added
- Storage trait abstraction with pluggable backends
- JSONL journal storage backend for persistent studies (behind `journal` feature)
- Differential Evolution sampler
- Gaussian Process sampler with Expected Improvement acquisition function
- NSGA-III sampler for many-objective optimization with reference-point decomposition
- MOEA/D sampler for decomposition-based multi-objective optimization
- `StudyBuilder` for fluent study construction
### Changed
- Move `next_trial_id` counter from `Study` into `Storage` trait
- Replace `rand 0.10` with `fastrand 2.3` for simpler, faster random number generation
- Remove `visualization` and `fanova` feature flags (now always available)
- Simplify CI test commands by removing feature matrix
### Fixed
- Seed importance tests to eliminate flakiness
- Use atomic counter for temp paths in journal tests
### Removed
- Reverted SQLite storage backend (may return in a future release)
## [0.8.1] - 2026-02-11
### Fixed
- Resolve broken rustdoc intra-doc links for `Error::NoCompletedTrials`
## [0.8.0] - 2026-02-11
### Added
- Multi-objective optimization with NSGA-II
- Multi-Objective TPE (MOTPE) sampler
- Pareto front analysis utilities
- CSV and JSON data export for visualization
- HTML visualization reports with Plotly.js charts
- fANOVA (functional ANOVA) parameter importance via random forest
### Fixed
- Rename variable to pass typos check
## [0.7.2] - 2026-02-11
### Changed
- Update `nalgebra` dependency to version 0.34
- Add advisory ignore for unmaintained transitive dependency
## [0.7.1] - 2026-02-11
### Changed
- Update `tracing` dependency to version 0.1.29
## [0.7.0] - 2026-02-11
### Added
- CMA-ES (Covariance Matrix Adaptation Evolution Strategy) sampler behind `cma-es` feature flag
- Sobol quasi-random sampler behind `sobol` feature flag
- BOHB (Bayesian Optimization with HyperBand) sampler for budget-aware optimization
- Parameter importance analysis via Spearman rank correlation
- Constraint handling with feasibility-aware trial ranking
- `optimize_with_retries()` for automatic retry of failed trials
- `optimize_with_checkpoint()` with atomic save writes for crash recovery
- `Study::summary()` and `Display` impl for study overview
- `IntoIterator` for `&Study` and `iter()` method
- Tracing integration behind `tracing` feature flag
- Serde serialization support behind `serde` feature flag
- Benchmark suite with criterion and standard test functions
## [0.6.0] - 2026-02-11
### Added
- Pruning system with `Pruner` trait and `NopPruner` default implementation
- Intermediate value reporting on trials for pruner integration
- `TrialPruned` error variant and `Pruned` trial state
- `ThresholdPruner` for fixed-bound trial pruning
- `MedianPruner` for statistics-based trial pruning
- `PercentilePruner` for configurable percentile-based trial pruning
- `PatientPruner` for patience-based trial pruning
- `SuccessiveHalvingPruner` for SHA-based trial pruning
- `HyperbandPruner` for multi-bracket trial pruning
- `WilcoxonPruner` for statistics-based trial pruning
- `Study::minimize()` and `Study::maximize()` constructor shortcuts
- `From<RangeInclusive>` for `FloatParam` and `IntParam`
- Timeout-based optimization with `optimize_until`
- `Study::top_trials(n)` for retrieving the best N trials
- Ask-and-tell interface with `ask()` and `tell()` methods
- Trial user attributes for logging and analysis
- `enqueue_trial()` for pre-specified parameter evaluation
## [0.5.1] - 2026-02-10
### Changed
- Update random number generation to use `rand::make_rng()` and upgrade `rand` to version 0.10
## [0.5.0] - 2026-02-06
### Added
- Typed parameter API with `FloatParam`, `IntParam`, `CategoricalParam`, `BoolParam`, and `EnumParam`
- `#[derive(Categorical)]` proc macro for deriving categorical parameters from enums (behind `derive` feature)
- `.name()` builder method on all parameter types for custom labels
- `CompletedTrial::get()` for typed parameter access
- `Display` impl on `ParamValue`
- Prelude module at `optimizer::prelude::*`
### Changed
- Reorganize sampler module structure and update imports
- Remove `log` dependency
## [0.4.0] - 2026-02-02
### Added
- Multivariate TPE sampler for correlated parameter search
- Gamma strategies for TPE sampler (linear, sqrt, fixed) with examples
- Example: async API parameter optimization
- Example: ML hyperparameter tuning
### Fixed
- Handle end value in `suggest` method to avoid panic on underflow
## [0.3.1] - 2026-01-31
### Added
- Grid search sampler for exhaustive parameter exploration
- `suggest_bool()` method for boolean parameter suggestion
- `SuggestableRange` trait and `suggest_range()` method for parameter suggestion from ranges
- Documentation, keywords, categories, and readme fields in `Cargo.toml`
### Changed
- Replace `TpeError` with unified `Error` type across the library
- Add permissions for read access to contents in CI and scheduled workflows
## [0.3.0] - 2026-01-30
### Added
- Cross-target compilation checks in CI
### Changed
- Internal code refactoring and cleanup
## [0.2.0] - 2026-01-30
### Added
- CI step for unused dependencies check with `cargo-machete`
- CI step to publish to crates.io
- Default typo extension for 'Tpe'
### Changed
- Remove `serde` dependency (was unused)
- Remove unused `ordered-float` dependency
## [0.1.1] - 2026-01-30
### Fixed
- Minor release fixes
## [0.1.0] - 2026-01-30
### Added
- Initial implementation of the optimization library
- TPE (Tree-structured Parzen Estimator) sampler with optional fixed bandwidth KDE
- Random sampler
- Async optimization support
- CI workflows for testing, coverage (Codecov), auditing, and publishing
- README with project overview, features, and quick start guide
---
## Release Workflow
This changelog is maintained automatically with [git-cliff](https://git-cliff.org/).
1. Write commits using [Conventional Commits](https://www.conventionalcommits.org/) (`feat:`, `fix:`, `refactor:`, etc.).
2. Tag the release: `git tag v0.X.0`
3. Regenerate the changelog: `make changelog`
4. Commit the updated `CHANGELOG.md` and push.
[0.9.1]: https://github.com/raimannma/rust-optimizer/compare/v0.9.0...v0.9.1
[0.9.0]: https://github.com/raimannma/rust-optimizer/compare/v0.8.1...v0.9.0
[0.8.1]: https://github.com/raimannma/rust-optimizer/compare/v0.8.0...v0.8.1
[0.8.0]: https://github.com/raimannma/rust-optimizer/compare/v0.7.2...v0.8.0
[0.7.2]: https://github.com/raimannma/rust-optimizer/compare/v0.7.1...v0.7.2
[0.7.1]: https://github.com/raimannma/rust-optimizer/compare/v0.7.0...v0.7.1
[0.7.0]: https://github.com/raimannma/rust-optimizer/compare/v0.6.0...v0.7.0
[0.6.0]: https://github.com/raimannma/rust-optimizer/compare/v0.5.1...v0.6.0
[0.5.1]: https://github.com/raimannma/rust-optimizer/compare/v0.5.0...v0.5.1
[0.5.0]: https://github.com/raimannma/rust-optimizer/compare/v0.4.0...v0.5.0
[0.4.0]: https://github.com/raimannma/rust-optimizer/compare/v0.3.1...v0.4.0
[0.3.1]: https://github.com/raimannma/rust-optimizer/compare/v0.3.0...v0.3.1
[0.3.0]: https://github.com/raimannma/rust-optimizer/compare/v0.2.0...v0.3.0
[0.2.0]: https://github.com/raimannma/rust-optimizer/compare/v0.1.1...v0.2.0
[0.1.1]: https://github.com/raimannma/rust-optimizer/compare/v0.1.0...v0.1.1
[0.1.0]: https://github.com/raimannma/rust-optimizer/releases/tag/v0.1.0
+3
View File
@@ -0,0 +1,3 @@
## Git Guidelines
- If the user asks to make a commit, do not add any co-author by claude or any other co-author note.
+79 -8
View File
@@ -1,26 +1,97 @@
[workspace]
members = ["optimizer-derive"]
[package]
name = "optimizer"
version = "0.3.0"
version = "1.0.0"
edition = "2024"
rust-version = "1.88"
rust-version = "1.89"
license = "MIT"
authors = ["Manuel Raimann <raimannma@outlook.de"]
description = "A Rust library for optimization algorithms."
authors = ["Manuel Raimann <raimannma@outlook.de>"]
description = "Bayesian and population-based optimization library with an Optuna-like API for hyperparameter tuning and black-box optimization"
repository = "https://github.com/raimannma/rust-optimizer"
documentation = "https://docs.rs/optimizer"
keywords = ["optimization", "hyperparameter", "tpe", "grid-search", "bayesian"]
categories = ["algorithm", "science", "data-structures"]
categories = ["algorithms", "science", "mathematics"]
readme = "README.md"
[dependencies]
rand = "0.9"
fastrand = "2.3"
thiserror = "2"
parking_lot = "0.12"
tokio = { version = "1", features = ["sync", "rt-multi-thread"], optional = true }
tokio = { version = "1.30", features = ["sync", "rt-multi-thread"], optional = true }
optimizer-derive = { version = "0.1.0", path = "optimizer-derive", optional = true }
serde = { version = "1", features = ["derive"], optional = true }
serde_json = { version = "1", optional = true }
tracing = { version = "0.1.29", optional = true }
sobol_burley = { version = "0.5", optional = true }
nalgebra = { version = "0.34", optional = true }
fs2 = { version = "0.4", optional = true }
[features]
default = []
async = ["dep:tokio"]
derive = ["dep:optimizer-derive"]
serde = ["dep:serde", "dep:serde_json"]
journal = ["dep:fs2", "serde"]
tracing = ["dep:tracing"]
sobol = ["dep:sobol_burley"]
cma-es = ["dep:nalgebra"]
gp = ["dep:nalgebra"]
[dev-dependencies]
tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
tokio = { version = "1.30", features = ["rt-multi-thread", "macros", "time"] }
optimizer-derive = { version = "0.1.0", path = "optimizer-derive" }
serde_json = "1"
criterion = { version = "0.8", features = ["html_reports"] }
[[bench]]
name = "samplers"
harness = false
[[bench]]
name = "optimization"
harness = false
[[example]]
name = "basic_optimization"
path = "examples/basic_optimization.rs"
[[example]]
name = "parameter_types"
path = "examples/parameter_types.rs"
required-features = ["derive"]
[[example]]
name = "sampler_comparison"
path = "examples/sampler_comparison.rs"
[[example]]
name = "pruning"
path = "examples/pruning.rs"
[[example]]
name = "early_stopping"
path = "examples/early_stopping.rs"
[[example]]
name = "async_parallel"
path = "examples/async_parallel.rs"
required-features = ["async"]
[[example]]
name = "journal_storage"
path = "examples/journal_storage.rs"
required-features = ["journal"]
[[example]]
name = "ask_and_tell"
path = "examples/ask_and_tell.rs"
[[example]]
name = "multi_objective"
path = "examples/multi_objective.rs"
[[test]]
name = "journal_tests"
required-features = ["journal"]
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 optimizer
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+6
View File
@@ -0,0 +1,6 @@
.PHONY: changelog
# Regenerate CHANGELOG.md from git tags using git-cliff.
# Install: cargo install git-cliff
changelog:
git-cliff --output CHANGELOG.md
+48 -68
View File
@@ -1,93 +1,73 @@
# optimizer
A Rust library for black-box optimization with multiple sampling strategies.
Bayesian and population-based optimization library with an Optuna-like API
for hyperparameter tuning and black-box optimization. Supports 12 samplers,
8 pruners, multi-objective optimization, async parallelism, and persistent storage.
[![Docs](https://docs.rs/optimizer/badge.svg)](https://docs.rs/optimizer)
[![Crates.io](https://img.shields.io/crates/v/optimizer.svg)](https://crates.io/crates/optimizer)
[![codecov](https://codecov.io/gh/raimannma/rust-optimizer/graph/badge.svg?token=WOE77XJ4M6)](https://codecov.io/gh/raimannma/rust-optimizer)
## Features
- Optuna-like API for hyperparameter optimization
- Multiple sampling strategies:
- **Random Search** - Simple random sampling for baseline comparisons
- **TPE (Tree-Parzen Estimator)** - Bayesian optimization for efficient search
- **Grid Search** - Exhaustive search over a specified parameter grid
- Float, integer, and categorical parameter types
- Log-scale and stepped parameter sampling
- Sync and async optimization with parallel trial evaluation
## Quick Start
```rust
use optimizer::{Direction, Study};
use optimizer::sampler::tpe::TpeSampler;
use optimizer::prelude::*;
let sampler = TpeSampler::builder().seed(42).build().unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(-10.0, 10.0).name("x");
study
.optimize_with_sampler(20, |trial| {
let x = trial.suggest_float("x", -10.0, 10.0)?;
Ok::<_, optimizer::Error>(x * x)
})
.unwrap();
study.optimize(50, |trial| {
let val = x.suggest(trial)?;
Ok::<_, Error>((val - 3.0).powi(2))
}).unwrap();
let best = study.best_trial().unwrap();
println!("Best value: {} at x={:?}", best.value, best.params);
println!("Best x = {:.4}, f(x) = {:.4}", best.get(&x).unwrap(), best.value);
```
## Samplers
## Features at a Glance
### Random Search
```rust
use optimizer::{Direction, Study};
use optimizer::sampler::random::RandomSampler;
let study: Study<f64> = Study::with_sampler(
Direction::Minimize,
RandomSampler::with_seed(42),
);
```
### TPE (Tree-Parzen Estimator)
```rust
use optimizer::{Direction, Study};
use optimizer::sampler::tpe::TpeSampler;
let sampler = TpeSampler::builder()
.gamma(0.15) // Quantile for good/bad split
.n_startup_trials(20) // Random trials before TPE kicks in
.n_ei_candidates(32) // Candidates to evaluate
.seed(42)
.build()
.unwrap();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
```
### Grid Search
```rust
use optimizer::{Direction, Study};
use optimizer::sampler::grid::GridSearchSampler;
let sampler = GridSearchSampler::builder()
.n_points_per_param(10) // Number of points per parameter dimension
.build();
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
```
- **[Samplers](https://docs.rs/optimizer/latest/optimizer/sampler/)** — Random, TPE, Multivariate TPE, Grid, Sobol, CMA-ES, Gaussian Process, Differential Evolution, BOHB, NSGA-II, NSGA-III, MOEA/D
- **[Pruners](https://docs.rs/optimizer/latest/optimizer/pruner/)** — Median, Percentile, Threshold, Patient, Hyperband, Successive Halving, Wilcoxon, Nop
- **[Parameters](https://docs.rs/optimizer/latest/optimizer/parameter/)** — Float, Int, Categorical, Bool, and Enum types with `.name()` labels and typed access
- **[Multi-objective](https://docs.rs/optimizer/latest/optimizer/multi_objective/)** — Pareto front extraction with NSGA-II/III and MOEA/D
- **[Async & parallel](https://docs.rs/optimizer/latest/optimizer/struct.Study.html#method.optimize_parallel)** — Concurrent trial evaluation with Tokio
- **[Storage backends](https://docs.rs/optimizer/latest/optimizer/storage/)** — In-memory (default) or JSONL journal for persistence and resumption
- **[Visualization](https://docs.rs/optimizer/latest/optimizer/fn.generate_html_report.html)** — HTML reports with optimization history and parameter importance
- **[Analysis](https://docs.rs/optimizer/latest/optimizer/struct.Study.html#method.fanova)** — fANOVA and Spearman correlation for parameter importance
## Feature Flags
- `async` - Enable async optimization methods (requires tokio)
| Flag | Enables | Default |
|------|---------|---------|
| `async` | Async/parallel optimization (Tokio) | No |
| `derive` | `#[derive(Categorical)]` for enum parameters | No |
| `serde` | Serialization of trials and parameters | No |
| `journal` | JSONL storage backend (implies `serde`) | No |
| `sobol` | Sobol quasi-random sampler | No |
| `cma-es` | CMA-ES sampler (requires `nalgebra`) | No |
| `gp` | Gaussian Process sampler (requires `nalgebra`) | No |
| `tracing` | Structured logging with `tracing` | No |
## Documentation
## Examples
Full API documentation is available at [docs.rs/optimizer](https://docs.rs/optimizer).
```sh
cargo run --example basic_optimization # Minimize a quadratic — simplest possible usage
cargo run --example parameter_types --features derive # All 5 param types + #[derive(Categorical)]
cargo run --example sampler_comparison # Compare Random, TPE, and Grid on the same problem
cargo run --example pruning # Trial pruning with MedianPruner
cargo run --example early_stopping # Halt a study when a target is reached
cargo run --example async_parallel --features async # Evaluate trials concurrently with tokio
cargo run --example journal_storage --features journal # Persist trials to disk and resume later
cargo run --example ask_and_tell # Decouple sampling from evaluation
cargo run --example multi_objective # Optimize competing objectives + Pareto front
```
## Learn More
- [API documentation](https://docs.rs/optimizer)
- [Changelog](CHANGELOG.md)
- [GitHub Issues](https://github.com/raimannma/rust-optimizer/issues)
## License
+112
View File
@@ -0,0 +1,112 @@
#[allow(dead_code)]
mod test_functions;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use optimizer::Study;
use optimizer::parameter::{FloatParam, Parameter};
use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
fn make_params(dims: usize) -> Vec<FloatParam> {
(0..dims)
.map(|i| FloatParam::new(-5.0, 5.0).name(format!("x{i}")))
.collect()
}
fn bench_tpe_sphere(c: &mut Criterion) {
let mut group = c.benchmark_group("tpe_sphere");
group.sample_size(10);
for dims in [2, 10, 50] {
let params = make_params(dims);
group.bench_with_input(BenchmarkId::new("dims", dims), &params, |b, params| {
b.iter(|| {
let study = Study::minimize(TpeSampler::builder().seed(42).build().unwrap());
study
.optimize(100, |trial: &mut optimizer::Trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(test_functions::sphere(&x))
})
.unwrap();
});
});
}
group.finish();
}
fn bench_tpe_rosenbrock(c: &mut Criterion) {
let mut group = c.benchmark_group("tpe_rosenbrock");
group.sample_size(10);
for dims in [2, 10] {
let params = make_params(dims);
group.bench_with_input(BenchmarkId::new("dims", dims), &params, |b, params| {
b.iter(|| {
let study = Study::minimize(TpeSampler::builder().seed(42).build().unwrap());
study
.optimize(100, |trial: &mut optimizer::Trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(test_functions::rosenbrock(&x))
})
.unwrap();
});
});
}
group.finish();
}
fn bench_random_vs_tpe(c: &mut Criterion) {
let mut group = c.benchmark_group("random_vs_tpe");
group.sample_size(10);
let params = make_params(5);
group.bench_function("random_5d", |b| {
b.iter(|| {
let study = Study::minimize(RandomSampler::with_seed(42));
study
.optimize(100, |trial: &mut optimizer::Trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(test_functions::sphere(&x))
})
.unwrap();
});
});
group.bench_function("tpe_5d", |b| {
b.iter(|| {
let study = Study::minimize(TpeSampler::builder().seed(42).build().unwrap());
study
.optimize(100, |trial: &mut optimizer::Trial| {
let x: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()
.unwrap();
Ok::<_, optimizer::Error>(test_functions::sphere(&x))
})
.unwrap();
});
});
group.finish();
}
criterion_group!(
benches,
bench_tpe_sphere,
bench_tpe_rosenbrock,
bench_random_vs_tpe
);
criterion_main!(benches);
+117
View File
@@ -0,0 +1,117 @@
use std::collections::HashMap;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use optimizer::parameter::{FloatParam, Parameter};
use optimizer::sampler::grid::GridSampler;
use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::sampler::{CompletedTrial, Sampler};
/// Build a synthetic history of `n` completed trials over `dims` float parameters.
fn build_history(n: usize, dims: usize) -> Vec<CompletedTrial<f64>> {
let params: Vec<FloatParam> = (0..dims)
.map(|i| FloatParam::new(-5.0, 5.0).name(format!("x{i}")))
.collect();
let mut history = Vec::with_capacity(n);
let sampler = RandomSampler::with_seed(42);
for trial_id in 0..n {
let id = trial_id as u64;
let mut param_values = HashMap::new();
let mut distributions = HashMap::new();
let mut param_labels = HashMap::new();
for p in &params {
let dist = p.distribution();
let val = sampler.sample(&dist, id, &history);
param_values.insert(p.id(), val);
distributions.insert(p.id(), dist);
param_labels.insert(p.id(), p.label());
}
// Use sphere function value as objective
let value: f64 = param_values
.values()
.map(|v| {
let optimizer::parameter::ParamValue::Float(f) = v else {
unreachable!()
};
f * f
})
.sum();
history.push(CompletedTrial::new(
id,
param_values,
distributions,
param_labels,
value,
));
}
history
}
fn bench_tpe_sample(c: &mut Criterion) {
let mut group = c.benchmark_group("tpe_sample");
let dist = FloatParam::new(-5.0, 5.0).distribution();
let tpe = TpeSampler::builder().seed(42).build().unwrap();
for history_size in [10, 100, 1000] {
let history = build_history(history_size, 2);
group.bench_with_input(
BenchmarkId::new("history", history_size),
&history,
|b, history| {
b.iter(|| tpe.sample(&dist, history.len() as u64, history));
},
);
}
group.finish();
}
fn bench_random_sample(c: &mut Criterion) {
let mut group = c.benchmark_group("random_sample");
let dist = FloatParam::new(-5.0, 5.0).distribution();
let sampler = RandomSampler::with_seed(42);
for history_size in [10, 100, 1000] {
let history = build_history(history_size, 2);
group.bench_with_input(
BenchmarkId::new("history", history_size),
&history,
|b, history| {
b.iter(|| sampler.sample(&dist, history.len() as u64, history));
},
);
}
group.finish();
}
fn bench_grid_sample(c: &mut Criterion) {
let mut group = c.benchmark_group("grid_sample");
let dist = FloatParam::new(-5.0, 5.0).distribution();
let history: Vec<CompletedTrial<f64>> = Vec::new();
for grid_points in [5, 10, 50] {
group.bench_with_input(
BenchmarkId::new("points", grid_points),
&grid_points,
|b, _| {
b.iter(|| {
// Fresh sampler each iteration since grid tracks used points
let sampler = GridSampler::builder()
.n_points_per_param(grid_points)
.build();
sampler.sample(&dist, 0, &history)
});
},
);
}
group.finish();
}
criterion_group!(
benches,
bench_tpe_sample,
bench_random_sample,
bench_grid_sample
);
criterion_main!(benches);
+84
View File
@@ -0,0 +1,84 @@
//! Standard optimization test functions for benchmarking.
/// Sphere function: unimodal, convex. Global minimum f(0,...,0) = 0.
pub fn sphere(x: &[f64]) -> f64 {
x.iter().map(|xi| xi * xi).sum()
}
/// Rosenbrock function: narrow valley. Global minimum f(1,...,1) = 0.
pub fn rosenbrock(x: &[f64]) -> f64 {
x.windows(2)
.map(|w| 100.0 * (w[1] - w[0] * w[0]).powi(2) + (1.0 - w[0]).powi(2))
.sum()
}
/// Rastrigin function: highly multimodal. Global minimum f(0,...,0) = 0.
pub fn rastrigin(x: &[f64]) -> f64 {
let n = x.len() as f64;
10.0 * n
+ x.iter()
.map(|xi| xi * xi - 10.0 * (2.0 * std::f64::consts::PI * xi).cos())
.sum::<f64>()
}
/// Ackley function: nearly flat with a deep well. Global minimum f(0,...,0) = 0.
pub fn ackley(x: &[f64]) -> f64 {
let n = x.len() as f64;
let sum_sq: f64 = x.iter().map(|xi| xi * xi).sum();
let sum_cos: f64 = x
.iter()
.map(|xi| (2.0 * std::f64::consts::PI * xi).cos())
.sum();
-20.0 * (-0.2 * (sum_sq / n).sqrt()).exp() - (sum_cos / n).exp() + 20.0 + std::f64::consts::E
}
/// Branin function (2D only). Three global minima with f* ≈ 0.397887.
///
/// # Panics
///
/// Panics if `x` does not have exactly 2 elements.
pub fn branin(x: &[f64]) -> f64 {
assert!(x.len() == 2, "Branin requires exactly 2 dimensions");
let (x1, x2) = (x[0], x[1]);
let pi = std::f64::consts::PI;
let a = 1.0;
let b = 5.1 / (4.0 * pi * pi);
let c = 5.0 / pi;
let r = 6.0;
let s = 10.0;
let t = 1.0 / (8.0 * pi);
a * (x2 - b * x1 * x1 + c * x1 - r).powi(2) + s * (1.0 - t) * x1.cos() + s
}
/// Hartmann 6D function. Global minimum f* ≈ -3.3224.
///
/// # Panics
///
/// Panics if `x` does not have exactly 6 elements.
pub fn hartmann6(x: &[f64]) -> f64 {
assert!(x.len() == 6, "Hartmann6 requires exactly 6 dimensions");
let alpha = [1.0, 1.2, 3.0, 3.2];
let a_matrix = [
[10.0, 3.0, 17.0, 3.5, 1.7, 8.0],
[0.05, 10.0, 17.0, 0.1, 8.0, 14.0],
[3.0, 3.5, 1.7, 10.0, 17.0, 8.0],
[17.0, 8.0, 0.05, 10.0, 0.1, 14.0],
];
let p_matrix = [
[0.1312, 0.1696, 0.5569, 0.0124, 0.8283, 0.5886],
[0.2329, 0.4135, 0.8307, 0.3736, 0.1004, 0.9991],
[0.2348, 0.1451, 0.3522, 0.2883, 0.3047, 0.6650],
[0.4047, 0.8828, 0.8732, 0.5743, 0.1091, 0.0381],
];
let mut result = 0.0;
for i in 0..4 {
let mut inner = 0.0;
for (j, xj) in x.iter().enumerate() {
inner += a_matrix[i][j] * (xj - p_matrix[i][j]).powi(2);
}
result -= alpha[i] * (-inner).exp();
}
result
}
+99
View File
@@ -0,0 +1,99 @@
# git-cliff configuration
# https://git-cliff.org/docs/configuration
[changelog]
# changelog header
header = """
# Changelog
All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).\n
"""
# template for the changelog body
# https://keats.github.io/tera/docs/#introduction
body = """
{%- macro remote_url() -%}
https://github.com/raimannma/rust-optimizer
{%- endmacro -%}
{% if version -%}
## [{{ version | trim_start_matches(pat="v") }}] - {{ timestamp | date(format="%Y-%m-%d") }}
{% else -%}
## [Unreleased]
{% endif -%}
{% for group, commits in commits | group_by(attribute="group") %}
### {{ group | striptags | trim }}
{% for commit in commits %}
- {{ commit.message | split(pat="\n") | first | trim }}\
{%- if commit.breaking %} (**BREAKING**){% endif -%}
{% endfor %}
{% endfor %}
"""
# template for the changelog footer
footer = """
{%- macro remote_url() -%}
https://github.com/raimannma/rust-optimizer
{%- endmacro -%}
{% for release in releases -%}
{% if release.version -%}
{% if release.previous.version -%}
[{{ release.version | trim_start_matches(pat="v") }}]: \
{{ self::remote_url() }}/compare/{{ release.previous.version }}...{{ release.version }}
{% else -%}
[{{ release.version | trim_start_matches(pat="v") }}]: \
{{ self::remote_url() }}/releases/tag/{{ release.version }}
{% endif -%}
{% else -%}
{% if release.previous.version -%}
[Unreleased]: {{ self::remote_url() }}/compare/{{ release.previous.version }}...HEAD
{% endif -%}
{% endif -%}
{% endfor %}
"""
# remove the leading and trailing whitespace from the templates
trim = true
[git]
# parse the commits based on https://www.conventionalcommits.org
conventional_commits = true
# filter out the commits that are not conventional
filter_unconventional = false
# process each line of a commit as an individual commit
split_commits = false
# regex for preprocessing the commit messages
commit_preprocessors = [
{ pattern = '\((\w+\s)?#([0-9]+)\)', replace = "([#${2}](https://github.com/raimannma/rust-optimizer/issues/${2}))" },
]
# regex for parsing and grouping commits
commit_parsers = [
{ message = "^feat", group = "Added" },
{ message = "^fix", group = "Fixed" },
{ message = "^refactor", group = "Changed" },
{ message = "^perf", group = "Changed" },
{ message = "^doc", group = "Changed" },
{ message = "^style", group = "Changed" },
{ message = "^ci", skip = true },
{ message = "^chore\\(release\\)", skip = true },
{ message = "^chore: release", skip = true },
{ message = "^chore", group = "Changed" },
{ message = "^revert", group = "Removed" },
{ body = ".*security", group = "Security" },
]
# protect breaking changes from being skipped due to matching a skipping commit_parser
protect_breaking_commits = true
# filter out the commits that are not matched by commit parsers
filter_commits = true
# regex for matching git tags
tag_pattern = "v[0-9].*"
# regex for skipping tags
skip_tags = ""
# regex for ignoring tags
ignore_tags = ""
# sort the tags topologically
topo_order = false
# sort the commits inside sections by oldest first
sort_commits = "oldest"
+6 -2
View File
@@ -1,7 +1,11 @@
[advisories]
version = 2
db-path = "~/.cargo/advisory-db"
ignore = []
ignore = [
# paste is unmaintained but it's a transitive dep via simba -> nalgebra
# with no fix available upstream yet
"RUSTSEC-2024-0436",
]
[licenses]
version = 2
@@ -9,8 +13,8 @@ allow = [
"MIT",
"Apache-2.0",
"Apache-2.0 WITH LLVM-exception",
"BSD-2-Clause",
"Unicode-3.0",
"Zlib",
]
confidence-threshold = 0.8
+51
View File
@@ -0,0 +1,51 @@
//! Ask-and-tell interface — decouple sampling from evaluation.
//!
//! Use `ask()` to get a trial with sampled parameters, evaluate it however
//! you like (workers, GPUs, external processes), then `tell()` the result.
//! This is useful for batch evaluation or custom scheduling.
//!
//! Run with: `cargo run --example ask_and_tell`
use optimizer::prelude::*;
fn main() -> optimizer::Result<()> {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(-5.0, 5.0).name("x");
let y = FloatParam::new(-5.0, 5.0).name("y");
for batch in 0..3 {
let batch_size = 5;
let mut trials = Vec::with_capacity(batch_size);
// ask() creates trials with sampled parameters
for _ in 0..batch_size {
let mut trial = study.ask();
let xv = x.suggest(&mut trial)?;
let yv = y.suggest(&mut trial)?;
trials.push((trial, xv, yv));
}
// Evaluate the batch (could be sent to workers, GPUs, etc.)
for (trial, xv, yv) in trials {
let value = xv * xv + yv * yv;
study.tell(trial, Ok::<_, &str>(value));
}
println!(
"Batch {}: evaluated {batch_size} trials (total: {})",
batch + 1,
study.n_trials(),
);
}
let best = study.best_trial()?;
println!(
"Best: f({:.3}, {:.3}) = {:.6}",
best.get(&x).unwrap(),
best.get(&y).unwrap(),
best.value,
);
Ok(())
}
+46
View File
@@ -0,0 +1,46 @@
//! Async parallel optimization — evaluate multiple trials concurrently.
//!
//! Uses `optimize_parallel` with tokio to run several trials at once,
//! reducing wall-clock time when the objective involves blocking work.
//! Each sync closure is internally wrapped in `spawn_blocking`.
//!
//! Run with: `cargo run --example async_parallel --features async`
use optimizer::prelude::*;
#[tokio::main]
async fn main() -> optimizer::Result<()> {
let study: Study<f64> = Study::minimize(TpeSampler::new());
let x = FloatParam::new(-5.0, 5.0).name("x");
let y = FloatParam::new(-5.0, 5.0).name("y");
let n_trials = 30;
let concurrency = 4;
println!("Running {n_trials} trials with {concurrency} concurrent workers...");
let xc = x.clone();
let yc = y.clone();
study
.optimize_parallel(
n_trials,
concurrency,
move |trial: &mut optimizer::Trial| {
let xv = xc.suggest(trial)?;
let yv = yc.suggest(trial)?;
Ok::<_, optimizer::Error>(xv * xv + yv * yv)
},
)
.await?;
let best = study.best_trial()?;
println!(
"Best: f({:.3}, {:.3}) = {:.6}",
best.get(&x).unwrap(),
best.get(&y).unwrap(),
best.value,
);
Ok(())
}
+32
View File
@@ -0,0 +1,32 @@
//! Basic optimization example — the "hello world" of the optimizer crate.
//!
//! Minimizes a simple quadratic function f(x) = (x - 3)² using the default
//! random sampler. No feature flags are required.
//!
//! Run with: `cargo run --example basic_optimization`
use optimizer::prelude::*;
fn main() {
// Create a study that minimizes the objective function.
// The default sampler is random; for smarter sampling, pass a TpeSampler.
let study: Study<f64> = Study::new(Direction::Minimize);
// Search for x in [-10, 10]. The optimizer will suggest values from this range.
let x = FloatParam::new(-10.0, 10.0).name("x");
// Run 50 trials, each evaluating f(x) = (x - 3)²
study
.optimize(50, |trial: &mut optimizer::Trial| {
let x_val = x.suggest(trial)?;
let value = (x_val - 3.0).powi(2);
Ok::<_, Error>(value)
})
.unwrap();
// Retrieve and display the best result
let best = study.best_trial().unwrap();
println!("Best trial #{}", best.id);
println!(" x = {:.4}", best.get(&x).unwrap());
println!(" f(x) = {:.4}", best.value);
}
+58
View File
@@ -0,0 +1,58 @@
//! Early stopping — halt an entire study once a target is reached.
//!
//! Implements the [`Objective`] trait on a custom struct and uses the
//! [`after_trial`](Objective::after_trial) hook to return
//! `ControlFlow::Break(())` when the best value drops below a threshold.
//!
//! Run with: `cargo run --example early_stopping`
use std::ops::ControlFlow;
use optimizer::prelude::*;
/// An objective that minimises `(x - 3)^2` and stops early once the
/// value drops below `target`.
struct EarlyStopObjective {
x: FloatParam,
target: f64,
}
impl Objective<f64> for EarlyStopObjective {
type Error = Error;
fn evaluate(&self, trial: &mut Trial) -> Result<f64> {
let v = self.x.suggest(trial)?;
Ok((v - 3.0).powi(2))
}
fn after_trial(&self, _study: &Study<f64>, trial: &CompletedTrial<f64>) -> ControlFlow<()> {
if trial.value < self.target {
println!("Target {} reached at trial #{}", self.target, trial.id);
ControlFlow::Break(())
} else {
ControlFlow::Continue(())
}
}
}
fn main() -> optimizer::Result<()> {
let study: Study<f64> = Study::new(Direction::Minimize);
let x = FloatParam::new(-10.0, 10.0).name("x");
let objective = EarlyStopObjective {
x: x.clone(),
target: 0.01,
};
study.optimize(100, objective)?;
let best = study.best_trial()?;
println!(
"Stopped after {} trials — best f({:.4}) = {:.6}",
study.n_trials(),
best.get(&x).unwrap(),
best.value,
);
Ok(())
}
+68
View File
@@ -0,0 +1,68 @@
//! Journal storage — persist trials to disk and resume later.
//!
//! `JournalStorage` writes every trial to a JSONL file so that a study can
//! be resumed after a crash or across separate runs.
//!
//! Run with: `cargo run --example journal_storage --features journal`
use optimizer::prelude::*;
fn main() -> optimizer::Result<()> {
let path = std::env::temp_dir().join("optimizer_journal_example.jsonl");
// Clean up from any previous run
let _ = std::fs::remove_file(&path);
let x = FloatParam::new(-5.0, 5.0).name("x");
// --- First run: optimize 20 trials and persist to disk ---
{
let storage = JournalStorage::<f64>::new(&path);
let study: Study<f64> = Study::builder()
.minimize()
.sampler(TpeSampler::new())
.storage(storage)
.build();
study.optimize(20, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(xv * xv)
})?;
println!(
"First run: {} trials saved to {}",
study.n_trials(),
path.display(),
);
}
// --- Second run: resume from the journal file ---
{
let storage = JournalStorage::<f64>::open(&path)?;
let study: Study<f64> = Study::builder()
.minimize()
.sampler(TpeSampler::new())
.storage(storage)
.build();
let before = study.n_trials();
study.optimize(10, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(xv * xv)
})?;
let best = study.best_trial()?;
println!(
"Resumed: {}{} trials, best f({:.4}) = {:.6}",
before,
study.n_trials(),
best.get(&x).unwrap(),
best.value,
);
}
// Clean up
let _ = std::fs::remove_file(&path);
Ok(())
}
+49
View File
@@ -0,0 +1,49 @@
//! Multi-objective optimization — optimize competing objectives simultaneously.
//!
//! `MultiObjectiveStudy` returns the Pareto front: the set of solutions where
//! no objective can be improved without worsening another.
//!
//! Run with: `cargo run --example multi_objective`
use optimizer::multi_objective::MultiObjectiveStudy;
use optimizer::prelude::*;
fn main() -> optimizer::Result<()> {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let x = FloatParam::new(0.0, 1.0).name("x");
// Classic bi-objective: f1(x) = x², f2(x) = (x-1)²
// The Pareto front is the curve where improving f1 worsens f2.
study.optimize(50, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let f1 = xv * xv;
let f2 = (xv - 1.0) * (xv - 1.0);
Ok::<_, optimizer::Error>(vec![f1, f2])
})?;
let front = study.pareto_front();
println!(
"Ran {} trials, Pareto front has {} solutions:",
study.n_trials(),
front.len(),
);
let mut sorted = front.clone();
sorted.sort_by(|a, b| a.values[0].partial_cmp(&b.values[0]).unwrap());
for (i, trial) in sorted.iter().take(5).enumerate() {
println!(
" {}: x={:.3}, f1={:.4}, f2={:.4}",
i + 1,
trial.get(&x).unwrap(),
trial.values[0],
trial.values[1],
);
}
if sorted.len() > 5 {
println!(" ... and {} more", sorted.len() - 5);
}
Ok(())
}
+73
View File
@@ -0,0 +1,73 @@
//! Parameter types example — demonstrates all five parameter types and the derive feature.
//!
//! Shows `FloatParam`, `IntParam`, `CategoricalParam`, `BoolParam`, and `EnumParam`
//! with `.name()` labels, `#[derive(Categorical)]` for enums, and typed access
//! to results via `CompletedTrial::get()`.
//!
//! Run with: `cargo run --example parameter_types --features derive`
use optimizer::prelude::*;
use optimizer_derive::Categorical;
/// Activation functions — `#[derive(Categorical)]` auto-generates the
/// `Categorical` trait, mapping each variant to a sequential index.
#[derive(Clone, Debug, Categorical)]
enum Activation {
Relu,
Sigmoid,
Tanh,
Gelu,
}
fn main() {
let study: Study<f64> = Study::new(Direction::Minimize);
// --- Define one of each parameter type, each with a human-readable .name() ---
// Float: learning rate on a log scale (common for ML hyperparameters)
let lr = FloatParam::new(1e-5, 1e-1).log_scale().name("lr");
// Int: number of hidden layers (stepped by 1, the default)
let n_layers = IntParam::new(1, 5).name("n_layers");
// Categorical: optimizer algorithm chosen from a list of strings
let optimizer = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]).name("optimizer");
// Bool: whether to apply dropout
let use_dropout = BoolParam::new().name("use_dropout");
// Enum: activation function — uses #[derive(Categorical)] above
let activation = EnumParam::<Activation>::new().name("activation");
// --- Run the optimization ---
study
.optimize(30, |trial: &mut optimizer::Trial| {
let lr_val = lr.suggest(trial)?;
let layers = n_layers.suggest(trial)?;
let opt = optimizer.suggest(trial)?;
let dropout = use_dropout.suggest(trial)?;
let act = activation.suggest(trial)?;
// Simulated loss that depends on all parameters
let loss = lr_val * f64::from(layers as i32)
+ if opt == "adam" { -0.05 } else { 0.0 }
+ if dropout { -0.02 } else { 0.0 }
+ match act {
Activation::Gelu => -0.03,
Activation::Relu => -0.01,
_ => 0.0,
};
Ok::<_, Error>(loss)
})
.unwrap();
// --- Retrieve best trial and read back each parameter with typed .get() ---
let best = study.best_trial().unwrap();
println!("Best trial #{} — loss = {:.6}", best.id, best.value);
println!(" lr = {:.6}", best.get(&lr).unwrap());
println!(" n_layers = {}", best.get(&n_layers).unwrap());
println!(" optimizer = {}", best.get(&optimizer).unwrap());
println!(" use_dropout = {}", best.get(&use_dropout).unwrap());
println!(" activation = {:?}", best.get(&activation).unwrap());
}
+67
View File
@@ -0,0 +1,67 @@
//! Trial pruning — stop unpromising trials early with `MedianPruner`.
//!
//! When your objective involves an iterative loop (e.g. training epochs),
//! the pruner compares intermediate values across trials and kills the
//! ones that fall below the median — saving compute on bad configurations.
//!
//! Run with: `cargo run --example pruning`
use optimizer::prelude::*;
fn main() -> optimizer::Result<()> {
// MedianPruner prunes trials whose intermediate value falls below the
// median of previously completed trials at the same step.
let study: Study<f64> = Study::builder()
.minimize()
.sampler(RandomSampler::with_seed(42))
.pruner(
MedianPruner::new(Direction::Minimize)
.n_warmup_steps(3) // run at least 3 epochs before pruning
.n_min_trials(3), // need 3 completed trials before pruning kicks in
)
.build();
let lr = FloatParam::new(1e-4, 1.0).name("learning_rate");
let momentum = FloatParam::new(0.0, 0.99).name("momentum");
let n_epochs: u64 = 20;
study.optimize(30, |trial: &mut optimizer::Trial| {
let lr_val = lr.suggest(trial)?;
let mom = momentum.suggest(trial)?;
// Simulated training loop — good hyperparameters converge to low loss,
// bad ones plateau high, giving the pruner something to cut.
let mut loss = 1.0;
for epoch in 0..n_epochs {
let lr_penalty = (lr_val.log10() - 0.01_f64.log10()).powi(2);
let mom_penalty = (mom - 0.8).powi(2);
let base_loss = 0.02 + 0.05 * lr_penalty + 1.5 * mom_penalty;
let progress = (epoch as f64 + 1.0) / n_epochs as f64;
loss = base_loss + (1.0 - base_loss) * (-3.5 * progress).exp();
// Report intermediate value so the pruner can evaluate this trial.
trial.report(epoch, loss);
// Check whether the pruner recommends stopping early.
if trial.should_prune() {
Err(TrialPruned)?;
}
}
Ok::<_, Error>(loss)
})?;
// --- Results ---
let best = study.best_trial()?;
println!(
"Recorded {} trials ({} pruned)",
study.n_trials(),
study.n_pruned_trials()
);
println!("Best trial #{}: loss = {:.6}", best.id, best.value);
println!(" learning_rate = {:.6}", best.get(&lr).unwrap());
println!(" momentum = {:.4}", best.get(&momentum).unwrap());
Ok(())
}
+78
View File
@@ -0,0 +1,78 @@
//! Sampler comparison example — benchmarks Random, TPE, and Grid samplers on the same problem.
//!
//! Runs the Sphere function f(x, y) = x² + y² with each sampler and compares the best
//! value found. This shows how sampler choice affects optimization quality.
//!
//! Run with: `cargo run --example sampler_comparison`
use optimizer::prelude::*;
/// Shared objective function: Sphere function with global minimum at (0, 0).
/// Simple enough to solve well, but 2-D so samplers have room to differ.
fn sphere(x: f64, y: f64) -> f64 {
x.powi(2) + y.powi(2)
}
/// Run an optimization study and return the best value found.
fn run_study(study: Study<f64>, n_trials: usize) -> f64 {
// Use asymmetric ranges so the Grid sampler tracks each parameter independently.
let x = FloatParam::new(-5.0, 5.0).name("x");
let y = FloatParam::new(-3.0, 3.0).name("y");
study
.optimize(n_trials, |trial: &mut optimizer::Trial| {
let x_val = x.suggest(trial)?;
let y_val = y.suggest(trial)?;
Ok::<_, Error>(sphere(x_val, y_val))
})
.unwrap();
let best = study.best_trial().unwrap();
println!(
" Best trial #{:>3}: x = {:>7.4}, y = {:>7.4}, f(x,y) = {:.6}",
best.id,
best.get(&x).unwrap(),
best.get(&y).unwrap(),
best.value,
);
best.value
}
fn main() {
let n_trials: usize = 100;
println!("Comparing samplers on Sphere(x, y) = x² + y² ({n_trials} trials each)");
println!();
// --- Random sampler (baseline) ---
// Pure random search: samples uniformly at random. Fast but not guided.
println!("1. Random sampler:");
let random_best = run_study(Study::minimize(RandomSampler::with_seed(42)), n_trials);
// --- TPE sampler (Bayesian) ---
// Tree-structured Parzen Estimator: builds a probabilistic model of good vs bad
// regions and focuses sampling where improvements are likely.
println!("\n2. TPE sampler (Bayesian):");
let tpe = TpeSampler::builder()
.n_startup_trials(10) // random exploration for the first 10 trials
.n_ei_candidates(24) // candidates evaluated per Expected Improvement step
.gamma(0.25) // top 25% of trials define the "good" distribution
.seed(42)
.build()
.unwrap();
let tpe_best = run_study(Study::minimize(tpe), n_trials);
// --- Grid sampler (exhaustive) ---
// Evaluates evenly spaced grid points. Each parameter gets its own grid that
// is sampled in order, so n_points_per_param must be >= n_trials.
println!("\n3. Grid sampler (exhaustive):");
let grid = GridSampler::builder()
.n_points_per_param(n_trials) // one grid point per trial per parameter
.build();
let grid_best = run_study(Study::minimize(grid), n_trials);
// --- Summary ---
println!("\n--- Summary ---");
println!(" Random : {random_best:.6}");
println!(" TPE : {tpe_best:.6}");
println!(" Grid : {grid_best:.6}");
}
+15
View File
@@ -0,0 +1,15 @@
[package]
name = "optimizer-derive"
version = "0.1.0"
edition = "2024"
rust-version = "1.89"
license = "MIT"
description = "Derive macros for the optimizer crate"
repository = "https://github.com/raimannma/rust-optimizer"
[lib]
proc-macro = true
[dependencies]
syn = { version = "2", features = ["full"] }
quote = "1"
+71
View File
@@ -0,0 +1,71 @@
use proc_macro::TokenStream;
use quote::quote;
use syn::{Data, DeriveInput, Fields, parse_macro_input};
/// Derive macro for the `Categorical` trait on fieldless enums.
///
/// Generates an implementation of `optimizer::parameter::Categorical` that maps
/// enum variants to/from sequential indices.
///
/// # Example
///
/// ```ignore
/// use optimizer::parameter::Categorical;
///
/// #[derive(Clone, Categorical)]
/// enum Color {
/// Red,
/// Green,
/// Blue,
/// }
/// ```
#[proc_macro_derive(Categorical)]
pub fn derive_categorical(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = &input.ident;
let Data::Enum(data_enum) = &input.data else {
return syn::Error::new_spanned(&input, "Categorical can only be derived for enums")
.to_compile_error()
.into();
};
// Validate all variants are fieldless
for variant in &data_enum.variants {
if !matches!(variant.fields, Fields::Unit) {
return syn::Error::new_spanned(
variant,
"Categorical can only be derived for enums with unit variants (no fields)",
)
.to_compile_error()
.into();
}
}
let n_choices = data_enum.variants.len();
let variant_names: Vec<_> = data_enum.variants.iter().map(|v| &v.ident).collect();
let indices: Vec<usize> = (0..n_choices).collect();
let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
let expanded = quote! {
impl #impl_generics optimizer::parameter::Categorical for #name #ty_generics #where_clause {
const N_CHOICES: usize = #n_choices;
fn from_index(index: usize) -> Self {
match index {
#(#indices => #name::#variant_names,)*
_ => panic!("invalid index {} for {} with {} variants", index, stringify!(#name), #n_choices),
}
}
fn to_index(&self) -> usize {
match self {
#(#name::#variant_names => #indices,)*
}
}
}
};
expanded.into()
}
+3 -3
View File
@@ -1,8 +1,8 @@
# Nightly Features
imports_granularity = "Module"
group_imports = "StdExternalCrate"
#imports_granularity = "Module"
#group_imports = "StdExternalCrate"
#format_strings = true
format_code_in_doc_comments = true
#format_code_in_doc_comments = true
# Stable Features
merge_derives = true
+4
View File
@@ -2,6 +2,7 @@
/// Distribution for floating-point parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FloatDistribution {
/// Lower bound (inclusive).
pub low: f64,
@@ -15,6 +16,7 @@ pub struct FloatDistribution {
/// Distribution for integer parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct IntDistribution {
/// Lower bound (inclusive).
pub low: i64,
@@ -28,6 +30,7 @@ pub struct IntDistribution {
/// Distribution for categorical parameters.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct CategoricalDistribution {
/// Number of choices available.
pub n_choices: usize,
@@ -35,6 +38,7 @@ pub struct CategoricalDistribution {
/// Enum wrapping all parameter distribution types.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Distribution {
/// A floating-point distribution.
Float(FloatDistribution),
+112 -11
View File
@@ -1,6 +1,22 @@
//! Error types for the optimizer crate.
//!
//! All fallible operations in the crate return [`Result<T>`], which is an
//! alias for `core::result::Result<T, Error>`. The [`Error`] enum covers
//! parameter validation, sampling conflicts, pruning signals, and
//! feature-gated I/O errors.
/// Errors returned by optimizer operations.
///
/// Most variants are returned during parameter validation or trial
/// management. The [`TrialPruned`](Error::TrialPruned) variant has special
/// significance — it signals early stopping and is typically raised via
/// the [`TrialPruned`](super::TrialPruned) convenience type.
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum Error {
/// Returned when the lower bound is greater than the upper bound.
/// The lower bound exceeds the upper bound in a
/// [`FloatParam`](crate::parameter::FloatParam) or
/// [`IntParam`](crate::parameter::IntParam).
#[error("invalid bounds: low ({low}) must be less than or equal to high ({high})")]
InvalidBounds {
/// The lower bound value.
@@ -9,19 +25,22 @@ pub enum Error {
high: f64,
},
/// Returned when log scale is used with non-positive bounds.
/// Log-scale is enabled but the lower bound is not positive (float) or
/// is less than 1 (integer).
#[error("invalid log bounds: low must be positive for log scale")]
InvalidLogBounds,
/// Returned when step size is not positive.
/// The step size provided to a parameter is not positive.
#[error("invalid step: step must be positive")]
InvalidStep,
/// Returned when categorical choices are empty.
/// A [`CategoricalParam`](crate::parameter::CategoricalParam) was created
/// with an empty choices vector.
#[error("categorical choices cannot be empty")]
EmptyChoices,
/// Returned when a parameter is suggested with a different configuration.
/// The same [`ParamId`](crate::parameter::ParamId) was suggested twice
/// with a different distribution configuration.
#[error("parameter conflict for '{name}': {reason}")]
ParameterConflict {
/// The name of the conflicting parameter.
@@ -30,30 +49,112 @@ pub enum Error {
reason: String,
},
/// Returned when requesting the best trial but no trials have completed.
/// [`Study::best_trial`](crate::Study::best_trial) or similar was called
/// before any trial completed successfully.
#[error("no completed trials available")]
NoCompletedTrials,
/// Returned when gamma is not in the valid range (0.0, 1.0).
/// The gamma value for TPE sampling is outside the open interval (0, 1).
#[error("invalid gamma: {0} must be in (0.0, 1.0)")]
InvalidGamma(f64),
/// Returned when bandwidth is not positive.
/// A KDE bandwidth value is not positive.
#[error("invalid bandwidth: {0} must be positive")]
InvalidBandwidth(f64),
/// Returned when KDE is created with empty samples.
/// A kernel density estimator was constructed with no samples.
#[error("KDE requires at least one sample")]
EmptySamples,
/// Returned when an internal invariant is violated.
/// Multivariate KDE samples have zero dimensions.
#[error("multivariate KDE samples must have at least one dimension")]
ZeroDimensions,
/// A sample in the multivariate KDE has a different number of dimensions
/// than the first sample.
#[error(
"dimension mismatch: expected {expected} dimensions but sample {sample_index} has {got}"
)]
DimensionMismatch {
/// The expected number of dimensions.
expected: usize,
/// The actual number of dimensions in the sample.
got: usize,
/// The index of the sample with mismatched dimensions.
sample_index: usize,
},
/// The bandwidth vector length does not match the number of KDE dimensions.
#[error("bandwidth dimension mismatch: expected {expected} bandwidths but got {got}")]
BandwidthDimensionMismatch {
/// The expected number of bandwidths.
expected: usize,
/// The actual number of bandwidths provided.
got: usize,
},
/// The objective signalled that this trial should be pruned (stopped
/// early). Typically raised via `Err(TrialPruned)?` inside the
/// objective closure.
#[error("trial was pruned")]
TrialPruned,
/// The multi-objective closure returned a different number of values
/// than the number of directions configured on the study.
#[error("objective dimension mismatch: expected {expected} values, got {got}")]
ObjectiveDimensionMismatch {
/// The expected number of objective values.
expected: usize,
/// The actual number of objective values returned.
got: usize,
},
/// An internal invariant was violated. This indicates a bug in the
/// library rather than a user error.
#[error("internal error: {0}")]
Internal(&'static str),
/// Returned when an async task fails.
/// An async worker task failed. Only available with the `async` feature.
#[cfg(feature = "async")]
#[error("async task error: {0}")]
TaskError(String),
/// A storage I/O operation failed. Only available with the `journal`
/// feature.
#[cfg(feature = "journal")]
#[error("storage error: {0}")]
Storage(String),
}
/// A convenience alias for `core::result::Result<T, Error>`.
pub type Result<T> = core::result::Result<T, Error>;
/// Convenience type for signalling a pruned trial from an objective function.
///
/// Implements `Into<Error>` so it can be used with `?` in objectives that
/// return `Result<V, Error>`.
///
/// # Examples
///
/// ```
/// use optimizer::{Error, TrialPruned};
///
/// fn objective_that_prunes() -> Result<f64, Error> {
/// // ... some computation ...
/// Err(TrialPruned)?
/// }
/// ```
#[derive(Debug)]
pub struct TrialPruned;
impl core::fmt::Display for TrialPruned {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "trial was pruned")
}
}
impl From<TrialPruned> for Error {
fn from(_: TrialPruned) -> Self {
Error::TrialPruned
}
}
+604
View File
@@ -0,0 +1,604 @@
//! fANOVA (functional ANOVA) parameter importance via random forest.
//!
//! fANOVA decomposes the variance of the objective function into
//! contributions from individual parameters (**main effects**) and
//! parameter pairs (**interaction effects**). This helps answer the
//! question: *"Which parameters matter most, and do any parameters
//! interact?"*
//!
//! # Algorithm
//!
//! 1. Fit a random forest to the mapping `(parameters) → objective`
//! 2. Apply functional ANOVA decomposition to the trained forest
//! 3. Compute main effects: the variance explained by each parameter alone
//! 4. Compute interaction effects: the additional variance explained by
//! pairs of parameters beyond their individual contributions
//! 5. Normalize so all importances sum to 1.0
//!
//! # When to use
//!
//! - **After optimization**: call [`Study::fanova()`](crate::Study::fanova)
//! or [`Study::fanova_with_config()`](crate::Study::fanova_with_config)
//! to identify which parameters had the most impact
//! - **Interaction detection**: unlike Spearman correlation
//! ([`Study::param_importance()`](crate::Study::param_importance)),
//! fANOVA can detect non-linear relationships and parameter interactions
//! - **Hyperparameter tuning**: focus tuning effort on high-importance
//! parameters and fix low-importance ones to reasonable defaults
//!
//! # Reference
//!
//! Hutter, F., Hoos, H. & Leyton-Brown, K. (2014). "An Efficient
//! Approach for Assessing Hyperparameter Importance." ICML 2014.
//!
//! # Example
//!
//! ```
//! use optimizer::prelude::*;
//!
//! let study: Study<f64> = Study::new(Direction::Minimize);
//! let x = FloatParam::new(0.0, 10.0).name("x");
//! let y = FloatParam::new(0.0, 10.0).name("y");
//!
//! study
//! .optimize(50, |trial: &mut optimizer::Trial| {
//! let xv = x.suggest(trial)?;
//! let yv = y.suggest(trial)?;
//! // x matters much more than y
//! Ok::<_, optimizer::Error>(3.0 * xv + 0.1 * yv)
//! })
//! .unwrap();
//!
//! let result = study.fanova().unwrap();
//! // Main effects sorted by descending importance
//! assert_eq!(result.main_effects[0].0, "x");
//! ```
/// Result of fANOVA analysis.
///
/// All importance values are fractions of total variance and sum to 1.0
/// across main effects and interactions combined.
#[derive(Debug, Clone)]
pub struct FanovaResult {
/// Per-parameter importance (fraction of total variance explained).
///
/// Sorted by descending importance. Each entry is
/// `(parameter_name, importance)` where importance is in `[0.0, 1.0]`.
pub main_effects: Vec<(String, f64)>,
/// Pairwise interaction importance (fraction of total variance explained).
///
/// Sorted by descending importance. Each entry is
/// `((param_a, param_b), importance)`. Only pairs with non-negligible
/// interaction (> 1e-10) are included.
pub interactions: Vec<((String, String), f64)>,
}
/// Configuration for fANOVA analysis.
///
/// Use [`Default::default()`] for reasonable settings, or customize
/// the random forest parameters for specific needs. Pass to
/// [`Study::fanova_with_config()`](crate::Study::fanova_with_config).
#[derive(Debug, Clone)]
pub struct FanovaConfig {
/// Number of trees in the random forest (default: 64).
pub n_trees: usize,
/// Maximum depth of each tree. `None` for unlimited (default: `None`).
pub max_depth: Option<usize>,
/// Minimum samples required to split a node (default: 2).
pub min_samples_split: usize,
/// Minimum samples required in a leaf node (default: 1).
pub min_samples_leaf: usize,
/// Random seed for reproducibility (default: `Some(42)`).
pub seed: Option<u64>,
}
impl Default for FanovaConfig {
fn default() -> Self {
Self {
n_trees: 64,
max_depth: None,
min_samples_split: 2,
min_samples_leaf: 1,
seed: Some(42),
}
}
}
// --- Decision Tree ---
/// A node in the regression tree (arena-allocated).
#[derive(Debug, Clone)]
enum TreeNode {
Leaf {
value: f64,
n_samples: usize,
},
Split {
feature: usize,
threshold: f64,
left: usize,
right: usize,
n_samples: usize,
},
}
/// A regression decision tree for fANOVA.
#[derive(Debug, Clone)]
struct DecisionTree {
nodes: Vec<TreeNode>,
}
impl DecisionTree {
/// Build a tree from the given data using the specified bootstrap indices.
fn build(
data: &[Vec<f64>],
targets: &[f64],
indices: &[usize],
config: &FanovaConfig,
rng: &mut fastrand::Rng,
) -> Self {
let mut tree = Self { nodes: Vec::new() };
tree.build_node(data, targets, indices, 0, config, rng);
tree
}
#[allow(clippy::cast_precision_loss)]
fn build_node(
&mut self,
data: &[Vec<f64>],
targets: &[f64],
indices: &[usize],
depth: usize,
config: &FanovaConfig,
rng: &mut fastrand::Rng,
) -> usize {
let n = indices.len();
let mean = indices.iter().map(|&i| targets[i]).sum::<f64>() / n as f64;
// Stopping conditions
if n < config.min_samples_split || config.max_depth.is_some_and(|d| depth >= d) {
let idx = self.nodes.len();
self.nodes.push(TreeNode::Leaf {
value: mean,
n_samples: n,
});
return idx;
}
// Pure node check (all targets identical)
#[allow(clippy::float_cmp)]
if indices.iter().all(|&i| targets[i] == targets[indices[0]]) {
let idx = self.nodes.len();
self.nodes.push(TreeNode::Leaf {
value: mean,
n_samples: n,
});
return idx;
}
let n_features = data[0].len();
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
let max_features = ((n_features as f64).sqrt().ceil() as usize)
.max(1)
.min(n_features);
let candidates = partial_shuffle(n_features, max_features, rng);
// Total variance at this node
let total_var: f64 = indices.iter().map(|&i| (targets[i] - mean).powi(2)).sum();
if total_var == 0.0 {
let idx = self.nodes.len();
self.nodes.push(TreeNode::Leaf {
value: mean,
n_samples: n,
});
return idx;
}
let mut best_score = f64::NEG_INFINITY;
let mut best_feature = 0;
let mut best_threshold = 0.0;
for &feat in &candidates {
let mut values: Vec<f64> = indices.iter().map(|&i| data[i][feat]).collect();
values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
values.dedup();
if values.len() < 2 {
continue;
}
for w in values.windows(2) {
let threshold = f64::midpoint(w[0], w[1]);
let (l_sum, l_sq, l_n, r_sum, r_sq, r_n) =
split_stats(data, targets, indices, feat, threshold);
if l_n < config.min_samples_leaf || r_n < config.min_samples_leaf {
continue;
}
let l_var = l_sq - l_sum * l_sum / l_n as f64;
let r_var = r_sq - r_sum * r_sum / r_n as f64;
let score = total_var - l_var - r_var;
if score > best_score {
best_score = score;
best_feature = feat;
best_threshold = threshold;
}
}
}
if best_score <= 0.0 {
let idx = self.nodes.len();
self.nodes.push(TreeNode::Leaf {
value: mean,
n_samples: n,
});
return idx;
}
let (left_indices, right_indices): (Vec<usize>, Vec<usize>) = indices
.iter()
.partition(|&&i| data[i][best_feature] <= best_threshold);
if left_indices.is_empty() || right_indices.is_empty() {
let idx = self.nodes.len();
self.nodes.push(TreeNode::Leaf {
value: mean,
n_samples: n,
});
return idx;
}
// Reserve slot for this split node (placeholder replaced below)
let node_idx = self.nodes.len();
self.nodes.push(TreeNode::Leaf {
value: 0.0,
n_samples: 0,
});
let left = self.build_node(data, targets, &left_indices, depth + 1, config, rng);
let right = self.build_node(data, targets, &right_indices, depth + 1, config, rng);
self.nodes[node_idx] = TreeNode::Split {
feature: best_feature,
threshold: best_threshold,
left,
right,
n_samples: n,
};
node_idx
}
/// Compute marginal prediction for a given feature subset.
///
/// Features in `subset` use values from `feature_values`.
/// Features not in `subset` are marginalized by weighting branches
/// proportionally to their training-data fractions.
fn marginal_predict(&self, subset: &[usize], feature_values: &[f64]) -> f64 {
self.marginal_predict_at(0, subset, feature_values)
}
#[allow(clippy::cast_precision_loss)]
fn marginal_predict_at(&self, idx: usize, subset: &[usize], vals: &[f64]) -> f64 {
match self.nodes[idx] {
TreeNode::Leaf { value, .. } => value,
TreeNode::Split {
feature,
threshold,
left,
right,
n_samples,
} => {
if subset.contains(&feature) {
if vals[feature] <= threshold {
self.marginal_predict_at(left, subset, vals)
} else {
self.marginal_predict_at(right, subset, vals)
}
} else {
let l_n = self.n_samples(left) as f64;
let r_n = self.n_samples(right) as f64;
let total = n_samples as f64;
(l_n / total) * self.marginal_predict_at(left, subset, vals)
+ (r_n / total) * self.marginal_predict_at(right, subset, vals)
}
}
}
}
fn n_samples(&self, idx: usize) -> usize {
match self.nodes[idx] {
TreeNode::Leaf { n_samples, .. } | TreeNode::Split { n_samples, .. } => n_samples,
}
}
}
// --- Helper Functions ---
/// Select `k` random indices from `0..n` using partial Fisher-Yates shuffle.
fn partial_shuffle(n: usize, k: usize, rng: &mut fastrand::Rng) -> Vec<usize> {
let mut indices: Vec<usize> = (0..n).collect();
let k = k.min(n);
for i in 0..k {
let j = rng.usize(i..n);
indices.swap(i, j);
}
indices.truncate(k);
indices
}
/// Compute left/right split statistics for variance reduction.
#[allow(clippy::cast_precision_loss)]
fn split_stats(
data: &[Vec<f64>],
targets: &[f64],
indices: &[usize],
feature: usize,
threshold: f64,
) -> (f64, f64, usize, f64, f64, usize) {
let (mut l_sum, mut l_sq, mut l_n) = (0.0, 0.0, 0usize);
let (mut r_sum, mut r_sq, mut r_n) = (0.0, 0.0, 0usize);
for &i in indices {
let y = targets[i];
if data[i][feature] <= threshold {
l_sum += y;
l_sq += y * y;
l_n += 1;
} else {
r_sum += y;
r_sq += y * y;
r_n += 1;
}
}
(l_sum, l_sq, l_n, r_sum, r_sq, r_n)
}
/// Population variance of a slice.
#[allow(clippy::cast_precision_loss)]
fn variance(values: &[f64]) -> f64 {
if values.is_empty() {
return 0.0;
}
let n = values.len() as f64;
let mean = values.iter().sum::<f64>() / n;
values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n
}
// --- Public API ---
/// Run fANOVA analysis on pre-processed numerical data.
///
/// `data` is `n_samples` rows, each with `n_features` columns.
/// `targets` has one entry per sample.
/// `feature_names` maps feature index to human-readable name.
#[allow(clippy::cast_precision_loss)]
pub(crate) fn compute_fanova(
data: &[Vec<f64>],
targets: &[f64],
feature_names: &[String],
config: &FanovaConfig,
) -> FanovaResult {
let n_samples = data.len();
let n_features = data[0].len();
let mut rng: fastrand::Rng = config
.seed
.map_or_else(fastrand::Rng::new, fastrand::Rng::with_seed);
// Build random forest with bootstrap sampling
let trees: Vec<DecisionTree> = (0..config.n_trees)
.map(|_| {
let bootstrap: Vec<usize> = (0..n_samples).map(|_| rng.usize(0..n_samples)).collect();
DecisionTree::build(data, targets, &bootstrap, config, &mut rng)
})
.collect();
// Compute main effects: V_j = Var[E[f | x_j]]
let main_var: Vec<f64> = (0..n_features)
.map(|j| {
let subset = [j];
let preds: Vec<f64> = (0..n_samples)
.map(|i| {
trees
.iter()
.map(|t| t.marginal_predict(&subset, &data[i]))
.sum::<f64>()
/ trees.len() as f64
})
.collect();
variance(&preds)
})
.collect();
// Compute pairwise interaction effects: V_{j,k} - V_j - V_k
let mut interactions: Vec<((String, String), f64)> = Vec::new();
for j in 0..n_features {
for k in (j + 1)..n_features {
let subset = [j, k];
let preds: Vec<f64> = (0..n_samples)
.map(|i| {
trees
.iter()
.map(|t| t.marginal_predict(&subset, &data[i]))
.sum::<f64>()
/ trees.len() as f64
})
.collect();
let joint = variance(&preds);
let interaction = (joint - main_var[j] - main_var[k]).max(0.0);
if interaction > 1e-10 {
interactions.push((
(feature_names[j].clone(), feature_names[k].clone()),
interaction,
));
}
}
}
// Normalize so all importances sum to 1.0
let total: f64 =
main_var.iter().sum::<f64>() + interactions.iter().map(|(_, v)| *v).sum::<f64>();
let mut main_effects: Vec<(String, f64)> = feature_names
.iter()
.zip(&main_var)
.map(|(name, &v)| (name.clone(), if total > 0.0 { v / total } else { 0.0 }))
.collect();
main_effects.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(core::cmp::Ordering::Equal));
if total > 0.0 {
for entry in &mut interactions {
entry.1 /= total;
}
}
interactions.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(core::cmp::Ordering::Equal));
FanovaResult {
main_effects,
interactions,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::rng_util;
#[test]
fn single_dominant_parameter() {
// f(x, y) = x — only x matters
let mut rng = fastrand::Rng::with_seed(0);
let n = 100;
let data: Vec<Vec<f64>> = (0..n)
.map(|_| {
vec![
rng_util::f64_range(&mut rng, 0.0, 10.0),
rng_util::f64_range(&mut rng, 0.0, 10.0),
]
})
.collect();
let targets: Vec<f64> = data.iter().map(|row| row[0]).collect();
let result = compute_fanova(
&data,
&targets,
&["x".into(), "y".into()],
&FanovaConfig::default(),
);
assert_eq!(result.main_effects[0].0, "x");
assert!(
result.main_effects[0].1 > 0.8,
"x importance = {}",
result.main_effects[0].1
);
}
#[test]
fn interaction_detection() {
// f(x, y) = x * y — both matter and interact
let mut rng = fastrand::Rng::with_seed(42);
let n = 200;
let data: Vec<Vec<f64>> = (0..n)
.map(|_| {
vec![
rng_util::f64_range(&mut rng, 0.0, 10.0),
rng_util::f64_range(&mut rng, 0.0, 10.0),
]
})
.collect();
let targets: Vec<f64> = data.iter().map(|row| row[0] * row[1]).collect();
let config = FanovaConfig {
n_trees: 128,
..FanovaConfig::default()
};
let result = compute_fanova(&data, &targets, &["x".into(), "y".into()], &config);
assert!(
!result.interactions.is_empty(),
"should detect x*y interaction"
);
assert!(
result.interactions[0].1 > 0.05,
"interaction importance = {}",
result.interactions[0].1
);
}
#[test]
fn variance_computation() {
assert!((variance(&[1.0, 2.0, 3.0, 4.0, 5.0]) - 2.0).abs() < 1e-10);
assert!(variance(&[5.0, 5.0, 5.0]).abs() < 1e-10);
assert!(variance(&[]).abs() < 1e-10);
}
#[test]
fn three_params_one_dominant() {
// f(x, y, z) = 3*x + 0.1*y + 0*z
let mut rng = fastrand::Rng::with_seed(7);
let n = 150;
let data: Vec<Vec<f64>> = (0..n)
.map(|_| {
vec![
rng_util::f64_range(&mut rng, 0.0, 10.0),
rng_util::f64_range(&mut rng, 0.0, 10.0),
rng_util::f64_range(&mut rng, 0.0, 10.0),
]
})
.collect();
let targets: Vec<f64> = data.iter().map(|r| 3.0 * r[0] + 0.1 * r[1]).collect();
let result = compute_fanova(
&data,
&targets,
&["x".into(), "y".into(), "z".into()],
&FanovaConfig::default(),
);
// x should be the most important
assert_eq!(result.main_effects[0].0, "x");
assert!(result.main_effects[0].1 > 0.5);
// z should have near-zero importance
let z_imp = result
.main_effects
.iter()
.find(|(name, _)| name == "z")
.map_or(0.0, |(_, v)| *v);
assert!(z_imp < 0.1, "z importance = {z_imp}");
}
#[test]
fn importances_sum_to_one() {
let mut rng = fastrand::Rng::with_seed(3);
let n = 100;
let data: Vec<Vec<f64>> = (0..n)
.map(|_| {
vec![
rng_util::f64_range(&mut rng, 0.0, 10.0),
rng_util::f64_range(&mut rng, 0.0, 10.0),
]
})
.collect();
let targets: Vec<f64> = data.iter().map(|r| r[0] + r[1]).collect();
let result = compute_fanova(
&data,
&targets,
&["x".into(), "y".into()],
&FanovaConfig::default(),
);
let total: f64 = result.main_effects.iter().map(|(_, v)| *v).sum::<f64>()
+ result.interactions.iter().map(|(_, v)| *v).sum::<f64>();
assert!(
(total - 1.0).abs() < 1e-10,
"importances should sum to 1.0, got {total}"
);
}
}
+115
View File
@@ -0,0 +1,115 @@
//! Parameter importance via Spearman rank correlation.
//!
//! Compute the absolute Spearman rank correlation between each parameter
//! and the objective value to estimate which parameters most influence
//! the outcome. This is a lightweight, non-parametric alternative to
//! [`fANOVA`](crate::fanova) that works well for monotonic relationships.
//!
//! # How it works
//!
//! 1. Rank parameter values and objective values independently
//! 2. Compute the Pearson correlation on the ranks (= Spearman ρ)
//! 3. Take the absolute value (direction of correlation is not relevant
//! for importance)
//!
//! # When to use
//!
//! - **Quick importance check**: call
//! [`Study::param_importance()`](crate::Study::param_importance) after
//! optimization for a fast, interpretable ranking
//! - **Monotonic relationships**: Spearman captures monotonic (not just
//! linear) correlations but may miss non-monotonic effects or interactions
//! - For interaction detection or non-linear importance, use
//! [`fANOVA`](crate::fanova) instead
/// Assign average ranks to a slice of `f64` values (handles ties).
#[allow(clippy::cast_precision_loss, clippy::float_cmp)]
pub(crate) fn rank(values: &[f64]) -> Vec<f64> {
let n = values.len();
let mut indexed: Vec<(usize, f64)> = values.iter().copied().enumerate().collect();
indexed.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
let mut ranks = vec![0.0; n];
let mut i = 0;
while i < n {
// Find the run of tied values.
let mut j = i + 1;
while j < n && indexed[j].1 == indexed[i].1 {
j += 1;
}
// Average rank for the tie group (1-based ranks).
let avg = (i + 1..=j).sum::<usize>() as f64 / (j - i) as f64;
for item in &indexed[i..j] {
ranks[item.0] = avg;
}
i = j;
}
ranks
}
/// Pearson correlation coefficient on two equal-length slices.
#[allow(clippy::cast_precision_loss)]
fn pearson(x: &[f64], y: &[f64]) -> f64 {
let n = x.len() as f64;
let mean_x = x.iter().sum::<f64>() / n;
let mean_y = y.iter().sum::<f64>() / n;
let mut cov = 0.0;
let mut var_x = 0.0;
let mut var_y = 0.0;
for (xi, yi) in x.iter().zip(y.iter()) {
let dx = xi - mean_x;
let dy = yi - mean_y;
cov += dx * dy;
var_x += dx * dx;
var_y += dy * dy;
}
let denom = (var_x * var_y).sqrt();
if denom == 0.0 { 0.0 } else { cov / denom }
}
/// Spearman rank correlation (Pearson on ranks).
pub(crate) fn spearman(x: &[f64], y: &[f64]) -> f64 {
pearson(&rank(x), &rank(y))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rank_no_ties() {
let ranks = rank(&[30.0, 10.0, 20.0]);
assert_eq!(ranks, vec![3.0, 1.0, 2.0]);
}
#[test]
fn rank_with_ties() {
let ranks = rank(&[10.0, 20.0, 20.0, 30.0]);
assert_eq!(ranks, vec![1.0, 2.5, 2.5, 4.0]);
}
#[test]
fn perfect_positive_correlation() {
let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
let y = vec![2.0, 4.0, 6.0, 8.0, 10.0];
let r = spearman(&x, &y);
assert!((r - 1.0).abs() < 1e-10);
}
#[test]
fn perfect_negative_correlation() {
let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
let y = vec![10.0, 8.0, 6.0, 4.0, 2.0];
let r = spearman(&x, &y);
assert!((r + 1.0).abs() < 1e-10);
}
#[test]
fn zero_variance_returns_zero() {
let x = vec![5.0, 5.0, 5.0];
let y = vec![1.0, 2.0, 3.0];
assert!(spearman(&x, &y).abs() < f64::EPSILON);
}
}
+13
View File
@@ -0,0 +1,13 @@
//! Kernel Density Estimation for parameter distributions.
//!
//! This module provides kernel density estimators used by TPE samplers
//! to model probability distributions over good and bad trial regions.
//!
//! - [`univariate`] - Univariate (single-parameter) KDE
//! - [`multivariate`] - Multivariate (joint-parameter) KDE for capturing parameter dependencies
mod multivariate;
mod univariate;
pub(crate) use multivariate::MultivariateKDE;
pub(crate) use univariate::KernelDensityEstimator;
+918
View File
@@ -0,0 +1,918 @@
//! Multivariate Kernel Density Estimation for joint parameter distributions.
//!
//! This module provides a multivariate kernel density estimator that captures
//! dependencies between parameters. Unlike the univariate KDE which models each
//! parameter independently, the multivariate KDE models the joint distribution
//! to better capture correlations between parameters.
use crate::error::{Error, Result};
/// A multivariate Gaussian kernel density estimator for joint distributions.
///
/// `MultivariateKDE` estimates a joint probability density function from a set
/// of multi-dimensional samples. This is used in multivariate TPE to model
/// the joint distributions l(x) and g(x) across multiple parameters simultaneously,
/// capturing their dependencies.
///
/// # Examples
///
/// ```ignore
/// use crate::kde::MultivariateKDE;
///
/// // Create samples: 3 samples with 2 dimensions each
/// let samples = vec![
/// vec![1.0, 2.0],
/// vec![1.5, 2.5],
/// vec![2.0, 3.0],
/// ];
/// let kde = MultivariateKDE::new(samples).unwrap();
///
/// // Get dimensionality
/// assert_eq!(kde.n_dims(), 2);
/// ```
#[derive(Clone, Debug)]
pub(crate) struct MultivariateKDE {
/// The sample points used to construct the KDE.
/// Each inner Vec is one sample with `n_dims` values.
samples: Vec<Vec<f64>>,
/// The bandwidth (standard deviation) for each dimension.
/// Uses a diagonal bandwidth matrix (independent bandwidths per dimension).
bandwidths: Vec<f64>,
/// The number of dimensions.
n_dims: usize,
}
impl MultivariateKDE {
/// Creates a new multivariate KDE with automatic bandwidth selection using Scott's rule.
///
/// Scott's rule for multivariate KDE sets bandwidth per dimension as:
/// `h_j = n^(-1/(d+4)) * sigma_j`
///
/// where n is the number of samples, d is the dimensionality, and `sigma_j` is the
/// standard deviation of the j-th dimension.
///
/// # Errors
///
/// Returns `Error::EmptySamples` if `samples` is empty.
/// Returns `Error::DimensionMismatch` if samples have inconsistent dimensions.
/// Returns `Error::ZeroDimensions` if samples have zero dimensions.
#[allow(clippy::cast_precision_loss)]
pub(crate) fn new(samples: Vec<Vec<f64>>) -> Result<Self> {
if samples.is_empty() {
return Err(Error::EmptySamples);
}
let n_dims = samples[0].len();
if n_dims == 0 {
return Err(Error::ZeroDimensions);
}
// Check all samples have the same dimensionality
for (i, sample) in samples.iter().enumerate() {
if sample.len() != n_dims {
return Err(Error::DimensionMismatch {
expected: n_dims,
got: sample.len(),
sample_index: i,
});
}
}
let bandwidths = Self::scotts_rule_multivariate(&samples, n_dims);
Ok(Self {
samples,
bandwidths,
n_dims,
})
}
/// Creates a new multivariate KDE with specified bandwidths.
///
/// Use this when you want explicit control over the smoothing parameters.
///
/// # Errors
///
/// Returns `Error::EmptySamples` if `samples` is empty.
/// Returns `Error::DimensionMismatch` if samples have inconsistent dimensions.
/// Returns `Error::ZeroDimensions` if samples have zero dimensions.
/// Returns `Error::BandwidthDimensionMismatch` if bandwidths length doesn't match dimensions.
/// Returns `Error::InvalidBandwidth` if any bandwidth is not positive.
#[cfg(test)]
pub(crate) fn with_bandwidths(samples: Vec<Vec<f64>>, bandwidths: Vec<f64>) -> Result<Self> {
if samples.is_empty() {
return Err(Error::EmptySamples);
}
let n_dims = samples[0].len();
if n_dims == 0 {
return Err(Error::ZeroDimensions);
}
// Check all samples have the same dimensionality
for (i, sample) in samples.iter().enumerate() {
if sample.len() != n_dims {
return Err(Error::DimensionMismatch {
expected: n_dims,
got: sample.len(),
sample_index: i,
});
}
}
// Check bandwidths length matches dimensions
if bandwidths.len() != n_dims {
return Err(Error::BandwidthDimensionMismatch {
expected: n_dims,
got: bandwidths.len(),
});
}
// Check all bandwidths are positive
for &bw in &bandwidths {
if bw <= 0.0 {
return Err(Error::InvalidBandwidth(bw));
}
}
Ok(Self {
samples,
bandwidths,
n_dims,
})
}
/// Computes bandwidths using Scott's rule for multivariate KDE.
///
/// Scott's rule for d dimensions: `h_j = n^(-1/(d+4)) * sigma_j`
#[allow(clippy::cast_precision_loss)]
fn scotts_rule_multivariate(samples: &[Vec<f64>], n_dims: usize) -> Vec<f64> {
let n = samples.len() as f64;
let d = n_dims as f64;
// Scott's rule exponent for multivariate: -1/(d+4)
let exponent = -1.0 / (d + 4.0);
let scale_factor = n.powf(exponent);
(0..n_dims)
.map(|dim| {
let std_dev = Self::dimension_std_dev(samples, dim);
// For degenerate case where all samples are identical in this dimension,
// use a small positive bandwidth
if std_dev < f64::EPSILON {
1.0
} else {
scale_factor * std_dev
}
})
.collect()
}
/// Computes the sample standard deviation for a single dimension.
#[allow(clippy::cast_precision_loss)]
fn dimension_std_dev(samples: &[Vec<f64>], dim: usize) -> f64 {
let n = samples.len() as f64;
let values: Vec<f64> = samples.iter().map(|s| s[dim]).collect();
let mean = values.iter().sum::<f64>() / n;
let variance = values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / n;
variance.sqrt()
}
/// Returns the number of dimensions.
#[cfg(test)]
pub(crate) fn n_dims(&self) -> usize {
self.n_dims
}
/// Returns the number of samples.
#[cfg(test)]
pub(crate) fn n_samples(&self) -> usize {
self.samples.len()
}
/// Returns the bandwidths for each dimension.
#[cfg(test)]
pub(crate) fn bandwidths(&self) -> &[f64] {
&self.bandwidths
}
/// Returns a reference to the samples.
#[cfg(test)]
pub(crate) fn samples(&self) -> &[Vec<f64>] {
&self.samples
}
/// Returns the log probability density at point `x`.
///
/// This computes the log-density for numerical stability, using the formula:
///
/// `log f(x) = log((1/n) * Σ_i Π_j K_hj((x_j - x_ij) / h_j))`
///
/// The computation uses the log-sum-exp trick for numerical stability:
///
/// `log(Σ exp(a_i)) = max(a) + log(Σ exp(a_i - max(a)))`
///
/// # Panics
///
/// Panics if `x.len() != self.n_dims`.
#[allow(clippy::cast_precision_loss)]
pub(crate) fn log_pdf(&self, x: &[f64]) -> f64 {
assert_eq!(
x.len(),
self.n_dims,
"Point dimension {} doesn't match KDE dimension {}",
x.len(),
self.n_dims
);
let n = self.samples.len() as f64;
// Precompute log normalization constant for each dimension
// For a Gaussian kernel: K_h(z) = (1/(h*sqrt(2*pi))) * exp(-0.5*z^2)
// log(K_h(z)) = -log(h) - 0.5*log(2*pi) - 0.5*z^2
let log_2pi = (2.0 * core::f64::consts::PI).ln();
let log_norm_per_dim: Vec<f64> = self
.bandwidths
.iter()
.map(|&h| -h.ln() - 0.5 * log_2pi)
.collect();
// Compute log of kernel contribution for each sample
// log(prod_j K_hj(z_j)) = sum_j log(K_hj(z_j))
let log_kernels: Vec<f64> = self
.samples
.iter()
.map(|sample| {
let mut log_kernel_sum = 0.0;
for j in 0..self.n_dims {
let z = (x[j] - sample[j]) / self.bandwidths[j];
log_kernel_sum += log_norm_per_dim[j] - 0.5 * z * z;
}
log_kernel_sum
})
.collect();
// Use log-sum-exp trick for numerical stability
// log(sum(exp(log_kernels))) = max + log(sum(exp(log_kernels - max)))
let max_log_kernel = log_kernels
.iter()
.copied()
.fold(f64::NEG_INFINITY, f64::max);
// Handle case where all log_kernels are -inf (extremely unlikely point)
if max_log_kernel.is_infinite() && max_log_kernel < 0.0 {
return f64::NEG_INFINITY;
}
let sum_exp: f64 = log_kernels
.iter()
.map(|&lk| (lk - max_log_kernel).exp())
.sum();
// log((1/n) * sum) = -log(n) + log(sum)
-n.ln() + max_log_kernel + sum_exp.ln()
}
/// Returns the probability density at point `x`.
///
/// The density is computed as the average of multivariate Gaussian kernels
/// centered at each sample point:
///
/// `f(x) = (1/n) * Σ_i Π_j K_hj((x_j - x_ij) / h_j)`
///
/// where `K_hj` is the univariate Gaussian kernel with bandwidth `h_j`.
///
/// This method computes in log-space for numerical stability and then
/// exponentiates the result.
///
/// # Panics
///
/// Panics if `x.len() != self.n_dims`.
#[cfg(test)]
pub(crate) fn pdf(&self, x: &[f64]) -> f64 {
self.log_pdf(x).exp()
}
/// Samples a point from the estimated joint density distribution.
///
/// Sampling works by:
/// 1. Uniformly selecting one of the kernel centers (samples)
/// 2. Adding independent Gaussian noise to each dimension with that
/// dimension's bandwidth as the standard deviation
///
/// This is equivalent to sampling from a mixture of multivariate Gaussians
/// with diagonal covariance matrices, where each mixture component is
/// centered at a sample point.
///
/// # Returns
///
/// A `Vec<f64>` of length `n_dims` representing a sample from the KDE.
pub(crate) fn sample(&self, rng: &mut fastrand::Rng) -> Vec<f64> {
// Select a random sample to center the kernel on
let idx = rng.usize(0..self.samples.len());
let center = &self.samples[idx];
// Add independent Gaussian noise to each dimension
// Using Box-Muller transform for Gaussian sampling
center
.iter()
.zip(self.bandwidths.iter())
.map(|(&center_j, &bandwidth_j)| {
let u1: f64 = rng.f64();
let u2: f64 = rng.f64();
// Box-Muller transform: generates standard normal variate
let z = (-2.0 * u1.ln()).sqrt() * (2.0 * core::f64::consts::PI * u2).cos();
// Scale by bandwidth and shift by center
center_j + z * bandwidth_j
})
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_multivariate_kde_new_basic() {
let samples = vec![vec![1.0, 2.0], vec![1.5, 2.5], vec![2.0, 3.0]];
let kde = MultivariateKDE::new(samples).unwrap();
assert_eq!(kde.n_dims(), 2);
assert_eq!(kde.n_samples(), 3);
assert_eq!(kde.bandwidths().len(), 2);
}
#[test]
fn test_multivariate_kde_new_single_sample() {
let samples = vec![vec![1.0, 2.0, 3.0]];
let kde = MultivariateKDE::new(samples).unwrap();
assert_eq!(kde.n_dims(), 3);
assert_eq!(kde.n_samples(), 1);
// With single sample, bandwidths should default to 1.0 (degenerate case)
for &bw in kde.bandwidths() {
assert!((bw - 1.0).abs() < f64::EPSILON);
}
}
#[test]
fn test_multivariate_kde_new_single_dimension() {
let samples = vec![vec![1.0], vec![2.0], vec![3.0], vec![4.0], vec![5.0]];
let kde = MultivariateKDE::new(samples).unwrap();
assert_eq!(kde.n_dims(), 1);
assert_eq!(kde.n_samples(), 5);
assert_eq!(kde.bandwidths().len(), 1);
// Bandwidth should be positive
assert!(kde.bandwidths()[0] > 0.0);
}
#[test]
fn test_multivariate_kde_scotts_rule() {
// Create samples with known statistics
// 10 samples, 2 dimensions
let samples: Vec<Vec<f64>> = (0..10)
.map(|i| {
let x = f64::from(i);
vec![x, x * 2.0] // Second dimension has 2x variance
})
.collect();
let kde = MultivariateKDE::new(samples).unwrap();
// Scott's rule: h = n^(-1/(d+4)) * sigma
// n=10, d=2: exponent = -1/6 ≈ -0.167
// 10^(-1/6) ≈ 0.681
// First dim std_dev ≈ 2.87, second ≈ 5.74
// Expected bandwidths: ~1.95 and ~3.91
let bw = kde.bandwidths();
assert!(
bw[0] > 1.0 && bw[0] < 3.0,
"First bandwidth {} unexpected",
bw[0]
);
assert!(
bw[1] > 2.0 && bw[1] < 6.0,
"Second bandwidth {} unexpected",
bw[1]
);
// Second bandwidth should be approximately 2x the first
assert!(
(bw[1] / bw[0] - 2.0).abs() < 0.1,
"Ratio {} not close to 2",
bw[1] / bw[0]
);
}
#[test]
fn test_multivariate_kde_with_bandwidths() {
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
let bandwidths = vec![0.5, 1.0];
let kde = MultivariateKDE::with_bandwidths(samples, bandwidths).unwrap();
assert_eq!(kde.n_dims(), 2);
assert!((kde.bandwidths()[0] - 0.5).abs() < f64::EPSILON);
assert!((kde.bandwidths()[1] - 1.0).abs() < f64::EPSILON);
}
#[test]
fn test_multivariate_kde_empty_samples() {
let samples: Vec<Vec<f64>> = vec![];
let result = MultivariateKDE::new(samples);
assert!(matches!(result, Err(Error::EmptySamples)));
}
#[test]
fn test_multivariate_kde_zero_dimensions() {
let samples = vec![vec![], vec![]];
let result = MultivariateKDE::new(samples);
assert!(matches!(result, Err(Error::ZeroDimensions)));
}
#[test]
fn test_multivariate_kde_dimension_mismatch() {
let samples = vec![vec![1.0, 2.0], vec![3.0]]; // Second sample has wrong dimensions
let result = MultivariateKDE::new(samples);
assert!(matches!(
result,
Err(Error::DimensionMismatch {
expected: 2,
got: 1,
sample_index: 1
})
));
}
#[test]
fn test_multivariate_kde_with_bandwidths_wrong_length() {
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
let bandwidths = vec![0.5]; // Only 1 bandwidth for 2 dimensions
let result = MultivariateKDE::with_bandwidths(samples, bandwidths);
assert!(matches!(
result,
Err(Error::BandwidthDimensionMismatch {
expected: 2,
got: 1
})
));
}
#[test]
fn test_multivariate_kde_with_bandwidths_zero() {
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
let bandwidths = vec![0.5, 0.0]; // Second bandwidth is zero
let result = MultivariateKDE::with_bandwidths(samples, bandwidths);
assert!(matches!(result, Err(Error::InvalidBandwidth(bw)) if bw == 0.0));
}
#[test]
fn test_multivariate_kde_with_bandwidths_negative() {
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
let bandwidths = vec![0.5, -1.0]; // Negative bandwidth
let result = MultivariateKDE::with_bandwidths(samples, bandwidths);
assert!(
matches!(result, Err(Error::InvalidBandwidth(bw)) if (bw - (-1.0)).abs() < f64::EPSILON)
);
}
#[test]
fn test_multivariate_kde_identical_samples() {
// All samples identical - should handle degenerate case
let samples = vec![vec![5.0, 10.0], vec![5.0, 10.0], vec![5.0, 10.0]];
let kde = MultivariateKDE::new(samples).unwrap();
// Bandwidths should default to 1.0 for degenerate dimensions
for &bw in kde.bandwidths() {
assert!((bw - 1.0).abs() < f64::EPSILON);
}
}
#[test]
fn test_multivariate_kde_high_dimensional() {
// Test with higher dimensions
let samples: Vec<Vec<f64>> = (0..20)
.map(|i| {
let x = f64::from(i);
vec![x, x * 0.5, x * 2.0, x * 0.1, x * 10.0]
})
.collect();
let kde = MultivariateKDE::new(samples).unwrap();
assert_eq!(kde.n_dims(), 5);
assert_eq!(kde.n_samples(), 20);
assert_eq!(kde.bandwidths().len(), 5);
// All bandwidths should be positive
for &bw in kde.bandwidths() {
assert!(bw > 0.0);
}
}
#[test]
fn test_multivariate_kde_samples_accessor() {
let samples = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
let kde = MultivariateKDE::new(samples.clone()).unwrap();
assert_eq!(kde.samples(), &samples);
}
// ==================== PDF Tests ====================
#[test]
fn test_multivariate_kde_pdf_basic() {
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]];
let kde = MultivariateKDE::new(samples).unwrap();
// Density should be positive everywhere
assert!(kde.pdf(&[0.0, 0.0]) > 0.0);
assert!(kde.pdf(&[1.0, 1.0]) > 0.0);
assert!(kde.pdf(&[2.0, 2.0]) > 0.0);
assert!(kde.pdf(&[0.5, 0.5]) > 0.0);
// Density should be higher near sample points
let near_density = kde.pdf(&[1.0, 1.0]);
let far_density = kde.pdf(&[10.0, 10.0]);
assert!(near_density > far_density);
}
#[test]
fn test_multivariate_kde_pdf_with_custom_bandwidths() {
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]];
let bandwidths = vec![0.5, 0.5];
let kde = MultivariateKDE::with_bandwidths(samples, bandwidths).unwrap();
// Density should be positive
assert!(kde.pdf(&[0.5, 0.5]) > 0.0);
assert!(kde.pdf(&[0.0, 0.0]) > 0.0);
}
#[test]
fn test_multivariate_kde_pdf_single_sample() {
let samples = vec![vec![5.0, 10.0]];
let kde = MultivariateKDE::new(samples).unwrap();
// Should have positive density near the sample
assert!(kde.pdf(&[5.0, 10.0]) > 0.0);
assert!(kde.pdf(&[4.5, 9.5]) > 0.0);
}
#[test]
fn test_multivariate_kde_log_pdf_consistency() {
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]];
let kde = MultivariateKDE::new(samples).unwrap();
// log_pdf and pdf should be consistent: exp(log_pdf(x)) == pdf(x)
let test_points = vec![
vec![0.0, 0.0],
vec![1.0, 1.0],
vec![0.5, 0.5],
vec![3.0, 3.0],
];
for point in test_points {
let log_p = kde.log_pdf(&point);
let p = kde.pdf(&point);
let p_from_log = log_p.exp();
assert!(
(p - p_from_log).abs() < 1e-10,
"pdf={p}, exp(log_pdf)={p_from_log}"
);
}
}
#[test]
fn test_multivariate_kde_pdf_integrates_to_one_1d() {
// Test with 1D case first (should match univariate KDE behavior)
let samples = vec![vec![0.0], vec![1.0], vec![2.0], vec![3.0], vec![4.0]];
let kde = MultivariateKDE::new(samples).unwrap();
// Numerical integration over a wide range
let n_points = 1000;
let low = -10.0;
let high = 15.0;
let dx = (high - low) / f64::from(n_points);
let integral: f64 = (0..n_points)
.map(|i| {
let x = low + (f64::from(i) + 0.5) * dx;
kde.pdf(&[x]) * dx
})
.sum();
// Should be approximately 1.0 (within numerical error)
assert!(
(integral - 1.0).abs() < 0.02,
"1D integral = {integral}, expected ~1.0"
);
}
#[test]
fn test_multivariate_kde_pdf_integrates_to_one_2d() {
// Test with 2D case
let samples = vec![
vec![0.0, 0.0],
vec![1.0, 0.0],
vec![0.0, 1.0],
vec![1.0, 1.0],
vec![0.5, 0.5],
];
let kde = MultivariateKDE::new(samples).unwrap();
// Numerical integration over a 2D grid
let n_points = 100; // 100x100 = 10000 points
let low = -5.0;
let high = 6.0;
let dx = (high - low) / f64::from(n_points);
let mut integral = 0.0;
for i in 0..n_points {
for j in 0..n_points {
let x = low + (f64::from(i) + 0.5) * dx;
let y = low + (f64::from(j) + 0.5) * dx;
integral += kde.pdf(&[x, y]) * dx * dx;
}
}
// Should be approximately 1.0 (within numerical error)
// 2D integration has more error, so use larger tolerance
assert!(
(integral - 1.0).abs() < 0.05,
"2D integral = {integral}, expected ~1.0"
);
}
#[test]
fn test_multivariate_kde_pdf_symmetry() {
// With symmetric samples, PDF should be symmetric
let samples = vec![
vec![1.0, 0.0],
vec![-1.0, 0.0],
vec![0.0, 1.0],
vec![0.0, -1.0],
];
let kde = MultivariateKDE::new(samples).unwrap();
// Density at symmetric points should be equal
let d1 = kde.pdf(&[0.5, 0.0]);
let d2 = kde.pdf(&[-0.5, 0.0]);
assert!(
(d1 - d2).abs() < 1e-10,
"Symmetric points have different densities: {d1} vs {d2}"
);
let d3 = kde.pdf(&[0.0, 0.5]);
let d4 = kde.pdf(&[0.0, -0.5]);
assert!(
(d3 - d4).abs() < 1e-10,
"Symmetric points have different densities: {d3} vs {d4}"
);
}
#[test]
fn test_multivariate_kde_pdf_high_dimensional() {
// Test with higher dimensions
let samples: Vec<Vec<f64>> = (0..10)
.map(|i| {
let x = f64::from(i) * 0.1;
vec![x, x, x, x, x] // 5D
})
.collect();
let kde = MultivariateKDE::new(samples).unwrap();
// Density should be positive
assert!(kde.pdf(&[0.5, 0.5, 0.5, 0.5, 0.5]) > 0.0);
assert!(kde.pdf(&[0.0, 0.0, 0.0, 0.0, 0.0]) > 0.0);
// Log PDF should be finite
let log_p = kde.log_pdf(&[0.5, 0.5, 0.5, 0.5, 0.5]);
assert!(log_p.is_finite());
}
#[test]
fn test_multivariate_kde_pdf_numerical_stability() {
// Test numerical stability with points far from samples
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]];
let kde = MultivariateKDE::new(samples).unwrap();
// Very far point should have very small but non-negative density
let far_pdf = kde.pdf(&[100.0, 100.0]);
assert!(far_pdf >= 0.0);
assert!(far_pdf.is_finite() || far_pdf == 0.0);
// Log PDF should be finite (or -inf for zero density)
let far_log_pdf = kde.log_pdf(&[100.0, 100.0]);
assert!(far_log_pdf.is_finite() || far_log_pdf.is_infinite());
}
#[test]
#[should_panic(expected = "Point dimension")]
fn test_multivariate_kde_pdf_wrong_dimension() {
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0]];
let kde = MultivariateKDE::new(samples).unwrap();
// Should panic with wrong dimension
let _ = kde.pdf(&[0.0]); // Only 1 value for 2D KDE
}
// ==================== Sampling Tests ====================
#[test]
fn test_multivariate_kde_sample_basic() {
let samples = vec![vec![0.0, 0.0], vec![1.0, 1.0], vec![2.0, 2.0]];
let kde = MultivariateKDE::new(samples).unwrap();
let mut rng = fastrand::Rng::new();
// Sample should have correct dimensionality
let sample = kde.sample(&mut rng);
assert_eq!(sample.len(), 2);
}
#[test]
fn test_multivariate_kde_sample_in_reasonable_range() {
let samples = vec![
vec![0.0, 0.0],
vec![1.0, 1.0],
vec![2.0, 2.0],
vec![3.0, 3.0],
vec![4.0, 4.0],
];
let kde = MultivariateKDE::new(samples).unwrap();
let mut rng = fastrand::Rng::new();
// 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[0] > -10.0 && s[0] < 15.0,
"Sample dimension 0: {} outside expected range",
s[0]
);
assert!(
s[1] > -10.0 && s[1] < 15.0,
"Sample dimension 1: {} outside expected range",
s[1]
);
}
}
#[test]
fn test_multivariate_kde_sample_single_sample() {
// When KDE has only one sample, all samples should be centered around it
let samples = vec![vec![5.0, 10.0]];
let kde = MultivariateKDE::new(samples).unwrap();
let mut rng = fastrand::Rng::new();
// Generate many samples and check they cluster around (5.0, 10.0)
let n_samples = 100;
let mut sum_x = 0.0;
let mut sum_y = 0.0;
for _ in 0..n_samples {
let s = kde.sample(&mut rng);
sum_x += s[0];
sum_y += s[1];
}
let mean_x = sum_x / f64::from(n_samples);
let mean_y = sum_y / f64::from(n_samples);
// Mean should be close to the single sample point
// With 100 samples and bandwidth=1.0, we expect mean within ~0.3 of center
assert!(
(mean_x - 5.0).abs() < 1.0,
"Mean x={mean_x}, expected close to 5.0"
);
assert!(
(mean_y - 10.0).abs() < 1.0,
"Mean y={mean_y}, expected close to 10.0"
);
}
#[test]
fn test_multivariate_kde_sample_high_dimensional() {
// Test sampling in higher dimensions
let samples: Vec<Vec<f64>> = (0..10)
.map(|i| {
let x = f64::from(i) * 0.5;
vec![x, x * 2.0, x * 0.5, x + 1.0, x - 1.0] // 5D
})
.collect();
let kde = MultivariateKDE::new(samples).unwrap();
let mut rng = fastrand::Rng::new();
// Sample should have correct dimensionality
for _ in 0..50 {
let sample = kde.sample(&mut rng);
assert_eq!(sample.len(), 5);
// All values should be finite
for &val in &sample {
assert!(val.is_finite(), "Sample value is not finite: {val}");
}
}
}
#[test]
#[allow(clippy::cast_precision_loss)]
fn test_multivariate_kde_sample_respects_bandwidth() {
// Create samples all at origin with large custom bandwidth
let data = vec![vec![0.0, 0.0], vec![0.0, 0.0], vec![0.0, 0.0]];
let bandwidths = vec![0.1, 10.0]; // Small bandwidth in x, large in y
let kde = MultivariateKDE::with_bandwidths(data, bandwidths).unwrap();
let mut rng = fastrand::Rng::new();
// Generate samples and check variance in each dimension
let n_samples = 1000;
let mut values_x: Vec<f64> = Vec::with_capacity(n_samples);
let mut values_y: Vec<f64> = Vec::with_capacity(n_samples);
for _ in 0..n_samples {
let s = kde.sample(&mut rng);
values_x.push(s[0]);
values_y.push(s[1]);
}
// Compute sample variances
let n = n_samples as f64;
let mean_x: f64 = values_x.iter().sum::<f64>() / n;
let mean_y: f64 = values_y.iter().sum::<f64>() / n;
let var_x: f64 = values_x.iter().map(|x| (x - mean_x).powi(2)).sum::<f64>() / n;
let var_y: f64 = values_y.iter().map(|y| (y - mean_y).powi(2)).sum::<f64>() / n;
// Variance should be approximately bandwidth^2
// x: bandwidth=0.1, expected variance ~0.01
// y: bandwidth=10.0, expected variance ~100.0
assert!(
var_x < 0.05,
"X variance {var_x} too large for bandwidth 0.1"
);
assert!(
var_y > 50.0 && var_y < 200.0,
"Y variance {var_y} unexpected for bandwidth 10.0"
);
}
#[test]
fn test_multivariate_kde_sample_distribution_shape() {
// Create samples along a diagonal line
let data = vec![
vec![0.0, 0.0],
vec![1.0, 1.0],
vec![2.0, 2.0],
vec![3.0, 3.0],
vec![4.0, 4.0],
];
let kde = MultivariateKDE::new(data).unwrap();
let mut rng = fastrand::Rng::new();
// Sample many points and verify the mean is near the center
let n_samples = 500;
let mut sum = [0.0, 0.0];
for _ in 0..n_samples {
let s = kde.sample(&mut rng);
sum[0] += s[0];
sum[1] += s[1];
}
let mean_x = sum[0] / f64::from(n_samples);
let mean_y = sum[1] / f64::from(n_samples);
// Mean should be close to (2.0, 2.0) - the center of the samples
assert!(
(mean_x - 2.0).abs() < 0.5,
"Mean x={mean_x}, expected close to 2.0"
);
assert!(
(mean_y - 2.0).abs() < 0.5,
"Mean y={mean_y}, expected close to 2.0"
);
}
#[test]
fn test_multivariate_kde_sample_deterministic_with_seeded_rng() {
let data = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
let kde = MultivariateKDE::new(data).unwrap();
// Use a seeded RNG for reproducibility
let mut rng1 = fastrand::Rng::with_seed(42);
let mut rng2 = fastrand::Rng::with_seed(42);
// Same seed should produce same samples
let result1 = kde.sample(&mut rng1);
let result2 = kde.sample(&mut rng2);
assert!(
(result1[0] - result2[0]).abs() < f64::EPSILON,
"Samples with same seed differ in dimension 0"
);
assert!(
(result1[1] - result2[1]).abs() < f64::EPSILON,
"Samples with same seed differ in dimension 1"
);
}
}
+6 -8
View File
@@ -3,8 +3,6 @@
//! 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;
use crate::error::{Error, Result};
/// A Gaussian kernel density estimator for continuous distributions.
@@ -26,7 +24,7 @@ use crate::error::{Error, Result};
/// assert!(density > 0.0);
///
/// // Sample from the estimated distribution
/// let mut rng = rand::rng();
/// let mut rng = fastrand::Rng::new();
/// let sample = kde.sample(&mut rng);
/// ```
#[derive(Clone, Debug)]
@@ -132,15 +130,15 @@ impl KernelDensityEstimator {
/// Sampling works by:
/// 1. Uniformly selecting one of the kernel centers (samples)
/// 2. Adding Gaussian noise with the bandwidth as standard deviation
pub(crate) fn sample<R: Rng>(&self, rng: &mut R) -> f64 {
pub(crate) fn sample(&self, rng: &mut fastrand::Rng) -> f64 {
// Select a random sample to center the kernel on
let idx = rng.random_range(0..self.samples.len());
let idx = rng.usize(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 u1: f64 = rng.f64();
let u2: f64 = rng.f64();
let z = (-2.0 * u1.ln()).sqrt() * (2.0 * core::f64::consts::PI * u2).cos();
center + z * self.bandwidth
@@ -211,7 +209,7 @@ mod tests {
fn test_kde_sample_in_reasonable_range() {
let samples = vec![0.0, 1.0, 2.0, 3.0, 4.0];
let kde = KernelDensityEstimator::new(samples).unwrap();
let mut rng = rand::rng();
let mut rng = fastrand::Rng::new();
// Samples should generally be in a reasonable range around the data
for _ in 0..100 {
+141 -139
View File
@@ -9,173 +9,175 @@
#![deny(clippy::pedantic)]
#![deny(clippy::std_instead_of_core)]
//! A black-box optimization library with multiple sampling strategies.
//! Bayesian and population-based optimization library with an Optuna-like API
//! for hyperparameter tuning and black-box optimization. It ships 12 samplers
//! (from random search to CMA-ES and NSGA-III), 8 pruners, async/parallel
//! evaluation, and optional journal-based persistence — all with zero required
//! feature flags for the common case.
//!
//! This library provides an Optuna-like API for hyperparameter optimization
//! with support for multiple sampling algorithms:
//! # Getting Started
//!
//! - **Random Search** - Simple random sampling for baseline comparisons
//! - **TPE (Tree-Parzen Estimator)** - Bayesian optimization for efficient search
//! - **Grid Search** - Exhaustive search over a specified parameter grid
//!
//! Additional features include:
//!
//! - Float, integer, and categorical parameter types
//! - Log-scale and stepped parameter sampling
//! - Synchronous and async optimization
//! - Parallel trial evaluation with bounded concurrency
//!
//! # Quick Start
//! Minimize a function in five lines — no feature flags needed:
//!
//! ```
//! use optimizer::sampler::tpe::TpeSampler;
//! use optimizer::{Direction, Study};
//! use optimizer::prelude::*;
//!
//! // Create a study with TPE sampler
//! let sampler = TpeSampler::builder().seed(42).build().unwrap();
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! let study: Study<f64> = Study::new(Direction::Minimize);
//! let x = FloatParam::new(-10.0, 10.0).name("x");
//!
//! // Optimize x^2 for 20 trials
//! study
//! .optimize_with_sampler(20, |trial| {
//! let x = trial.suggest_float("x", -10.0, 10.0)?;
//! Ok::<_, optimizer::Error>(x * x)
//! .optimize(50, |trial: &mut optimizer::Trial| {
//! let v = x.suggest(trial)?;
//! Ok::<_, Error>((v - 3.0).powi(2))
//! })
//! .unwrap();
//!
//! // Get the best result
//! let best = study.best_trial().unwrap();
//! println!("Best value: {} at x={:?}", best.value, best.params);
//! println!("x = {:.4}, f(x) = {:.4}", best.get(&x).unwrap(), best.value);
//! ```
//!
//! # Creating a Study
//! # Core Concepts
//!
//! A [`Study`] manages optimization trials. Create one with an optimization direction:
//! | Type | Role |
//! |------|------|
//! | [`Study`] | Drive an optimization loop: create trials, record results, track the best. |
//! | [`Trial`] | A single evaluation of the objective function, carrying suggested parameter values. |
//! | [`Parameter`](parameter::Parameter) | Define the search space — [`FloatParam`](parameter::FloatParam), [`IntParam`](parameter::IntParam), [`CategoricalParam`](parameter::CategoricalParam), [`BoolParam`](parameter::BoolParam), [`EnumParam`](parameter::EnumParam). |
//! | [`Sampler`](sampler::Sampler) | Strategy for choosing the next point to evaluate (TPE, CMA-ES, random, etc.). |
//! | [`Direction`] | Whether the study minimizes or maximizes the objective value. |
//!
//! ```
//! use optimizer::sampler::random::RandomSampler;
//! use optimizer::sampler::tpe::TpeSampler;
//! use optimizer::{Direction, Study};
//! # Sampler Guide
//!
//! // Minimize with default random sampler
//! let study: Study<f64> = Study::new(Direction::Minimize);
//! ## Single-objective samplers
//!
//! // Maximize with TPE sampler
//! let study: Study<f64> = Study::with_sampler(Direction::Maximize, TpeSampler::new());
//! | Sampler | Algorithm | Best for | Feature flag |
//! |---------|-----------|----------|--------------|
//! | [`RandomSampler`](sampler::RandomSampler) | Uniform random | Baselines, high-dimensional | — |
//! | [`TpeSampler`](sampler::TpeSampler) | Tree-Parzen Estimator | General-purpose Bayesian | — |
//! | [`GridSearchSampler`](sampler::GridSampler) | Exhaustive grid | Small, discrete spaces | — |
//! | `SobolSampler` | Sobol quasi-random sequence | Space-filling, low dimensions | `sobol` |
//! | `CmaEsSampler` | CMA-ES | Continuous, moderate dimensions | `cma-es` |
//! | `GpSampler` | Gaussian Process + EI | Expensive objectives, few trials | `gp` |
//! | [`DESampler`](sampler::DESampler) | Differential Evolution | Non-convex, population-based | — |
//! | [`BohbSampler`](sampler::BohbSampler) | BOHB (TPE + `HyperBand`) | Budget-aware early stopping | — |
//!
//! // With seeded sampler for reproducibility
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
//! ```
//! ## Multi-objective samplers
//!
//! # Suggesting Parameters
//!
//! Within the objective function, use [`Trial`] to suggest parameter values:
//!
//! ```
//! use optimizer::{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::<_, optimizer::Error>(x * n as f64)
//! })
//! .unwrap();
//! ```
//!
//! # Available Samplers
//!
//! ## Random Search
//!
//! The simplest sampling strategy, useful for baselines:
//!
//! ```
//! use optimizer::sampler::random::RandomSampler;
//! use optimizer::{Direction, Study};
//!
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
//! ```
//!
//! ## TPE (Tree-Parzen Estimator)
//!
//! Bayesian optimization that learns from previous trials:
//!
//! ```
//! use optimizer::sampler::tpe::TpeSampler;
//!
//! let sampler = TpeSampler::builder()
//! .gamma(0.15) // Quantile for good/bad split
//! .n_startup_trials(20) // Random trials before TPE
//! .n_ei_candidates(32) // Candidates to evaluate
//! .seed(42) // Reproducibility
//! .build()
//! .unwrap();
//! ```
//!
//! ## Grid Search
//!
//! Exhaustive search over a discretized parameter space:
//!
//! ```
//! use optimizer::sampler::grid::GridSearchSampler;
//! use optimizer::{Direction, Study};
//!
//! let sampler = GridSearchSampler::builder()
//! .n_points_per_param(10) // Points per parameter dimension
//! .build();
//!
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! ```
//!
//! # Async and Parallel Optimization
//!
//! With the `async` feature enabled, you can run trials asynchronously:
//!
//! ```ignore
//! use optimizer::{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?;
//! ```
//! | Sampler | Algorithm | Best for | Feature flag |
//! |---------|-----------|----------|--------------|
//! | [`Nsga2Sampler`](sampler::Nsga2Sampler) | NSGA-II | 2-3 objectives | — |
//! | [`Nsga3Sampler`](sampler::Nsga3Sampler) | NSGA-III (reference-point) | 3+ objectives | — |
//! | [`MoeadSampler`](sampler::MoeadSampler) | MOEA/D (decomposition) | Many objectives, structured fronts | — |
//! | [`MotpeSampler`](sampler::MotpeSampler) | Multi-Objective TPE | Bayesian multi-objective | — |
//!
//! # Feature Flags
//!
//! - `async`: Enable async optimization methods (requires tokio)
//! | Flag | What it enables | Default |
//! |------|----------------|---------|
//! | `async` | Async/parallel optimization via tokio (`Study::optimize_async`, `Study::optimize_parallel`) | off |
//! | `derive` | `#[derive(Categorical)]` for enum parameters | off |
//! | `serde` | `Serialize`/`Deserialize` on public types, `Study::save`/`Study::load` | off |
//! | `journal` | `JournalStorage` — JSONL persistence with file locking (enables `serde`) | off |
//! | `sobol` | `SobolSampler` — quasi-random low-discrepancy sequences | off |
//! | `cma-es` | `CmaEsSampler` — Covariance Matrix Adaptation Evolution Strategy | off |
//! | `gp` | `GpSampler` — Gaussian Process surrogate with Expected Improvement | off |
//! | `tracing` | Structured log events via [`tracing`](https://docs.rs/tracing) at key optimization points | off |
mod distribution;
/// Emit a `tracing::info!` event when the `tracing` feature is enabled.
/// No-op otherwise.
#[cfg(feature = "tracing")]
macro_rules! trace_info {
($($arg:tt)*) => { tracing::info!($($arg)*) };
}
#[cfg(not(feature = "tracing"))]
macro_rules! trace_info {
($($arg:tt)*) => {};
}
/// Emit a `tracing::debug!` event when the `tracing` feature is enabled.
/// No-op otherwise.
#[cfg(feature = "tracing")]
macro_rules! trace_debug {
($($arg:tt)*) => { tracing::debug!($($arg)*) };
}
#[cfg(not(feature = "tracing"))]
macro_rules! trace_debug {
($($arg:tt)*) => {};
}
pub mod distribution;
mod error;
mod fanova;
mod importance;
mod kde;
mod param;
pub mod multi_objective;
pub mod objective;
pub mod param;
pub mod parameter;
pub mod pareto;
pub mod pruner;
mod rng_util;
pub mod sampler;
pub mod storage;
mod study;
mod trial;
mod types;
mod visualization;
pub use error::{Error, Result};
pub use param::ParamValue;
pub use study::Study;
pub use trial::Trial;
pub use error::{Error, Result, TrialPruned};
pub use fanova::{FanovaConfig, FanovaResult};
pub use objective::Objective;
#[cfg(feature = "derive")]
pub use optimizer_derive::Categorical;
#[cfg(feature = "serde")]
pub use study::StudySnapshot;
pub use study::{Study, StudyBuilder};
pub use trial::{AttrValue, Trial};
pub use types::{Direction, TrialState};
pub use visualization::generate_html_report;
/// Convenient wildcard import for the most common types.
///
/// ```
/// use optimizer::prelude::*;
/// ```
pub mod prelude {
#[cfg(feature = "derive")]
pub use optimizer_derive::Categorical as DeriveCategory;
pub use crate::error::{Error, Result, TrialPruned};
pub use crate::fanova::{FanovaConfig, FanovaResult};
pub use crate::multi_objective::{
MultiObjectiveSampler, MultiObjectiveStudy, MultiObjectiveTrial,
};
pub use crate::objective::Objective;
pub use crate::parameter::{
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, ParamValue,
Parameter,
};
pub use crate::pruner::{
HyperbandPruner, MedianPruner, NopPruner, PatientPruner, PercentilePruner, Pruner,
SuccessiveHalvingPruner, ThresholdPruner, WilcoxonPruner,
};
#[cfg(feature = "cma-es")]
pub use crate::sampler::CmaEsSampler;
#[cfg(feature = "gp")]
pub use crate::sampler::GpSampler;
#[cfg(feature = "sobol")]
pub use crate::sampler::SobolSampler;
pub use crate::sampler::{
BohbSampler, CompletedTrial, DESampler, DEStrategy, Decomposition, GridSampler,
MoeadSampler, MotpeSampler, Nsga2Sampler, Nsga3Sampler, RandomSampler, TpeSampler,
};
#[cfg(feature = "journal")]
pub use crate::storage::JournalStorage;
pub use crate::storage::{MemoryStorage, Storage};
#[cfg(feature = "serde")]
pub use crate::study::StudySnapshot;
pub use crate::study::{Study, StudyBuilder};
pub use crate::trial::{AttrValue, Trial};
pub use crate::types::Direction;
pub use crate::visualization::generate_html_report;
}
+413
View File
@@ -0,0 +1,413 @@
//! Multi-objective optimization via a dedicated study type.
//!
//! [`MultiObjectiveStudy`] manages trials that return **multiple** objective
//! values simultaneously. It supports arbitrary numbers of objectives with
//! per-objective directions (minimize or maximize).
//!
//! # Key concepts
//!
//! In multi-objective optimization there is usually no single best solution.
//! Instead, there is a **Pareto front** — the set of solutions where no
//! objective can be improved without worsening another. Use
//! [`pareto_front()`](MultiObjectiveStudy::pareto_front) to retrieve these
//! non-dominated solutions after optimization.
//!
//! A solution **dominates** another if it is at least as good in all
//! objectives and strictly better in at least one. Solutions that are not
//! dominated by any other are called **Pareto-optimal**.
//!
//! # Samplers
//!
//! By default a random sampler is used. For smarter search, pass a
//! [`MultiObjectiveSampler`] such as [`Nsga2Sampler`](crate::sampler::Nsga2Sampler),
//! [`Nsga3Sampler`](crate::sampler::Nsga3Sampler), or
//! [`MoeadSampler`](crate::sampler::MoeadSampler) via
//! [`MultiObjectiveStudy::with_sampler`].
//!
//! # Examples
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::multi_objective::MultiObjectiveStudy;
//! use optimizer::parameter::{FloatParam, Parameter};
//!
//! let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
//! let x = FloatParam::new(0.0, 1.0);
//!
//! study
//! .optimize(20, |trial: &mut optimizer::Trial| {
//! let xv = x.suggest(trial)?;
//! Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
//! })
//! .unwrap();
//!
//! let front = study.pareto_front();
//! assert!(!front.is_empty());
//! ```
use core::sync::atomic::{AtomicU64, Ordering};
use std::collections::HashMap;
use std::sync::Arc;
use parking_lot::RwLock;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::parameter::{ParamId, Parameter};
use crate::pruner::NopPruner;
use crate::sampler::random::RandomMultiObjectiveSampler;
use crate::sampler::{CompletedTrial, Sampler};
use crate::trial::{AttrValue, Trial};
use crate::types::{Direction, TrialState};
// ---------------------------------------------------------------------------
// MultiObjectiveTrial
// ---------------------------------------------------------------------------
/// A completed trial with multiple objective values.
///
/// Each trial stores its sampled parameter values, the vector of
/// objective values (one per objective), and optional constraint values.
/// Retrieve typed parameter values with [`get()`](Self::get) and check
/// constraint feasibility with [`is_feasible()`](Self::is_feasible).
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct MultiObjectiveTrial {
/// The unique identifier for this trial.
pub id: u64,
/// The sampled parameter values, keyed by parameter id.
pub params: HashMap<ParamId, ParamValue>,
/// The parameter distributions used, keyed by parameter id.
pub distributions: HashMap<ParamId, Distribution>,
/// Human-readable labels for parameters, keyed by parameter id.
pub param_labels: HashMap<ParamId, String>,
/// The objective values (one per objective).
pub values: Vec<f64>,
/// The state of the trial.
pub state: TrialState,
/// User-defined attributes stored during the trial.
pub user_attrs: HashMap<String, AttrValue>,
/// Constraint values for this trial (<=0.0 means feasible).
#[cfg_attr(feature = "serde", serde(default))]
pub constraints: Vec<f64>,
}
impl MultiObjectiveTrial {
/// Returns the typed value for the given parameter.
///
/// Returns `None` if the parameter was not used in this trial.
///
/// # Panics
///
/// Panics if the stored value is incompatible with the parameter type.
pub fn get<P: Parameter>(&self, param: &P) -> Option<P::Value> {
self.params.get(&param.id()).map(|v| {
param
.cast_param_value(v)
.expect("parameter type mismatch: stored value incompatible with parameter")
})
}
/// Returns `true` if all constraints are satisfied (values <= 0.0).
///
/// A trial with no constraints is considered feasible.
#[must_use]
pub fn is_feasible(&self) -> bool {
self.constraints.iter().all(|&c| c <= 0.0)
}
/// Gets a user attribute by key.
#[must_use]
pub fn user_attr(&self, key: &str) -> Option<&AttrValue> {
self.user_attrs.get(key)
}
/// Returns all user attributes.
#[must_use]
pub fn user_attrs(&self) -> &HashMap<String, AttrValue> {
&self.user_attrs
}
}
// ---------------------------------------------------------------------------
// MultiObjectiveSampler trait
// ---------------------------------------------------------------------------
/// Trait for samplers aware of multi-objective history.
///
/// Separate from [`Sampler`] because multi-objective algorithms (e.g.,
/// NSGA-II) need access to the full vector of objective values per trial
/// (`&[MultiObjectiveTrial]`) and the per-objective directions
/// (`&[Direction]`).
///
/// Implementations include [`Nsga2Sampler`](crate::sampler::Nsga2Sampler),
/// [`Nsga3Sampler`](crate::sampler::Nsga3Sampler), and
/// [`MoeadSampler`](crate::sampler::MoeadSampler).
pub trait MultiObjectiveSampler: Send + Sync {
/// Samples a parameter value from the given distribution.
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[MultiObjectiveTrial],
directions: &[Direction],
) -> ParamValue;
}
// ---------------------------------------------------------------------------
// MoSamplerBridge — bridges MultiObjectiveSampler to Sampler trait
// ---------------------------------------------------------------------------
/// Bridges a [`MultiObjectiveSampler`] to the [`Sampler`] trait so that
/// `Trial::with_sampler()` can use it.
struct MoSamplerBridge {
inner: Arc<dyn MultiObjectiveSampler>,
history: Arc<RwLock<Vec<MultiObjectiveTrial>>>,
directions: Vec<Direction>,
}
impl Sampler for MoSamplerBridge {
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
_history: &[CompletedTrial],
) -> ParamValue {
let mo_history = self.history.read();
self.inner
.sample(distribution, trial_id, &mo_history, &self.directions)
}
}
// ---------------------------------------------------------------------------
// MultiObjectiveStudy
// ---------------------------------------------------------------------------
/// A study for multi-objective optimization.
///
/// Manage trials that return multiple objective values. Supports
/// arbitrary numbers of objectives with independent minimize/maximize
/// directions. After optimization, call [`pareto_front()`](Self::pareto_front)
/// to retrieve the non-dominated solutions.
///
/// For single-objective optimization, use [`Study`](crate::Study) instead.
///
/// # Examples
///
/// ```
/// use optimizer::Direction;
/// use optimizer::multi_objective::MultiObjectiveStudy;
/// use optimizer::parameter::{FloatParam, Parameter};
///
/// // Bi-objective: minimize both
/// let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
/// let x = FloatParam::new(0.0, 1.0);
///
/// study
/// .optimize(30, |trial: &mut optimizer::Trial| {
/// let xv = x.suggest(trial)?;
/// Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
/// })
/// .unwrap();
///
/// let front = study.pareto_front();
/// assert!(!front.is_empty());
/// ```
pub struct MultiObjectiveStudy {
directions: Vec<Direction>,
sampler: Arc<dyn MultiObjectiveSampler>,
completed_trials: Arc<RwLock<Vec<MultiObjectiveTrial>>>,
next_trial_id: AtomicU64,
}
impl MultiObjectiveStudy {
/// Creates a new multi-objective study with the given directions.
///
/// Uses a random sampler by default.
///
/// # Arguments
///
/// * `directions` - One direction per objective (minimize or maximize).
#[must_use]
pub fn new(directions: Vec<Direction>) -> Self {
Self {
directions,
sampler: Arc::new(RandomMultiObjectiveSampler::new()),
completed_trials: Arc::new(RwLock::new(Vec::new())),
next_trial_id: AtomicU64::new(0),
}
}
/// Creates a new study with a custom multi-objective sampler.
#[must_use]
pub fn with_sampler(
directions: Vec<Direction>,
sampler: impl MultiObjectiveSampler + 'static,
) -> Self {
Self {
directions,
sampler: Arc::new(sampler),
completed_trials: Arc::new(RwLock::new(Vec::new())),
next_trial_id: AtomicU64::new(0),
}
}
/// Returns the optimization directions.
#[must_use]
pub fn directions(&self) -> &[Direction] {
&self.directions
}
/// Returns the number of objectives.
#[must_use]
pub fn n_objectives(&self) -> usize {
self.directions.len()
}
/// Returns the number of completed trials.
#[must_use]
pub fn n_trials(&self) -> usize {
self.completed_trials.read().len()
}
/// Returns all completed trials.
#[must_use]
pub fn trials(&self) -> Vec<MultiObjectiveTrial> {
self.completed_trials.read().clone()
}
/// Return the Pareto-optimal trials (the non-dominated front).
///
/// Uses fast non-dominated sorting (Deb et al., 2002) from the
/// [`pareto`](crate::pareto) module. Returns an empty vec if no
/// trials have completed.
#[must_use]
pub fn pareto_front(&self) -> Vec<MultiObjectiveTrial> {
let trials = self.completed_trials.read();
let complete: Vec<_> = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.collect();
if complete.is_empty() {
return Vec::new();
}
let values: Vec<Vec<f64>> = complete.iter().map(|t| t.values.clone()).collect();
let fronts = crate::pareto::fast_non_dominated_sort(&values, &self.directions);
if fronts.is_empty() {
return Vec::new();
}
fronts[0].iter().map(|&i| complete[i].clone()).collect()
}
/// Creates a new trial wired to the study's MO sampler.
fn create_trial(&self) -> Trial {
let id = self.next_trial_id.fetch_add(1, Ordering::SeqCst);
let bridge: Arc<dyn Sampler> = Arc::new(MoSamplerBridge {
inner: Arc::clone(&self.sampler),
history: Arc::clone(&self.completed_trials),
directions: self.directions.clone(),
});
// Dummy f64 history — the bridge ignores it.
let dummy_history: Arc<RwLock<Vec<CompletedTrial<f64>>>> =
Arc::new(RwLock::new(Vec::new()));
Trial::with_sampler(id, bridge, dummy_history, Arc::new(NopPruner))
}
/// Records a completed trial.
fn complete_trial(&self, trial: Trial, values: Vec<f64>) {
let mo_trial = trial.into_multi_objective_trial(values, TrialState::Complete);
self.completed_trials.write().push(mo_trial);
}
/// Records a failed trial (not stored in history).
fn fail_trial(trial: &mut Trial) {
trial.set_failed();
}
/// Request a new trial for the ask/tell interface.
///
/// After creating the trial, suggest parameters on it, evaluate your
/// objective externally, then pass the trial back to [`tell()`](Self::tell).
pub fn ask(&self) -> Trial {
self.create_trial()
}
/// Report the result of a trial obtained from [`ask()`](Self::ask).
///
/// Pass `Ok(values)` for a successful evaluation or `Err(reason)` for a failure.
///
/// # Errors
///
/// Returns `ObjectiveDimensionMismatch` if the number of values doesn't
/// match the number of directions.
pub fn tell(
&self,
mut trial: Trial,
result: core::result::Result<Vec<f64>, impl ToString>,
) -> crate::Result<()> {
if let Ok(values) = result {
if values.len() != self.directions.len() {
return Err(crate::Error::ObjectiveDimensionMismatch {
expected: self.directions.len(),
got: values.len(),
});
}
self.complete_trial(trial, values);
} else {
Self::fail_trial(&mut trial);
}
Ok(())
}
/// Runs multi-objective optimization for `n_trials` trials.
///
/// The objective function must return a `Vec<f64>` with one value per
/// objective.
///
/// # Errors
///
/// Returns `ObjectiveDimensionMismatch` if the objective returns the wrong
/// number of values. Returns `NoCompletedTrials` if all trials fail.
pub fn optimize<F, E>(&self, n_trials: usize, mut objective: F) -> crate::Result<()>
where
F: FnMut(&mut Trial) -> core::result::Result<Vec<f64>, E>,
E: ToString,
{
for _ in 0..n_trials {
let mut trial = self.create_trial();
match objective(&mut trial) {
Ok(values) => {
if values.len() != self.directions.len() {
return Err(crate::Error::ObjectiveDimensionMismatch {
expected: self.directions.len(),
got: values.len(),
});
}
self.complete_trial(trial, values);
}
Err(_) => {
Self::fail_trial(&mut trial);
}
}
}
let has_complete = self
.completed_trials
.read()
.iter()
.any(|t| t.state == TrialState::Complete);
if !has_complete {
return Err(crate::Error::NoCompletedTrials);
}
Ok(())
}
}
+146
View File
@@ -0,0 +1,146 @@
//! The [`Objective`] trait defines what gets optimized.
//!
//! # Closures work directly
//!
//! Any `Fn(&mut Trial) -> Result<V, E>` closure automatically implements
//! [`Objective`], so you can pass closures straight to
//! [`Study::optimize`](crate::Study::optimize):
//!
//! ```
//! use optimizer::prelude::*;
//!
//! let study: Study<f64> = Study::new(Direction::Minimize);
//! let x = FloatParam::new(-10.0, 10.0).name("x");
//!
//! study
//! .optimize(50, |trial: &mut optimizer::Trial| {
//! let v = x.suggest(trial)?;
//! Ok::<_, Error>((v - 3.0).powi(2))
//! })
//! .unwrap();
//! ```
//!
//! # Structs for lifecycle hooks
//!
//! For richer control — early stopping or per-trial logging — implement
//! [`Objective`] on a struct and pass it to the same
//! [`Study::optimize`](crate::Study::optimize) method:
//!
//! ```
//! use std::ops::ControlFlow;
//!
//! use optimizer::Objective;
//! use optimizer::prelude::*;
//!
//! struct QuadraticWithEarlyStopping {
//! x: FloatParam,
//! target: f64,
//! }
//!
//! impl Objective<f64> for QuadraticWithEarlyStopping {
//! type Error = Error;
//!
//! fn evaluate(&self, trial: &mut Trial) -> Result<f64> {
//! let v = self.x.suggest(trial)?;
//! Ok((v - 3.0).powi(2))
//! }
//!
//! fn after_trial(&self, _study: &Study<f64>, trial: &CompletedTrial<f64>) -> ControlFlow<()> {
//! if trial.value < self.target {
//! ControlFlow::Break(())
//! } else {
//! ControlFlow::Continue(())
//! }
//! }
//! }
//!
//! let study: Study<f64> = Study::new(Direction::Minimize);
//! let obj = QuadraticWithEarlyStopping {
//! x: FloatParam::new(-10.0, 10.0).name("x"),
//! target: 1.0,
//! };
//! study.optimize(200, obj).unwrap();
//! assert!(study.best_value().unwrap() < 1.0);
//! ```
use core::ops::ControlFlow;
use crate::sampler::CompletedTrial;
use crate::study::Study;
use crate::trial::Trial;
/// Defines an objective function with lifecycle hooks for optimization.
///
/// The only required method is [`evaluate`](Objective::evaluate), which
/// computes the objective value for a given trial. Optional hooks provide
/// early stopping ([`before_trial`](Objective::before_trial),
/// [`after_trial`](Objective::after_trial)).
///
/// # Closures implement `Objective` automatically
///
/// A blanket implementation covers all `Fn(&mut Trial) -> Result<V, E>`
/// closures, so you can pass closures directly to
/// [`Study::optimize`](crate::Study::optimize) without wrapping them.
///
/// # Thread safety
///
/// The async optimization methods (`optimize_async`, `optimize_parallel`)
/// additionally require `Send + Sync + 'static` on the objective. The
/// sync `optimize` method has no thread-safety requirements.
pub trait Objective<V: PartialOrd = f64> {
/// The error type returned by [`evaluate`](Objective::evaluate).
type Error: ToString + 'static;
/// Evaluate the objective function for a single trial.
///
/// Sample parameters from `trial` via
/// [`Parameter::suggest`](crate::parameter::Parameter::suggest) and
/// return the objective value. Return `Err(TrialPruned)` to prune a
/// trial early.
///
/// # Errors
///
/// Any error whose type implements `ToString`. Pruning errors
/// (`Error::TrialPruned` or `TrialPruned`) are handled specially —
/// the trial is recorded as pruned rather than failed.
fn evaluate(&self, trial: &mut Trial) -> Result<V, Self::Error>;
/// Called before each trial is created.
///
/// Return `ControlFlow::Break(())` to stop the optimization loop
/// before the next trial starts.
///
/// Default: always continues.
fn before_trial(&self, _study: &Study<V>) -> ControlFlow<()> {
ControlFlow::Continue(())
}
/// Called after each **completed** trial (not failed or pruned).
///
/// The trial is passed directly as the argument *before* it is pushed
/// to storage, so `study.n_trials()` and `study.trials()` do not yet
/// include this trial. The trial is always pushed to storage after this
/// callback returns, regardless of the return value.
///
/// Return `ControlFlow::Break(())` to stop the optimization loop.
///
/// Default: always continues.
fn after_trial(&self, _study: &Study<V>, _trial: &CompletedTrial<V>) -> ControlFlow<()> {
ControlFlow::Continue(())
}
}
/// Blanket implementation: any `Fn(&mut Trial) -> Result<V, E>` is an
/// `Objective` with no lifecycle hooks.
impl<F, V, E> Objective<V> for F
where
F: Fn(&mut Trial) -> Result<V, E>,
V: PartialOrd,
E: ToString + 'static,
{
type Error = E;
fn evaluate(&self, trial: &mut Trial) -> Result<V, E> {
self(trial)
}
}
+36 -8
View File
@@ -1,16 +1,44 @@
//! Parameter value storage types.
//! Raw parameter value storage.
//!
//! [`ParamValue`] is the type-erased representation of a sampled parameter.
//! Users rarely construct `ParamValue` directly — the
//! [`Parameter::suggest`](crate::parameter::Parameter::suggest) method returns
//! the already-typed value (e.g., `f64` for [`FloatParam`](crate::parameter::FloatParam)).
//!
//! `ParamValue` is useful when inspecting raw trial data via
//! [`Trial::params`](crate::Trial::params) or
//! [`CompletedTrial::params`](crate::sampler::CompletedTrial).
/// Represents a sampled parameter value.
/// A type-erased sampled parameter value.
///
/// This enum stores different parameter value types uniformly.
/// For categorical parameters, the `Categorical` variant stores
/// the index into the choices array.
/// Stores float, integer, or categorical (index) values uniformly.
/// For categorical parameters the `Categorical` variant stores the
/// zero-based index into the choices array, not the choice itself.
///
/// # Display
///
/// `ParamValue` implements [`Display`](core::fmt::Display): floats and
/// integers print their numeric value, and categoricals print `category(i)`.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum ParamValue {
/// A floating-point parameter value.
/// A floating-point parameter value (from [`FloatParam`](crate::parameter::FloatParam)).
Float(f64),
/// An integer parameter value.
/// An integer parameter value (from [`IntParam`](crate::parameter::IntParam)).
Int(i64),
/// A categorical parameter value, stored as an index into the choices array.
/// A categorical index into the choices array (from
/// [`CategoricalParam`](crate::parameter::CategoricalParam),
/// [`BoolParam`](crate::parameter::BoolParam), or
/// [`EnumParam`](crate::parameter::EnumParam)).
Categorical(usize),
}
impl core::fmt::Display for ParamValue {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Float(v) => write!(f, "{v}"),
Self::Int(v) => write!(f, "{v}"),
Self::Categorical(v) => write!(f, "category({v})"),
}
}
}
+1156
View File
File diff suppressed because it is too large Load Diff
+660
View File
@@ -0,0 +1,660 @@
//! Pareto front analysis utilities for multi-objective optimization.
//!
//! In multi-objective optimization there is generally no single best
//! solution. Instead, the goal is to find the **Pareto front** — the set
//! of solutions where no objective can be improved without worsening
//! another. This module provides tools for computing and analyzing Pareto
//! fronts.
//!
//! # Available functions
//!
//! | Function | Purpose |
//! |---|---|
//! | [`hypervolume`] | Measure the quality of a Pareto front (volume of dominated space) |
//! | [`non_dominated_sort`] | Rank solutions into successive fronts (front 0, 1, …) |
//! | [`pareto_front_indices`] | Filter to non-dominated (Pareto-optimal) solutions only |
//! | [`crowding_distance`] | Measure diversity/spread within a single front |
//!
//! # When to use
//!
//! - **Evaluating front quality**: Use [`hypervolume`] to compare two
//! Pareto fronts — a higher hypervolume indicates a better-quality front.
//! - **Ranking all solutions**: Use [`non_dominated_sort`] to partition
//! solutions into successive fronts, useful for selection in evolutionary
//! algorithms.
//! - **Extracting the best solutions**: Use [`pareto_front_indices`] to get
//! only the non-dominated set.
//! - **Diversity measurement**: Use [`crowding_distance`] to quantify how
//! spread out solutions are within a front, which helps maintain diversity.
//!
//! Internally, this module also provides the fast non-dominated sorting
//! algorithm (Deb et al., 2002) used by
//! [`MultiObjectiveStudy::pareto_front()`](crate::multi_objective::MultiObjectiveStudy::pareto_front)
//! and [`Nsga2Sampler`](crate::sampler::Nsga2Sampler).
//!
//! # Example
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::pareto::{
//! crowding_distance, hypervolume, non_dominated_sort, pareto_front_indices,
//! };
//!
//! let solutions = vec![
//! vec![1.0, 5.0], // Pareto-optimal
//! vec![5.0, 1.0], // Pareto-optimal
//! vec![3.0, 3.0], // Pareto-optimal
//! vec![4.0, 4.0], // Dominated by (3, 3)
//! ];
//! let dirs = [Direction::Minimize, Direction::Minimize];
//!
//! // Non-dominated sorting: front 0 has indices {0, 1, 2}
//! let fronts = non_dominated_sort(&solutions, &dirs);
//! assert_eq!(fronts.len(), 2);
//!
//! // Pareto front indices (shortcut for fronts[0])
//! let mut front = pareto_front_indices(&solutions, &dirs);
//! front.sort();
//! assert_eq!(front, vec![0, 1, 2]);
//!
//! // Hypervolume with reference point (6, 6)
//! let front_values: Vec<_> = front.iter().map(|&i| solutions[i].clone()).collect();
//! let hv = hypervolume(&front_values, &[6.0, 6.0], &dirs);
//! assert!(hv > 0.0);
//!
//! // Crowding distance for diversity analysis
//! let cd = crowding_distance(&front_values, &dirs);
//! assert!(cd[0].is_infinite()); // boundary solution
//! ```
use crate::types::Direction;
/// Returns `true` if solution `a` Pareto-dominates solution `b`.
///
/// A solution dominates another if it is at least as good in all objectives
/// and strictly better in at least one, respecting the given directions.
#[allow(clippy::module_name_repetitions)]
pub(crate) fn dominates(a: &[f64], b: &[f64], directions: &[Direction]) -> bool {
debug_assert_eq!(a.len(), b.len());
debug_assert_eq!(a.len(), directions.len());
let mut strictly_better = false;
for ((&av, &bv), dir) in a.iter().zip(b.iter()).zip(directions.iter()) {
let better = match dir {
Direction::Minimize => av < bv,
Direction::Maximize => av > bv,
};
let worse = match dir {
Direction::Minimize => av > bv,
Direction::Maximize => av < bv,
};
if worse {
return false;
}
if better {
strictly_better = true;
}
}
strictly_better
}
/// Constrained dominance: feasible beats infeasible, among infeasible
/// prefer lower total constraint violation, among feasible use Pareto dominance.
pub(crate) fn constrained_dominates(
a_values: &[f64],
b_values: &[f64],
a_constraints: &[f64],
b_constraints: &[f64],
directions: &[Direction],
) -> bool {
let a_feasible = a_constraints.iter().all(|&c| c <= 0.0);
let b_feasible = b_constraints.iter().all(|&c| c <= 0.0);
match (a_feasible, b_feasible) {
(true, false) => true,
(false, true) => false,
(false, false) => {
let a_violation: f64 = a_constraints.iter().map(|c| c.max(0.0)).sum();
let b_violation: f64 = b_constraints.iter().map(|c| c.max(0.0)).sum();
a_violation < b_violation
}
(true, true) => dominates(a_values, b_values, directions),
}
}
/// Fast non-dominated sorting (Deb et al., 2002).
///
/// Returns `Vec<Vec<usize>>` where `fronts[0]` is the Pareto front,
/// each inner vec contains indices into `values`.
///
/// Complexity: O(M * N^2) where M = objectives, N = solutions.
#[allow(clippy::cast_possible_truncation)]
pub(crate) fn fast_non_dominated_sort(
values: &[Vec<f64>],
directions: &[Direction],
) -> Vec<Vec<usize>> {
fast_non_dominated_sort_constrained(values, directions, &[])
}
/// Fast non-dominated sorting with constraint support.
///
/// `constraints` is either empty (no constraints) or has the same length
/// as `values`, where each entry is the constraint vector for that solution.
#[allow(clippy::cast_possible_truncation)]
pub(crate) fn fast_non_dominated_sort_constrained(
values: &[Vec<f64>],
directions: &[Direction],
constraints: &[Vec<f64>],
) -> Vec<Vec<usize>> {
let n = values.len();
if n == 0 {
return Vec::new();
}
let has_constraints = !constraints.is_empty();
let empty_constraints: Vec<f64> = Vec::new();
// S_p: set of solutions dominated by p
let mut dominated_by: Vec<Vec<usize>> = vec![Vec::new(); n];
// n_p: domination count for p
let mut domination_count: Vec<usize> = vec![0; n];
for i in 0..n {
for j in (i + 1)..n {
let (a_c, b_c) = if has_constraints {
(&constraints[i], &constraints[j])
} else {
(&empty_constraints, &empty_constraints)
};
let i_dom_j = if has_constraints {
constrained_dominates(&values[i], &values[j], a_c, b_c, directions)
} else {
dominates(&values[i], &values[j], directions)
};
let j_dom_i = if has_constraints {
constrained_dominates(&values[j], &values[i], b_c, a_c, directions)
} else {
dominates(&values[j], &values[i], directions)
};
if i_dom_j {
dominated_by[i].push(j);
domination_count[j] += 1;
} else if j_dom_i {
dominated_by[j].push(i);
domination_count[i] += 1;
}
}
}
let mut fronts: Vec<Vec<usize>> = Vec::new();
let mut current_front: Vec<usize> = (0..n).filter(|&i| domination_count[i] == 0).collect();
while !current_front.is_empty() {
let mut next_front: Vec<usize> = Vec::new();
for &p in &current_front {
for &q in &dominated_by[p] {
domination_count[q] -= 1;
if domination_count[q] == 0 {
next_front.push(q);
}
}
}
fronts.push(current_front);
current_front = next_front;
}
fronts
}
/// Crowding distance for one front (index-based, internal API).
///
/// Boundary solutions get `f64::INFINITY`. Returns one distance value per
/// solution in the front, in the same order as `front_indices`.
#[allow(clippy::cast_precision_loss)]
pub(crate) fn crowding_distance_indexed(front_indices: &[usize], values: &[Vec<f64>]) -> Vec<f64> {
let n = front_indices.len();
if n <= 2 {
return vec![f64::INFINITY; n];
}
let m = values[front_indices[0]].len(); // number of objectives
let mut distances = vec![0.0_f64; n];
// Helper to look up objective value for a front member.
let val = |front_pos: usize, obj: usize| -> f64 { values[front_indices[front_pos]][obj] };
for obj in 0..m {
// Sort front positions by this objective
let mut sorted: Vec<usize> = (0..n).collect();
sorted.sort_by(|&a, &b| {
val(a, obj)
.partial_cmp(&val(b, obj))
.unwrap_or(core::cmp::Ordering::Equal)
});
// Boundary solutions get infinity
distances[sorted[0]] = f64::INFINITY;
distances[sorted[n - 1]] = f64::INFINITY;
let range = val(sorted[n - 1], obj) - val(sorted[0], obj);
if range > 0.0 {
for i in 1..(n - 1) {
distances[sorted[i]] += (val(sorted[i + 1], obj) - val(sorted[i - 1], obj)) / range;
}
}
}
distances
}
// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
/// Compute the hypervolume indicator of a Pareto front.
///
/// The hypervolume is the volume of the objective space dominated by
/// the Pareto front and bounded by a reference point. A **higher**
/// hypervolume indicates a better front (closer to the ideal and more
/// spread out).
///
/// Each entry in `front` is one solution's objective values.
/// `reference_point` should be worse than all front members in every
/// objective (e.g., the worst acceptable values). Solutions that do
/// not strictly dominate the reference point are ignored.
///
/// Uses recursive slicing for dimensions > 1. Complexity grows with
/// the number of objectives and front size.
///
/// # Panics
///
/// Panics (in debug) if dimensions of `front`, `reference_point`, and
/// `directions` are inconsistent.
#[must_use]
#[allow(clippy::cast_precision_loss)]
pub fn hypervolume(front: &[Vec<f64>], reference_point: &[f64], directions: &[Direction]) -> f64 {
if front.is_empty() {
return 0.0;
}
let d = reference_point.len();
debug_assert!(front.iter().all(|p| p.len() == d));
debug_assert_eq!(d, directions.len());
// Normalize to minimize-space (negate maximized objectives).
let normalized: Vec<Vec<f64>> = front
.iter()
.map(|p| {
p.iter()
.zip(directions)
.map(|(&v, dir)| match dir {
Direction::Minimize => v,
Direction::Maximize => -v,
})
.collect()
})
.collect();
let ref_norm: Vec<f64> = reference_point
.iter()
.zip(directions)
.map(|(&v, dir)| match dir {
Direction::Minimize => v,
Direction::Maximize => -v,
})
.collect();
// Keep only points strictly dominated by the reference point.
let filtered: Vec<Vec<f64>> = normalized
.into_iter()
.filter(|p| p.iter().zip(&ref_norm).all(|(&pv, &rv)| pv < rv))
.collect();
if filtered.is_empty() {
return 0.0;
}
hv_recursive(&filtered, &ref_norm)
}
/// Recursive hypervolume via slicing on the last objective.
///
/// All points are in minimize-space and dominated by `reference`.
#[allow(clippy::cast_precision_loss)]
fn hv_recursive(points: &[Vec<f64>], reference: &[f64]) -> f64 {
let d = reference.len();
// Base case: 1-D hypervolume is just the gap from the best point to ref.
if d == 1 {
let min_val = points.iter().map(|p| p[0]).fold(f64::INFINITY, f64::min);
return (reference[0] - min_val).max(0.0);
}
// Single point: hypervolume is the product of gaps.
if points.len() == 1 {
return points[0]
.iter()
.zip(reference)
.map(|(&p, &r)| (r - p).max(0.0))
.product();
}
// Sort by last objective ascending.
let mut sorted: Vec<&Vec<f64>> = points.iter().collect();
sorted.sort_by(|a, b| {
a[d - 1]
.partial_cmp(&b[d - 1])
.unwrap_or(core::cmp::Ordering::Equal)
});
let sub_ref: Vec<f64> = reference[..d - 1].to_vec();
let mut result = 0.0;
for i in 0..sorted.len() {
let height = if i + 1 < sorted.len() {
sorted[i + 1][d - 1] - sorted[i][d - 1]
} else {
reference[d - 1] - sorted[i][d - 1]
};
if height <= 0.0 {
continue;
}
// Project points[0..=i] onto the first d-1 dimensions and
// keep only the non-dominated subset.
let projected: Vec<Vec<f64>> = sorted[..=i].iter().map(|p| p[..d - 1].to_vec()).collect();
let non_dom = non_dominated_minimize(&projected);
if !non_dom.is_empty() {
result += height * hv_recursive(&non_dom, &sub_ref);
}
}
result
}
/// Return the non-dominated subset of `points` in minimize-space.
fn non_dominated_minimize(points: &[Vec<f64>]) -> Vec<Vec<f64>> {
let mut result = Vec::new();
'outer: for (i, p) in points.iter().enumerate() {
for (j, q) in points.iter().enumerate() {
if i == j {
continue;
}
// Check if q dominates p (all <=, at least one <).
let mut all_leq = true;
let mut any_lt = false;
for (&qv, &pv) in q.iter().zip(p.iter()) {
if qv > pv {
all_leq = false;
break;
}
if qv < pv {
any_lt = true;
}
}
if all_leq && any_lt {
continue 'outer;
}
}
result.push(p.clone());
}
result
}
/// Compute non-dominated sorting of a set of solutions.
///
/// Return a vec of fronts, where `fronts[0]` is the Pareto front
/// (non-dominated solutions), `fronts[1]` is the next-best front
/// (dominated only by front 0), and so on. Each inner vec contains
/// indices into the original `solutions` slice.
///
/// Use the fast non-dominated sorting algorithm from
/// Deb et al. (2002) with O(M × N²) complexity, where M is the
/// number of objectives and N is the number of solutions.
#[must_use]
pub fn non_dominated_sort(solutions: &[Vec<f64>], directions: &[Direction]) -> Vec<Vec<usize>> {
fast_non_dominated_sort(solutions, directions)
}
/// Filter solutions to return only non-dominated (Pareto-optimal) indices.
///
/// Equivalent to `non_dominated_sort(solutions, directions)[0]` but
/// communicates the intent more clearly. Use this when you only need
/// the Pareto front and not the full ranking.
#[must_use]
pub fn pareto_front_indices(solutions: &[Vec<f64>], directions: &[Direction]) -> Vec<usize> {
let fronts = fast_non_dominated_sort(solutions, directions);
fronts.into_iter().next().unwrap_or_default()
}
/// Compute crowding distance for diversity measurement.
///
/// Return one distance value per solution in `front` (same order).
/// Boundary solutions (best/worst in any objective) receive
/// [`f64::INFINITY`]. Interior solutions get a finite positive value
/// proportional to the gap between their neighbors in each objective.
///
/// Crowding distance is used by NSGA-II to prefer well-spread
/// solutions when two solutions are in the same front.
///
/// `directions` is accepted for API consistency but does not affect
/// the result, since crowding distance measures spacing regardless of
/// optimization direction.
#[must_use]
#[allow(clippy::cast_precision_loss, clippy::needless_range_loop)]
pub fn crowding_distance(front: &[Vec<f64>], _directions: &[Direction]) -> Vec<f64> {
let n = front.len();
if n <= 2 {
return vec![f64::INFINITY; n];
}
let m = front[0].len();
let mut distances = vec![0.0_f64; n];
for obj in 0..m {
let mut sorted: Vec<usize> = (0..n).collect();
sorted.sort_by(|&a, &b| {
front[a][obj]
.partial_cmp(&front[b][obj])
.unwrap_or(core::cmp::Ordering::Equal)
});
distances[sorted[0]] = f64::INFINITY;
distances[sorted[n - 1]] = f64::INFINITY;
let range = front[sorted[n - 1]][obj] - front[sorted[0]][obj];
if range > 0.0 {
for i in 1..(n - 1) {
distances[sorted[i]] +=
(front[sorted[i + 1]][obj] - front[sorted[i - 1]][obj]) / range;
}
}
}
distances
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_dominates_basic() {
let dirs = [Direction::Minimize, Direction::Minimize];
assert!(dominates(&[1.0, 1.0], &[2.0, 2.0], &dirs));
assert!(!dominates(&[2.0, 2.0], &[1.0, 1.0], &dirs));
// Equal does not dominate
assert!(!dominates(&[1.0, 1.0], &[1.0, 1.0], &dirs));
}
#[test]
fn test_dominates_incomparable() {
let dirs = [Direction::Minimize, Direction::Minimize];
assert!(!dominates(&[1.0, 3.0], &[3.0, 1.0], &dirs));
assert!(!dominates(&[3.0, 1.0], &[1.0, 3.0], &dirs));
}
#[test]
fn test_dominates_maximize() {
let dirs = [Direction::Maximize, Direction::Minimize];
// a = (5, 1) vs b = (3, 2): a is better in both
assert!(dominates(&[5.0, 1.0], &[3.0, 2.0], &dirs));
assert!(!dominates(&[3.0, 2.0], &[5.0, 1.0], &dirs));
}
#[test]
fn test_nds_known() {
let values = vec![
vec![1.0, 5.0], // front 0
vec![5.0, 1.0], // front 0
vec![3.0, 3.0], // front 0 (non-dominated)
vec![4.0, 4.0], // front 1 (dominated by #2)
vec![6.0, 6.0], // front 2
];
let dirs = [Direction::Minimize, Direction::Minimize];
let fronts = fast_non_dominated_sort(&values, &dirs);
assert_eq!(fronts.len(), 3);
let mut f0 = fronts[0].clone();
f0.sort_unstable();
assert_eq!(f0, vec![0, 1, 2]);
assert_eq!(fronts[1], vec![3]);
assert_eq!(fronts[2], vec![4]);
}
#[test]
fn test_crowding_indexed_boundaries() {
let values = vec![vec![1.0, 5.0], vec![3.0, 3.0], vec![5.0, 1.0]];
let front = vec![0, 1, 2];
let cd = crowding_distance_indexed(&front, &values);
assert!(cd[0].is_infinite());
assert!(cd[2].is_infinite());
assert!(cd[1].is_finite());
assert!(cd[1] > 0.0);
}
// ---- Public API tests ----
#[test]
fn test_hypervolume_2d_minimize() {
// Front: (1,3), (2,2), (3,1) with ref (4,4) — all minimize
let front = vec![vec![1.0, 3.0], vec![2.0, 2.0], vec![3.0, 1.0]];
let dirs = [Direction::Minimize, Direction::Minimize];
let hv = hypervolume(&front, &[4.0, 4.0], &dirs);
// Strip 1: x=[1,2), h=4-3=1 → area=1
// Strip 2: x=[2,3), h=4-2=2 → area=2
// Strip 3: x=[3,4], h=4-1=3 → area=3
// Total = 6
assert!((hv - 6.0).abs() < 1e-10);
}
#[test]
fn test_hypervolume_2d_maximize() {
// Front: (3,1), (2,2), (1,3) with ref (0,0) — all maximize
let front = vec![vec![3.0, 1.0], vec![2.0, 2.0], vec![1.0, 3.0]];
let dirs = [Direction::Maximize, Direction::Maximize];
let hv = hypervolume(&front, &[0.0, 0.0], &dirs);
// In negate-space: points become (-3,-1),(-2,-2),(-1,-3), ref=(0,0)
// Same geometry as minimize test above → area = 6
assert!((hv - 6.0).abs() < 1e-10);
}
#[test]
fn test_hypervolume_single_point() {
let front = vec![vec![1.0, 1.0]];
let dirs = [Direction::Minimize, Direction::Minimize];
let hv = hypervolume(&front, &[3.0, 3.0], &dirs);
// Rectangle: (3-1) * (3-1) = 4
assert!((hv - 4.0).abs() < 1e-10);
}
#[test]
fn test_hypervolume_empty_front() {
let front: Vec<Vec<f64>> = vec![];
let dirs = [Direction::Minimize];
assert!(hypervolume(&front, &[1.0], &dirs).abs() < f64::EPSILON);
}
#[test]
fn test_hypervolume_point_at_ref() {
// Point not strictly better than ref → contributes nothing
let front = vec![vec![5.0, 5.0]];
let dirs = [Direction::Minimize, Direction::Minimize];
let hv = hypervolume(&front, &[5.0, 5.0], &dirs);
assert!(hv.abs() < f64::EPSILON);
}
#[test]
fn test_hypervolume_3d() {
// Single point in 3D: (1,1,1) with ref (2,2,2)
let front = vec![vec![1.0, 1.0, 1.0]];
let dirs = [
Direction::Minimize,
Direction::Minimize,
Direction::Minimize,
];
let hv = hypervolume(&front, &[2.0, 2.0, 2.0], &dirs);
assert!((hv - 1.0).abs() < 1e-10);
}
#[test]
fn test_non_dominated_sort_public() {
let values = vec![
vec![1.0, 5.0],
vec![5.0, 1.0],
vec![3.0, 3.0],
vec![4.0, 4.0],
];
let dirs = [Direction::Minimize, Direction::Minimize];
let fronts = non_dominated_sort(&values, &dirs);
assert_eq!(fronts.len(), 2);
let mut f0 = fronts[0].clone();
f0.sort_unstable();
assert_eq!(f0, vec![0, 1, 2]);
assert_eq!(fronts[1], vec![3]);
}
#[test]
fn test_pareto_front_indices_basic() {
let values = vec![
vec![1.0, 5.0],
vec![5.0, 1.0],
vec![3.0, 3.0],
vec![4.0, 4.0],
];
let dirs = [Direction::Minimize, Direction::Minimize];
let mut idx = pareto_front_indices(&values, &dirs);
idx.sort_unstable();
assert_eq!(idx, vec![0, 1, 2]);
}
#[test]
fn test_pareto_front_indices_empty() {
let values: Vec<Vec<f64>> = vec![];
let dirs = [Direction::Minimize];
assert!(pareto_front_indices(&values, &dirs).is_empty());
}
#[test]
fn test_crowding_distance_public() {
let front = vec![vec![1.0, 5.0], vec![3.0, 3.0], vec![5.0, 1.0]];
let dirs = [Direction::Minimize, Direction::Minimize];
let cd = crowding_distance(&front, &dirs);
assert!(cd[0].is_infinite());
assert!(cd[2].is_infinite());
assert!(cd[1].is_finite());
assert!(cd[1] > 0.0);
}
#[test]
fn test_crowding_distance_single_point() {
let front = vec![vec![2.0, 3.0]];
let dirs = [Direction::Minimize, Direction::Minimize];
let cd = crowding_distance(&front, &dirs);
assert_eq!(cd.len(), 1);
assert!(cd[0].is_infinite());
}
}
+603
View File
@@ -0,0 +1,603 @@
//! `HyperBand` pruner — adaptive budget scheduling with multiple SHA brackets.
//!
//! `HyperBand` addresses the main weakness of
//! [`SuccessiveHalvingPruner`](super::SuccessiveHalvingPruner): sensitivity to
//! the `min_resource` setting. It runs multiple Successive Halving brackets in
//! parallel, each with a different trade-off between the number of trials and
//! the starting budget:
//!
//! - **Bracket 0**: many trials, small starting budget (aggressive early pruning)
//! - **Bracket `s_max`**: few trials, full budget (no pruning)
//!
//! Trials are assigned to brackets in round-robin order. This ensures that
//! the overall search is robust regardless of how informative early steps are.
//!
//! # When to use
//!
//! - When you don't know how many epochs/steps are needed before performance
//! becomes predictive
//! - As a drop-in upgrade over [`SuccessiveHalvingPruner`](super::SuccessiveHalvingPruner)
//! when you can afford more total trials
//! - For large-scale hyperparameter searches where compute savings matter most
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `min_resource` | 1 | Smallest budget for the most aggressive bracket |
//! | `max_resource` | 81 | Full budget (last rung in every bracket) |
//! | `reduction_factor` | 3 | At each rung, keep top 1/η trials |
//! | `direction` | `Minimize` | Optimization direction |
//!
//! # Example
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::pruner::HyperbandPruner;
//!
//! let pruner = HyperbandPruner::new()
//! .min_resource(1)
//! .max_resource(81)
//! .reduction_factor(3)
//! .direction(Direction::Minimize);
//! ```
use core::sync::atomic::{AtomicU64, Ordering};
use std::collections::HashMap;
use std::sync::Mutex;
use super::Pruner;
use crate::sampler::CompletedTrial;
use crate::types::{Direction, TrialState};
/// `HyperBand` pruner that manages multiple Successive Halving brackets.
///
/// Hyperband addresses SHA's sensitivity to the `min_resource` choice by
/// running multiple brackets, each with a different tradeoff between the
/// number of configurations and the starting budget:
///
/// - Bracket 0: many trials, very small starting budget (aggressive pruning)
/// - Bracket 1: fewer trials, larger starting budget (moderate pruning)
/// - ...
/// - Bracket `s_max`: few trials, full budget (no pruning)
///
/// Trials are assigned to brackets in round-robin fashion. Each bracket
/// runs SHA with its own `min_resource` and rung schedule.
///
/// # Examples
///
/// ```
/// use optimizer::Direction;
/// use optimizer::pruner::HyperbandPruner;
///
/// let pruner = HyperbandPruner::new()
/// .min_resource(1)
/// .max_resource(81)
/// .reduction_factor(3)
/// .direction(Direction::Minimize);
/// ```
pub struct HyperbandPruner {
min_resource: u64,
max_resource: u64,
reduction_factor: u64,
direction: Direction,
/// Tracks which bracket each trial belongs to.
trial_brackets: Mutex<HashMap<u64, usize>>,
/// Counter for round-robin bracket assignment.
next_bracket: AtomicU64,
}
impl HyperbandPruner {
/// Create a new `HyperbandPruner` with default parameters.
///
/// Defaults: `min_resource=1`, `max_resource=81`, `reduction_factor=3`,
/// `direction=Minimize`.
#[must_use]
pub fn new() -> Self {
Self {
min_resource: 1,
max_resource: 81,
reduction_factor: 3,
direction: Direction::Minimize,
trial_brackets: Mutex::new(HashMap::new()),
next_bracket: AtomicU64::new(0),
}
}
/// Set the minimum resource (budget) per trial.
///
/// # Panics
///
/// Panics if `r` is 0.
#[must_use]
pub fn min_resource(mut self, r: u64) -> Self {
assert!(r > 0, "min_resource must be > 0, got {r}");
self.min_resource = r;
self
}
/// Set the maximum resource (budget) per trial.
///
/// # Panics
///
/// Panics if `r` is 0.
#[must_use]
pub fn max_resource(mut self, r: u64) -> Self {
assert!(r > 0, "max_resource must be > 0, got {r}");
self.max_resource = r;
self
}
/// Set the reduction factor (eta). At each rung, the top 1/eta trials survive.
///
/// # Panics
///
/// Panics if `eta` is less than 2.
#[must_use]
pub fn reduction_factor(mut self, eta: u64) -> Self {
assert!(eta >= 2, "reduction_factor must be >= 2, got {eta}");
self.reduction_factor = eta;
self
}
/// Set the optimization direction.
#[must_use]
pub fn direction(mut self, d: Direction) -> Self {
self.direction = d;
self
}
/// Compute `s_max = floor(log(max_resource / min_resource) / log(eta))`.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
fn s_max(&self) -> u64 {
let eta = self.reduction_factor as f64;
let ratio = self.max_resource as f64 / self.min_resource as f64;
(ratio.ln() / eta.ln()).floor() as u64
}
/// Compute the rung steps for a given bracket `s`.
///
/// For bracket `s`, the starting resource is `max_resource / eta^(s_max - s)`,
/// and rungs are spaced at powers of eta from there up to `max_resource`.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
fn rung_steps_for_bracket(&self, bracket: usize) -> Vec<u64> {
let s_max = self.s_max();
let eta = self.reduction_factor as f64;
// Starting resource for this bracket
let exponent = s_max.saturating_sub(bracket as u64);
let min_resource_bracket =
(self.max_resource as f64 / eta.powi(exponent as i32)).ceil() as u64;
let mut steps = Vec::new();
let mut rung: u32 = 0;
while let Some(power) = self.reduction_factor.checked_pow(rung) {
let step = min_resource_bracket.saturating_mul(power);
if step > self.max_resource {
break;
}
steps.push(step);
rung += 1;
}
steps
}
/// Assign a trial to a bracket (round-robin) and return the bracket index.
#[allow(clippy::cast_possible_truncation)]
fn assign_bracket(&self, trial_id: u64) -> usize {
let n_brackets = (self.s_max() + 1) as usize;
let mut map = self.trial_brackets.lock().expect("lock poisoned");
*map.entry(trial_id).or_insert_with(|| {
let idx = self.next_bracket.fetch_add(1, Ordering::Relaxed);
(idx as usize) % n_brackets
})
}
}
impl Default for HyperbandPruner {
fn default() -> Self {
Self::new()
}
}
#[allow(clippy::cast_precision_loss)]
impl Pruner for HyperbandPruner {
fn should_prune(
&self,
trial_id: u64,
step: u64,
intermediate_values: &[(u64, f64)],
completed_trials: &[CompletedTrial],
) -> bool {
let bracket = self.assign_bracket(trial_id);
let rungs = self.rung_steps_for_bracket(bracket);
// Find the highest rung step <= current step
let Some(&rung_step) = rungs.iter().rev().find(|&&r| r <= step) else {
return false;
};
// Never prune at the last rung (full budget)
if rung_step >= self.max_resource {
return false;
}
// Get the current trial's value at this rung step
let current_value =
if let Some(&(_, v)) = intermediate_values.iter().find(|(s, _)| *s == rung_step) {
v
} else if let Some(&(_, v)) = intermediate_values
.iter()
.rev()
.find(|(s, _)| *s <= rung_step)
{
v
} else {
return false;
};
self.is_pruned_at_rung(current_value, rung_step, bracket, completed_trials)
}
}
impl HyperbandPruner {
/// Determine whether a trial should be pruned at the given rung within its bracket.
///
/// Only compares against other trials in the same bracket.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
fn is_pruned_at_rung(
&self,
current_value: f64,
rung_step: u64,
bracket: usize,
completed_trials: &[CompletedTrial],
) -> bool {
let eta = self.reduction_factor as usize;
// Collect values at this rung step from trials in the same bracket
let map = self.trial_brackets.lock().expect("lock poisoned");
let mut values_at_rung: Vec<f64> = completed_trials
.iter()
.filter(|t| t.state == TrialState::Complete || t.state == TrialState::Pruned)
.filter(|t| map.get(&t.id).copied() == Some(bracket))
.filter_map(|t| {
t.intermediate_values
.iter()
.find(|(s, _)| *s == rung_step)
.map(|(_, v)| *v)
})
.collect();
drop(map);
// Need at least eta trials to make a meaningful comparison
if values_at_rung.len() < eta {
return false;
}
values_at_rung.push(current_value);
values_at_rung
.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
if self.direction == Direction::Maximize {
values_at_rung.reverse();
}
let n_keep = (values_at_rung.len() as f64 / eta as f64).ceil() as usize;
let threshold_idx = n_keep.max(1) - 1;
let threshold = values_at_rung[threshold_idx];
match self.direction {
Direction::Minimize => current_value > threshold,
Direction::Maximize => current_value < threshold,
}
}
}
#[cfg(test)]
#[allow(clippy::cast_precision_loss)]
mod tests {
use super::*;
fn make_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
use std::collections::HashMap;
use crate::parameter::ParamId;
CompletedTrial::with_intermediate_values(
id,
HashMap::<ParamId, crate::parameter::ParamValue>::new(),
HashMap::new(),
HashMap::new(),
0.0,
values.to_vec(),
HashMap::new(),
)
}
fn make_pruned_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
let mut t = make_trial(id, values);
t.state = TrialState::Pruned;
t
}
#[test]
fn s_max_default() {
let pruner = HyperbandPruner::new();
// s_max = floor(ln(81/1) / ln(3)) = floor(4.0) = 4
assert_eq!(pruner.s_max(), 4);
}
#[test]
fn s_max_custom() {
let pruner = HyperbandPruner::new()
.min_resource(1)
.max_resource(16)
.reduction_factor(2);
// s_max = floor(ln(16) / ln(2)) = floor(4.0) = 4
assert_eq!(pruner.s_max(), 4);
}
#[test]
fn bracket_count() {
let pruner = HyperbandPruner::new();
// s_max=4, so brackets 0..=4 → 5 brackets
assert_eq!(pruner.s_max() + 1, 5);
}
#[test]
fn rung_steps_bracket_0_default() {
let pruner = HyperbandPruner::new();
// Bracket 0: min_resource_bracket = ceil(81 / 3^4) = ceil(81/81) = 1
// Rungs: 1, 3, 9, 27, 81
assert_eq!(pruner.rung_steps_for_bracket(0), vec![1, 3, 9, 27, 81]);
}
#[test]
fn rung_steps_bracket_2_default() {
let pruner = HyperbandPruner::new();
// Bracket 2: min_resource_bracket = ceil(81 / 3^(4-2)) = ceil(81/9) = 9
// Rungs: 9, 27, 81
assert_eq!(pruner.rung_steps_for_bracket(2), vec![9, 27, 81]);
}
#[test]
fn rung_steps_bracket_4_default() {
let pruner = HyperbandPruner::new();
// Bracket 4 (s_max): min_resource_bracket = ceil(81 / 3^0) = 81
// Rungs: 81 only (no pruning, full budget)
assert_eq!(pruner.rung_steps_for_bracket(4), vec![81]);
}
#[test]
fn rung_steps_eta2() {
let pruner = HyperbandPruner::new()
.min_resource(1)
.max_resource(16)
.reduction_factor(2);
// s_max = 4
// Bracket 0: min=ceil(16/2^4)=1, rungs: 1,2,4,8,16
assert_eq!(pruner.rung_steps_for_bracket(0), vec![1, 2, 4, 8, 16]);
// Bracket 2: min=ceil(16/2^2)=4, rungs: 4,8,16
assert_eq!(pruner.rung_steps_for_bracket(2), vec![4, 8, 16]);
// Bracket 4: min=16, rungs: 16
assert_eq!(pruner.rung_steps_for_bracket(4), vec![16]);
}
#[test]
fn round_robin_bracket_assignment() {
let pruner = HyperbandPruner::new(); // 5 brackets (0..=4)
// Trials get assigned in round-robin: 0→0, 1→1, 2→2, 3→3, 4→4, 5→0, ...
assert_eq!(pruner.assign_bracket(100), 0);
assert_eq!(pruner.assign_bracket(101), 1);
assert_eq!(pruner.assign_bracket(102), 2);
assert_eq!(pruner.assign_bracket(103), 3);
assert_eq!(pruner.assign_bracket(104), 4);
assert_eq!(pruner.assign_bracket(105), 0); // wraps around
// Repeated calls for same trial return same bracket
assert_eq!(pruner.assign_bracket(100), 0);
assert_eq!(pruner.assign_bracket(103), 3);
}
#[test]
fn no_prune_before_first_rung() {
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
// Assign trial 0 to bracket 0 (rungs: 1, 3, 9, 27, 81)
pruner.assign_bracket(0);
// Register completed trials in bracket 0
let mut completed = Vec::new();
for i in 1..=9 {
pruner.assign_bracket(i);
completed.push(make_trial(i, &[(1, i as f64)]));
}
// Trial at step 0 (before rung 1) → don't prune
assert!(!pruner.should_prune(0, 0, &[(0, 100.0)], &completed));
}
#[test]
fn no_prune_at_max_resource() {
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
// Put all trials in bracket 0
let mut completed = Vec::new();
for i in 0..9 {
pruner.assign_bracket(i);
completed.push(make_trial(i, &[(81, (i + 1) as f64)]));
}
let trial_id = 9;
pruner.assign_bracket(trial_id);
// At max_resource (81), never prune
assert!(!pruner.should_prune(trial_id, 81, &[(81, 100.0)], &completed));
}
#[test]
fn prune_worst_in_bracket_minimize() {
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
// Force all trials into bracket 0 by assigning sequentially
// With 5 brackets, trials 0,5,10,... go to bracket 0
let bracket_0_ids: Vec<u64> = (0..5).map(|i| i * 5).collect();
// Assign all 25 trial IDs to fill brackets
for i in 0..25 {
pruner.assign_bracket(i);
}
// Create 9 completed trials in bracket 0 at rung step=1
let completed: Vec<_> = bracket_0_ids
.iter()
.take(3)
.enumerate()
.map(|(idx, &id)| make_trial(id, &[(1, (idx + 1) as f64)]))
.collect();
// Trial 25 → bracket 0 (25 % 5 == 0)
let test_id = 25;
pruner.assign_bracket(test_id);
assert_eq!(pruner.assign_bracket(test_id), 0);
// 3 completed + 1 current = 4. eta=3. ceil(4/3)=2. Threshold = 2.0
// Value 2.0 → keep
assert!(!pruner.should_prune(test_id, 1, &[(1, 2.0)], &completed));
// Value 3.0 → prune
assert!(pruner.should_prune(test_id, 1, &[(1, 3.0)], &completed));
}
#[test]
fn prune_worst_in_bracket_maximize() {
let pruner = HyperbandPruner::new().direction(Direction::Maximize);
// Assign trials so they end up in bracket 0
for i in 0..25 {
pruner.assign_bracket(i);
}
let completed: Vec<_> = [0u64, 5, 10]
.iter()
.enumerate()
.map(|(idx, &id)| make_trial(id, &[(1, (idx + 1) as f64)]))
.collect();
let test_id = 25;
pruner.assign_bracket(test_id);
// For maximize, best = highest. Values: 1,2,3 + current
// Value 2.0 → keep (threshold = 2.0 when sorted desc: 3,2,current,1)
assert!(!pruner.should_prune(test_id, 1, &[(1, 2.0)], &completed));
// Value 1.0 → prune
assert!(pruner.should_prune(test_id, 1, &[(1, 0.5)], &completed));
}
#[test]
fn different_brackets_have_different_aggressiveness() {
let pruner = HyperbandPruner::new()
.min_resource(1)
.max_resource(81)
.reduction_factor(3)
.direction(Direction::Minimize);
let rungs_0 = pruner.rung_steps_for_bracket(0);
let rungs_2 = pruner.rung_steps_for_bracket(2);
let rungs_4 = pruner.rung_steps_for_bracket(4);
// Bracket 0 has the most rungs (most aggressive)
assert!(rungs_0.len() > rungs_2.len());
// Bracket 4 has just 1 rung (no pruning)
assert_eq!(rungs_4.len(), 1);
// Bracket 0 starts earliest
assert!(rungs_0[0] < rungs_2[0]);
}
#[test]
fn trials_in_different_brackets_independent() {
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
// Assign trials: bracket 0 gets IDs 0,5,10,15,20
for i in 0..25 {
pruner.assign_bracket(i);
}
// Bracket 0 trials: bad values at rung step=1
let bracket_0_trials: Vec<_> = [0u64, 5, 10]
.iter()
.map(|&id| make_trial(id, &[(1, 100.0)]))
.collect();
// Bracket 1 trials: good values at rung step=1
let bracket_1_trials: Vec<_> = [1u64, 6, 11]
.iter()
.map(|&id| make_trial(id, &[(1, 1.0)]))
.collect();
let mut all_trials = bracket_0_trials;
all_trials.extend(bracket_1_trials);
// A new bracket-0 trial with value 50 should be compared against
// bracket-0 peers (100,100,100), not bracket-1 peers (1,1,1)
let test_id = 25; // bracket 0
pruner.assign_bracket(test_id);
// 3 peers at 100.0 + current at 50.0. ceil(4/3)=2. Sorted: 50,100,100,100. Threshold=100.0
// Value 50.0 < 100.0 → keep
assert!(!pruner.should_prune(test_id, 1, &[(1, 50.0)], &all_trials));
}
#[test]
fn includes_pruned_trials() {
let pruner = HyperbandPruner::new().direction(Direction::Minimize);
for i in 0..25 {
pruner.assign_bracket(i);
}
let completed = vec![
make_trial(0, &[(1, 1.0)]),
make_pruned_trial(5, &[(1, 8.0)]),
make_pruned_trial(10, &[(1, 9.0)]),
];
let test_id = 25;
pruner.assign_bracket(test_id);
// Values: 1.0, 8.0, 9.0 + current. eta=3.
// Value 1.0 → keep
assert!(!pruner.should_prune(test_id, 1, &[(1, 1.0)], &completed));
// Value 5.0 → prune (sorted: 1,5,8,9 → keep ceil(4/3)=2 → threshold=5.0, 5.0 not > 5.0 → keep)
assert!(!pruner.should_prune(test_id, 1, &[(1, 5.0)], &completed));
// Value 6.0 → prune (sorted: 1,6,8,9 → threshold=6.0, 6.0 not > 6.0 → keep)
assert!(!pruner.should_prune(test_id, 1, &[(1, 6.0)], &completed));
// Value 9.5 → prune
assert!(pruner.should_prune(test_id, 1, &[(1, 9.5)], &completed));
}
#[test]
#[should_panic(expected = "min_resource must be > 0")]
fn rejects_zero_min_resource() {
let _ = HyperbandPruner::new().min_resource(0);
}
#[test]
#[should_panic(expected = "max_resource must be > 0")]
fn rejects_zero_max_resource() {
let _ = HyperbandPruner::new().max_resource(0);
}
#[test]
#[should_panic(expected = "reduction_factor must be >= 2")]
fn rejects_reduction_factor_one() {
let _ = HyperbandPruner::new().reduction_factor(1);
}
}
+168
View File
@@ -0,0 +1,168 @@
//! Median pruner — the recommended default pruner for most use cases.
//!
//! At each step, the current trial's intermediate value is compared against
//! the median of all completed trials' values at the same step. Trials
//! performing worse than the median are pruned.
//!
//! This is a convenience wrapper around [`PercentilePruner`](super::PercentilePruner)
//! with a fixed percentile of 50%.
//!
//! # When to use
//!
//! - **Default choice** for any iterative objective (e.g., neural network training)
//! - Works well when intermediate values are a reasonable proxy for final performance
//! - Prunes roughly half of unpromising trials, giving a good speed/accuracy balance
//!
//! If your intermediate values are noisy, consider [`WilcoxonPruner`](super::WilcoxonPruner)
//! or wrapping this pruner in a [`PatientPruner`](super::PatientPruner).
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `n_warmup_steps` | 0 | Skip pruning in the first N steps |
//! | `n_min_trials` | 1 | Require at least N completed trials before pruning |
//!
//! # Example
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::pruner::MedianPruner;
//!
//! let pruner = MedianPruner::new(Direction::Minimize)
//! .n_warmup_steps(5)
//! .n_min_trials(3);
//! ```
use super::Pruner;
use super::percentile::compute_percentile;
use crate::sampler::CompletedTrial;
use crate::types::{Direction, TrialState};
/// Prune trials that are performing worse than the median of completed trials
/// at the same step.
///
/// This is the most commonly used pruner. It compares the current trial's
/// intermediate value at each step with the median of all completed trials'
/// values at that same step.
///
/// Equivalent to `PercentilePruner::new(50.0, direction)`.
///
/// # Examples
///
/// ```
/// use optimizer::Direction;
/// use optimizer::pruner::MedianPruner;
///
/// // Prune trials worse than median when minimizing, after 5 warmup steps
/// let pruner = MedianPruner::new(Direction::Minimize)
/// .n_warmup_steps(5)
/// .n_min_trials(3);
/// ```
pub struct MedianPruner {
/// The optimization direction.
direction: Direction,
/// Don't prune in the first N steps (let the trial warm up).
n_warmup_steps: u64,
/// Require at least N completed trials before pruning.
n_min_trials: usize,
}
impl MedianPruner {
/// Create a new `MedianPruner` for the given optimization direction.
///
/// By default, `n_warmup_steps` is 0 and `n_min_trials` is 1.
#[must_use]
pub fn new(direction: Direction) -> Self {
Self {
direction,
n_warmup_steps: 0,
n_min_trials: 1,
}
}
/// Set the number of warmup steps. No pruning occurs before this step.
#[must_use]
pub fn n_warmup_steps(mut self, n: u64) -> Self {
self.n_warmup_steps = n;
self
}
/// Set the minimum number of completed trials required before pruning.
///
/// # Panics
///
/// Panics if `n` is 0.
#[must_use]
pub fn n_min_trials(mut self, n: usize) -> Self {
assert!(n >= 1, "n_min_trials must be >= 1, got {n}");
self.n_min_trials = n;
self
}
}
impl Pruner for MedianPruner {
fn should_prune(
&self,
_trial_id: u64,
step: u64,
intermediate_values: &[(u64, f64)],
completed_trials: &[CompletedTrial],
) -> bool {
// 1. Don't prune during warmup
if step < self.n_warmup_steps {
return false;
}
// Get the current trial's latest value
let Some(&(_, current_value)) = intermediate_values.last() else {
return false;
};
// 2. Collect values at this step from completed (non-pruned) trials
let mut values_at_step: Vec<f64> = completed_trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.filter_map(|t| {
t.intermediate_values
.iter()
.find(|(s, _)| *s == step)
.map(|(_, v)| *v)
})
.collect();
// 3. Not enough trials
if values_at_step.len() < self.n_min_trials {
return false;
}
// 4. Compute median (50th percentile)
let median = compute_percentile(&mut values_at_step, 50.0);
// 5. Compare against median based on direction
match self.direction {
Direction::Minimize => current_value > median,
Direction::Maximize => current_value < median,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn compute_median_odd() {
assert!((compute_percentile(&mut [3.0, 1.0, 2.0], 50.0) - 2.0).abs() < f64::EPSILON);
}
#[test]
fn compute_median_even() {
assert!((compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 50.0) - 2.5).abs() < f64::EPSILON);
}
#[test]
fn compute_median_single() {
assert!((compute_percentile(&mut [5.0], 50.0) - 5.0).abs() < f64::EPSILON);
}
}
+215
View File
@@ -0,0 +1,215 @@
//! Pruner trait and implementations for trial pruning.
//!
//! Pruners decide whether to stop (prune) a trial early based on its
//! intermediate values compared to other trials. This is useful for
//! discarding unpromising trials before they complete, saving compute.
//!
//! # How pruning works
//!
//! During optimization, each trial reports intermediate values at discrete
//! steps (e.g., validation loss after each training epoch). A pruner inspects
//! these values and compares them against completed trials to decide whether
//! the current trial should be stopped early.
//!
//! The typical flow is:
//!
//! 1. Call [`Trial::report`](crate::Trial::report) to record an intermediate value.
//! 2. Call [`Trial::should_prune`](crate::Trial::should_prune) to check the pruner's decision.
//! 3. If the pruner says prune, return [`TrialPruned`](crate::TrialPruned) from the objective.
//!
//! # Available pruners
//!
//! | Pruner | Algorithm | Best for |
//! |--------|-----------|----------|
//! | [`MedianPruner`] | Prune below median at each step | General-purpose default |
//! | [`PercentilePruner`] | Prune below configurable percentile | Tunable aggressiveness |
//! | [`ThresholdPruner`] | Prune outside fixed bounds | Known divergence limits |
//! | [`PatientPruner`] | Require N consecutive prune signals | Noisy intermediate values |
//! | [`SuccessiveHalvingPruner`] | Keep top 1/η fraction at each rung | Budget-aware pruning |
//! | [`HyperbandPruner`] | Multiple SHA brackets with different budgets | Robust to budget choice |
//! | [`WilcoxonPruner`] | Statistical signed-rank test vs. best trial | Rigorous noisy pruning |
//! | [`NopPruner`] | Never prune | Disabling pruning explicitly |
//!
//! # When to use pruning
//!
//! Pruning is most beneficial when:
//!
//! - The objective function has a natural notion of "steps" (e.g., training epochs)
//! - Early steps are informative about final performance
//! - Trials are expensive enough that stopping bad ones early saves significant time
//!
//! 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;
mod nop;
mod patient;
pub(crate) mod percentile;
mod successive_halving;
mod threshold;
mod wilcoxon;
pub use hyperband::HyperbandPruner;
pub use median::MedianPruner;
pub use nop::NopPruner;
pub use patient::PatientPruner;
pub use percentile::PercentilePruner;
pub use successive_halving::SuccessiveHalvingPruner;
pub use threshold::ThresholdPruner;
pub use wilcoxon::WilcoxonPruner;
use crate::sampler::CompletedTrial;
/// Trait for pluggable trial pruning strategies.
///
/// Pruners are consulted after each intermediate value is reported to
/// decide whether the trial should be stopped early. The trait requires
/// `Send + Sync` to support concurrent and async optimization.
///
/// # Implementing a custom pruner
///
/// ```
/// use optimizer::pruner::Pruner;
/// use optimizer::sampler::CompletedTrial;
///
/// struct MyPruner {
/// threshold: f64,
/// }
///
/// impl Pruner for MyPruner {
/// fn should_prune(
/// &self,
/// _trial_id: u64,
/// _step: u64,
/// intermediate_values: &[(u64, f64)],
/// _completed_trials: &[CompletedTrial],
/// ) -> bool {
/// // Prune if the latest value exceeds the threshold
/// intermediate_values
/// .last()
/// .is_some_and(|&(_, v)| v > self.threshold)
/// }
/// }
/// ```
///
/// 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.
///
/// # Arguments
///
/// * `trial_id` - The current trial's ID.
/// * `step` - The step at which the intermediate value was reported.
/// * `intermediate_values` - All `(step, value)` pairs reported so far for this trial.
/// * `completed_trials` - History of all completed trials (for comparison).
fn should_prune(
&self,
trial_id: u64,
step: u64,
intermediate_values: &[(u64, f64)],
completed_trials: &[CompletedTrial],
) -> bool;
}
+29
View File
@@ -0,0 +1,29 @@
//! No-op pruner — never prune any trial.
//!
//! This is the default pruner used when no pruner is configured on a
//! [`Study`](crate::Study). It unconditionally returns `false` for every
//! pruning decision, allowing all trials to run to completion.
//!
//! # When to use
//!
//! - When you want to explicitly disable pruning
//! - As a baseline to compare against other pruners
//! - Already used by default — you rarely need to configure this manually
use super::Pruner;
use crate::sampler::CompletedTrial;
/// A pruner that never prunes. This is the default when no pruner is configured.
pub struct NopPruner;
impl Pruner for NopPruner {
fn should_prune(
&self,
_trial_id: u64,
_step: u64,
_intermediate_values: &[(u64, f64)],
_completed_trials: &[CompletedTrial],
) -> bool {
false
}
}
+191
View File
@@ -0,0 +1,191 @@
//! Patient pruner — require consecutive prune signals before actually pruning.
//!
//! Wraps any other pruner and adds a patience window: the inner pruner
//! must recommend pruning for `patience` consecutive steps before the
//! trial is actually pruned. This prevents premature pruning when
//! intermediate values are noisy and a single bad step doesn't indicate
//! a truly bad trial.
//!
//! # When to use
//!
//! - When your intermediate values have high variance (e.g., mini-batch loss)
//! - When the inner pruner is too aggressive on its own
//! - To add robustness to any statistical pruner without changing its threshold
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `inner` | *(required)* | The underlying pruner to wrap |
//! | `patience` | *(required)* | Number of consecutive prune signals required |
//!
//! # Example
//!
//! ```
//! use optimizer::pruner::{PatientPruner, ThresholdPruner};
//!
//! // Only prune after the threshold is exceeded 3 times in a row
//! let inner = ThresholdPruner::new().upper(100.0);
//! let pruner = PatientPruner::new(inner, 3);
//! ```
use std::collections::HashMap;
use std::sync::Mutex;
use super::Pruner;
use crate::sampler::CompletedTrial;
/// Wraps another pruner and adds a patience window.
///
/// The inner pruner must recommend pruning for `patience` consecutive
/// steps before this pruner actually prunes the trial. This is useful
/// to prevent premature pruning when intermediate values are noisy.
///
/// # Examples
///
/// ```
/// use optimizer::pruner::{PatientPruner, ThresholdPruner};
///
/// // Only prune after the threshold pruner recommends pruning 3 times in a row
/// let inner = ThresholdPruner::new().upper(100.0);
/// let pruner = PatientPruner::new(inner, 3);
/// ```
pub struct PatientPruner {
inner: Box<dyn Pruner>,
patience: u64,
/// Track consecutive prune recommendations per trial.
consecutive_counts: Mutex<HashMap<u64, u64>>,
}
impl PatientPruner {
/// Create a new `PatientPruner` wrapping the given inner pruner.
///
/// The inner pruner must recommend pruning for `patience` consecutive
/// calls before this pruner returns `true`.
pub fn new(inner: impl Pruner + 'static, patience: u64) -> Self {
Self {
inner: Box::new(inner),
patience,
consecutive_counts: Mutex::new(HashMap::new()),
}
}
}
impl Pruner for PatientPruner {
fn should_prune(
&self,
trial_id: u64,
step: u64,
intermediate_values: &[(u64, f64)],
completed_trials: &[CompletedTrial],
) -> bool {
let inner_says_prune =
self.inner
.should_prune(trial_id, step, intermediate_values, completed_trials);
let mut counts = self.consecutive_counts.lock().expect("lock poisoned");
let count = counts.entry(trial_id).or_insert(0);
if inner_says_prune {
*count += 1;
*count >= self.patience
} else {
*count = 0;
false
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pruner::ThresholdPruner;
/// A test pruner that always returns the given value.
struct ConstPruner(bool);
impl Pruner for ConstPruner {
fn should_prune(
&self,
_trial_id: u64,
_step: u64,
_intermediate_values: &[(u64, f64)],
_completed_trials: &[CompletedTrial],
) -> bool {
self.0
}
}
/// A pruner that returns values from a sequence.
struct SequencePruner(Mutex<Vec<bool>>);
impl Pruner for SequencePruner {
fn should_prune(
&self,
_trial_id: u64,
_step: u64,
_intermediate_values: &[(u64, f64)],
_completed_trials: &[CompletedTrial],
) -> bool {
self.0.lock().expect("lock poisoned").remove(0)
}
}
fn call(pruner: &PatientPruner, trial_id: u64, step: u64) -> bool {
pruner.should_prune(trial_id, step, &[(step, 0.0)], &[])
}
#[test]
fn patience_1_behaves_like_inner() {
let pruner = PatientPruner::new(ConstPruner(true), 1);
assert!(call(&pruner, 0, 0));
assert!(call(&pruner, 0, 1));
let pruner = PatientPruner::new(ConstPruner(false), 1);
assert!(!call(&pruner, 0, 0));
assert!(!call(&pruner, 0, 1));
}
#[test]
fn patience_3_requires_consecutive_recommendations() {
let pruner = PatientPruner::new(ConstPruner(true), 3);
assert!(!call(&pruner, 0, 0)); // count=1
assert!(!call(&pruner, 0, 1)); // count=2
assert!(call(&pruner, 0, 2)); // count=3 → prune
}
#[test]
fn counter_resets_on_no_prune() {
// Sequence: prune, prune, no-prune, prune, prune, prune
let seq = vec![true, true, false, true, true, true];
let pruner = PatientPruner::new(SequencePruner(Mutex::new(seq)), 3);
assert!(!call(&pruner, 0, 0)); // count=1
assert!(!call(&pruner, 0, 1)); // count=2
assert!(!call(&pruner, 0, 2)); // reset → count=0
assert!(!call(&pruner, 0, 3)); // count=1
assert!(!call(&pruner, 0, 4)); // count=2
assert!(call(&pruner, 0, 5)); // count=3 → prune
}
#[test]
fn independent_per_trial() {
let pruner = PatientPruner::new(ConstPruner(true), 2);
assert!(!call(&pruner, 0, 0)); // trial 0: count=1
assert!(!call(&pruner, 1, 0)); // trial 1: count=1
assert!(call(&pruner, 0, 1)); // trial 0: count=2 → prune
assert!(!call(&pruner, 2, 0)); // trial 2: count=1
assert!(call(&pruner, 1, 1)); // trial 1: count=2 → prune
}
#[test]
fn works_with_threshold_pruner() {
let inner = ThresholdPruner::new().upper(10.0);
let pruner = PatientPruner::new(inner, 2);
// Value below threshold → inner says no
assert!(!pruner.should_prune(0, 0, &[(0, 5.0)], &[]));
// Value above threshold → inner says yes, count=1
assert!(!pruner.should_prune(0, 1, &[(0, 5.0), (1, 15.0)], &[]));
// Value above threshold again → count=2 → prune
assert!(pruner.should_prune(0, 2, &[(0, 5.0), (1, 15.0), (2, 20.0)], &[]));
}
}
+334
View File
@@ -0,0 +1,334 @@
//! Percentile pruner — prune trials outside the top N% at each step.
//!
//! A generalization of [`MedianPruner`](super::MedianPruner) that lets you
//! control how aggressively to prune. At each step, the current trial's
//! intermediate value is compared against the given percentile of all
//! completed trials' values at the same step.
//!
//! # When to use
//!
//! - When you want finer control over pruning aggressiveness than median pruning
//! - Lower percentiles (e.g., 25%) are more aggressive — only keep the best quarter
//! - Higher percentiles (e.g., 75%) are more lenient — keep the top three quarters
//! - Percentile 50% is equivalent to [`MedianPruner`](super::MedianPruner)
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `percentile` | *(required)* | Keep trials in the top N% — range `(0, 100)` |
//! | `n_warmup_steps` | 0 | Skip pruning in the first N steps |
//! | `n_min_trials` | 1 | Require at least N completed trials before pruning |
//!
//! # Example
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::pruner::PercentilePruner;
//!
//! // Keep only the top 25% of trials (aggressive pruning)
//! let pruner = PercentilePruner::new(25.0, Direction::Minimize)
//! .n_warmup_steps(5)
//! .n_min_trials(3);
//! ```
use super::Pruner;
use crate::sampler::CompletedTrial;
use crate::types::{Direction, TrialState};
/// Prune trials that are not in the top `percentile`% of completed trials
/// at the same training step.
///
/// `PercentilePruner::new(50.0, direction)` is equivalent to `MedianPruner`.
/// `PercentilePruner::new(25.0, direction)` keeps only the top 25% of trials.
///
/// # Examples
///
/// ```
/// use optimizer::Direction;
/// use optimizer::pruner::PercentilePruner;
///
/// // Keep only the top 25% of trials (aggressive pruning)
/// let pruner = PercentilePruner::new(25.0, Direction::Minimize)
/// .n_warmup_steps(5)
/// .n_min_trials(3);
/// ```
pub struct PercentilePruner {
/// Keep trials in the top `percentile`%. Range: (0.0, 100.0).
percentile: f64,
/// Don't prune in the first N steps (let the trial warm up).
n_warmup_steps: u64,
/// Require at least N completed trials before pruning.
n_min_trials: usize,
/// The optimization direction.
direction: Direction,
}
impl PercentilePruner {
/// Create a new `PercentilePruner` for the given percentile and direction.
///
/// The `percentile` value must be in `(0.0, 100.0)`.
/// A percentile of 50.0 is equivalent to median pruning.
///
/// # Panics
///
/// Panics if `percentile` is not in `(0.0, 100.0)`.
#[must_use]
pub fn new(percentile: f64, direction: Direction) -> Self {
assert!(
percentile > 0.0 && percentile < 100.0,
"percentile must be in (0.0, 100.0), got {percentile}"
);
Self {
percentile,
n_warmup_steps: 0,
n_min_trials: 1,
direction,
}
}
/// Set the number of warmup steps. No pruning occurs before this step.
#[must_use]
pub fn n_warmup_steps(mut self, n: u64) -> Self {
self.n_warmup_steps = n;
self
}
/// Set the minimum number of completed trials required before pruning.
///
/// # Panics
///
/// Panics if `n` is 0.
#[must_use]
pub fn n_min_trials(mut self, n: usize) -> Self {
assert!(n >= 1, "n_min_trials must be >= 1, got {n}");
self.n_min_trials = n;
self
}
}
impl Pruner for PercentilePruner {
fn should_prune(
&self,
_trial_id: u64,
step: u64,
intermediate_values: &[(u64, f64)],
completed_trials: &[CompletedTrial],
) -> bool {
// 1. Don't prune during warmup
if step < self.n_warmup_steps {
return false;
}
// Get the current trial's latest value
let Some(&(_, current_value)) = intermediate_values.last() else {
return false;
};
// 2. Collect values at this step from completed (non-pruned) trials
let mut values_at_step: Vec<f64> = completed_trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.filter_map(|t| {
t.intermediate_values
.iter()
.find(|(s, _)| *s == step)
.map(|(_, v)| *v)
})
.collect();
// 3. Not enough trials
if values_at_step.len() < self.n_min_trials {
return false;
}
// 4. Compute percentile threshold
let threshold = compute_percentile(&mut values_at_step, self.percentile);
// 5. Compare against threshold based on direction
match self.direction {
Direction::Minimize => current_value > threshold,
Direction::Maximize => current_value < threshold,
}
}
}
/// Compute the given percentile of a non-empty slice. Sorts the slice in place.
///
/// Uses linear interpolation between the two nearest ranks.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
pub(crate) fn compute_percentile(values: &mut [f64], percentile: f64) -> f64 {
assert!(!values.is_empty(), "compute_percentile: empty input");
values.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
let len = values.len();
if len == 1 {
return values[0];
}
// Rank in [0, len-1] range
let rank = percentile / 100.0 * (len - 1) as f64;
let lower = rank.floor() as usize;
let upper = rank.ceil() as usize;
if lower == upper {
values[lower]
} else {
let frac = rank - lower as f64;
values[lower] * (1.0 - frac) + values[upper] * frac
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn compute_percentile_median_odd() {
// Percentile 50 on odd-length slice = median
let val = compute_percentile(&mut [3.0, 1.0, 2.0], 50.0);
assert!((val - 2.0).abs() < f64::EPSILON);
}
#[test]
fn compute_percentile_median_even() {
// Percentile 50 on even-length slice = median (interpolated)
let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 50.0);
assert!((val - 2.5).abs() < f64::EPSILON);
}
#[test]
fn compute_percentile_25() {
// [1.0, 2.0, 3.0, 4.0], rank = 0.25 * 3 = 0.75
// interpolate: 1.0 * 0.25 + 2.0 * 0.75 = 1.75
let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 25.0);
assert!((val - 1.75).abs() < f64::EPSILON);
}
#[test]
fn compute_percentile_75() {
// [1.0, 2.0, 3.0, 4.0], rank = 0.75 * 3 = 2.25
// interpolate: 3.0 * 0.75 + 4.0 * 0.25 = 3.25
let val = compute_percentile(&mut [4.0, 1.0, 3.0, 2.0], 75.0);
assert!((val - 3.25).abs() < f64::EPSILON);
}
#[test]
fn compute_percentile_single() {
let val = compute_percentile(&mut [5.0], 50.0);
assert!((val - 5.0).abs() < f64::EPSILON);
}
#[test]
#[should_panic(expected = "percentile must be in (0.0, 100.0)")]
fn new_rejects_zero() {
let _ = PercentilePruner::new(0.0, Direction::Minimize);
}
#[test]
#[should_panic(expected = "percentile must be in (0.0, 100.0)")]
fn new_rejects_hundred() {
let _ = PercentilePruner::new(100.0, Direction::Minimize);
}
fn make_completed_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
use std::collections::HashMap;
use crate::parameter::ParamId;
CompletedTrial::with_intermediate_values(
id,
HashMap::<ParamId, crate::parameter::ParamValue>::new(),
HashMap::new(),
HashMap::new(),
0.0,
values.to_vec(),
HashMap::new(),
)
}
#[test]
fn percentile_50_matches_median_behavior() {
let pruner = PercentilePruner::new(50.0, Direction::Minimize);
let completed = vec![
make_completed_trial(0, &[(0, 1.0), (1, 2.0)]),
make_completed_trial(1, &[(0, 3.0), (1, 4.0)]),
make_completed_trial(2, &[(0, 5.0), (1, 6.0)]),
];
// Median at step 1 is 4.0
// Value 5.0 > 4.0 → prune
assert!(pruner.should_prune(3, 1, &[(0, 3.0), (1, 5.0)], &completed));
// Value 3.0 < 4.0 → keep
assert!(!pruner.should_prune(3, 1, &[(0, 3.0), (1, 3.0)], &completed));
}
#[test]
fn percentile_25_is_more_aggressive() {
let pruner_25 = PercentilePruner::new(25.0, Direction::Minimize);
let pruner_75 = PercentilePruner::new(75.0, Direction::Minimize);
let completed = vec![
make_completed_trial(0, &[(0, 1.0)]),
make_completed_trial(1, &[(0, 2.0)]),
make_completed_trial(2, &[(0, 3.0)]),
make_completed_trial(3, &[(0, 4.0)]),
];
// 25th percentile at step 0: 1.75
// 75th percentile at step 0: 3.25
// Value 2.5: above 25th (prune), below 75th (keep)
assert!(pruner_25.should_prune(4, 0, &[(0, 2.5)], &completed));
assert!(!pruner_75.should_prune(4, 0, &[(0, 2.5)], &completed));
}
#[test]
fn warmup_prevents_pruning() {
let pruner = PercentilePruner::new(50.0, Direction::Minimize).n_warmup_steps(5);
let completed = vec![make_completed_trial(0, &[(0, 1.0)])];
// Step 3 < warmup 5 → no prune even with bad value
assert!(!pruner.should_prune(1, 3, &[(3, 100.0)], &completed));
}
#[test]
fn n_min_trials_prevents_pruning() {
let pruner = PercentilePruner::new(50.0, Direction::Minimize).n_min_trials(5);
let completed = vec![
make_completed_trial(0, &[(0, 1.0)]),
make_completed_trial(1, &[(0, 2.0)]),
];
// Only 2 trials, need 5 → no prune
assert!(!pruner.should_prune(2, 0, &[(0, 100.0)], &completed));
}
#[test]
fn maximize_direction() {
let pruner = PercentilePruner::new(50.0, Direction::Maximize);
let completed = vec![
make_completed_trial(0, &[(0, 1.0)]),
make_completed_trial(1, &[(0, 3.0)]),
make_completed_trial(2, &[(0, 5.0)]),
];
// Median at step 0 is 3.0
// Value 2.0 < 3.0 → prune (maximize wants higher)
assert!(pruner.should_prune(3, 0, &[(0, 2.0)], &completed));
// Value 4.0 > 3.0 → keep
assert!(!pruner.should_prune(3, 0, &[(0, 4.0)], &completed));
}
#[test]
fn near_boundary_percentiles() {
let pruner_low = PercentilePruner::new(1.0, Direction::Minimize);
let pruner_high = PercentilePruner::new(99.0, Direction::Minimize);
let completed = vec![
make_completed_trial(0, &[(0, 1.0)]),
make_completed_trial(1, &[(0, 2.0)]),
make_completed_trial(2, &[(0, 3.0)]),
make_completed_trial(3, &[(0, 100.0)]),
];
// Percentile 1 is very aggressive (threshold near 1.0)
// Value 1.5 should be pruned
assert!(pruner_low.should_prune(4, 0, &[(0, 1.5)], &completed));
// Percentile 99 is very lenient (threshold near 100.0)
// Value 50.0 should not be pruned
assert!(!pruner_high.should_prune(4, 0, &[(0, 50.0)], &completed));
}
}
+516
View File
@@ -0,0 +1,516 @@
//! Successive Halving (SHA) pruner — budget-aware pruning at exponential rungs.
//!
//! Trials are evaluated at exponentially-spaced "rungs" (checkpoints). At each
//! rung, only the top 1/η fraction of trials survive to the next rung. This
//! is a principled way to allocate compute budget: give many trials a small
//! budget, then progressively invest more in the best ones.
//!
//! For example, with `min_resource=1`, `max_resource=81`, `reduction_factor=3`:
//!
//! | Rung | Step | Survivors |
//! |------|------|-----------|
//! | 0 | 1 | top 1/3 |
//! | 1 | 3 | top 1/3 |
//! | 2 | 9 | top 1/3 |
//! | 3 | 27 | top 1/3 |
//! | 4 | 81 | all (full budget) |
//!
//! # When to use
//!
//! - When your objective has a natural "budget" dimension (epochs, iterations)
//! - When early performance is a reasonable predictor of final performance
//! - When you want a principled alternative to median pruning
//!
//! If you're unsure about the right `min_resource`, consider
//! [`HyperbandPruner`](super::HyperbandPruner) which runs multiple brackets
//! to hedge against that choice.
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `min_resource` | 1 | Step at which the first rung is placed |
//! | `max_resource` | 81 | Full budget (final rung, no pruning) |
//! | `reduction_factor` | 3 | At each rung, keep top 1/η trials |
//! | `min_early_stopping_rate` | 0 | Skip the first N rungs |
//! | `direction` | `Minimize` | Optimization direction |
//!
//! # Example
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::pruner::SuccessiveHalvingPruner;
//!
//! let pruner = SuccessiveHalvingPruner::new()
//! .min_resource(1)
//! .max_resource(81)
//! .reduction_factor(3)
//! .direction(Direction::Minimize);
//! ```
use super::Pruner;
use crate::sampler::CompletedTrial;
use crate::types::{Direction, TrialState};
/// Successive Halving pruner based on the SHA algorithm.
///
/// Trials are evaluated at exponentially-spaced "rungs". At each rung,
/// only the top 1/eta fraction of trials survive to the next rung.
///
/// For example, with `min_resource=1`, `max_resource=81`, `reduction_factor=3`:
/// - Rung 0: evaluate at step 1, keep top 1/3
/// - Rung 1: evaluate at step 3, keep top 1/3
/// - Rung 2: evaluate at step 9, keep top 1/3
/// - Rung 3: evaluate at step 27, keep top 1/3
/// - Rung 4: evaluate at step 81 (full budget)
///
/// # Examples
///
/// ```
/// use optimizer::Direction;
/// use optimizer::pruner::SuccessiveHalvingPruner;
///
/// let pruner = SuccessiveHalvingPruner::new()
/// .min_resource(1)
/// .max_resource(81)
/// .reduction_factor(3)
/// .direction(Direction::Minimize);
/// ```
pub struct SuccessiveHalvingPruner {
min_resource: u64,
max_resource: u64,
reduction_factor: u64,
min_early_stopping_rate: u64,
direction: Direction,
}
impl SuccessiveHalvingPruner {
/// Create a new `SuccessiveHalvingPruner` with default parameters.
///
/// Defaults: `min_resource=1`, `max_resource=81`, `reduction_factor=3`,
/// `min_early_stopping_rate=0`, `direction=Minimize`.
#[must_use]
pub fn new() -> Self {
Self {
min_resource: 1,
max_resource: 81,
reduction_factor: 3,
min_early_stopping_rate: 0,
direction: Direction::Minimize,
}
}
/// Set the minimum resource (budget) per trial.
///
/// # Panics
///
/// Panics if `r` is 0.
#[must_use]
pub fn min_resource(mut self, r: u64) -> Self {
assert!(r > 0, "min_resource must be > 0, got {r}");
self.min_resource = r;
self
}
/// Set the maximum resource (budget) per trial.
///
/// # Panics
///
/// Panics if `r` is 0.
#[must_use]
pub fn max_resource(mut self, r: u64) -> Self {
assert!(r > 0, "max_resource must be > 0, got {r}");
self.max_resource = r;
self
}
/// Set the reduction factor (eta). At each rung, the top 1/eta trials survive.
///
/// # Panics
///
/// Panics if `eta` is less than 2.
#[must_use]
pub fn reduction_factor(mut self, eta: u64) -> Self {
assert!(eta >= 2, "reduction_factor must be >= 2, got {eta}");
self.reduction_factor = eta;
self
}
/// Set the minimum early stopping rate. Skips the first N rungs.
#[must_use]
pub fn min_early_stopping_rate(mut self, n: u64) -> Self {
self.min_early_stopping_rate = n;
self
}
/// Set the optimization direction.
#[must_use]
pub fn direction(mut self, d: Direction) -> Self {
self.direction = d;
self
}
/// Compute the rung steps: `[min_resource * eta^(s), ...]` up to `max_resource`,
/// skipping the first `min_early_stopping_rate` rungs.
fn rung_steps(&self) -> Vec<u64> {
let eta = self.reduction_factor;
let mut steps = Vec::new();
let mut rung: u32 = 0;
while let Some(power) = eta.checked_pow(rung) {
let step = self.min_resource.saturating_mul(power);
if step > self.max_resource {
break;
}
if u64::from(rung) >= self.min_early_stopping_rate {
steps.push(step);
}
rung += 1;
}
steps
}
}
impl Default for SuccessiveHalvingPruner {
fn default() -> Self {
Self::new()
}
}
#[allow(clippy::cast_precision_loss)]
impl Pruner for SuccessiveHalvingPruner {
fn should_prune(
&self,
_trial_id: u64,
step: u64,
intermediate_values: &[(u64, f64)],
completed_trials: &[CompletedTrial],
) -> bool {
let rungs = self.rung_steps();
// Find the highest rung step <= current step
let Some(&rung_step) = rungs.iter().rev().find(|&&r| r <= step) else {
// No rung matches (before the first rung) → don't prune
return false;
};
// If this is the last rung (full budget), don't prune
if rung_step >= self.max_resource {
return false;
}
// Get the current trial's value at this rung step
let Some(&(_, current_value)) = intermediate_values.iter().find(|(s, _)| *s == rung_step)
else {
// Trial hasn't reported a value at this exact rung step.
// Use the latest intermediate value at or before the rung step instead.
let Some(&(_, current_value)) = intermediate_values
.iter()
.rev()
.find(|(s, _)| *s <= rung_step)
else {
return false;
};
return self.is_pruned_at_rung(current_value, rung_step, completed_trials);
};
self.is_pruned_at_rung(current_value, rung_step, completed_trials)
}
}
impl SuccessiveHalvingPruner {
/// Determine whether a trial with `current_value` should be pruned at the given rung.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
fn is_pruned_at_rung(
&self,
current_value: f64,
rung_step: u64,
completed_trials: &[CompletedTrial],
) -> bool {
let eta = self.reduction_factor as usize;
// Collect values at this rung step from all trials that reached it
let mut values_at_rung: Vec<f64> = completed_trials
.iter()
.filter(|t| t.state == TrialState::Complete || t.state == TrialState::Pruned)
.filter_map(|t| {
t.intermediate_values
.iter()
.find(|(s, _)| *s == rung_step)
.map(|(_, v)| *v)
})
.collect();
// Need at least eta trials to make a meaningful comparison
// (with fewer trials, we can't determine the top 1/eta fraction)
if values_at_rung.len() < eta {
return false;
}
// Include the current trial's value for ranking
values_at_rung.push(current_value);
// Sort based on direction: best values first
values_at_rung
.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(core::cmp::Ordering::Equal));
if self.direction == Direction::Maximize {
values_at_rung.reverse();
}
// Keep top 1/eta fraction
let n_keep = (values_at_rung.len() as f64 / eta as f64).ceil() as usize;
let threshold_idx = n_keep.max(1) - 1;
let threshold = values_at_rung[threshold_idx];
// Prune if current value is worse than the threshold
match self.direction {
Direction::Minimize => current_value > threshold,
Direction::Maximize => current_value < threshold,
}
}
}
#[cfg(test)]
#[allow(clippy::cast_precision_loss)]
mod tests {
use super::*;
fn make_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
use std::collections::HashMap;
use crate::parameter::ParamId;
CompletedTrial::with_intermediate_values(
id,
HashMap::<ParamId, crate::parameter::ParamValue>::new(),
HashMap::new(),
HashMap::new(),
0.0,
values.to_vec(),
HashMap::new(),
)
}
fn make_pruned_trial(id: u64, values: &[(u64, f64)]) -> CompletedTrial {
let mut t = make_trial(id, values);
t.state = TrialState::Pruned;
t
}
#[test]
fn rung_steps_default() {
let pruner = SuccessiveHalvingPruner::new();
let rungs = pruner.rung_steps();
// min=1, max=81, eta=3 → 1, 3, 9, 27, 81
assert_eq!(rungs, vec![1, 3, 9, 27, 81]);
}
#[test]
fn rung_steps_custom() {
let pruner = SuccessiveHalvingPruner::new()
.min_resource(2)
.max_resource(32)
.reduction_factor(2);
let rungs = pruner.rung_steps();
// 2, 4, 8, 16, 32
assert_eq!(rungs, vec![2, 4, 8, 16, 32]);
}
#[test]
fn rung_steps_with_early_stopping_rate() {
let pruner = SuccessiveHalvingPruner::new().min_early_stopping_rate(2);
let rungs = pruner.rung_steps();
// Skip rung 0 (step=1) and rung 1 (step=3), keep rung 2+ (9, 27, 81)
assert_eq!(rungs, vec![9, 27, 81]);
}
#[test]
fn no_prune_before_first_rung() {
let pruner = SuccessiveHalvingPruner::new()
.min_resource(10)
.max_resource(100)
.reduction_factor(3);
let completed = vec![
make_trial(0, &[(5, 1.0)]),
make_trial(1, &[(5, 2.0)]),
make_trial(2, &[(5, 3.0)]),
];
// Step 5 is before the first rung (10)
assert!(!pruner.should_prune(3, 5, &[(5, 100.0)], &completed));
}
#[test]
fn no_prune_with_single_trial() {
let pruner = SuccessiveHalvingPruner::new();
let completed = vec![make_trial(0, &[(1, 5.0)]), make_trial(1, &[(1, 3.0)])];
// Only 2 completed trials at rung + 1 current = 3 total, threshold = ceil(3/3) = 1
// With eta=3, we need at least 3 completed trials
assert!(!pruner.should_prune(2, 1, &[(1, 10.0)], &completed));
}
#[test]
fn prune_worst_trials_at_rung() {
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
// 9 completed trials at rung step=1, with values 1..=9
let completed: Vec<_> = (0..9)
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
.collect();
// With eta=3, keep top 1/3. 10 total values → ceil(10/3) = 4 kept
// Best 4 values: 1, 2, 3, 4. Threshold = 4.0
// Value 3.0 → keep (in top 1/3)
assert!(!pruner.should_prune(9, 1, &[(1, 3.0)], &completed));
// Value 5.0 → prune (not in top 1/3)
assert!(pruner.should_prune(9, 1, &[(1, 5.0)], &completed));
}
#[test]
fn top_fraction_survives() {
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
// 6 completed trials at step=1
let completed: Vec<_> = (0..6)
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
.collect();
// 7 total (6 + current). ceil(7/3) = 3 keep. Threshold = 3.0
// Value 2.0 → keep
assert!(!pruner.should_prune(6, 1, &[(1, 2.0)], &completed));
// Value 3.0 → keep (at threshold)
assert!(!pruner.should_prune(6, 1, &[(1, 3.0)], &completed));
// Value 4.0 → prune
assert!(pruner.should_prune(6, 1, &[(1, 4.0)], &completed));
}
#[test]
fn maximize_direction() {
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Maximize);
let completed: Vec<_> = (0..6)
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
.collect();
// 7 total. For maximize, best = highest. ceil(7/3)=3. Top 3: 6,5,4. Threshold=4.0
// Value 5.0 → keep
assert!(!pruner.should_prune(6, 1, &[(1, 5.0)], &completed));
// Value 4.0 → keep (at threshold)
assert!(!pruner.should_prune(6, 1, &[(1, 4.0)], &completed));
// Value 3.0 → prune
assert!(pruner.should_prune(6, 1, &[(1, 3.0)], &completed));
}
#[test]
fn reduction_factor_2() {
let pruner = SuccessiveHalvingPruner::new()
.reduction_factor(2)
.min_resource(1)
.max_resource(16)
.direction(Direction::Minimize);
// Rungs: 1, 2, 4, 8, 16
assert_eq!(pruner.rung_steps(), vec![1, 2, 4, 8, 16]);
// 4 completed trials at rung step=1
let completed: Vec<_> = (0..4)
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
.collect();
// With eta=2, 5 total. ceil(5/2) = 3 keep. Threshold = 3.0
// Value 3.0 → keep
assert!(!pruner.should_prune(4, 1, &[(1, 3.0)], &completed));
// Value 4.0 → prune
assert!(pruner.should_prune(4, 1, &[(1, 4.0)], &completed));
}
#[test]
fn reduction_factor_4() {
let pruner = SuccessiveHalvingPruner::new()
.reduction_factor(4)
.min_resource(1)
.max_resource(64)
.direction(Direction::Minimize);
// Rungs: 1, 4, 16, 64
assert_eq!(pruner.rung_steps(), vec![1, 4, 16, 64]);
// 12 completed trials at rung step=1
let completed: Vec<_> = (0..12)
.map(|i| make_trial(i, &[(1, (i + 1) as f64)]))
.collect();
// With eta=4, 13 total. ceil(13/4) = 4 keep. Threshold = 4.0
// Value 4.0 → keep
assert!(!pruner.should_prune(12, 1, &[(1, 4.0)], &completed));
// Value 5.0 → prune
assert!(pruner.should_prune(12, 1, &[(1, 5.0)], &completed));
}
#[test]
fn non_contiguous_steps() {
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
// Trials reporting at rung step=3 (not step=1)
let completed: Vec<_> = (0..6)
.map(|i| make_trial(i, &[(3, (i + 1) as f64)]))
.collect();
// Current trial reports at step 5 (between rung 3 and rung 9)
// Highest rung <= 5 is 3. Use value at rung step 3.
// Trial has value at step 3 → use it
assert!(!pruner.should_prune(6, 5, &[(3, 2.0)], &completed));
assert!(pruner.should_prune(6, 5, &[(3, 5.0)], &completed));
}
#[test]
fn no_prune_at_max_resource() {
let pruner = SuccessiveHalvingPruner::new();
let completed: Vec<_> = (0..9)
.map(|i| make_trial(i, &[(81, (i + 1) as f64)]))
.collect();
// At the max resource rung, never prune (trial should complete)
assert!(!pruner.should_prune(9, 81, &[(81, 100.0)], &completed));
}
#[test]
fn includes_pruned_trials_in_comparison() {
let pruner = SuccessiveHalvingPruner::new().direction(Direction::Minimize);
// Mix of completed and pruned trials at rung step=1
let completed = vec![
make_trial(0, &[(1, 1.0)]),
make_trial(1, &[(1, 2.0)]),
make_pruned_trial(2, &[(1, 8.0)]),
make_pruned_trial(3, &[(1, 9.0)]),
make_pruned_trial(4, &[(1, 10.0)]),
];
// 6 total. ceil(6/3) = 2 keep. Threshold = 2.0
// Value 2.0 → keep
assert!(!pruner.should_prune(5, 1, &[(1, 2.0)], &completed));
// Value 3.0 → prune
assert!(pruner.should_prune(5, 1, &[(1, 3.0)], &completed));
}
#[test]
#[should_panic(expected = "min_resource must be > 0")]
fn rejects_zero_min_resource() {
let _ = SuccessiveHalvingPruner::new().min_resource(0);
}
#[test]
#[should_panic(expected = "max_resource must be > 0")]
fn rejects_zero_max_resource() {
let _ = SuccessiveHalvingPruner::new().max_resource(0);
}
#[test]
#[should_panic(expected = "reduction_factor must be >= 2")]
fn rejects_reduction_factor_one() {
let _ = SuccessiveHalvingPruner::new().reduction_factor(1);
}
}
+113
View File
@@ -0,0 +1,113 @@
//! Threshold pruner — prune trials whose values fall outside fixed bounds.
//!
//! Unlike statistical pruners that compare against other trials, the
//! threshold pruner uses absolute bounds. Any trial whose latest
//! intermediate value exceeds the upper bound or falls below the lower
//! bound is pruned immediately.
//!
//! # When to use
//!
//! - When you know hard limits for valid intermediate values (e.g., loss should
//! never exceed 100.0)
//! - To catch diverging or NaN-producing trials early
//! - Often combined with other pruners via [`PatientPruner`](super::PatientPruner)
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `upper` | `None` | Prune if value exceeds this bound |
//! | `lower` | `None` | Prune if value falls below this bound |
//!
//! # Example
//!
//! ```
//! use optimizer::pruner::ThresholdPruner;
//!
//! // Prune if loss exceeds 100.0 or accuracy drops below 0.0
//! let pruner = ThresholdPruner::new().upper(100.0).lower(0.0);
//! ```
use super::Pruner;
use crate::sampler::CompletedTrial;
/// Prune trials whose intermediate values exceed fixed thresholds.
///
/// Useful for cutting off trials that are clearly diverging or stuck
/// at bad values early in training.
///
/// # Examples
///
/// ```
/// use optimizer::pruner::ThresholdPruner;
///
/// // Prune if the intermediate value exceeds 100.0 or falls below 0.0
/// let pruner = ThresholdPruner::new().upper(100.0).lower(0.0);
/// ```
pub struct ThresholdPruner {
/// Prune if intermediate value is greater than this. `None` = no upper bound.
upper: Option<f64>,
/// Prune if intermediate value is less than this. `None` = no lower bound.
lower: Option<f64>,
}
impl ThresholdPruner {
/// Create a new `ThresholdPruner` with no thresholds set.
///
/// By default, no pruning occurs. Use [`upper`](Self::upper) and
/// [`lower`](Self::lower) to set bounds.
#[must_use]
pub fn new() -> Self {
Self {
upper: None,
lower: None,
}
}
/// Set the upper threshold. Trials with intermediate values above this
/// will be pruned.
#[must_use]
pub fn upper(mut self, threshold: f64) -> Self {
self.upper = Some(threshold);
self
}
/// Set the lower threshold. Trials with intermediate values below this
/// will be pruned.
#[must_use]
pub fn lower(mut self, threshold: f64) -> Self {
self.lower = Some(threshold);
self
}
}
impl Default for ThresholdPruner {
fn default() -> Self {
Self::new()
}
}
impl Pruner for ThresholdPruner {
fn should_prune(
&self,
_trial_id: u64,
_step: u64,
intermediate_values: &[(u64, f64)],
_completed_trials: &[CompletedTrial],
) -> bool {
let Some(&(_, latest_value)) = intermediate_values.last() else {
return false;
};
if let Some(upper) = self.upper
&& latest_value > upper
{
return true;
}
if let Some(lower) = self.lower
&& latest_value < lower
{
return true;
}
false
}
}
+572
View File
@@ -0,0 +1,572 @@
//! Wilcoxon pruner — statistically rigorous pruning for noisy objectives.
//!
//! Uses the Wilcoxon signed-rank test to compare the current trial's
//! intermediate values against the best completed trial at matching steps.
//! The test accounts for the paired, step-aligned nature of the comparison
//! and only prunes when the difference is statistically significant.
//!
//! This is more principled than [`MedianPruner`](super::MedianPruner) for
//! noisy objectives because a single bad step won't trigger pruning — the
//! test considers the full distribution of paired differences.
//!
//! # When to use
//!
//! - When intermediate values have high variance (e.g., mini-batch loss,
//! stochastic reward signals)
//! - When you want a statistical guarantee that pruned trials are truly worse
//! - When you have enough steps (at least 6) for a meaningful test
//!
//! For less noisy objectives, [`MedianPruner`](super::MedianPruner) is simpler
//! and often sufficient.
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `p_value_threshold` | 0.05 | Significance level — lower is more conservative |
//! | `n_warmup_steps` | 0 | Skip pruning in the first N steps |
//! | `n_min_trials` | 1 | Require at least N completed trials before pruning |
//!
//! # Example
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::pruner::WilcoxonPruner;
//!
//! let pruner = WilcoxonPruner::new(Direction::Minimize)
//! .p_value_threshold(0.05)
//! .n_warmup_steps(5)
//! .n_min_trials(1);
//! ```
use core::cmp::Ordering;
use super::Pruner;
use crate::sampler::CompletedTrial;
use crate::types::{Direction, TrialState};
/// Prune trials using a Wilcoxon signed-rank test comparing intermediate
/// values against the best completed trial.
///
/// More principled than `MedianPruner` for noisy objectives — it accounts
/// for the paired nature of step-aligned comparisons and doesn't prune
/// on random fluctuations.
///
/// The test compares intermediate values at matching steps between the
/// current trial and the best completed trial. If the current trial is
/// statistically significantly worse (p < threshold), it is pruned.
///
/// # Examples
///
/// ```
/// use optimizer::Direction;
/// use optimizer::pruner::WilcoxonPruner;
///
/// let pruner = WilcoxonPruner::new(Direction::Minimize)
/// .p_value_threshold(0.05)
/// .n_warmup_steps(5)
/// .n_min_trials(1);
/// ```
pub struct WilcoxonPruner {
/// Significance level (default 0.05). Lower = more conservative.
p_value_threshold: f64,
/// Don't prune in the first N steps (let the trial warm up).
n_warmup_steps: u64,
/// Require at least N completed trials before pruning.
n_min_trials: usize,
/// The optimization direction.
direction: Direction,
}
impl WilcoxonPruner {
/// Create a new `WilcoxonPruner` for the given optimization direction.
///
/// By default, `p_value_threshold` is 0.05, `n_warmup_steps` is 0,
/// and `n_min_trials` is 1.
#[must_use]
pub fn new(direction: Direction) -> Self {
Self {
p_value_threshold: 0.05,
n_warmup_steps: 0,
n_min_trials: 1,
direction,
}
}
/// Set the p-value threshold for significance.
///
/// Must be in (0.0, 1.0). Lower values are more conservative (harder to prune).
///
/// # Panics
///
/// Panics if `p` is not in the open interval (0.0, 1.0).
#[must_use]
pub fn p_value_threshold(mut self, p: f64) -> Self {
assert!(
p > 0.0 && p < 1.0,
"p_value_threshold must be in (0.0, 1.0)"
);
self.p_value_threshold = p;
self
}
/// Set the number of warmup steps. No pruning occurs before this step.
#[must_use]
pub fn n_warmup_steps(mut self, n: u64) -> Self {
self.n_warmup_steps = n;
self
}
/// Set the minimum number of completed trials required before pruning.
#[must_use]
pub fn n_min_trials(mut self, n: usize) -> Self {
self.n_min_trials = n;
self
}
}
impl Pruner for WilcoxonPruner {
fn should_prune(
&self,
_trial_id: u64,
step: u64,
intermediate_values: &[(u64, f64)],
completed_trials: &[CompletedTrial],
) -> bool {
if step < self.n_warmup_steps {
return false;
}
let completed: Vec<&CompletedTrial> = completed_trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.collect();
if completed.len() < self.n_min_trials {
return false;
}
// Find the best completed trial by final objective value.
let best = match self.direction {
Direction::Minimize => completed
.iter()
.min_by(|a, b| a.value.partial_cmp(&b.value).unwrap_or(Ordering::Equal)),
Direction::Maximize => completed
.iter()
.max_by(|a, b| a.value.partial_cmp(&b.value).unwrap_or(Ordering::Equal)),
};
let Some(best) = best else {
return false;
};
// Pair intermediate values at matching steps.
let pairs: Vec<(f64, f64)> = intermediate_values
.iter()
.filter_map(|&(s, current_v)| {
best.intermediate_values
.iter()
.find(|(bs, _)| *bs == s)
.map(|&(_, best_v)| (current_v, best_v))
})
.collect();
// Need at least 6 pairs for a meaningful test.
if pairs.len() < 6 {
return false;
}
// Compute signed differences: current - best.
// For minimization: positive diff means current is worse.
// For maximization: negative diff means current is worse.
let differences: Vec<f64> = pairs
.iter()
.map(|&(current, best_v)| current - best_v)
.collect();
// Run the Wilcoxon signed-rank test.
let p_value = wilcoxon_signed_rank_test(&differences, self.direction);
p_value < self.p_value_threshold
}
}
/// Perform a one-sided Wilcoxon signed-rank test.
///
/// Tests whether the values tend to be worse than zero (positive for
/// minimization, negative for maximization).
///
/// Returns a p-value. Small p-values indicate the current trial is
/// significantly worse.
fn wilcoxon_signed_rank_test(differences: &[f64], direction: Direction) -> f64 {
// 1. Remove zero differences.
let nonzero: Vec<f64> = differences.iter().copied().filter(|d| *d != 0.0).collect();
let n = nonzero.len();
if n < 6 {
return 1.0; // Not enough data
}
// 2. Rank by absolute value.
let mut abs_ranked: Vec<(usize, f64, f64)> = nonzero
.iter()
.enumerate()
.map(|(i, &d)| (i, d.abs(), d))
.collect();
abs_ranked.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
// 3. Assign ranks with tie correction.
let ranks = assign_ranks(&abs_ranked);
// 4. Compute W+ (sum of ranks for positive differences) and
// W- (sum of ranks for negative differences).
let mut w_plus = 0.0;
let mut w_minus = 0.0;
for (i, &(_, _, orig)) in abs_ranked.iter().enumerate() {
if orig > 0.0 {
w_plus += ranks[i];
} else {
w_minus += ranks[i];
}
}
// For one-sided test:
// - Minimization: we want to detect positive diffs (current worse).
// Large W+ means significantly worse. Test statistic = W-.
// - Maximization: we want to detect negative diffs (current worse).
// Large W- means significantly worse. Test statistic = W+.
let w = match direction {
Direction::Minimize => w_minus,
Direction::Maximize => w_plus,
};
// 5. Normal approximation for the p-value.
#[allow(clippy::cast_precision_loss)]
let n_f = n as f64;
let mean = n_f * (n_f + 1.0) / 4.0;
let variance = n_f * (n_f + 1.0) * (2.0 * n_f + 1.0) / 24.0;
// Tie correction for variance.
let tie_correction = compute_tie_correction(&ranks);
let adjusted_variance = variance - tie_correction;
if adjusted_variance <= 0.0 {
return 1.0;
}
let std_dev = adjusted_variance.sqrt();
// Continuity correction: shift W by 0.5 towards mean.
let continuity = if w < mean { 0.5 } else { -0.5 };
let z = (w + continuity - mean) / std_dev;
// One-sided p-value (lower tail): probability that the test statistic
// is this small or smaller under H0.
normal_cdf(z)
}
/// Assign average ranks, handling ties.
fn assign_ranks(sorted: &[(usize, f64, f64)]) -> Vec<f64> {
let n = sorted.len();
let mut ranks = vec![0.0; n];
let mut i = 0;
while i < n {
let mut j = i;
// Find all items tied with sorted[i].
while j < n
&& (sorted[j].1 - sorted[i].1).abs() < f64::EPSILON * sorted[i].1.max(1.0) * 100.0
{
j += 1;
}
// Average rank for the tie group. Ranks are 1-based.
#[allow(clippy::cast_precision_loss)]
let avg_rank = (i + 1 + j) as f64 / 2.0;
for rank in ranks.iter_mut().take(j).skip(i) {
*rank = avg_rank;
}
i = j;
}
ranks
}
/// Compute the tie correction term for the variance.
/// For each tie group of size t, subtract t^3 - t from the sum,
/// then divide by 48.
fn compute_tie_correction(ranks: &[f64]) -> f64 {
let mut correction = 0.0;
let mut i = 0;
while i < ranks.len() {
let mut j = i;
while j < ranks.len() && (ranks[j] - ranks[i]).abs() < f64::EPSILON {
j += 1;
}
#[allow(clippy::cast_precision_loss)]
let t = (j - i) as f64;
if t > 1.0 {
correction += t * t * t - t;
}
i = j;
}
correction / 48.0
}
/// Standard normal CDF using an approximation (Abramowitz & Stegun).
fn normal_cdf(x: f64) -> f64 {
// Use the complementary error function relationship:
// Φ(x) = 0.5 * erfc(-x / √2)
0.5 * erfc(-x / core::f64::consts::SQRT_2)
}
/// Complementary error function approximation.
/// Maximum error: 1.5 × 10⁻⁷ (Abramowitz & Stegun formula 7.1.26).
fn erfc(x: f64) -> f64 {
let t = 1.0 / (1.0 + 0.327_591_1 * x.abs());
let poly = t
* (0.254_829_592
+ t * (-0.284_496_736
+ t * (1.421_413_741 + t * (-1.453_152_027 + t * 1.061_405_429))));
let result = poly * (-x * x).exp();
if x >= 0.0 { result } else { 2.0 - result }
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use super::*;
fn trial_with_values(
id: u64,
value: f64,
intermediate_values: Vec<(u64, f64)>,
) -> CompletedTrial {
CompletedTrial::with_intermediate_values(
id,
HashMap::new(),
HashMap::new(),
HashMap::new(),
value,
intermediate_values,
HashMap::new(),
)
}
#[test]
fn no_prune_during_warmup() {
let pruner = WilcoxonPruner::new(Direction::Minimize).n_warmup_steps(10);
let completed = vec![trial_with_values(
0,
0.1,
(0..20).map(|s| (s, 0.1)).collect(),
)];
let current: Vec<(u64, f64)> = (0..8).map(|s| (s, 100.0)).collect();
assert!(!pruner.should_prune(1, 7, &current, &completed));
}
#[test]
fn no_prune_with_insufficient_trials() {
let pruner = WilcoxonPruner::new(Direction::Minimize).n_min_trials(5);
let completed = vec![trial_with_values(
0,
0.1,
(0..20).map(|s| (s, 0.1)).collect(),
)];
let current: Vec<(u64, f64)> = (0..10).map(|s| (s, 100.0)).collect();
assert!(!pruner.should_prune(1, 9, &current, &completed));
}
#[test]
fn no_prune_with_fewer_than_6_pairs() {
let pruner = WilcoxonPruner::new(Direction::Minimize);
let completed = vec![trial_with_values(
0,
0.1,
(0..5).map(|s| (s, 0.1)).collect(),
)];
// Only 5 matching steps
let current: Vec<(u64, f64)> = (0..5).map(|s| (s, 100.0)).collect();
assert!(!pruner.should_prune(1, 4, &current, &completed));
}
#[test]
fn prune_when_consistently_worse_minimize() {
let pruner = WilcoxonPruner::new(Direction::Minimize);
// Best trial has low values.
let best_values: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
let completed = vec![trial_with_values(0, 0.1, best_values)];
// Current trial is consistently much worse.
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 10.0)).collect();
assert!(pruner.should_prune(1, 19, &current, &completed));
}
#[test]
fn prune_when_consistently_worse_maximize() {
let pruner = WilcoxonPruner::new(Direction::Maximize);
// Best trial has high values.
let best_values: Vec<(u64, f64)> = (0..20).map(|s| (s, 10.0)).collect();
let completed = vec![trial_with_values(0, 10.0, best_values)];
// Current trial is consistently much worse.
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
assert!(pruner.should_prune(1, 19, &current, &completed));
}
#[test]
fn no_prune_when_statistically_similar() {
let pruner = WilcoxonPruner::new(Direction::Minimize);
// Best trial and current trial have very similar values with noise.
let best_values: Vec<(u64, f64)> = (0..20_u64)
.map(|s| {
let noise = if s.is_multiple_of(2) { 0.01 } else { -0.01 };
(s, 1.0 + noise)
})
.collect();
let completed = vec![trial_with_values(0, 1.0, best_values)];
// Current trial is similar — alternating above/below.
let current: Vec<(u64, f64)> = (0..20_u64)
.map(|s| {
let noise = if s.is_multiple_of(2) { -0.01 } else { 0.01 };
(s, 1.0 + noise)
})
.collect();
assert!(!pruner.should_prune(1, 19, &current, &completed));
}
#[test]
fn selects_best_trial_minimize() {
let pruner = WilcoxonPruner::new(Direction::Minimize);
// Two completed trials: trial 0 is better (lower).
let completed = vec![
trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect()),
trial_with_values(1, 5.0, (0..20).map(|s| (s, 5.0)).collect()),
];
// Current trial is worse than the best but similar to the second.
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 5.0)).collect();
assert!(pruner.should_prune(2, 19, &current, &completed));
}
#[test]
fn selects_best_trial_maximize() {
let pruner = WilcoxonPruner::new(Direction::Maximize);
// Two completed trials: trial 1 is better (higher).
let completed = vec![
trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect()),
trial_with_values(1, 10.0, (0..20).map(|s| (s, 10.0)).collect()),
];
// Current trial is worse than the best.
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 0.1)).collect();
assert!(pruner.should_prune(2, 19, &current, &completed));
}
#[test]
fn ignores_pruned_trials() {
let pruner = WilcoxonPruner::new(Direction::Minimize);
// Only a pruned trial — no complete trials.
let mut trial = trial_with_values(0, 0.1, (0..20).map(|s| (s, 0.1)).collect());
trial.state = TrialState::Pruned;
let completed = vec![trial];
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 100.0)).collect();
assert!(!pruner.should_prune(1, 19, &current, &completed));
}
#[test]
fn lower_p_value_is_more_conservative() {
let strict = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.001);
let lenient = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.1);
let completed = vec![trial_with_values(
0,
0.1,
(0..20).map(|s| (s, 0.1)).collect(),
)];
// Moderately worse — should pass lenient but maybe not strict.
let current: Vec<(u64, f64)> = (0..20)
.map(|s| if s < 15 { (s, 0.2) } else { (s, 0.15) })
.collect();
let lenient_prunes = lenient.should_prune(1, 19, &current, &completed);
let strict_prunes = strict.should_prune(1, 19, &current, &completed);
// A stricter threshold should never prune when a lenient one doesn't.
if !lenient_prunes {
assert!(!strict_prunes);
}
}
#[test]
#[should_panic(expected = "p_value_threshold must be in (0.0, 1.0)")]
fn panics_on_zero_p_value() {
let _ = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(0.0);
}
#[test]
#[should_panic(expected = "p_value_threshold must be in (0.0, 1.0)")]
fn panics_on_one_p_value() {
let _ = WilcoxonPruner::new(Direction::Minimize).p_value_threshold(1.0);
}
#[test]
fn correct_signed_rank_statistic() {
// Known example: differences [1, 2, 3, 4, 5, 6] (all positive).
// Ranks: 1, 2, 3, 4, 5, 6. W+ = 21, W- = 0.
// For minimization (testing if positive = worse), W- = 0.
// This should give a very small p-value.
let diffs = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
let p = wilcoxon_signed_rank_test(&diffs, Direction::Minimize);
assert!(
p < 0.05,
"p-value {p} should be < 0.05 for all-positive diffs"
);
}
#[test]
fn symmetric_differences_not_significant() {
// Balanced differences: half positive, half negative.
let diffs = vec![1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0, -4.0];
let p = wilcoxon_signed_rank_test(&diffs, Direction::Minimize);
assert!(p > 0.05, "p-value {p} should be > 0.05 for symmetric diffs");
}
#[test]
fn normal_cdf_known_values() {
assert!((normal_cdf(0.0) - 0.5).abs() < 1e-6);
assert!(normal_cdf(-10.0) < 1e-6);
assert!((normal_cdf(10.0) - 1.0).abs() < 1e-6);
assert!((normal_cdf(-1.96) - 0.025).abs() < 0.001);
}
#[test]
fn no_intermediate_values() {
let pruner = WilcoxonPruner::new(Direction::Minimize);
let completed = vec![trial_with_values(
0,
0.1,
(0..20).map(|s| (s, 0.1)).collect(),
)];
assert!(!pruner.should_prune(1, 0, &[], &completed));
}
#[test]
fn no_completed_trials() {
let pruner = WilcoxonPruner::new(Direction::Minimize);
let current: Vec<(u64, f64)> = (0..20).map(|s| (s, 1.0)).collect();
assert!(!pruner.should_prune(1, 19, &current, &[]));
}
}
+58
View File
@@ -0,0 +1,58 @@
use crate::distribution::Distribution;
/// Generate a random `f64` in the range `[low, high)`.
#[inline]
pub(crate) fn f64_range(rng: &mut fastrand::Rng, low: f64, high: f64) -> f64 {
low + rng.f64() * (high - low)
}
/// Combine a base seed, trial id, and distribution fingerprint into a
/// deterministic per-call seed using `MurmurHash3`'s 64-bit finalizer.
#[inline]
pub(crate) fn mix_seed(base: u64, trial_id: u64, dist_fingerprint: u64) -> u64 {
let mut h = base
.wrapping_mul(0xff51_afd7_ed55_8ccd)
.wrapping_add(trial_id)
.wrapping_mul(0xc4ce_b9fe_1a85_ec53)
.wrapping_add(dist_fingerprint);
h ^= h >> 33;
h = h.wrapping_mul(0xff51_afd7_ed55_8ccd);
h ^= h >> 33;
h = h.wrapping_mul(0xc4ce_b9fe_1a85_ec53);
h ^= h >> 33;
h
}
/// Stable `u64` fingerprint for a [`Distribution`], using variant tags and
/// `f64::to_bits()` for float fields so that distinct distributions within
/// the same trial produce different RNG streams.
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
pub(crate) fn distribution_fingerprint(distribution: &Distribution) -> u64 {
match distribution {
Distribution::Float(d) => {
let mut h: u64 = 1;
h = h.wrapping_mul(31).wrapping_add(d.low.to_bits());
h = h.wrapping_mul(31).wrapping_add(d.high.to_bits());
h = h.wrapping_mul(31).wrapping_add(u64::from(d.log_scale));
if let Some(step) = d.step {
h = h.wrapping_mul(31).wrapping_add(step.to_bits());
}
h
}
Distribution::Int(d) => {
let mut h: u64 = 2;
h = h.wrapping_mul(31).wrapping_add(d.low as u64);
h = h.wrapping_mul(31).wrapping_add(d.high as u64);
h = h.wrapping_mul(31).wrapping_add(u64::from(d.log_scale));
if let Some(step) = d.step {
h = h.wrapping_mul(31).wrapping_add(step as u64);
}
h
}
Distribution::Categorical(d) => {
let mut h: u64 = 3;
h = h.wrapping_mul(31).wrapping_add(d.n_choices as u64);
h
}
}
}
+737
View File
@@ -0,0 +1,737 @@
//! BOHB (Bayesian Optimization + `HyperBand`) sampler.
//!
//! BOHB combines TPE's model-guided sampling with `HyperBand`'s budget-aware
//! evaluation. Instead of building one global TPE model, BOHB conditions
//! its TPE model on trials evaluated at a specific budget level, giving
//! better-calibrated proposals for each rung of the `HyperBand` schedule.
//!
//! # How it works
//!
//! 1. Compute all `HyperBand` rung steps (budget levels) from the config.
//! 2. On each `sample()` call, scan the history's `intermediate_values`
//! to find the **largest budget level** with enough observations
//! (`>= min_points_in_model`).
//! 3. Build a filtered history where each trial's `value` is replaced
//! with its intermediate value at that budget level.
//! 4. Delegate to an internal [`TpeSampler`] for the actual sampling.
//! 5. Fall back to random sampling if no budget level has enough data.
//!
//! # When to use
//!
//! - You are tuning hyperparameters for models that support early stopping
//! (e.g., neural networks with configurable epoch counts).
//! - You want to combine model-guided search with aggressive pruning of
//! unpromising configurations.
//! - Your objective has a natural "budget" axis (epochs, iterations, data
//! fraction) reported via [`Trial::report`](crate::Trial::report).
//!
//! Pair `BohbSampler` with [`matching_pruner`](BohbSampler::matching_pruner)
//! to get a `HyperBandPruner` whose budget schedule is consistent with
//! the sampler's conditioning levels.
//!
//! # Configuration
//!
//! - `min_resource` / `max_resource` — budget range (default: 1 … 81)
//! - `reduction_factor` (η) — successive halving factor (default: 3)
//! - `min_points_in_model` — observations needed before TPE replaces random (default: 10)
//! - All [`TpeSamplerBuilder`](super::tpe::TpeSamplerBuilder) options (gamma, seed, etc.)
//!
//! # Examples
//!
//! ```
//! use optimizer::sampler::bohb::BohbSampler;
//! use optimizer::{Direction, Study};
//!
//! let bohb = BohbSampler::new();
//! let pruner = bohb.matching_pruner(Direction::Minimize);
//! let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
//! ```
//!
//! Custom configuration via builder:
//!
//! ```
//! use optimizer::sampler::bohb::BohbSampler;
//!
//! let bohb = BohbSampler::builder()
//! .min_resource(1)
//! .max_resource(81)
//! .reduction_factor(3)
//! .min_points_in_model(10)
//! .seed(42)
//! .build()
//! .unwrap();
//! ```
use crate::distribution::Distribution;
use crate::error::Result;
use crate::param::ParamValue;
use crate::pruner::HyperbandPruner;
use crate::sampler::tpe::TpeSampler;
use crate::sampler::{CompletedTrial, Sampler};
use crate::types::Direction;
/// A BOHB sampler that combines TPE with `HyperBand` budget awareness.
///
/// BOHB filters trial history by budget level before delegating to TPE,
/// so the surrogate model is conditioned on trials evaluated at the same
/// resource level. This produces better-calibrated parameter proposals
/// than using a single global model across all budgets.
///
/// Use [`BohbSampler::matching_pruner`] to create a [`HyperbandPruner`]
/// with matching `HyperBand` parameters.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::bohb::BohbSampler;
/// use optimizer::{Direction, Study};
///
/// let bohb = BohbSampler::builder()
/// .min_resource(1)
/// .max_resource(27)
/// .reduction_factor(3)
/// .seed(42)
/// .build()
/// .unwrap();
/// let pruner = bohb.matching_pruner(Direction::Minimize);
/// let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
/// ```
pub struct BohbSampler {
min_resource: u64,
max_resource: u64,
reduction_factor: u64,
min_points_in_model: usize,
tpe: TpeSampler,
}
impl BohbSampler {
/// Creates a new BOHB sampler with default settings.
///
/// Defaults:
/// - `min_resource`: 1
/// - `max_resource`: 81
/// - `reduction_factor`: 3
/// - `min_points_in_model`: 10
/// - TPE: default settings
#[must_use]
pub fn new() -> Self {
Self {
min_resource: 1,
max_resource: 81,
reduction_factor: 3,
min_points_in_model: 10,
tpe: TpeSampler::new(),
}
}
/// Creates a builder for configuring a BOHB sampler.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::bohb::BohbSampler;
///
/// let sampler = BohbSampler::builder()
/// .min_resource(1)
/// .max_resource(27)
/// .reduction_factor(3)
/// .min_points_in_model(5)
/// .seed(42)
/// .build()
/// .unwrap();
/// ```
#[must_use]
pub fn builder() -> BohbSamplerBuilder {
BohbSamplerBuilder::new()
}
/// Creates a [`HyperbandPruner`] with matching Hyperband parameters.
///
/// This ensures the pruner's budget schedule is consistent with the
/// budget levels used by BOHB for model conditioning.
#[must_use]
pub fn matching_pruner(&self, direction: Direction) -> HyperbandPruner {
HyperbandPruner::new()
.min_resource(self.min_resource)
.max_resource(self.max_resource)
.reduction_factor(self.reduction_factor)
.direction(direction)
}
/// Compute all unique budget levels (rung steps) across all Hyperband brackets.
///
/// Returns sorted ascending.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
fn all_budget_levels(&self) -> Vec<u64> {
let eta = self.reduction_factor as f64;
let ratio = self.max_resource as f64 / self.min_resource as f64;
let s_max = (ratio.ln() / eta.ln()).floor() as u64;
let mut levels = Vec::new();
for bracket in 0..=s_max {
let exponent = s_max.saturating_sub(bracket);
let min_resource_bracket =
(self.max_resource as f64 / eta.powi(exponent as i32)).ceil() as u64;
let mut rung: u32 = 0;
while let Some(power) = self.reduction_factor.checked_pow(rung) {
let step = min_resource_bracket.saturating_mul(power);
if step > self.max_resource {
break;
}
levels.push(step);
rung += 1;
}
}
levels.sort_unstable();
levels.dedup();
levels
}
/// Build a filtered history for a specific budget level.
///
/// For each trial that has an intermediate value at the given budget step,
/// creates a new `CompletedTrial` with `value` replaced by the intermediate
/// value at that step.
fn filter_history_for_budget(history: &[CompletedTrial], budget: u64) -> Vec<CompletedTrial> {
history
.iter()
.filter_map(|trial| {
trial
.intermediate_values
.iter()
.find(|(step, _)| *step == budget)
.map(|(_, iv)| CompletedTrial {
id: trial.id,
params: trial.params.clone(),
distributions: trial.distributions.clone(),
param_labels: trial.param_labels.clone(),
value: *iv,
intermediate_values: trial.intermediate_values.clone(),
state: trial.state,
user_attrs: trial.user_attrs.clone(),
constraints: trial.constraints.clone(),
})
})
.collect()
}
}
impl Default for BohbSampler {
fn default() -> Self {
Self::new()
}
}
impl Sampler for BohbSampler {
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue {
// Find the largest budget level with enough observations
let levels = self.all_budget_levels();
for &budget in levels.iter().rev() {
let count = history
.iter()
.filter(|t| {
t.intermediate_values
.iter()
.any(|(step, _)| *step == budget)
})
.count();
if count >= self.min_points_in_model {
let filtered = Self::filter_history_for_budget(history, budget);
return self.tpe.sample(distribution, trial_id, &filtered);
}
}
// Not enough data at any budget level: delegate to TPE with empty history
// which triggers its uniform-random startup behavior.
self.tpe.sample(distribution, trial_id, &[])
}
}
/// Builder for configuring a [`BohbSampler`].
///
/// # Defaults
///
/// - `min_resource`: 1
/// - `max_resource`: 81
/// - `reduction_factor`: 3 (η)
/// - `min_points_in_model`: 10
/// - TPE: default settings (gamma = 0.25, etc.)
///
/// # Examples
///
/// ```
/// use optimizer::sampler::bohb::BohbSamplerBuilder;
///
/// let sampler = BohbSamplerBuilder::new()
/// .min_resource(1)
/// .max_resource(81)
/// .reduction_factor(3)
/// .gamma(0.15)
/// .seed(42)
/// .build()
/// .unwrap();
/// ```
pub struct BohbSamplerBuilder {
min_resource: u64,
max_resource: u64,
reduction_factor: u64,
min_points_in_model: usize,
tpe_builder: crate::sampler::tpe::TpeSamplerBuilder,
}
impl BohbSamplerBuilder {
/// Creates a new builder with default settings.
#[must_use]
pub fn new() -> Self {
Self {
min_resource: 1,
max_resource: 81,
reduction_factor: 3,
min_points_in_model: 10,
tpe_builder: crate::sampler::tpe::TpeSamplerBuilder::new(),
}
}
/// Sets the minimum resource (budget) per trial.
///
/// # Panics
///
/// Panics if `r` is 0.
#[must_use]
pub fn min_resource(mut self, r: u64) -> Self {
assert!(r > 0, "min_resource must be > 0, got {r}");
self.min_resource = r;
self
}
/// Sets the maximum resource (budget) per trial.
///
/// # Panics
///
/// Panics if `r` is 0.
#[must_use]
pub fn max_resource(mut self, r: u64) -> Self {
assert!(r > 0, "max_resource must be > 0, got {r}");
self.max_resource = r;
self
}
/// Sets the reduction factor (eta).
///
/// # Panics
///
/// Panics if `eta` is less than 2.
#[must_use]
pub fn reduction_factor(mut self, eta: u64) -> Self {
assert!(eta >= 2, "reduction_factor must be >= 2, got {eta}");
self.reduction_factor = eta;
self
}
/// Sets the minimum number of observations at a budget level before
/// BOHB uses TPE instead of random sampling.
#[must_use]
pub fn min_points_in_model(mut self, n: usize) -> Self {
self.min_points_in_model = n;
self
}
/// Sets a fixed gamma value for the internal TPE sampler.
#[must_use]
pub fn gamma(mut self, gamma: f64) -> Self {
self.tpe_builder = self.tpe_builder.gamma(gamma);
self
}
/// Sets a custom gamma strategy for the internal TPE sampler.
#[must_use]
pub fn gamma_strategy<G: crate::sampler::tpe::GammaStrategy + 'static>(
mut self,
strategy: G,
) -> Self {
self.tpe_builder = self.tpe_builder.gamma_strategy(strategy);
self
}
/// Sets the number of EI candidates for the internal TPE sampler.
#[must_use]
pub fn n_ei_candidates(mut self, n: usize) -> Self {
self.tpe_builder = self.tpe_builder.n_ei_candidates(n);
self
}
/// Sets a fixed KDE bandwidth for the internal TPE sampler.
#[must_use]
pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self {
self.tpe_builder = self.tpe_builder.kde_bandwidth(bandwidth);
self
}
/// Sets a seed for reproducible sampling.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.tpe_builder = self.tpe_builder.seed(seed);
self
}
/// Builds the configured [`BohbSampler`].
///
/// # Errors
///
/// Returns an error if the TPE configuration is invalid (e.g. gamma
/// not in (0, 1) or bandwidth not positive).
pub fn build(self) -> Result<BohbSampler> {
let tpe = self.tpe_builder.build()?;
Ok(BohbSampler {
min_resource: self.min_resource,
max_resource: self.max_resource,
reduction_factor: self.reduction_factor,
min_points_in_model: self.min_points_in_model,
tpe,
})
}
}
impl Default for BohbSamplerBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
#[allow(clippy::cast_precision_loss)]
mod tests {
use std::collections::HashMap;
use super::*;
use crate::distribution::{FloatDistribution, IntDistribution};
use crate::parameter::ParamId;
use crate::types::TrialState;
fn make_trial_with_intermediates(
id: u64,
value: f64,
params: Vec<(ParamId, ParamValue, Distribution)>,
intermediate_values: Vec<(u64, f64)>,
) -> CompletedTrial {
let mut param_map = HashMap::new();
let mut dist_map = HashMap::new();
for (param_id, pv, dist) in params {
param_map.insert(param_id, pv);
dist_map.insert(param_id, dist);
}
CompletedTrial {
id,
params: param_map,
distributions: dist_map,
param_labels: HashMap::new(),
value,
intermediate_values,
state: TrialState::Complete,
user_attrs: HashMap::new(),
constraints: Vec::new(),
}
}
#[test]
fn budget_levels_default() {
let bohb = BohbSampler::new();
let levels = bohb.all_budget_levels();
// With min=1, max=81, eta=3:
// bracket 0: 1, 3, 9, 27, 81
// bracket 1: 3, 9, 27, 81
// bracket 2: 9, 27, 81
// bracket 3: 27, 81
// bracket 4: 81
// Unique sorted: [1, 3, 9, 27, 81]
assert_eq!(levels, vec![1, 3, 9, 27, 81]);
}
#[test]
fn budget_levels_eta2() {
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(16)
.reduction_factor(2)
.build()
.unwrap();
let levels = bohb.all_budget_levels();
// s_max = floor(ln(16)/ln(2)) = 4
// bracket 0: 1, 2, 4, 8, 16
// bracket 1: 2, 4, 8, 16
// bracket 2: 4, 8, 16
// bracket 3: 8, 16
// bracket 4: 16
// Unique sorted: [1, 2, 4, 8, 16]
assert_eq!(levels, vec![1, 2, 4, 8, 16]);
}
#[test]
fn filter_history_selects_correct_budget() {
let x_id = ParamId::new();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let history = vec![
make_trial_with_intermediates(
0,
0.5,
vec![(x_id, ParamValue::Float(0.3), dist.clone())],
vec![(1, 0.9), (3, 0.7), (9, 0.5)],
),
make_trial_with_intermediates(
1,
0.4,
vec![(x_id, ParamValue::Float(0.6), dist.clone())],
vec![(1, 0.8), (3, 0.4)],
),
make_trial_with_intermediates(
2,
0.3,
vec![(x_id, ParamValue::Float(0.1), dist.clone())],
vec![(1, 0.7)],
),
];
// Budget 3: trials 0 and 1 have intermediate values at step 3
let filtered = BohbSampler::filter_history_for_budget(&history, 3);
assert_eq!(filtered.len(), 2);
assert!((filtered[0].value - 0.7).abs() < f64::EPSILON);
assert!((filtered[1].value - 0.4).abs() < f64::EPSILON);
// Budget 9: only trial 0
let filtered = BohbSampler::filter_history_for_budget(&history, 9);
assert_eq!(filtered.len(), 1);
assert!((filtered[0].value - 0.5).abs() < f64::EPSILON);
// Budget 27: nobody
let filtered = BohbSampler::filter_history_for_budget(&history, 27);
assert!(filtered.is_empty());
}
#[test]
fn matching_pruner_has_same_params() {
let bohb = BohbSampler::builder()
.min_resource(2)
.max_resource(64)
.reduction_factor(4)
.build()
.unwrap();
let pruner = bohb.matching_pruner(Direction::Minimize);
// We can't directly inspect HyperbandPruner fields, but we can
// verify it was created without panicking with the same params.
// The pruner's rung steps should match BOHB's budget levels.
// Just verify it doesn't panic.
drop(pruner);
}
#[test]
fn fallback_to_random_when_insufficient_data() {
let bohb = BohbSampler::builder()
.min_points_in_model(10)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
// Only 3 trials with intermediate values (< min_points_in_model=10)
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..3)
.map(|i| {
make_trial_with_intermediates(
i,
i as f64,
vec![(x_id, ParamValue::Float(i as f64 / 3.0), dist.clone())],
vec![(1, i as f64)],
)
})
.collect();
// Should not panic, should sample within bounds
for trial_id in 0..20 {
let val = bohb.sample(&dist, trial_id, &history);
if let ParamValue::Float(v) = val {
assert!((0.0..=1.0).contains(&v));
} else {
panic!("Expected Float");
}
}
}
#[test]
fn uses_budget_level_when_enough_data() {
let bohb = BohbSampler::builder()
.min_points_in_model(5)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 10.0,
log_scale: false,
step: None,
});
// Create 20 trials with intermediate values at budget 1.
// Good trials have x near 2.0, bad trials have x far from 2.0.
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..20)
.map(|i| {
let x = i as f64 / 2.0;
let iv_at_1 = (x - 2.0).powi(2);
make_trial_with_intermediates(
i,
iv_at_1, // final value same as intermediate for simplicity
vec![(x_id, ParamValue::Float(x), dist.clone())],
vec![(1, iv_at_1)],
)
})
.collect();
// Should use TPE on filtered history at budget 1
let val = bohb.sample(&dist, 100, &history);
if let ParamValue::Float(v) = val {
assert!((0.0..=10.0).contains(&v), "Value {v} out of bounds");
} else {
panic!("Expected Float");
}
}
#[test]
fn prefers_largest_budget_level() {
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(9)
.reduction_factor(3)
.min_points_in_model(3)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 10.0,
log_scale: false,
step: None,
});
// Budget levels: [1, 3, 9]
assert_eq!(bohb.all_budget_levels(), vec![1, 3, 9]);
// Create 5 trials with intermediates at budget 1 and 3
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..5)
.map(|i| {
let x = i as f64;
make_trial_with_intermediates(
i,
x,
vec![(x_id, ParamValue::Float(x), dist.clone())],
vec![(1, x * 2.0), (3, x)],
)
})
.collect();
// Budget 3 has 5 observations (>= 3), budget 9 has 0.
// BOHB should pick budget 3 (largest with enough data).
// The filtered history at budget 3 has values [0, 1, 2, 3, 4].
let filtered_3 = BohbSampler::filter_history_for_budget(&history, 3);
assert_eq!(filtered_3.len(), 5);
let filtered_9 = BohbSampler::filter_history_for_budget(&history, 9);
assert_eq!(filtered_9.len(), 0);
// Should sample successfully
let val = bohb.sample(&dist, 100, &history);
assert!(matches!(val, ParamValue::Float(_)));
}
#[test]
fn builder_validates_tpe_params() {
// Invalid gamma
let result = BohbSampler::builder().gamma(1.5).build();
assert!(result.is_err());
// Invalid bandwidth
let result = BohbSampler::builder().kde_bandwidth(-1.0).build();
assert!(result.is_err());
}
#[test]
#[should_panic(expected = "min_resource must be > 0")]
fn builder_rejects_zero_min_resource() {
let _ = BohbSampler::builder().min_resource(0);
}
#[test]
#[should_panic(expected = "max_resource must be > 0")]
fn builder_rejects_zero_max_resource() {
let _ = BohbSampler::builder().max_resource(0);
}
#[test]
#[should_panic(expected = "reduction_factor must be >= 2")]
fn builder_rejects_small_reduction_factor() {
let _ = BohbSampler::builder().reduction_factor(1);
}
#[test]
fn int_distribution_works() {
let bohb = BohbSampler::builder()
.min_points_in_model(3)
.seed(42)
.build()
.unwrap();
let dist = Distribution::Int(IntDistribution {
low: 0,
high: 100,
log_scale: false,
step: None,
});
let x_id = ParamId::new();
let history: Vec<CompletedTrial> = (0..10)
.map(|i| {
make_trial_with_intermediates(
i,
i as f64,
vec![(x_id, ParamValue::Int(i.cast_signed() * 10), dist.clone())],
vec![(1, i as f64)],
)
})
.collect();
let val = bohb.sample(&dist, 100, &history);
if let ParamValue::Int(v) = val {
assert!((0..=100).contains(&v));
} else {
panic!("Expected Int");
}
}
}
+954
View File
@@ -0,0 +1,954 @@
//! CMA-ES (Covariance Matrix Adaptation Evolution Strategy) sampler.
//!
//! CMA-ES is a stochastic, population-based optimizer that maintains a
//! multivariate Gaussian distribution over continuous parameters and adapts
//! its **mean**, **covariance matrix**, and **step-size** (σ) based on
//! trial rankings. It is widely regarded as one of the most effective
//! derivative-free optimizers for continuous search spaces.
//!
//! # Algorithm overview
//!
//! Each generation:
//! 1. **Sample** λ (population size) candidates from N(m, σ²C).
//! 2. **Evaluate** and **rank** the candidates by objective value.
//! 3. **Update** the mean toward the best μ candidates (weighted recombination).
//! 4. **Adapt** the covariance matrix C via rank-one and rank-μ updates, and
//! adapt σ via cumulative step-size adaptation (CSA).
//!
//! Over time the search distribution narrows and rotates to align with the
//! landscape, efficiently exploiting structure in the objective function.
//!
//! # When to use
//!
//! - **Continuous parameters only** (float/int). Categorical parameters are
//! sampled uniformly at random and do not participate in the CMA-ES model.
//! - **Moderate dimensionality** — works well up to ~100 continuous dimensions.
//! Beyond that, the O(n²) covariance matrix becomes expensive to maintain.
//! - **Non-separable objectives** — CMA-ES learns parameter correlations
//! through the covariance matrix, making it especially effective on
//! rotated or ill-conditioned landscapes.
//! - **Moderate evaluation budgets** — typically needs ≈10×n to 100×n
//! evaluations to converge, where n is the number of continuous dimensions.
//!
//! For very cheap evaluations in low dimensions (d ≤ 20), consider
//! [`GpSampler`](super::gp::GpSampler) instead. For high-dimensional
//! separable problems, TPE may be more efficient.
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `sigma0` | `avg_range / 4` | Initial step size — controls exploration breadth |
//! | `population_size` | `4 + ⌊3 ln n⌋` | Candidates per generation (λ) |
//! | `seed` | random | RNG seed for reproducibility |
//!
//! # Feature flag
//!
//! Requires the **`cma-es`** feature (adds the `nalgebra` dependency):
//!
//! ```toml
//! [dependencies]
//! optimizer = { version = "...", features = ["cma-es"] }
//! ```
//!
//! # Examples
//!
//! ```
//! use optimizer::sampler::cma_es::CmaEsSampler;
//! use optimizer::{Direction, Study};
//!
//! // Minimize a 2-D sphere function with CMA-ES
//! let sampler = CmaEsSampler::builder()
//! .sigma0(1.0)
//! .population_size(10)
//! .seed(42)
//! .build();
//!
//! let mut study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! ```
use std::collections::HashMap;
use nalgebra::{DMatrix, DVector};
use parking_lot::Mutex;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::rng_util;
use crate::sampler::{CompletedTrial, Sampler};
use super::common::{from_internal, internal_bounds, sample_random};
/// CMA-ES sampler for continuous optimization.
///
/// Adapts a multivariate Gaussian to concentrate around promising regions
/// of the search space. Best suited for continuous (float/int) parameters
/// in moderate dimensions (up to ~100).
///
/// # Examples
///
/// ```
/// use optimizer::sampler::cma_es::CmaEsSampler;
/// use optimizer::{Direction, Study};
///
/// // Default configuration
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, CmaEsSampler::new());
///
/// // With seed for reproducibility
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, CmaEsSampler::with_seed(42));
///
/// // Custom configuration via builder
/// let sampler = CmaEsSampler::builder()
/// .sigma0(0.5)
/// .population_size(20)
/// .seed(42)
/// .build();
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
/// ```
pub struct CmaEsSampler {
state: Mutex<CmaEsState>,
}
impl CmaEsSampler {
/// Creates a new CMA-ES sampler with a random seed.
#[must_use]
pub fn new() -> Self {
Self {
state: Mutex::new(CmaEsState::new(None, None, None)),
}
}
/// Creates a new CMA-ES sampler with a fixed seed for reproducibility.
#[must_use]
pub fn with_seed(seed: u64) -> Self {
Self {
state: Mutex::new(CmaEsState::new(None, None, Some(seed))),
}
}
/// Creates a builder for configuring a `CmaEsSampler`.
#[must_use]
pub fn builder() -> CmaEsSamplerBuilder {
CmaEsSamplerBuilder::new()
}
}
impl Default for CmaEsSampler {
fn default() -> Self {
Self::new()
}
}
/// Builder for configuring a [`CmaEsSampler`].
///
/// All options have sensible defaults:
/// - `sigma0`: auto-computed as average range / 4
/// - `population_size`: `4 + floor(3 * ln(n))`
/// - `seed`: random
///
/// # Examples
///
/// ```
/// use optimizer::sampler::cma_es::CmaEsSamplerBuilder;
///
/// let sampler = CmaEsSamplerBuilder::new()
/// .sigma0(0.3)
/// .population_size(10)
/// .seed(42)
/// .build();
/// ```
#[derive(Debug, Clone, Default)]
pub struct CmaEsSamplerBuilder {
sigma0: Option<f64>,
population_size: Option<usize>,
seed: Option<u64>,
}
impl CmaEsSamplerBuilder {
/// Creates a new builder with default settings.
#[must_use]
pub fn new() -> Self {
Self::default()
}
/// Sets the initial step size (sigma).
///
/// Controls the initial spread of the search distribution. Larger values
/// explore more broadly; smaller values search more locally.
///
/// Default: `average_range / 4` (auto-computed from parameter bounds).
#[must_use]
pub fn sigma0(mut self, sigma0: f64) -> Self {
self.sigma0 = Some(sigma0);
self
}
/// Sets the population size (lambda).
///
/// Number of candidate solutions evaluated per generation.
/// Larger populations improve robustness but increase the number of
/// trials per generation.
///
/// Default: `4 + floor(3 * ln(n))` where `n` is the number of continuous dimensions.
#[must_use]
pub fn population_size(mut self, population_size: usize) -> Self {
self.population_size = Some(population_size);
self
}
/// Sets the random seed for reproducibility.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.seed = Some(seed);
self
}
/// Builds the configured [`CmaEsSampler`].
#[must_use]
pub fn build(self) -> CmaEsSampler {
CmaEsSampler {
state: Mutex::new(CmaEsState::new(
self.sigma0,
self.population_size,
self.seed,
)),
}
}
}
// ---------------------------------------------------------------------------
// Internal types
// ---------------------------------------------------------------------------
/// Describes how a parameter dimension maps into the CMA-ES internal vector.
#[derive(Clone, Debug)]
struct DimensionInfo {
/// The distribution for this dimension (stored for decoding).
distribution: Distribution,
/// Whether this dimension participates in CMA-ES (Float/Int = true, Categorical = false).
is_continuous: bool,
/// Internal-space bounds for continuous dimensions: `(low, high)`.
/// For log-scale parameters these are in log-space.
bounds: Option<(f64, f64)>,
}
/// A candidate solution produced by the CMA-ES distribution.
#[derive(Clone, Debug)]
struct Candidate {
/// Internal-space vector (only continuous dimensions).
x: DVector<f64>,
/// Values for categorical dimensions (index in `dimensions` → categorical index).
categorical_values: HashMap<usize, usize>,
}
/// Tracks per-trial sampling progress.
#[derive(Clone, Debug)]
struct TrialProgress {
/// Index of the candidate assigned to this trial.
candidate_idx: usize,
/// Next dimension to return for this trial.
next_dim: usize,
}
/// The CMA-ES algorithm constants, derived from dimension count.
#[derive(Clone, Debug)]
struct CmaEsConstants {
/// Dimension of the continuous search space.
n: usize,
/// Population size (lambda).
lambda: usize,
/// Parent count (mu = lambda / 2).
mu: usize,
/// Recombination weights (length mu).
weights: Vec<f64>,
/// Variance effective selection mass.
mu_eff: f64,
/// Learning rate for the cumulation of the step-size control.
c_sigma: f64,
/// Damping for sigma.
d_sigma: f64,
/// Learning rate for the cumulation of the rank-one update.
c_c: f64,
/// Learning rate for the rank-one update of C.
c_1: f64,
/// Learning rate for the rank-mu update of C.
c_mu: f64,
/// Expected norm of N(0, I) in n dimensions.
chi_n: f64,
}
/// The mutable CMA-ES algorithm state.
struct CmaEsAlgorithm {
/// Distribution mean.
mean: DVector<f64>,
/// Step size.
sigma: f64,
/// Covariance matrix.
c: DMatrix<f64>,
/// Evolution path for sigma.
p_sigma: DVector<f64>,
/// Evolution path for rank-one update.
p_c: DVector<f64>,
/// Eigenvectors of C (columns of B).
b: DMatrix<f64>,
/// Sqrt of eigenvalues of C (diagonal of D).
d: DVector<f64>,
/// C^{-1/2} for sigma path update.
inv_sqrt_c: DMatrix<f64>,
/// Generation counter (for eigendecomposition scheduling).
generation: usize,
/// Last generation at which eigendecomposition was performed.
last_eigen_generation: usize,
/// Algorithm constants.
constants: CmaEsConstants,
}
/// Phase of the CMA-ES state machine.
enum Phase {
/// Discovering the search space structure (first trial).
Discovery,
/// Steady-state sampling and updating.
Active(Box<CmaEsAlgorithm>),
}
/// Top-level mutable state behind the `Mutex`.
struct CmaEsState {
/// The RNG used for sampling.
rng: fastrand::Rng,
/// User-provided initial sigma (None = auto).
sigma0: Option<f64>,
/// User-provided population size (None = auto).
user_lambda: Option<usize>,
/// Current phase.
phase: Phase,
/// Discovered dimension info (populated during discovery).
dimensions: Vec<DimensionInfo>,
/// Current generation's candidates.
candidates: Vec<Candidate>,
/// Mapping from `trial_id` to its progress.
trial_progress: HashMap<u64, TrialProgress>,
/// Number of candidates assigned so far in the current generation.
assigned_count: usize,
/// Trial IDs assigned in the current generation (for tracking completion).
generation_trial_ids: Vec<u64>,
/// Last `trial_id` seen during discovery (to detect when first trial ends).
discovery_trial_id: Option<u64>,
}
impl CmaEsState {
fn new(sigma0: Option<f64>, user_lambda: Option<usize>, seed: Option<u64>) -> Self {
let rng = seed.map_or_else(fastrand::Rng::new, fastrand::Rng::with_seed);
Self {
rng,
sigma0,
user_lambda,
phase: Phase::Discovery,
dimensions: Vec::new(),
candidates: Vec::new(),
trial_progress: HashMap::new(),
assigned_count: 0,
generation_trial_ids: Vec::new(),
discovery_trial_id: None,
}
}
}
// ---------------------------------------------------------------------------
// CMA-ES algorithm helpers
// ---------------------------------------------------------------------------
impl CmaEsConstants {
/// Compute all CMA-ES constants from dimension count and population size.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
fn new(n: usize, user_lambda: Option<usize>) -> Self {
let n_f = n as f64;
// Population size
let lambda = user_lambda.unwrap_or_else(|| 4 + (3.0 * n_f.ln()).max(0.0).floor() as usize);
let lambda = lambda.max(4); // ensure at least 4
// Parent count
let mu = lambda / 2;
// Recombination weights (log-proportional)
let log_half_lambda = f64::midpoint(lambda as f64, 1.0).ln();
let raw_weights: Vec<f64> = (0..mu)
.map(|i| log_half_lambda - ((i + 1) as f64).ln())
.collect();
let w_sum: f64 = raw_weights.iter().sum();
let weights: Vec<f64> = raw_weights.iter().map(|w| w / w_sum).collect();
// Variance effective selection mass
let w_sq_sum: f64 = weights.iter().map(|w| w * w).sum();
let mu_eff = 1.0 / w_sq_sum;
// Learning rates
let c_sigma = (mu_eff + 2.0) / (n_f + mu_eff + 5.0);
let d_sigma = 1.0 + 2.0 * (((mu_eff - 1.0) / (n_f + 1.0)).sqrt() - 1.0).max(0.0) + c_sigma;
let c_c = (4.0 + mu_eff / n_f) / (n_f + 4.0 + 2.0 * mu_eff / n_f);
let c_1 = 2.0 / ((n_f + 1.3).powi(2) + mu_eff);
let c_mu_raw = (2.0 * (mu_eff - 2.0 + 1.0 / mu_eff)) / ((n_f + 2.0).powi(2) + mu_eff);
let c_mu = c_mu_raw.min(1.0 - c_1);
// Expected norm of N(0, I)
let chi_n = n_f.sqrt() * (1.0 - 1.0 / (4.0 * n_f) + 1.0 / (21.0 * n_f * n_f));
Self {
n,
lambda,
mu,
weights,
mu_eff,
c_sigma,
d_sigma,
c_c,
c_1,
c_mu,
chi_n,
}
}
}
impl CmaEsAlgorithm {
/// Initialize the CMA-ES algorithm state.
#[allow(clippy::cast_precision_loss)]
fn new(dimensions: &[DimensionInfo], sigma0: Option<f64>, user_lambda: Option<usize>) -> Self {
let n = dimensions.iter().filter(|d| d.is_continuous).count();
let constants = CmaEsConstants::new(n, user_lambda);
// Compute initial mean (center of bounds) and auto sigma
let mut mean = DVector::zeros(n);
let mut total_range = 0.0;
let mut ci = 0;
for dim in dimensions {
if dim.is_continuous {
if let Some((lo, hi)) = dim.bounds {
mean[ci] = f64::midpoint(lo, hi);
total_range += hi - lo;
}
ci += 1;
}
}
let sigma = sigma0.unwrap_or_else(|| {
if n > 0 {
(total_range / n as f64) / 4.0
} else {
1.0
}
});
let c = DMatrix::identity(n, n);
let p_sigma = DVector::zeros(n);
let p_c = DVector::zeros(n);
let b = DMatrix::identity(n, n);
let d = DVector::from_element(n, 1.0);
let inv_sqrt_c = DMatrix::identity(n, n);
Self {
mean,
sigma,
c,
p_sigma,
p_c,
b,
d,
inv_sqrt_c,
generation: 0,
last_eigen_generation: 0,
constants,
}
}
/// Generate `lambda` candidate vectors from the current distribution.
fn generate_candidates(
&self,
rng: &mut fastrand::Rng,
dimensions: &[DimensionInfo],
) -> Vec<Candidate> {
let n = self.constants.n;
let lambda = self.constants.lambda;
let mut candidates = Vec::with_capacity(lambda);
for _ in 0..lambda {
let candidate = self.generate_single_candidate(rng, dimensions, n);
candidates.push(candidate);
}
candidates
}
/// Generate a single candidate from the current distribution.
fn generate_single_candidate(
&self,
rng: &mut fastrand::Rng,
dimensions: &[DimensionInfo],
n: usize,
) -> Candidate {
// x = mean + sigma * B * D * z where z ~ N(0, I)
let x = self.sample_with_rejection(rng, dimensions, n);
// Sample categorical dimensions randomly
let mut categorical_values = HashMap::new();
for (i, dim) in dimensions.iter().enumerate() {
if !dim.is_continuous
&& let Distribution::Categorical(cat) = &dim.distribution
{
categorical_values.insert(i, rng.usize(0..cat.n_choices));
}
}
Candidate {
x,
categorical_values,
}
}
/// Sample a candidate vector with rejection sampling for bounds.
fn sample_with_rejection(
&self,
rng: &mut fastrand::Rng,
dimensions: &[DimensionInfo],
n: usize,
) -> DVector<f64> {
let max_attempts = 100;
for _ in 0..max_attempts {
let z = DVector::from_fn(n, |_, _| sample_standard_normal(rng));
let x = &self.mean + self.sigma * (&self.b * self.d.component_mul(&z));
if is_within_bounds(&x, dimensions) {
return x;
}
}
// Fallback: clip to bounds
let z = DVector::from_fn(n, |_, _| sample_standard_normal(rng));
let mut x = &self.mean + self.sigma * (&self.b * self.d.component_mul(&z));
clip_to_bounds(&mut x, dimensions);
x
}
/// Run the CMA-ES update step given ranked candidates.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_possible_wrap
)]
fn update(&mut self, ranked_candidates: &[&DVector<f64>]) {
let n = self.constants.n;
let mu = self.constants.mu;
let sigma = self.sigma;
// New mean = weighted sum of top-mu candidates
let mut new_mean = DVector::zeros(n);
for (i, &x) in ranked_candidates.iter().take(mu).enumerate() {
new_mean += self.constants.weights[i] * x;
}
// Displacement in mean
let mean_diff = &new_mean - &self.mean;
// Update p_sigma (cumulation for sigma control)
let inv_sqrt_c_times_diff = &self.inv_sqrt_c * &mean_diff / sigma;
self.p_sigma = (1.0 - self.constants.c_sigma) * &self.p_sigma
+ (self.constants.c_sigma * (2.0 - self.constants.c_sigma) * self.constants.mu_eff)
.sqrt()
* &inv_sqrt_c_times_diff;
// h_sigma: stall indicator
let p_sigma_norm = self.p_sigma.norm();
let threshold =
(1.0 - (1.0 - self.constants.c_sigma).powi(2 * (self.generation as i32 + 1))).sqrt()
* (1.4 + 2.0 / (n as f64 + 1.0))
* self.constants.chi_n;
let h_sigma = if p_sigma_norm < threshold { 1.0 } else { 0.0 };
// Update p_c (cumulation for rank-one update)
self.p_c = (1.0 - self.constants.c_c) * &self.p_c
+ h_sigma
* (self.constants.c_c * (2.0 - self.constants.c_c) * self.constants.mu_eff).sqrt()
* &mean_diff
/ sigma;
// Rank-one and rank-mu update of C
let delta_h = (1.0 - h_sigma) * self.constants.c_c * (2.0 - self.constants.c_c);
let old_c_weight =
1.0 - self.constants.c_1 - self.constants.c_mu + self.constants.c_1 * delta_h;
// Rank-one term
let rank_one = self.constants.c_1 * &self.p_c * self.p_c.transpose();
// Rank-mu term
let mut rank_mu = DMatrix::zeros(n, n);
for (i, &x) in ranked_candidates.iter().take(mu).enumerate() {
let y = (x - &self.mean) / sigma;
rank_mu += self.constants.weights[i] * &y * y.transpose();
}
let rank_mu = self.constants.c_mu * rank_mu;
self.c = old_c_weight * &self.c + rank_one + rank_mu;
// Update sigma via CSA
self.sigma *= ((self.constants.c_sigma / self.constants.d_sigma)
* (p_sigma_norm / self.constants.chi_n - 1.0))
.exp();
self.sigma = self.sigma.clamp(1e-20, 1e10);
// Update mean
self.mean = new_mean;
self.generation += 1;
// Eigendecomposition (every n/10 generations, minimum every generation for small n)
let eigen_interval = (n / 10).max(1);
if self.generation - self.last_eigen_generation >= eigen_interval {
self.update_eigen();
}
}
/// Perform eigendecomposition of C and update B, D, `inv_sqrt_c`.
fn update_eigen(&mut self) {
let n = self.constants.n;
// Enforce symmetry
self.c = (&self.c + self.c.transpose()) / 2.0;
// Eigendecomposition
let eigen = self.c.clone().symmetric_eigen();
let eigenvalues = &eigen.eigenvalues;
let eigenvectors = &eigen.eigenvectors;
// Clamp eigenvalues for numerical stability
let mut d_vec = DVector::zeros(n);
for i in 0..n {
d_vec[i] = eigenvalues[i].max(1e-20).sqrt();
}
self.b = eigenvectors.clone();
self.d = d_vec;
// Compute C^{-1/2} = B * D^{-1} * B^T
let d_inv = DVector::from_fn(n, |i, _| 1.0 / self.d[i]);
let d_inv_diag = DMatrix::from_diagonal(&d_inv);
self.inv_sqrt_c = &self.b * d_inv_diag * self.b.transpose();
self.last_eigen_generation = self.generation;
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/// Check whether a continuous candidate vector is within bounds.
fn is_within_bounds(x: &DVector<f64>, dimensions: &[DimensionInfo]) -> bool {
let mut ci = 0;
for dim in dimensions {
if dim.is_continuous {
if let Some((lo, hi)) = dim.bounds
&& (x[ci] < lo || x[ci] > hi)
{
return false;
}
ci += 1;
}
}
true
}
/// Clip a continuous candidate vector to bounds.
fn clip_to_bounds(x: &mut DVector<f64>, dimensions: &[DimensionInfo]) {
let mut ci = 0;
for dim in dimensions {
if dim.is_continuous {
if let Some((lo, hi)) = dim.bounds {
x[ci] = x[ci].clamp(lo, hi);
}
ci += 1;
}
}
}
/// Sample a value from the standard normal distribution using Box-Muller transform.
fn sample_standard_normal(rng: &mut fastrand::Rng) -> f64 {
// Box-Muller transform
let u1: f64 = rng_util::f64_range(rng, f64::EPSILON, 1.0);
let u2: f64 = rng_util::f64_range(rng, 0.0_f64, core::f64::consts::TAU);
(-2.0 * u1.ln()).sqrt() * u2.cos()
}
// ---------------------------------------------------------------------------
// Sampler trait implementation
// ---------------------------------------------------------------------------
impl Sampler for CmaEsSampler {
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue {
let mut state = self.state.lock();
match &state.phase {
Phase::Discovery => sample_discovery(&mut state, distribution, trial_id),
Phase::Active(_) => sample_active(&mut state, distribution, trial_id, history),
}
}
}
/// Handle sampling during the discovery phase.
fn sample_discovery(
state: &mut CmaEsState,
distribution: &Distribution,
trial_id: u64,
) -> ParamValue {
// Check if this is a new trial (discovery phase ended for previous trial)
if let Some(prev_id) = state.discovery_trial_id
&& trial_id != prev_id
{
// First trial is done; we know the search space. Initialize CMA-ES.
finalize_discovery(state);
// Now delegate to active sampling
return sample_active(state, distribution, trial_id, &[]);
}
// Record this trial_id
state.discovery_trial_id = Some(trial_id);
// Record this dimension
let is_continuous = !matches!(distribution, Distribution::Categorical(_));
let bounds = internal_bounds(distribution);
state.dimensions.push(DimensionInfo {
distribution: distribution.clone(),
is_continuous,
bounds,
});
// Sample randomly for the discovery trial
sample_random(&mut state.rng, distribution)
}
/// Finalize discovery and transition to the active phase.
fn finalize_discovery(state: &mut CmaEsState) {
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
let algo = CmaEsAlgorithm::new(&state.dimensions, state.sigma0, state.user_lambda);
// Generate first batch of candidates
let candidates = if n_continuous > 0 {
algo.generate_candidates(&mut state.rng, &state.dimensions)
} else {
// Pure categorical: generate random categorical candidates
generate_pure_categorical_candidates(
&mut state.rng,
&state.dimensions,
algo.constants.lambda,
)
};
state.candidates = candidates;
state.assigned_count = 0;
state.generation_trial_ids.clear();
state.trial_progress.clear();
state.phase = Phase::Active(Box::new(algo));
}
/// Generate candidates that are purely categorical (no continuous dims).
fn generate_pure_categorical_candidates(
rng: &mut fastrand::Rng,
dimensions: &[DimensionInfo],
lambda: usize,
) -> Vec<Candidate> {
(0..lambda)
.map(|_| {
let mut categorical_values = HashMap::new();
for (i, dim) in dimensions.iter().enumerate() {
if let Distribution::Categorical(cat) = &dim.distribution {
categorical_values.insert(i, rng.usize(0..cat.n_choices));
}
}
Candidate {
x: DVector::zeros(0),
categorical_values,
}
})
.collect()
}
/// Handle sampling during the active phase.
fn sample_active(
state: &mut CmaEsState,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue {
// Check if we need to update (all generation candidates assigned and completed)
maybe_update_generation(state, history);
// Assign a candidate to this trial if not yet done
if !state.trial_progress.contains_key(&trial_id) {
assign_candidate(state, trial_id);
}
let progress = state.trial_progress.get_mut(&trial_id).unwrap();
let dim_idx = progress.next_dim;
progress.next_dim += 1;
// Safety check: dim_idx should be within dimensions
if dim_idx >= state.dimensions.len() {
// Extra dimension not seen during discovery; sample randomly
return sample_random(&mut state.rng, distribution);
}
let candidate = &state.candidates[progress.candidate_idx];
let dim_info = &state.dimensions[dim_idx];
if dim_info.is_continuous {
// Map from internal continuous index to the candidate's x vector
let ci = state.dimensions[..dim_idx]
.iter()
.filter(|d| d.is_continuous)
.count();
if ci < candidate.x.len() {
from_internal(candidate.x[ci], &dim_info.distribution)
} else {
sample_random(&mut state.rng, distribution)
}
} else {
// Categorical: use pre-sampled value
if let Some(&cat_idx) = candidate.categorical_values.get(&dim_idx) {
ParamValue::Categorical(cat_idx)
} else {
sample_random(&mut state.rng, distribution)
}
}
}
/// Assign a candidate to a trial.
fn assign_candidate(state: &mut CmaEsState, trial_id: u64) {
let candidate_idx = if state.assigned_count < state.candidates.len() {
let idx = state.assigned_count;
state.assigned_count += 1;
idx
} else {
// Overflow: generate an extra candidate from current distribution
let extra = generate_overflow_candidate(state);
state.candidates.push(extra);
let idx = state.candidates.len() - 1;
state.assigned_count = state.candidates.len();
idx
};
state.trial_progress.insert(
trial_id,
TrialProgress {
candidate_idx,
next_dim: 0,
},
);
state.generation_trial_ids.push(trial_id);
}
/// Generate an overflow candidate from the current distribution.
fn generate_overflow_candidate(state: &mut CmaEsState) -> Candidate {
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
if n_continuous == 0 {
let mut categorical_values = HashMap::new();
for (i, dim) in state.dimensions.iter().enumerate() {
if let Distribution::Categorical(cat) = &dim.distribution {
categorical_values.insert(i, state.rng.usize(0..cat.n_choices));
}
}
return Candidate {
x: DVector::zeros(0),
categorical_values,
};
}
match &state.phase {
Phase::Active(algo) => {
algo.generate_single_candidate(&mut state.rng, &state.dimensions, n_continuous)
}
Phase::Discovery => {
// Should not happen, but handle gracefully
Candidate {
x: DVector::zeros(n_continuous),
categorical_values: HashMap::new(),
}
}
}
}
/// Check if we should run the CMA-ES update and generate a new generation.
fn maybe_update_generation(state: &mut CmaEsState, history: &[CompletedTrial]) {
let Phase::Active(algo) = &state.phase else {
return;
};
let lambda = algo.constants.lambda;
let n_continuous = algo.constants.n;
// Only update when at least lambda candidates have been assigned
if state.generation_trial_ids.len() < lambda {
return;
}
// Check if the first lambda trial IDs are all completed
let trial_ids: Vec<u64> = state
.generation_trial_ids
.iter()
.take(lambda)
.copied()
.collect();
let history_ids: HashMap<u64, f64> = history.iter().map(|t| (t.id, t.value)).collect();
let all_completed = trial_ids.iter().all(|id| history_ids.contains_key(id));
if !all_completed {
return;
}
// No continuous dimensions: just regenerate random candidates
if n_continuous == 0 {
state.candidates =
generate_pure_categorical_candidates(&mut state.rng, &state.dimensions, lambda);
state.assigned_count = 0;
state.generation_trial_ids.clear();
state.trial_progress.clear();
return;
}
// Collect (candidate_x, value) for the generation
let mut ranked: Vec<(&DVector<f64>, f64)> = trial_ids
.iter()
.filter_map(|id| {
let progress = state.trial_progress.get(id)?;
let value = *history_ids.get(id)?;
Some((&state.candidates[progress.candidate_idx].x, value))
})
.collect();
// Sort ascending (minimize)
ranked.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
let ranked_xs: Vec<&DVector<f64>> = ranked.iter().map(|(x, _)| *x).collect();
// Run update
let Phase::Active(algo) = &mut state.phase else {
return;
};
algo.update(&ranked_xs);
// Generate new candidates
let new_candidates = algo.generate_candidates(&mut state.rng, &state.dimensions);
state.candidates = new_candidates;
state.assigned_count = 0;
state.generation_trial_ids.clear();
state.trial_progress.clear();
}
+134
View File
@@ -0,0 +1,134 @@
//! Shared distribution-level utilities used across multiple samplers.
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::rng_util;
/// Compute internal-space bounds for a distribution.
#[allow(clippy::cast_precision_loss)]
pub(crate) fn internal_bounds(distribution: &Distribution) -> Option<(f64, f64)> {
match distribution {
Distribution::Float(d) => {
if d.log_scale {
if d.low <= 0.0 || d.high <= 0.0 {
return None;
}
Some((d.low.ln(), d.high.ln()))
} else {
Some((d.low, d.high))
}
}
Distribution::Int(d) => {
if d.log_scale {
if d.low < 1 {
return None;
}
Some(((d.low as f64).ln(), (d.high as f64).ln()))
} else {
Some((d.low as f64, d.high as f64))
}
}
Distribution::Categorical(_) => None,
}
}
/// Convert an internal-space value back to a `ParamValue`.
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
pub(crate) fn from_internal(value: f64, distribution: &Distribution) -> ParamValue {
match distribution {
Distribution::Float(d) => {
let v = if d.log_scale { value.exp() } else { value };
let v = if let Some(step) = d.step {
let k = ((v - d.low) / step).round();
d.low + k * step
} else {
v
};
ParamValue::Float(v.clamp(d.low, d.high))
}
Distribution::Int(d) => {
let v = if d.log_scale { value.exp() } else { value };
let v = if let Some(step) = d.step {
let k = ((v - d.low as f64) / step as f64).round() as i64;
d.low.saturating_add(k.saturating_mul(step))
} else {
v.round() as i64
};
ParamValue::Int(v.clamp(d.low, d.high))
}
Distribution::Categorical(_) => {
unreachable!("from_internal should not be called for categorical distributions")
}
}
}
/// Convert a `ParamValue` to its internal-space representation.
#[allow(clippy::cast_precision_loss, dead_code)]
pub(crate) fn to_internal(value: &ParamValue, distribution: &Distribution) -> f64 {
match (value, distribution) {
(ParamValue::Float(v), Distribution::Float(d)) => {
if d.log_scale {
v.ln()
} else {
*v
}
}
(ParamValue::Int(v), Distribution::Int(d)) => {
if d.log_scale {
(*v as f64).ln()
} else {
*v as f64
}
}
_ => 0.0,
}
}
/// Sample a random value for any distribution.
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
pub(crate) fn sample_random(rng: &mut fastrand::Rng, distribution: &Distribution) -> ParamValue {
match distribution {
Distribution::Float(d) => {
let value = if d.log_scale {
let log_low = d.low.ln();
let log_high = d.high.ln();
let v = rng_util::f64_range(rng, log_low, log_high).exp();
if let Some(step) = d.step {
let k = ((v - d.low) / step).round();
(d.low + k * step).clamp(d.low, d.high)
} else {
v
}
} else if let Some(step) = d.step {
let n_steps = ((d.high - d.low) / step).floor() as i64;
let k = rng.i64(0..=n_steps);
d.low + (k as f64) * step
} else {
rng_util::f64_range(rng, d.low, d.high)
};
ParamValue::Float(value)
}
Distribution::Int(d) => {
let value = if d.log_scale {
let log_low = (d.low as f64).ln();
let log_high = (d.high as f64).ln();
let v = rng_util::f64_range(rng, log_low, log_high).exp();
let raw = if let Some(step) = d.step {
let k = ((v - d.low as f64) / step as f64).round() as i64;
d.low.saturating_add(k.saturating_mul(step))
} else {
v.round() as i64
};
raw.clamp(d.low, d.high)
} else if let Some(step) = d.step {
let n_steps = (d.high - d.low) / step;
let k = rng.i64(0..=n_steps);
d.low + k * step
} else {
rng.i64(d.low..=d.high)
};
ParamValue::Int(value)
}
Distribution::Categorical(d) => ParamValue::Categorical(rng.usize(0..d.n_choices)),
}
}
+926
View File
@@ -0,0 +1,926 @@
//! Differential Evolution (DE) sampler.
//!
//! DE is a population-based metaheuristic that maintains a pool of candidate
//! solutions and creates new candidates through **mutation** (combining
//! difference vectors of existing members) and **binomial crossover**. A
//! trial vector replaces its parent only if it achieves a better objective
//! value, guaranteeing monotonic improvement of the population.
//!
//! # Algorithm overview
//!
//! Each generation, for every population member *xᵢ*:
//! 1. **Mutation** — create a mutant vector *v* from other population
//! members using the selected [`DEStrategy`]:
//! - `Rand1`: `v = x_r1 + F * (x_r2 - x_r3)`
//! - `Best1`: `v = x_best + F * (x_r1 - x_r2)`
//! - `CurrentToBest1`: `v = x_i + F * (x_best - x_i) + F * (x_r1 - x_r2)`
//! 2. **Crossover** — create a trial vector *u* by mixing *v* and *xᵢ*
//! dimension-by-dimension with probability CR.
//! 3. **Selection** — replace *xᵢ* with *u* if `f(u) ≤ f(xᵢ)`.
//!
//! # When to use
//!
//! - **Continuous parameters** (float/int). Categorical parameters are
//! sampled uniformly at random and do not participate in DE.
//! - **Moderate to large search spaces** — DE scales better than GP-based
//! methods to higher dimensions, though it may need more evaluations.
//! - **Multi-modal landscapes** — the `Rand1` strategy maintains diversity
//! and avoids premature convergence.
//! - **No feature flags required** — DE is available with default features.
//!
//! For non-separable problems in moderate dimensions, consider
//! `CmaEsSampler` which learns parameter
//! correlations. For expensive functions with few dimensions, consider
//! `GpSampler`.
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `population_size` | `max(10n, 15)` | Candidates per generation |
//! | `mutation_factor` (F) | 0.8 | Differential amplification — higher = more exploration |
//! | `crossover_rate` (CR) | 0.9 | Probability of taking a dimension from the mutant |
//! | `strategy` | `Rand1` | Mutation strategy (see [`DEStrategy`]) |
//! | `seed` | random | RNG seed for reproducibility |
//!
//! # Examples
//!
//! ```
//! use optimizer::sampler::de::{DESampler, DEStrategy};
//! use optimizer::{Direction, Study};
//!
//! // Minimize with DE using the Best1 strategy for faster convergence
//! let sampler = DESampler::builder()
//! .mutation_factor(0.7)
//! .crossover_rate(0.9)
//! .strategy(DEStrategy::Best1)
//! .population_size(20)
//! .seed(42)
//! .build();
//!
//! let mut study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! ```
use std::collections::HashMap;
use parking_lot::Mutex;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::rng_util;
use crate::sampler::{CompletedTrial, Sampler};
use super::common::{from_internal, internal_bounds, sample_random};
/// Differential Evolution mutation strategy.
///
/// Controls how mutant vectors are created from the current population.
#[derive(Clone, Copy, Debug, Default)]
pub enum DEStrategy {
/// DE/rand/1: `v = x_r1 + F * (x_r2 - x_r3)`
///
/// The most robust strategy. Uses three random population members.
#[default]
Rand1,
/// DE/best/1: `v = x_best + F * (x_r1 - x_r2)`
///
/// Greedier strategy that biases toward the current best solution.
Best1,
/// DE/current-to-best/1: `v = x_i + F * (x_best - x_i) + F * (x_r1 - x_r2)`
///
/// Balances exploration and exploitation by blending the current
/// individual with the best.
CurrentToBest1,
}
/// Differential Evolution sampler for continuous global optimization.
///
/// Maintains a population of candidate solutions. New candidates are
/// created by combining (mutating + crossing over) existing members.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::de::DESampler;
/// use optimizer::{Direction, Study};
///
/// // Default configuration
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, DESampler::new());
///
/// // With seed for reproducibility
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, DESampler::with_seed(42));
///
/// // Custom configuration via builder
/// use optimizer::sampler::de::DEStrategy;
/// let sampler = DESampler::builder()
/// .mutation_factor(0.8)
/// .crossover_rate(0.9)
/// .strategy(DEStrategy::Best1)
/// .population_size(30)
/// .seed(42)
/// .build();
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
/// ```
pub struct DESampler {
state: Mutex<State>,
}
impl DESampler {
/// Creates a new DE sampler with default settings and a random seed.
#[must_use]
pub fn new() -> Self {
Self {
state: Mutex::new(State::new(None, 0.8, 0.9, DEStrategy::Rand1, None)),
}
}
/// Creates a new DE sampler with a fixed seed for reproducibility.
#[must_use]
pub fn with_seed(seed: u64) -> Self {
Self {
state: Mutex::new(State::new(None, 0.8, 0.9, DEStrategy::Rand1, Some(seed))),
}
}
/// Creates a builder for configuring a `DESampler`.
#[must_use]
pub fn builder() -> DESamplerBuilder {
DESamplerBuilder::new()
}
}
impl Default for DESampler {
fn default() -> Self {
Self::new()
}
}
/// Builder for configuring a [`DESampler`].
///
/// All options have sensible defaults:
/// - `population_size`: `max(10 * n_dims, 15)` (auto-computed from parameter count)
/// - `mutation_factor` (F): 0.8
/// - `crossover_rate` (CR): 0.9
/// - `strategy`: `Rand1`
/// - `seed`: random
///
/// # Examples
///
/// ```
/// use optimizer::sampler::de::{DESamplerBuilder, DEStrategy};
///
/// let sampler = DESamplerBuilder::new()
/// .mutation_factor(0.5)
/// .crossover_rate(0.7)
/// .strategy(DEStrategy::CurrentToBest1)
/// .population_size(20)
/// .seed(42)
/// .build();
/// ```
#[derive(Debug, Clone)]
pub struct DESamplerBuilder {
population_size: Option<usize>,
mutation_factor: f64,
crossover_rate: f64,
strategy: DEStrategy,
seed: Option<u64>,
}
impl Default for DESamplerBuilder {
fn default() -> Self {
Self::new()
}
}
impl DESamplerBuilder {
/// Creates a new builder with default settings.
#[must_use]
pub fn new() -> Self {
Self {
population_size: None,
mutation_factor: 0.8,
crossover_rate: 0.9,
strategy: DEStrategy::Rand1,
seed: None,
}
}
/// Sets the population size.
///
/// Number of candidate solutions maintained across generations.
/// Larger populations improve robustness but require more evaluations
/// per generation.
///
/// Default: `max(10 * n_continuous_dims, 15)`.
#[must_use]
pub fn population_size(mut self, size: usize) -> Self {
self.population_size = Some(size);
self
}
/// Sets the mutation factor (F).
///
/// Controls the amplification of differential variation.
/// Typical values are in `[0.5, 1.0]`. Higher values increase
/// exploration; lower values favor exploitation.
///
/// Default: 0.8.
#[must_use]
pub fn mutation_factor(mut self, f: f64) -> Self {
self.mutation_factor = f;
self
}
/// Sets the crossover rate (CR).
///
/// Probability of each dimension being taken from the mutant vector
/// rather than the parent. Typical values are in `[0.7, 1.0]`.
///
/// Default: 0.9.
#[must_use]
pub fn crossover_rate(mut self, cr: f64) -> Self {
self.crossover_rate = cr;
self
}
/// Sets the mutation strategy.
///
/// Default: [`DEStrategy::Rand1`].
#[must_use]
pub fn strategy(mut self, strategy: DEStrategy) -> Self {
self.strategy = strategy;
self
}
/// Sets the random seed for reproducibility.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.seed = Some(seed);
self
}
/// Builds the configured [`DESampler`].
#[must_use]
pub fn build(self) -> DESampler {
DESampler {
state: Mutex::new(State::new(
self.population_size,
self.mutation_factor,
self.crossover_rate,
self.strategy,
self.seed,
)),
}
}
}
// ---------------------------------------------------------------------------
// Internal types
// ---------------------------------------------------------------------------
/// Describes how a parameter dimension maps into the DE internal vector.
#[derive(Clone, Debug)]
struct DimensionInfo {
/// The distribution for this dimension (stored for decoding).
distribution: Distribution,
/// Whether this dimension participates in DE (Float/Int = true, Categorical = false).
is_continuous: bool,
/// Internal-space bounds for continuous dimensions: `(low, high)`.
/// For log-scale parameters these are in log-space.
bounds: Option<(f64, f64)>,
}
/// A candidate solution produced by mutation + crossover.
#[derive(Clone, Debug)]
struct Candidate {
/// Internal-space vector (only continuous dimensions).
x: Vec<f64>,
/// Values for categorical dimensions (index in `dimensions` -> categorical index).
categorical_values: HashMap<usize, usize>,
/// Index of the population member this candidate competes against.
target_idx: usize,
}
/// Tracks per-trial sampling progress.
#[derive(Clone, Debug)]
struct TrialProgress {
/// Index of the candidate assigned to this trial.
candidate_idx: usize,
/// Next dimension to return for this trial.
next_dim: usize,
}
/// Phase of the DE state machine.
enum Phase {
/// Discovering the search space structure (first trial).
Discovery,
/// Active sampling and evolving.
Active,
}
/// Top-level mutable state behind the `Mutex`.
struct State {
/// The RNG used for sampling.
rng: fastrand::Rng,
/// User-provided population size (None = auto).
user_population_size: Option<usize>,
/// Mutation factor (F).
mutation_factor: f64,
/// Crossover rate (CR).
crossover_rate: f64,
/// Mutation strategy.
strategy: DEStrategy,
/// Current phase.
phase: Phase,
/// Discovered dimension info (populated during discovery).
dimensions: Vec<DimensionInfo>,
/// Last `trial_id` seen during discovery.
discovery_trial_id: Option<u64>,
// --- Population state ---
/// Current population (internal-space vectors, continuous dims only).
population: Vec<Vec<f64>>,
/// Categorical values for each population member.
population_categorical: Vec<HashMap<usize, usize>>,
/// Objective values for the current population.
population_values: Vec<f64>,
/// Index of the best population member.
best_idx: usize,
/// Whether the initial population has been evaluated.
initialized: bool,
/// Effective population size (resolved after discovery).
population_size: usize,
// --- Current generation ---
/// Current generation's candidates.
candidates: Vec<Candidate>,
/// Mapping from `trial_id` to its progress.
trial_progress: HashMap<u64, TrialProgress>,
/// Number of candidates assigned so far in the current generation.
assigned_count: usize,
/// Trial IDs assigned in the current generation.
generation_trial_ids: Vec<u64>,
}
impl State {
fn new(
user_population_size: Option<usize>,
mutation_factor: f64,
crossover_rate: f64,
strategy: DEStrategy,
seed: Option<u64>,
) -> Self {
let rng = seed.map_or_else(fastrand::Rng::new, fastrand::Rng::with_seed);
Self {
rng,
user_population_size,
mutation_factor,
crossover_rate,
strategy,
phase: Phase::Discovery,
dimensions: Vec::new(),
discovery_trial_id: None,
population: Vec::new(),
population_categorical: Vec::new(),
population_values: Vec::new(),
best_idx: 0,
initialized: false,
population_size: 0,
candidates: Vec::new(),
trial_progress: HashMap::new(),
assigned_count: 0,
generation_trial_ids: Vec::new(),
}
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/// Sample a random value in internal space for a continuous dimension.
fn sample_random_internal(rng: &mut fastrand::Rng, bounds: (f64, f64)) -> f64 {
rng_util::f64_range(rng, bounds.0, bounds.1)
}
/// Clamp a value to the given bounds.
fn clamp_to_bounds(value: f64, bounds: Option<(f64, f64)>) -> f64 {
if let Some((lo, hi)) = bounds {
value.clamp(lo, hi)
} else {
value
}
}
// ---------------------------------------------------------------------------
// DE algorithm
// ---------------------------------------------------------------------------
/// Select `count` distinct random indices from `0..n`, all different from `exclude`.
fn select_random_indices(
rng: &mut fastrand::Rng,
n: usize,
count: usize,
exclude: &[usize],
) -> Vec<usize> {
let mut selected = Vec::with_capacity(count);
while selected.len() < count {
let idx = rng.usize(0..n);
if !exclude.contains(&idx) && !selected.contains(&idx) {
selected.push(idx);
}
}
selected
}
/// Generate trial vectors (mutation + crossover) for the current population.
fn generate_trial_vectors(state: &mut State) -> Vec<Candidate> {
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
let pop_size = state.population_size;
let mut candidates = Vec::with_capacity(pop_size);
for i in 0..pop_size {
// Mutation
let mutant = create_mutant_with_rng(state, i, n_continuous);
// Crossover (binomial)
let j_rand = state.rng.usize(0..n_continuous.max(1));
let trial_x: Vec<f64> = if n_continuous > 0 {
(0..n_continuous)
.map(|j| {
let use_mutant = j == j_rand || state.rng.f64() < state.crossover_rate;
let val = if use_mutant {
mutant[j]
} else {
state.population[i][j]
};
// Clamp to bounds
let dim_bounds = continuous_dim_bounds(&state.dimensions, j);
clamp_to_bounds(val, dim_bounds)
})
.collect()
} else {
Vec::new()
};
// Categorical: randomly sample (DE doesn't optimize categoricals)
let mut categorical_values = HashMap::new();
for (dim_idx, dim) in state.dimensions.iter().enumerate() {
if !dim.is_continuous
&& let Distribution::Categorical(cat) = &dim.distribution
{
categorical_values.insert(dim_idx, state.rng.usize(0..cat.n_choices));
}
}
candidates.push(Candidate {
x: trial_x,
categorical_values,
target_idx: i,
});
}
candidates
}
/// Create a mutant vector, consuming RNG from state.
fn create_mutant_with_rng(state: &mut State, target_idx: usize, n_continuous: usize) -> Vec<f64> {
if n_continuous == 0 {
return Vec::new();
}
let pop = &state.population;
let best_idx = state.best_idx;
let f = state.mutation_factor;
let pop_size = state.population_size;
match state.strategy {
DEStrategy::Rand1 => {
let indices = select_random_indices(&mut state.rng, pop_size, 3, &[target_idx]);
let (r1, r2, r3) = (indices[0], indices[1], indices[2]);
(0..n_continuous)
.map(|j| pop[r1][j] + f * (pop[r2][j] - pop[r3][j]))
.collect()
}
DEStrategy::Best1 => {
let indices = select_random_indices(&mut state.rng, pop_size, 2, &[target_idx]);
let (r1, r2) = (indices[0], indices[1]);
(0..n_continuous)
.map(|j| pop[best_idx][j] + f * (pop[r1][j] - pop[r2][j]))
.collect()
}
DEStrategy::CurrentToBest1 => {
let indices = select_random_indices(&mut state.rng, pop_size, 2, &[target_idx]);
let (r1, r2) = (indices[0], indices[1]);
(0..n_continuous)
.map(|j| {
pop[target_idx][j]
+ f * (pop[best_idx][j] - pop[target_idx][j])
+ f * (pop[r1][j] - pop[r2][j])
})
.collect()
}
}
}
/// Get the bounds for the j-th continuous dimension.
fn continuous_dim_bounds(
dimensions: &[DimensionInfo],
continuous_idx: usize,
) -> Option<(f64, f64)> {
let mut ci = 0;
for dim in dimensions {
if dim.is_continuous {
if ci == continuous_idx {
return dim.bounds;
}
ci += 1;
}
}
None
}
/// Generate the initial random population.
fn generate_initial_population(state: &mut State) -> Vec<Candidate> {
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
let mut candidates = Vec::with_capacity(state.population_size);
for i in 0..state.population_size {
let x: Vec<f64> = if n_continuous > 0 {
let mut v = Vec::with_capacity(n_continuous);
for dim in &state.dimensions {
if dim.is_continuous {
let val = if let Some(bounds) = dim.bounds {
sample_random_internal(&mut state.rng, bounds)
} else {
0.0
};
v.push(val);
}
}
v
} else {
Vec::new()
};
let mut categorical_values = HashMap::new();
for (dim_idx, dim) in state.dimensions.iter().enumerate() {
if !dim.is_continuous
&& let Distribution::Categorical(cat) = &dim.distribution
{
categorical_values.insert(dim_idx, state.rng.usize(0..cat.n_choices));
}
}
candidates.push(Candidate {
x,
categorical_values,
target_idx: i,
});
}
candidates
}
// ---------------------------------------------------------------------------
// Sampler trait implementation
// ---------------------------------------------------------------------------
impl Sampler for DESampler {
#[allow(clippy::cast_precision_loss)]
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue {
let mut state = self.state.lock();
match &state.phase {
Phase::Discovery => sample_discovery(&mut state, distribution, trial_id),
Phase::Active => sample_active(&mut state, distribution, trial_id, history),
}
}
}
/// Handle sampling during the discovery phase.
fn sample_discovery(state: &mut State, distribution: &Distribution, trial_id: u64) -> ParamValue {
// Check if this is a new trial (discovery phase ended for previous trial)
if let Some(prev_id) = state.discovery_trial_id
&& trial_id != prev_id
{
// First trial is done; we know the search space. Initialize DE.
finalize_discovery(state);
return sample_active(state, distribution, trial_id, &[]);
}
// Record this trial_id
state.discovery_trial_id = Some(trial_id);
// Record this dimension
let is_continuous = !matches!(distribution, Distribution::Categorical(_));
let bounds = internal_bounds(distribution);
state.dimensions.push(DimensionInfo {
distribution: distribution.clone(),
is_continuous,
bounds,
});
// Sample randomly for the discovery trial
sample_random(&mut state.rng, distribution)
}
/// Finalize discovery and transition to the active phase.
#[allow(clippy::cast_precision_loss)]
fn finalize_discovery(state: &mut State) {
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
// Resolve population size
state.population_size = state
.user_population_size
.unwrap_or_else(|| (10 * n_continuous).max(15));
// Ensure population size is at least 4 (DE needs distinct random indices)
state.population_size = state.population_size.max(4);
// Generate initial random population
state.candidates = generate_initial_population(state);
state.assigned_count = 0;
state.generation_trial_ids.clear();
state.trial_progress.clear();
state.phase = Phase::Active;
}
/// Handle sampling during the active phase.
fn sample_active(
state: &mut State,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue {
// Check if we need to process completed trials and start a new generation
maybe_update_generation(state, history);
// Assign a candidate to this trial if not yet done
if !state.trial_progress.contains_key(&trial_id) {
assign_candidate(state, trial_id);
}
let progress = state.trial_progress.get_mut(&trial_id).unwrap();
let dim_idx = progress.next_dim;
progress.next_dim += 1;
// Safety check
if dim_idx >= state.dimensions.len() {
return sample_random(&mut state.rng, distribution);
}
let candidate = &state.candidates[progress.candidate_idx];
let dim_info = &state.dimensions[dim_idx];
if dim_info.is_continuous {
// Map from overall dimension index to continuous index
let ci = state.dimensions[..dim_idx]
.iter()
.filter(|d| d.is_continuous)
.count();
if ci < candidate.x.len() {
from_internal(candidate.x[ci], &dim_info.distribution)
} else {
sample_random(&mut state.rng, distribution)
}
} else {
// Categorical: use pre-sampled value
if let Some(&cat_idx) = candidate.categorical_values.get(&dim_idx) {
ParamValue::Categorical(cat_idx)
} else {
sample_random(&mut state.rng, distribution)
}
}
}
/// Assign a candidate to a trial.
fn assign_candidate(state: &mut State, trial_id: u64) {
let candidate_idx = if state.assigned_count < state.candidates.len() {
let idx = state.assigned_count;
state.assigned_count += 1;
idx
} else {
// Overflow: generate an extra random candidate
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
let x: Vec<f64> = (0..n_continuous)
.map(|j| {
let bounds = continuous_dim_bounds(&state.dimensions, j);
if let Some(b) = bounds {
sample_random_internal(&mut state.rng, b)
} else {
0.0
}
})
.collect();
let mut categorical_values = HashMap::new();
for (dim_idx, dim) in state.dimensions.iter().enumerate() {
if !dim.is_continuous
&& let Distribution::Categorical(cat) = &dim.distribution
{
categorical_values.insert(dim_idx, state.rng.usize(0..cat.n_choices));
}
}
state.candidates.push(Candidate {
x,
categorical_values,
target_idx: 0, // overflow candidates don't compete
});
let idx = state.candidates.len() - 1;
state.assigned_count = state.candidates.len();
idx
};
state.trial_progress.insert(
trial_id,
TrialProgress {
candidate_idx,
next_dim: 0,
},
);
state.generation_trial_ids.push(trial_id);
}
/// Check if we should process completed trials and start a new generation.
fn maybe_update_generation(state: &mut State, history: &[CompletedTrial]) {
let pop_size = state.population_size;
// Only update when at least pop_size candidates have been assigned
if state.generation_trial_ids.len() < pop_size {
return;
}
// Check if the first pop_size trial IDs are all completed
let trial_ids: Vec<u64> = state
.generation_trial_ids
.iter()
.take(pop_size)
.copied()
.collect();
let history_map: HashMap<u64, f64> = history.iter().map(|t| (t.id, t.value)).collect();
let all_completed = trial_ids.iter().all(|id| history_map.contains_key(id));
if !all_completed {
return;
}
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
if state.initialized {
// Subsequent generations: selection
perform_selection(state, &trial_ids, &history_map);
} else {
// First generation: store as initial population
initialize_population(state, &trial_ids, &history_map, n_continuous);
}
// Generate next generation's trial vectors
state.candidates = if state.initialized && n_continuous > 0 {
generate_trial_vectors(state)
} else {
generate_initial_population(state)
};
state.assigned_count = 0;
state.generation_trial_ids.clear();
state.trial_progress.clear();
}
/// Initialize the population from the first generation's results.
fn initialize_population(
state: &mut State,
trial_ids: &[u64],
history_map: &HashMap<u64, f64>,
_n_continuous: usize,
) {
state.population.clear();
state.population_categorical.clear();
state.population_values.clear();
let mut best_value = f64::INFINITY;
let mut best_idx = 0;
for (i, &trial_id) in trial_ids.iter().enumerate() {
let progress = &state.trial_progress[&trial_id];
let candidate = &state.candidates[progress.candidate_idx];
let value = history_map[&trial_id];
state.population.push(candidate.x.clone());
state
.population_categorical
.push(candidate.categorical_values.clone());
state.population_values.push(value);
if value < best_value {
best_value = value;
best_idx = i;
}
}
state.best_idx = best_idx;
state.initialized = true;
}
/// Perform DE selection: replace parent if trial vector is better.
fn perform_selection(state: &mut State, trial_ids: &[u64], history_map: &HashMap<u64, f64>) {
for &trial_id in trial_ids {
let progress = &state.trial_progress[&trial_id];
let candidate = &state.candidates[progress.candidate_idx];
let trial_value = history_map[&trial_id];
let target_idx = candidate.target_idx;
if target_idx < state.population_size && trial_value <= state.population_values[target_idx]
{
state.population[target_idx] = candidate.x.clone();
state.population_categorical[target_idx] = candidate.categorical_values.clone();
state.population_values[target_idx] = trial_value;
}
}
// Update best index
let mut best_value = f64::INFINITY;
let mut best_idx = 0;
for (i, &val) in state.population_values.iter().enumerate() {
if val < best_value {
best_value = val;
best_idx = i;
}
}
state.best_idx = best_idx;
}
#[cfg(test)]
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
mod tests {
use super::*;
use crate::distribution::FloatDistribution;
#[test]
fn test_de_sampler_basic_float() {
let sampler = DESampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: -5.0,
high: 5.0,
log_scale: false,
step: None,
});
// Sample many values and check bounds
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value {
assert!(
(-5.0..=5.0).contains(&v),
"value {v} out of bounds at trial {i}"
);
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn test_de_sampler_reproducibility() {
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let sample_values = |seed: u64| {
let sampler = DESampler::with_seed(seed);
(0..20)
.map(|i| sampler.sample(&dist, i, &[]))
.collect::<Vec<_>>()
};
let v1 = sample_values(42);
let v2 = sample_values(42);
assert_eq!(v1, v2, "same seed should produce same results");
let v3 = sample_values(99);
assert_ne!(v1, v3, "different seeds should produce different results");
}
#[test]
fn test_de_strategy_default() {
assert!(matches!(DEStrategy::default(), DEStrategy::Rand1));
}
#[test]
fn test_builder_defaults() {
let builder = DESamplerBuilder::new();
assert!(builder.population_size.is_none());
assert!((builder.mutation_factor - 0.8).abs() < f64::EPSILON);
assert!((builder.crossover_rate - 0.9).abs() < f64::EPSILON);
assert!(matches!(builder.strategy, DEStrategy::Rand1));
assert!(builder.seed.is_none());
}
}
+537
View File
@@ -0,0 +1,537 @@
//! Shared types and genetic operators for evolutionary multi-objective samplers.
//!
//! This module extracts common functionality used by NSGA-II, NSGA-III, and MOEA/D:
//! candidate management, discovery/active phase logic, SBX crossover,
//! polynomial mutation, and Das-Dennis reference point generation.
use std::collections::HashMap;
use crate::distribution::Distribution;
use crate::multi_objective::MultiObjectiveTrial;
use crate::param::ParamValue;
use crate::rng_util;
/// Describes a parameter dimension discovered during the first trial.
#[derive(Clone, Debug)]
pub(crate) struct DimensionInfo {
pub distribution: Distribution,
}
/// A candidate solution: one value per dimension.
#[derive(Clone, Debug)]
pub(crate) struct Candidate {
pub params: Vec<ParamValue>,
}
/// Tracks per-trial sampling progress (which candidate, which dimension next).
#[derive(Clone, Debug)]
pub(crate) struct TrialProgress {
pub candidate_idx: usize,
pub next_dim: usize,
}
/// Phase of an evolutionary sampler.
pub(crate) enum Phase {
/// First trial reveals parameter dimensions.
Discovery,
/// Evolutionary optimisation.
Active,
}
/// Common state shared by all evolutionary multi-objective samplers.
pub(crate) struct EvolutionaryState {
pub rng: fastrand::Rng,
pub phase: Phase,
pub dimensions: Vec<DimensionInfo>,
pub population_size: usize,
pub candidates: Vec<Candidate>,
pub trial_progress: HashMap<u64, TrialProgress>,
pub assigned_count: usize,
pub generation_trial_ids: Vec<u64>,
pub discovery_trial_id: Option<u64>,
pub generation: usize,
}
impl EvolutionaryState {
pub(crate) fn new(seed: Option<u64>) -> Self {
let rng = seed.map_or_else(fastrand::Rng::new, fastrand::Rng::with_seed);
Self {
rng,
phase: Phase::Discovery,
dimensions: Vec::new(),
population_size: 4,
candidates: Vec::new(),
trial_progress: HashMap::new(),
assigned_count: 0,
generation_trial_ids: Vec::new(),
discovery_trial_id: None,
generation: 0,
}
}
}
// ---------------------------------------------------------------------------
// Discovery phase helpers
// ---------------------------------------------------------------------------
/// Handle sampling during the discovery phase.
///
/// Returns `Some(value)` if the discovery phase handled the sample,
/// or `None` if it transitioned to active phase and the caller should
/// generate candidates and sample from them.
pub(crate) fn sample_discovery(
evo: &mut EvolutionaryState,
distribution: &Distribution,
trial_id: u64,
) -> Option<ParamValue> {
if let Some(prev_id) = evo.discovery_trial_id
&& trial_id != prev_id
{
// A new trial arrived — transition to active phase
return None;
}
evo.discovery_trial_id = Some(trial_id);
evo.dimensions.push(DimensionInfo {
distribution: distribution.clone(),
});
Some(sample_random(&mut evo.rng, distribution))
}
/// Compute population size from dimensions and optional user override.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
pub(crate) fn compute_population_size(
n_dims: usize,
user_pop_size: Option<usize>,
minimum: usize,
) -> usize {
user_pop_size
.unwrap_or_else(|| (4.0 + 3.0 * (n_dims as f64).ln().max(0.0)).floor() as usize)
.max(minimum)
}
/// Transition from discovery to active phase.
pub(crate) fn finalize_discovery(evo: &mut EvolutionaryState, user_pop_size: Option<usize>) {
evo.population_size = compute_population_size(evo.dimensions.len(), user_pop_size, 4);
evo.phase = Phase::Active;
}
/// Generate `population_size` random candidates.
pub(crate) fn generate_random_candidates(evo: &mut EvolutionaryState) {
let pop = evo.population_size;
evo.candidates = (0..pop)
.map(|_| {
let params: Vec<ParamValue> = evo
.dimensions
.iter()
.map(|d| sample_random(&mut evo.rng, &d.distribution))
.collect();
Candidate { params }
})
.collect();
evo.assigned_count = 0;
evo.generation_trial_ids.clear();
evo.trial_progress.clear();
}
/// Assign a candidate to a trial and return the next dimension value.
pub(crate) fn sample_from_candidate(evo: &mut EvolutionaryState, trial_id: u64) -> ParamValue {
if !evo.trial_progress.contains_key(&trial_id) {
let candidate_idx = if evo.assigned_count < evo.candidates.len() {
let idx = evo.assigned_count;
evo.assigned_count += 1;
idx
} else {
// Overflow: generate a random candidate
let params: Vec<ParamValue> = evo
.dimensions
.iter()
.map(|d| sample_random(&mut evo.rng, &d.distribution))
.collect();
evo.candidates.push(Candidate { params });
let idx = evo.candidates.len() - 1;
evo.assigned_count = evo.candidates.len();
idx
};
evo.trial_progress.insert(
trial_id,
TrialProgress {
candidate_idx,
next_dim: 0,
},
);
evo.generation_trial_ids.push(trial_id);
}
let progress = evo.trial_progress.get_mut(&trial_id).unwrap();
let dim_idx = progress.next_dim;
progress.next_dim += 1;
if dim_idx >= evo.dimensions.len() {
return sample_random(&mut evo.rng, &evo.dimensions.last().unwrap().distribution);
}
evo.candidates[progress.candidate_idx].params[dim_idx].clone()
}
/// Extract parameter values from a trial, ordered by dimension index.
pub(crate) fn extract_trial_params(
trial: &MultiObjectiveTrial,
dimensions: &[DimensionInfo],
rng: &mut fastrand::Rng,
) -> Vec<ParamValue> {
let mut param_pairs: Vec<_> = trial.params.iter().collect();
param_pairs.sort_by_key(|(id, _)| *id);
dimensions
.iter()
.enumerate()
.map(|(dim_idx, dim_info)| {
if dim_idx < param_pairs.len() {
param_pairs[dim_idx].1.clone()
} else {
sample_random(rng, &dim_info.distribution)
}
})
.collect()
}
/// Install new offspring as the next generation's candidates.
pub(crate) fn advance_generation(evo: &mut EvolutionaryState, offspring: Vec<Candidate>) {
evo.candidates = offspring;
evo.assigned_count = 0;
evo.generation_trial_ids.clear();
evo.trial_progress.clear();
evo.generation += 1;
}
/// Check if the current generation is fully evaluated and return the
/// evaluated trials if so.
pub(crate) fn collect_evaluated_generation<'a>(
evo: &EvolutionaryState,
history: &'a [MultiObjectiveTrial],
) -> Option<Vec<&'a MultiObjectiveTrial>> {
let pop_size = evo.population_size;
if evo.generation_trial_ids.len() < pop_size {
return None;
}
let gen_ids: Vec<u64> = evo
.generation_trial_ids
.iter()
.take(pop_size)
.copied()
.collect();
let history_map: HashMap<u64, &MultiObjectiveTrial> =
history.iter().map(|t| (t.id, t)).collect();
if !gen_ids.iter().all(|id| history_map.contains_key(id)) {
return None;
}
Some(
gen_ids
.iter()
.filter_map(|id| history_map.get(id).copied())
.collect(),
)
}
// ---------------------------------------------------------------------------
// Genetic operators
// ---------------------------------------------------------------------------
/// SBX crossover for continuous params, uniform crossover for categorical.
pub(crate) fn crossover(
rng: &mut fastrand::Rng,
parent1: &[ParamValue],
parent2: &[ParamValue],
dimensions: &[DimensionInfo],
crossover_prob: f64,
eta: f64,
) -> (Vec<ParamValue>, Vec<ParamValue>) {
let n = parent1.len();
let mut child1 = parent1.to_vec();
let mut child2 = parent2.to_vec();
let u: f64 = rng_util::f64_range(rng, 0.0, 1.0);
if u > crossover_prob {
return (child1, child2);
}
for i in 0..n {
match (&parent1[i], &parent2[i], &dimensions[i].distribution) {
(ParamValue::Float(p1), ParamValue::Float(p2), Distribution::Float(d)) => {
if (p1 - p2).abs() < 1e-14 {
continue;
}
let (c1, c2) = sbx_crossover_f64(rng, *p1, *p2, d.low, d.high, eta);
child1[i] = ParamValue::Float(c1);
child2[i] = ParamValue::Float(c2);
}
(ParamValue::Int(p1), ParamValue::Int(p2), Distribution::Int(d)) => {
if p1 == p2 {
continue;
}
#[allow(clippy::cast_precision_loss)]
let (c1, c2) = sbx_crossover_f64(
rng,
*p1 as f64,
*p2 as f64,
d.low as f64,
d.high as f64,
eta,
);
#[allow(clippy::cast_possible_truncation)]
{
child1[i] = ParamValue::Int((c1.round() as i64).clamp(d.low, d.high));
child2[i] = ParamValue::Int((c2.round() as i64).clamp(d.low, d.high));
}
}
(ParamValue::Categorical(_), ParamValue::Categorical(_), _) => {
if rng_util::f64_range(rng, 0.0, 1.0) < 0.5 {
core::mem::swap(&mut child1[i], &mut child2[i]);
}
}
_ => {}
}
}
(child1, child2)
}
/// SBX crossover for a single float dimension.
pub(crate) fn sbx_crossover_f64(
rng: &mut fastrand::Rng,
p1: f64,
p2: f64,
low: f64,
high: f64,
eta: f64,
) -> (f64, f64) {
let u: f64 = rng_util::f64_range(rng, 0.0, 1.0);
let beta = if u <= 0.5 {
(2.0 * u).powf(1.0 / (eta + 1.0))
} else {
(1.0 / (2.0 * (1.0 - u))).powf(1.0 / (eta + 1.0))
};
let c1 = 0.5 * ((1.0 + beta) * p1 + (1.0 - beta) * p2);
let c2 = 0.5 * ((1.0 - beta) * p1 + (1.0 + beta) * p2);
(c1.clamp(low, high), c2.clamp(low, high))
}
/// Polynomial mutation for each dimension.
#[allow(clippy::cast_precision_loss)]
pub(crate) fn mutate(
rng: &mut fastrand::Rng,
individual: &mut [ParamValue],
dimensions: &[DimensionInfo],
eta: f64,
) {
let n = individual.len();
if n == 0 {
return;
}
let mutation_prob = 1.0 / n as f64;
for (i, value) in individual.iter_mut().enumerate() {
if rng_util::f64_range(rng, 0.0, 1.0) >= mutation_prob {
continue;
}
match (value, &dimensions[i].distribution) {
(v @ ParamValue::Float(_), Distribution::Float(d)) => {
let ParamValue::Float(x) = *v else {
unreachable!();
};
let mutated = polynomial_mutation_f64(rng, x, d.low, d.high, eta);
*v = ParamValue::Float(mutated);
}
(v @ ParamValue::Int(_), Distribution::Int(d)) => {
let ParamValue::Int(x) = *v else {
unreachable!();
};
#[allow(clippy::cast_possible_truncation)]
{
let mutated =
polynomial_mutation_f64(rng, x as f64, d.low as f64, d.high as f64, eta);
*v = ParamValue::Int((mutated.round() as i64).clamp(d.low, d.high));
}
}
(v @ ParamValue::Categorical(_), Distribution::Categorical(d)) => {
*v = ParamValue::Categorical(rng.usize(0..d.n_choices));
}
_ => {}
}
}
}
/// Polynomial mutation for a single float value.
pub(crate) fn polynomial_mutation_f64(
rng: &mut fastrand::Rng,
x: f64,
low: f64,
high: f64,
eta: f64,
) -> f64 {
let u: f64 = rng_util::f64_range(rng, 0.0, 1.0);
let range = high - low;
if range <= 0.0 {
return x;
}
let delta1 = (x - low) / range;
let delta2 = (high - x) / range;
let delta_q = if u < 0.5 {
let xy = 1.0 - delta1;
let val = 2.0 * u + (1.0 - 2.0 * u) * xy.powf(eta + 1.0);
val.powf(1.0 / (eta + 1.0)) - 1.0
} else {
let xy = 1.0 - delta2;
let val = 2.0 * (1.0 - u) + 2.0 * (u - 0.5) * xy.powf(eta + 1.0);
1.0 - val.powf(1.0 / (eta + 1.0))
};
(x + delta_q * range).clamp(low, high)
}
pub(crate) use super::common::sample_random;
// ---------------------------------------------------------------------------
// Das-Dennis reference point generation
// ---------------------------------------------------------------------------
/// Generate Das-Dennis (simplex-lattice) reference points.
///
/// Returns `C(H + M - 1, M - 1)` uniformly spaced points on the
/// `M`-dimensional unit simplex, where `M = n_objectives` and
/// `H = divisions`.
pub(crate) fn das_dennis(n_objectives: usize, divisions: usize) -> Vec<Vec<f64>> {
let mut points = Vec::new();
let mut point = vec![0.0_f64; n_objectives];
das_dennis_recursive(
n_objectives,
divisions,
0,
divisions,
&mut point,
&mut points,
);
points
}
#[allow(clippy::cast_precision_loss)]
fn das_dennis_recursive(
n_objectives: usize,
divisions: usize,
depth: usize,
remaining: usize,
current: &mut Vec<f64>,
result: &mut Vec<Vec<f64>>,
) {
if depth == n_objectives - 1 {
current[depth] = remaining as f64 / divisions as f64;
result.push(current.clone());
return;
}
for i in 0..=remaining {
current[depth] = i as f64 / divisions as f64;
das_dennis_recursive(
n_objectives,
divisions,
depth + 1,
remaining - i,
current,
result,
);
}
}
/// Choose the number of divisions for Das-Dennis to get close to a target
/// population size.
///
/// The number of reference points is `C(H + M - 1, M - 1)`. This function
/// finds the smallest `H` such that the number of points >= `target_pop`.
pub(crate) fn auto_divisions(n_objectives: usize, target_pop: usize) -> usize {
let m = n_objectives;
for h in 1..200 {
let n_points = n_combinations(h + m - 1, m - 1);
if n_points >= target_pop {
return h;
}
}
12
}
/// Compute `C(n, k)` = n! / (k! * (n-k)!).
fn n_combinations(n: usize, k: usize) -> usize {
if k > n {
return 0;
}
let k = k.min(n - k);
let mut result: usize = 1;
for i in 0..k {
result = result.saturating_mul(n - i) / (i + 1);
}
result
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_das_dennis_2d() {
let points = das_dennis(2, 4);
// C(4+1, 1) = 5 points
assert_eq!(points.len(), 5);
for p in &points {
let sum: f64 = p.iter().sum();
assert!((sum - 1.0).abs() < 1e-10, "point {p:?} doesn't sum to 1");
}
}
#[test]
fn test_das_dennis_3d() {
let points = das_dennis(3, 4);
// C(4+2, 2) = 15 points
assert_eq!(points.len(), 15);
for p in &points {
let sum: f64 = p.iter().sum();
assert!((sum - 1.0).abs() < 1e-10);
}
}
#[test]
fn test_auto_divisions() {
// For 2 objectives targeting 10 points: H=9 gives C(10,1)=10
let h = auto_divisions(2, 10);
let n = n_combinations(h + 1, 1);
assert!(n >= 10);
// For 3 objectives targeting ~91 points: H=12 gives C(14,2)=91
let h3 = auto_divisions(3, 91);
let n3 = n_combinations(h3 + 2, 2);
assert!(n3 >= 91);
}
#[test]
fn test_n_combinations() {
assert_eq!(n_combinations(5, 2), 10);
assert_eq!(n_combinations(4, 0), 1);
assert_eq!(n_combinations(4, 4), 1);
assert_eq!(n_combinations(6, 3), 20);
}
}
+790
View File
@@ -0,0 +1,790 @@
//! Gaussian Process (GP) sampler with Expected Improvement acquisition.
//!
//! A classical Bayesian optimization sampler that builds a Gaussian Process
//! surrogate model with a **Matérn 5/2 kernel** (with ARD lengthscales) and
//! selects the next trial by maximizing the **Expected Improvement (EI)**
//! acquisition function. Best suited for small, expensive evaluations in
//! low-dimensional continuous spaces (d ≤ 20).
//!
//! # Algorithm overview
//!
//! 1. **Startup phase** — the first `n_startup_trials` trials are sampled
//! uniformly at random to build an initial dataset.
//! 2. **Fit GP** — training observations are standardized (zero mean, unit
//! variance) and a GP with Matérn 5/2 kernel is fitted via Cholesky
//! decomposition. ARD lengthscales are set to the per-dimension standard
//! deviation of the training inputs.
//! 3. **Maximize EI** — `n_candidates` random points are evaluated under
//! the GP posterior and the point with the highest Expected Improvement
//! is returned as the next trial.
//!
//! The GP uses at most 100 training points (the most recent ones) to keep
//! the O(n³) fitting cost manageable.
//!
//! # When to use
//!
//! - **Expensive objective functions** where every evaluation is costly
//! (e.g. physical experiments, large simulations). The GP surrogate
//! amortizes this cost by making fewer evaluations.
//! - **Low-dimensional continuous spaces** — typically d ≤ 20. Beyond that,
//! the GP becomes unreliable and alternatives like
//! [`CmaEsSampler`](super::cma_es::CmaEsSampler) or
//! [`TpeSampler`](super::tpe::TpeSampler) are preferable.
//! - **Smooth, low-noise objectives** — the GP assumes smoothness through
//! the Matérn 5/2 kernel. Very noisy objectives require increasing
//! `noise_variance`.
//!
//! Categorical parameters are sampled uniformly at random and do not
//! participate in the GP model. If all parameters are categorical, the
//! sampler falls back to pure random sampling.
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |--------|---------|-------------|
//! | `n_startup_trials` | 10 | Random trials before GP-guided sampling begins |
//! | `n_candidates` | 1000 | Random candidates for EI maximization |
//! | `noise_variance` | 1e-6 | Observation noise added to kernel diagonal |
//! | `seed` | random | RNG seed for reproducibility |
//!
//! # Feature flag
//!
//! Requires the **`gp`** feature (adds the `nalgebra` dependency):
//!
//! ```toml
//! [dependencies]
//! optimizer = { version = "...", features = ["gp"] }
//! ```
//!
//! # Examples
//!
//! ```
//! use optimizer::sampler::gp::GpSampler;
//! use optimizer::{Direction, Study};
//!
//! // Minimize an expensive function with GP-based Bayesian optimization
//! let sampler = GpSampler::builder()
//! .n_startup_trials(5)
//! .n_candidates(500)
//! .seed(42)
//! .build();
//!
//! let mut study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! ```
use std::collections::HashMap;
use nalgebra::DMatrix;
use parking_lot::Mutex;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::rng_util;
use crate::sampler::{CompletedTrial, Sampler};
use super::common::{from_internal, internal_bounds, sample_random, to_internal};
// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
/// Gaussian Process sampler for Bayesian optimization.
///
/// Uses a GP surrogate with Matérn 5/2 kernel and Expected Improvement
/// acquisition to guide sampling toward promising regions of the search
/// space. Best suited for continuous (float/int) parameters in low
/// dimensions (up to ~20).
///
/// # Examples
///
/// ```
/// use optimizer::sampler::gp::GpSampler;
/// use optimizer::{Direction, Study};
///
/// // Default configuration
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, GpSampler::new());
///
/// // With seed for reproducibility
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, GpSampler::with_seed(42));
///
/// // Custom configuration via builder
/// let sampler = GpSampler::builder()
/// .n_startup_trials(15)
/// .n_candidates(2000)
/// .noise_variance(1e-4)
/// .seed(42)
/// .build();
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
/// ```
pub struct GpSampler {
state: Mutex<GpState>,
}
impl GpSampler {
/// Creates a new GP sampler with a random seed.
#[must_use]
pub fn new() -> Self {
Self {
state: Mutex::new(GpState::new(None, None, None, None)),
}
}
/// Creates a new GP sampler with a fixed seed for reproducibility.
#[must_use]
pub fn with_seed(seed: u64) -> Self {
Self {
state: Mutex::new(GpState::new(None, None, None, Some(seed))),
}
}
/// Creates a builder for configuring a `GpSampler`.
#[must_use]
pub fn builder() -> GpSamplerBuilder {
GpSamplerBuilder::new()
}
}
impl Default for GpSampler {
fn default() -> Self {
Self::new()
}
}
/// Builder for configuring a [`GpSampler`].
///
/// All options have sensible defaults:
/// - `n_startup_trials`: 10
/// - `n_candidates`: 1000
/// - `noise_variance`: 1e-6
/// - `seed`: random
///
/// # Examples
///
/// ```
/// use optimizer::sampler::gp::GpSamplerBuilder;
///
/// let sampler = GpSamplerBuilder::new()
/// .n_startup_trials(15)
/// .n_candidates(2000)
/// .noise_variance(1e-4)
/// .seed(42)
/// .build();
/// ```
#[derive(Debug, Clone, Default)]
pub struct GpSamplerBuilder {
n_startup_trials: Option<usize>,
n_candidates: Option<usize>,
noise_variance: Option<f64>,
seed: Option<u64>,
}
impl GpSamplerBuilder {
/// Creates a new builder with default settings.
#[must_use]
pub fn new() -> Self {
Self::default()
}
/// Sets the number of random trials before GP-guided sampling begins.
///
/// Default: 10.
#[must_use]
pub fn n_startup_trials(mut self, n: usize) -> Self {
self.n_startup_trials = Some(n);
self
}
/// Sets the number of random candidate points for acquisition optimization.
///
/// More candidates improve the quality of the acquisition maximum
/// at the cost of more GP predictions per trial.
///
/// Default: 1000.
#[must_use]
pub fn n_candidates(mut self, n: usize) -> Self {
self.n_candidates = Some(n);
self
}
/// Sets the observation noise variance added to the kernel diagonal.
///
/// Controls the assumed noise level. Larger values make the GP smoother.
///
/// Default: 1e-6 (near-noiseless).
#[must_use]
pub fn noise_variance(mut self, v: f64) -> Self {
self.noise_variance = Some(v);
self
}
/// Sets the random seed for reproducibility.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.seed = Some(seed);
self
}
/// Builds the configured [`GpSampler`].
#[must_use]
pub fn build(self) -> GpSampler {
GpSampler {
state: Mutex::new(GpState::new(
self.n_startup_trials,
self.n_candidates,
self.noise_variance,
self.seed,
)),
}
}
}
// ---------------------------------------------------------------------------
// Internal types
// ---------------------------------------------------------------------------
/// Default number of random startup trials before GP kicks in.
const DEFAULT_N_STARTUP: usize = 10;
/// Default number of candidate points for EI optimization.
const DEFAULT_N_CANDIDATES: usize = 1000;
/// Default observation noise variance.
const DEFAULT_NOISE_VAR: f64 = 1e-6;
/// Describes how a parameter dimension maps into the GP internal vector.
#[derive(Clone, Debug)]
struct DimensionInfo {
distribution: Distribution,
is_continuous: bool,
bounds: Option<(f64, f64)>,
}
/// Tracks per-trial sampling progress.
#[derive(Clone, Debug)]
struct TrialProgress {
/// The candidate values for each dimension.
values: Vec<ParamValue>,
/// Next dimension to return.
next_dim: usize,
}
/// Phase of the GP state machine.
enum GpPhase {
/// Discovering the search space (first trial).
Discovery,
/// Steady-state sampling.
Active,
}
/// A fitted GP model ready for predictions.
struct GpModel {
/// Cholesky factor L of K + σ²I.
cholesky: nalgebra::linalg::Cholesky<f64, nalgebra::Dyn>,
/// α = (K + σ²I)^{-1} y.
alpha: nalgebra::DVector<f64>,
/// Training inputs (each row is a data point, normalized to [0, 1]).
x_train: Vec<Vec<f64>>,
/// ARD lengthscales per dimension.
lengthscales: Vec<f64>,
/// Signal variance.
signal_var: f64,
/// Mean of original y values (for un-standardization, unused but kept for diagnostics).
_y_mean: f64,
/// Std dev of original y values (unused but kept for diagnostics).
_y_std: f64,
/// Best observed (standardized) y.
f_best: f64,
}
/// Top-level mutable state behind the `Mutex`.
struct GpState {
rng: fastrand::Rng,
n_startup_trials: usize,
n_candidates: usize,
noise_variance: f64,
phase: GpPhase,
dimensions: Vec<DimensionInfo>,
trial_progress: HashMap<u64, TrialProgress>,
discovery_trial_id: Option<u64>,
}
impl GpState {
fn new(
n_startup: Option<usize>,
n_candidates: Option<usize>,
noise_var: Option<f64>,
seed: Option<u64>,
) -> Self {
let rng = seed.map_or_else(fastrand::Rng::new, fastrand::Rng::with_seed);
Self {
rng,
n_startup_trials: n_startup.unwrap_or(DEFAULT_N_STARTUP),
n_candidates: n_candidates.unwrap_or(DEFAULT_N_CANDIDATES),
noise_variance: noise_var.unwrap_or(DEFAULT_NOISE_VAR),
phase: GpPhase::Discovery,
dimensions: Vec::new(),
trial_progress: HashMap::new(),
discovery_trial_id: None,
}
}
}
// ---------------------------------------------------------------------------
// Matérn 5/2 kernel
// ---------------------------------------------------------------------------
/// Matérn 5/2 kernel with ARD lengthscales.
///
/// `k(x1, x2) = σ² (1 + √5 r + 5/3 r²) exp(-√5 r)`
/// where `r = sqrt(Σ ((x1_i - x2_i) / l_i)²)`
fn matern52(x1: &[f64], x2: &[f64], lengthscales: &[f64], signal_var: f64) -> f64 {
let mut r_sq = 0.0;
for i in 0..x1.len() {
let diff = (x1[i] - x2[i]) / lengthscales[i];
r_sq += diff * diff;
}
let r = r_sq.sqrt();
let sqrt5_r = SQRT_5 * r;
signal_var * (1.0 + sqrt5_r + 5.0 / 3.0 * r_sq) * (-sqrt5_r).exp()
}
/// Build the kernel matrix `K + σ²I`.
fn kernel_matrix(
x: &[Vec<f64>],
lengthscales: &[f64],
signal_var: f64,
noise_var: f64,
) -> DMatrix<f64> {
let n = x.len();
DMatrix::from_fn(n, n, |i, j| {
let k = matern52(&x[i], &x[j], lengthscales, signal_var);
if i == j { k + noise_var } else { k }
})
}
/// Compute the kernel vector k(x*, X) for a test point.
fn kernel_vector(
x_star: &[f64],
x_train: &[Vec<f64>],
lengthscales: &[f64],
signal_var: f64,
) -> nalgebra::DVector<f64> {
nalgebra::DVector::from_fn(x_train.len(), |i, _| {
matern52(x_star, &x_train[i], lengthscales, signal_var)
})
}
/// Precomputed √5 constant.
const SQRT_5: f64 = 2.236_213_562_373_095;
// ---------------------------------------------------------------------------
// GP fitting and prediction
// ---------------------------------------------------------------------------
/// Fit a GP model to the training data.
///
/// Returns `None` if fitting fails (e.g. Cholesky decomposition failure).
#[allow(clippy::cast_precision_loss)]
fn fit_gp(x_train: &[Vec<f64>], y_train: &[f64], noise_var: f64) -> Option<GpModel> {
let n = y_train.len();
if n == 0 {
return None;
}
// Standardize y
let y_mean = y_train.iter().sum::<f64>() / n as f64;
let y_var = if n > 1 {
y_train.iter().map(|&y| (y - y_mean).powi(2)).sum::<f64>() / (n - 1) as f64
} else {
1.0
};
let y_std = y_var.sqrt().max(1e-10);
let y_standardized: Vec<f64> = y_train.iter().map(|&y| (y - y_mean) / y_std).collect();
let f_best = y_standardized.iter().copied().fold(f64::INFINITY, f64::min);
// ARD lengthscales: per-dimension std dev of training X, clamped
let d = if x_train.is_empty() {
0
} else {
x_train[0].len()
};
let lengthscales: Vec<f64> = (0..d)
.map(|j| {
let vals: Vec<f64> = x_train.iter().map(|x| x[j]).collect();
let mean_j = vals.iter().sum::<f64>() / n as f64;
let var_j = vals.iter().map(|&v| (v - mean_j).powi(2)).sum::<f64>() / n as f64;
var_j.sqrt().max(0.01)
})
.collect();
// Signal variance = 1.0 (data is standardized)
let signal_var = 1.0;
let k = kernel_matrix(x_train, &lengthscales, signal_var, noise_var);
let cholesky = nalgebra::linalg::Cholesky::new(k)?;
// α = (K + σ²I)^{-1} y
let y_vec = nalgebra::DVector::from_column_slice(&y_standardized);
let alpha = cholesky.solve(&y_vec);
Some(GpModel {
cholesky,
alpha,
x_train: x_train.to_vec(),
lengthscales,
signal_var,
_y_mean: y_mean,
_y_std: y_std,
f_best,
})
}
/// Predict mean and standard deviation at a test point.
fn predict(model: &GpModel, x: &[f64]) -> (f64, f64) {
let k_star = kernel_vector(x, &model.x_train, &model.lengthscales, model.signal_var);
// Mean: k*^T α
let mean = k_star.dot(&model.alpha);
// Variance: k(x*, x*) - k*^T (K + σ²I)^{-1} k*
let k_self = model.signal_var;
let v = model.cholesky.solve(&k_star);
let var = (k_self - k_star.dot(&v)).max(0.0);
(mean, var.sqrt())
}
// ---------------------------------------------------------------------------
// Normal distribution helpers (Abramowitz-Stegun approximation)
// ---------------------------------------------------------------------------
/// Standard normal PDF.
fn norm_pdf(x: f64) -> f64 {
const INV_SQRT_2PI: f64 = 0.398_942_280_401_432_7;
INV_SQRT_2PI * (-0.5 * x * x).exp()
}
/// Standard normal CDF (Abramowitz-Stegun rational approximation).
fn norm_cdf(x: f64) -> f64 {
// Hart approximation (higher precision than basic A&S)
if x < -8.0 {
return 0.0;
}
if x > 8.0 {
return 1.0;
}
let abs_x = x.abs();
let t = 1.0 / (1.0 + 0.231_641_9 * abs_x);
let t2 = t * t;
let t3 = t2 * t;
let t4 = t3 * t;
let t5 = t4 * t;
let poly = 0.319_381_530 * t - 0.356_563_782 * t2 + 1.781_477_937 * t3 - 1.821_255_978 * t4
+ 1.330_274_429 * t5;
let pdf = norm_pdf(abs_x);
let cdf = 1.0 - pdf * poly;
if x >= 0.0 { cdf } else { 1.0 - cdf }
}
// ---------------------------------------------------------------------------
// Expected Improvement
// ---------------------------------------------------------------------------
/// Compute Expected Improvement at a point.
///
/// `EI(x) = (f_best - mean) Φ(z) + std φ(z)`
/// where `z = (f_best - mean) / std`
fn expected_improvement(mean: f64, std: f64, f_best: f64) -> f64 {
if std < 1e-12 {
return (f_best - mean).max(0.0);
}
let z = (f_best - mean) / std;
let improvement = (f_best - mean) * norm_cdf(z) + std * norm_pdf(z);
improvement.max(0.0)
}
// ---------------------------------------------------------------------------
// Acquisition optimization
// ---------------------------------------------------------------------------
/// Find the point in [0, 1]^d that maximizes EI via multi-start random search.
fn optimize_acquisition(
model: &GpModel,
n_dims: usize,
n_candidates: usize,
rng: &mut fastrand::Rng,
) -> Vec<f64> {
let mut best_ei = f64::NEG_INFINITY;
let mut best_x = vec![0.5; n_dims];
for _ in 0..n_candidates {
let x: Vec<f64> = (0..n_dims)
.map(|_| rng_util::f64_range(rng, 0.0, 1.0))
.collect();
let (mean, std) = predict(model, &x);
let ei = expected_improvement(mean, std, model.f_best);
if ei > best_ei {
best_ei = ei;
best_x = x;
}
}
best_x
}
// ---------------------------------------------------------------------------
// Data preprocessing helpers
// ---------------------------------------------------------------------------
/// Convert an internal-space value to normalized [0, 1] using bounds.
fn to_normalized(value: f64, lo: f64, hi: f64) -> f64 {
if (hi - lo).abs() < 1e-15 {
0.5
} else {
(value - lo) / (hi - lo)
}
}
/// Convert a normalized [0, 1] value back to internal space.
fn from_normalized(value: f64, lo: f64, hi: f64) -> f64 {
lo + value * (hi - lo)
}
// ---------------------------------------------------------------------------
// Extract training data from history
// ---------------------------------------------------------------------------
/// Maximum number of training points to use for the GP.
/// Caps computational cost at O(`MAX_TRAIN_POINTS`^3) per trial.
const MAX_TRAIN_POINTS: usize = 100;
/// Establish a deterministic mapping from dimension index to `ParamId`
/// using the first trial in history.
///
/// Matches dimensions to params by distribution equality, consuming
/// matched params to correctly handle duplicate distributions.
fn establish_param_mapping(
trial: &CompletedTrial,
dimensions: &[DimensionInfo],
) -> Vec<Option<crate::parameter::ParamId>> {
use crate::parameter::ParamId;
let mut available: Vec<(ParamId, &Distribution)> =
trial.distributions.iter().map(|(id, d)| (*id, d)).collect();
// Sort for deterministic matching order
available.sort_by_key(|(id, _)| *id);
let mut mapping = Vec::with_capacity(dimensions.len());
for dim in dimensions {
let pos = available.iter().position(|(_, d)| **d == dim.distribution);
if let Some(pos) = pos {
mapping.push(Some(available.remove(pos).0));
} else {
mapping.push(None);
}
}
mapping
}
/// Build normalized training data from completed trials.
///
/// Returns `(x_train, y_train)` where x values are normalized to [0, 1]
/// per dimension using the bounds from `dimensions`. Only continuous
/// dimensions are included. Uses at most [`MAX_TRAIN_POINTS`] most recent
/// trials.
#[allow(clippy::cast_precision_loss)]
fn build_training_data(
history: &[CompletedTrial],
dimensions: &[DimensionInfo],
) -> (Vec<Vec<f64>>, Vec<f64>) {
if history.is_empty() {
return (Vec::new(), Vec::new());
}
// Use only the most recent trials to cap GP fitting cost
let start = history.len().saturating_sub(MAX_TRAIN_POINTS);
let recent = &history[start..];
// Establish dimension → ParamId mapping from the first trial
let param_mapping = establish_param_mapping(&recent[0], dimensions);
let continuous_indices: Vec<usize> = dimensions
.iter()
.enumerate()
.filter(|(_, d)| d.is_continuous)
.map(|(i, _)| i)
.collect();
let mut x_train = Vec::with_capacity(recent.len());
let mut y_train = Vec::with_capacity(recent.len());
for trial in recent {
let mut x_row = Vec::with_capacity(continuous_indices.len());
let mut valid = true;
for &dim_idx in &continuous_indices {
let dim_info = &dimensions[dim_idx];
if let Some(param_id) = param_mapping[dim_idx] {
if let Some(param_val) = trial.params.get(&param_id) {
let internal = to_internal(param_val, &dim_info.distribution);
let (lo, hi) = dim_info.bounds.unwrap_or((0.0, 1.0));
x_row.push(to_normalized(internal, lo, hi));
} else {
valid = false;
break;
}
} else {
valid = false;
break;
}
}
if valid && x_row.len() == continuous_indices.len() {
x_train.push(x_row);
y_train.push(trial.value);
}
}
(x_train, y_train)
}
// ---------------------------------------------------------------------------
// Sampler trait implementation
// ---------------------------------------------------------------------------
impl Sampler for GpSampler {
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue {
let mut state = self.state.lock();
match &state.phase {
GpPhase::Discovery => sample_discovery(&mut state, distribution, trial_id),
GpPhase::Active => sample_active(&mut state, distribution, trial_id, history),
}
}
}
/// Handle sampling during the discovery phase.
fn sample_discovery(state: &mut GpState, distribution: &Distribution, trial_id: u64) -> ParamValue {
// A new trial_id means discovery is done
if let Some(prev_id) = state.discovery_trial_id
&& trial_id != prev_id
{
finalize_discovery(state);
return sample_active(state, distribution, trial_id, &[]);
}
state.discovery_trial_id = Some(trial_id);
let is_continuous = !matches!(distribution, Distribution::Categorical(_));
let bounds = internal_bounds(distribution);
state.dimensions.push(DimensionInfo {
distribution: distribution.clone(),
is_continuous,
bounds,
});
sample_random(&mut state.rng, distribution)
}
/// Finalize discovery and transition to the active phase.
fn finalize_discovery(state: &mut GpState) {
state.phase = GpPhase::Active;
state.trial_progress.clear();
}
/// Handle sampling during the active phase.
fn sample_active(
state: &mut GpState,
distribution: &Distribution,
trial_id: u64,
history: &[CompletedTrial],
) -> ParamValue {
// If this trial already has progress, return the next pre-computed value
if let Some(progress) = state.trial_progress.get_mut(&trial_id) {
let dim_idx = progress.next_dim;
progress.next_dim += 1;
if dim_idx < progress.values.len() {
return progress.values[dim_idx].clone();
}
// Extra dimension not seen during discovery
return sample_random(&mut state.rng, distribution);
}
// New trial: compute all dimension values at once
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
let use_gp = n_continuous > 0 && history.len() >= state.n_startup_trials;
let values = if use_gp {
compute_gp_candidate(state, history)
} else {
compute_random_candidate(state)
};
let first_value = values
.first()
.cloned()
.unwrap_or_else(|| sample_random(&mut state.rng, distribution));
state.trial_progress.insert(
trial_id,
TrialProgress {
values,
next_dim: 1,
},
);
first_value
}
/// Compute a candidate using the GP model.
fn compute_gp_candidate(state: &mut GpState, history: &[CompletedTrial]) -> Vec<ParamValue> {
let (x_train, y_train) = build_training_data(history, &state.dimensions);
// Try to fit GP; fall back to random if it fails
let model = fit_gp(&x_train, &y_train, state.noise_variance);
let n_continuous = state.dimensions.iter().filter(|d| d.is_continuous).count();
let normalized_candidate = if let Some(ref model) = model {
optimize_acquisition(model, n_continuous, state.n_candidates, &mut state.rng)
} else {
// GP fitting failed; use random
(0..n_continuous)
.map(|_| rng_util::f64_range(&mut state.rng, 0.0, 1.0))
.collect()
};
// Convert normalized candidate back to parameter values
let mut values = Vec::with_capacity(state.dimensions.len());
let mut ci = 0; // continuous dimension index
for dim in &state.dimensions {
if dim.is_continuous {
let (lo, hi) = dim.bounds.unwrap_or((0.0, 1.0));
let internal_val = from_normalized(normalized_candidate[ci], lo, hi);
values.push(from_internal(internal_val, &dim.distribution));
ci += 1;
} else {
values.push(sample_random(&mut state.rng, &dim.distribution));
}
}
values
}
/// Compute a random candidate for all dimensions.
fn compute_random_candidate(state: &mut GpState) -> Vec<ParamValue> {
state
.dimensions
.iter()
.map(|dim| sample_random(&mut state.rng, &dim.distribution))
.collect()
}
+91 -59
View File
@@ -1,7 +1,39 @@
//! Grid search sampler implementation.
//! Grid search sampler — exhaustive evaluation of discretized parameter spaces.
//!
//! `GridSearchSampler` performs exhaustive grid search over the parameter space,
//! systematically evaluating all combinations of discretized parameter values.
//! [`GridSampler`] divides each parameter range into a fixed number of
//! evenly spaced points (or uses the explicit step size when defined) and
//! evaluates them sequentially. This guarantees complete coverage of the
//! search grid at the cost of scaling exponentially with the number of
//! parameters.
//!
//! # When to use
//!
//! - **Small, discrete spaces** — when you have a handful of categorical or
//! integer parameters and want to evaluate every combination.
//! - **Reproducibility** — grid search is fully deterministic with no random
//! component.
//! - **Benchmarking** — compare grid search results against adaptive samplers
//! to measure their benefit.
//!
//! Avoid grid search for high-dimensional or large continuous spaces;
//! prefer [`TpeSampler`](super::tpe::TpeSampler) or
//! [`RandomSampler`](super::random::RandomSampler) instead.
//!
//! # Configuration
//!
//! | Option | Default | Description |
//! |---|---|---|
//! | `n_points_per_param` | 10 | Points per continuous parameter (ignored when `step` is set) |
//!
//! # Example
//!
//! ```
//! use optimizer::prelude::*;
//! use optimizer::sampler::grid::GridSampler;
//!
//! let sampler = GridSampler::builder().n_points_per_param(5).build();
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! ```
use std::collections::HashMap;
@@ -295,48 +327,48 @@ struct GridState {
grids: HashMap<String, CachedGrid>,
}
/// A grid search sampler that exhaustively evaluates all grid points.
/// Exhaustive grid search sampler.
///
/// `GridSearchSampler` divides the parameter space into a grid and systematically
/// samples each point. This is useful when you want to evaluate all combinations
/// of parameter values, especially for discrete or small parameter spaces.
/// Divide each parameter range into evenly spaced points and evaluate them
/// sequentially. When a parameter has an explicit `step` size, the step grid
/// is used instead of auto-discretization.
///
/// # Grid Exhaustion
/// Grid state is tracked **per distribution key** (bounds + step + log-scale).
/// Parameters with identical distributions share the same grid counter, so
/// use distinct ranges when multiple parameters span the same domain.
///
/// The sampler tracks its position in the grid for each distribution independently.
/// When all grid points for a distribution have been sampled, subsequent calls to
/// `sample()` for that distribution will **panic** with the message:
/// `"GridSearchSampler: all grid points exhausted"`.
/// # Grid exhaustion
///
/// To avoid panics, use [`is_exhausted()`](Self::is_exhausted) to check if all
/// points have been sampled before calling `sample()`. You can also use
/// [`grid_size()`](Self::grid_size) to determine the total number of grid points
/// that will be sampled.
/// When all grid points for a distribution have been sampled, the next
/// `sample()` call for that distribution **panics**. Use
/// [`is_exhausted()`](Self::is_exhausted) to check before sampling, or
/// set `n_points_per_param` high enough to cover the planned number of
/// trials.
///
/// # Thread Safety
/// # Thread safety
///
/// `GridSearchSampler` is thread-safe (`Send + Sync`) and uses internal locking
/// to ensure safe concurrent access to grid state.
/// `GridSearchSampler` is `Send + Sync` and uses internal locking for
/// safe concurrent access.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::grid::GridSearchSampler;
/// use optimizer::sampler::grid::GridSampler;
///
/// // Create with default settings (10 points per parameter)
/// let sampler = GridSearchSampler::new();
/// // Default: 10 points per parameter
/// let sampler = GridSampler::new();
///
/// // Create with custom settings using the builder
/// let sampler = GridSearchSampler::builder().n_points_per_param(20).build();
/// // Custom grid density
/// let sampler = GridSampler::builder().n_points_per_param(20).build();
/// ```
pub struct GridSearchSampler {
pub struct GridSampler {
/// Number of grid points per parameter (used when auto-discretizing).
n_points_per_param: usize,
/// Thread-safe internal state for tracking grid positions.
state: Mutex<GridState>,
}
impl GridSearchSampler {
impl GridSampler {
/// Creates a new grid search sampler with default settings.
///
/// Default settings:
@@ -354,9 +386,9 @@ impl GridSearchSampler {
/// # Examples
///
/// ```
/// use optimizer::sampler::grid::GridSearchSampler;
/// use optimizer::sampler::grid::GridSampler;
///
/// let sampler = GridSearchSampler::builder().n_points_per_param(20).build();
/// let sampler = GridSampler::builder().n_points_per_param(20).build();
/// ```
#[must_use]
pub fn builder() -> GridSearchSamplerBuilder {
@@ -364,13 +396,13 @@ impl GridSearchSampler {
}
}
impl Default for GridSearchSampler {
impl Default for GridSampler {
fn default() -> Self {
Self::new()
}
}
impl GridSearchSampler {
impl GridSampler {
/// Returns `true` if all grid points for all tracked distributions have been sampled.
///
/// A distribution is considered exhausted when its `current_index` equals the number
@@ -384,9 +416,9 @@ impl GridSearchSampler {
/// # Examples
///
/// ```
/// use optimizer::sampler::grid::GridSearchSampler;
/// use optimizer::sampler::grid::GridSampler;
///
/// let sampler = GridSearchSampler::new();
/// let sampler = GridSampler::new();
/// // Initially exhausted (no distributions tracked yet)
/// assert!(sampler.is_exhausted());
/// ```
@@ -414,9 +446,9 @@ impl GridSearchSampler {
/// # Examples
///
/// ```
/// use optimizer::sampler::grid::GridSearchSampler;
/// use optimizer::sampler::grid::GridSampler;
///
/// let sampler = GridSearchSampler::new();
/// let sampler = GridSampler::new();
/// // No distributions tracked yet
/// assert_eq!(sampler.grid_size(), 0);
/// ```
@@ -427,7 +459,7 @@ impl GridSearchSampler {
}
}
/// Builder for configuring a [`GridSearchSampler`].
/// Builder for configuring a [`GridSampler`].
///
/// # Examples
///
@@ -479,7 +511,7 @@ impl GridSearchSamplerBuilder {
self
}
/// Builds the configured [`GridSearchSampler`].
/// Builds the configured [`GridSampler`].
///
/// # Examples
///
@@ -491,8 +523,8 @@ impl GridSearchSamplerBuilder {
/// .build();
/// ```
#[must_use]
pub fn build(self) -> GridSearchSampler {
GridSearchSampler {
pub fn build(self) -> GridSampler {
GridSampler {
n_points_per_param: self.n_points_per_param,
state: Mutex::new(GridState::default()),
}
@@ -534,7 +566,7 @@ fn distribution_key(dist: &Distribution) -> String {
}
}
impl Sampler for GridSearchSampler {
impl Sampler for GridSampler {
fn sample(
&self,
distribution: &Distribution,
@@ -569,7 +601,7 @@ impl Sampler for GridSearchSampler {
// Check if all points have been exhausted
assert!(
cached.current_index < cached.points.len(),
"GridSearchSampler: all grid points exhausted"
"GridSampler: all grid points exhausted"
);
// Get the current grid point and advance the index
@@ -844,7 +876,7 @@ mod tests {
#[test]
fn test_sampler_exhausts_after_expected_samples() {
let sampler = GridSearchSampler::new();
let sampler = GridSampler::new();
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 3 });
// Sample all 3 points
@@ -858,7 +890,7 @@ mod tests {
#[test]
fn test_sampler_exhaustion_with_int_distribution() {
let sampler = GridSearchSampler::builder().n_points_per_param(5).build();
let sampler = GridSampler::builder().n_points_per_param(5).build();
let dist = Distribution::Int(IntDistribution {
low: 0,
high: 100,
@@ -876,9 +908,9 @@ mod tests {
}
#[test]
#[should_panic(expected = "GridSearchSampler: all grid points exhausted")]
#[should_panic(expected = "GridSampler: all grid points exhausted")]
fn test_sampler_panics_after_exhaustion() {
let sampler = GridSearchSampler::new();
let sampler = GridSampler::new();
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 2 });
// Sample all 2 points
@@ -893,14 +925,14 @@ mod tests {
#[test]
fn test_is_exhausted_before_sampling() {
let sampler = GridSearchSampler::new();
let sampler = GridSampler::new();
// Newly created sampler is vacuously exhausted (no distributions tracked)
assert!(sampler.is_exhausted());
}
#[test]
fn test_is_exhausted_during_sampling() {
let sampler = GridSearchSampler::new();
let sampler = GridSampler::new();
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 3 });
// After first sample, not exhausted
@@ -918,7 +950,7 @@ mod tests {
#[test]
fn test_is_exhausted_multiple_distributions() {
let sampler = GridSearchSampler::new();
let sampler = GridSampler::new();
// Use different n_choices so they have different distribution keys
let dist1 = Distribution::Categorical(CategoricalDistribution { n_choices: 2 });
let dist2 = Distribution::Categorical(CategoricalDistribution { n_choices: 3 });
@@ -944,7 +976,7 @@ mod tests {
#[test]
fn test_builder_default() {
let sampler = GridSearchSampler::builder().build();
let sampler = GridSampler::builder().build();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -961,7 +993,7 @@ mod tests {
#[test]
fn test_builder_custom_n_points() {
let sampler = GridSearchSampler::builder().n_points_per_param(3).build();
let sampler = GridSampler::builder().n_points_per_param(3).build();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -979,7 +1011,7 @@ mod tests {
#[test]
fn test_new_default() {
let sampler = GridSearchSampler::new();
let sampler = GridSampler::new();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -999,8 +1031,8 @@ mod tests {
#[test]
fn test_reproducibility_same_grid_order() {
// Two samplers with the same configuration should produce the same grid order
let sampler1 = GridSearchSampler::builder().n_points_per_param(5).build();
let sampler2 = GridSearchSampler::builder().n_points_per_param(5).build();
let sampler1 = GridSampler::builder().n_points_per_param(5).build();
let sampler2 = GridSampler::builder().n_points_per_param(5).build();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
@@ -1019,8 +1051,8 @@ mod tests {
#[test]
fn test_reproducibility_int_distribution() {
let sampler1 = GridSearchSampler::new();
let sampler2 = GridSearchSampler::new();
let sampler1 = GridSampler::new();
let sampler2 = GridSampler::new();
let dist = Distribution::Int(IntDistribution {
low: 0,
@@ -1041,8 +1073,8 @@ mod tests {
#[test]
fn test_reproducibility_categorical() {
let sampler1 = GridSearchSampler::new();
let sampler2 = GridSearchSampler::new();
let sampler1 = GridSampler::new();
let sampler2 = GridSampler::new();
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 4 });
@@ -1059,13 +1091,13 @@ mod tests {
#[test]
fn test_grid_size_empty() {
let sampler = GridSearchSampler::new();
let sampler = GridSampler::new();
assert_eq!(sampler.grid_size(), 0);
}
#[test]
fn test_grid_size_single_distribution() {
let sampler = GridSearchSampler::builder().n_points_per_param(5).build();
let sampler = GridSampler::builder().n_points_per_param(5).build();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
@@ -1083,7 +1115,7 @@ mod tests {
#[test]
fn test_grid_size_multiple_distributions() {
let sampler = GridSearchSampler::builder().n_points_per_param(3).build();
let sampler = GridSampler::builder().n_points_per_param(3).build();
let dist1 = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
+444 -15
View File
@@ -1,13 +1,221 @@
//! 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")]
pub mod cma_es;
pub(crate) mod common;
pub mod de;
pub(crate) mod genetic;
#[cfg(feature = "gp")]
pub mod gp;
pub mod grid;
pub mod moead;
pub mod motpe;
pub mod nsga2;
pub mod nsga3;
pub mod random;
#[cfg(feature = "sobol")]
pub mod sobol;
pub mod tpe;
use std::collections::HashMap;
pub use bohb::BohbSampler;
#[cfg(feature = "cma-es")]
pub use cma_es::CmaEsSampler;
pub use de::{DESampler, DEStrategy};
#[cfg(feature = "gp")]
pub use gp::GpSampler;
pub use grid::GridSampler;
pub use moead::{Decomposition, MoeadSampler};
pub use motpe::MotpeSampler;
pub use nsga2::Nsga2Sampler;
pub use nsga3::Nsga3Sampler;
pub use random::RandomSampler;
#[cfg(feature = "sobol")]
pub use sobol::SobolSampler;
pub use tpe::TpeSampler;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::parameter::{ParamId, Parameter};
use crate::trial::AttrValue;
use crate::types::TrialState;
/// A completed trial with its parameters, distributions, and objective value.
///
@@ -15,30 +223,229 @@ use crate::param::ParamValue;
/// parameter values, their distributions, and the objective value returned
/// by the objective function.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::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 sampled parameter values, keyed by parameter id.
pub params: HashMap<ParamId, ParamValue>,
/// The parameter distributions used, keyed by parameter id.
pub distributions: HashMap<ParamId, Distribution>,
/// Human-readable labels for parameters, keyed by parameter id.
pub param_labels: HashMap<ParamId, String>,
/// The objective value returned by the objective function.
pub value: V,
/// Intermediate objective values reported during the trial.
pub intermediate_values: Vec<(u64, f64)>,
/// The state of the trial (Complete, Pruned, or Failed).
pub state: TrialState,
/// User-defined attributes stored during the trial.
pub user_attrs: HashMap<String, AttrValue>,
/// Constraint values for this trial (<=0.0 means feasible).
#[cfg_attr(feature = "serde", serde(default))]
pub constraints: Vec<f64>,
}
impl<V> CompletedTrial<V> {
/// Creates a new completed trial.
pub fn new(
id: u64,
params: HashMap<String, ParamValue>,
distributions: HashMap<String, Distribution>,
params: HashMap<ParamId, ParamValue>,
distributions: HashMap<ParamId, Distribution>,
param_labels: HashMap<ParamId, String>,
value: V,
) -> Self {
Self {
id,
params,
distributions,
param_labels,
value,
intermediate_values: Vec::new(),
state: TrialState::Complete,
user_attrs: HashMap::new(),
constraints: Vec::new(),
}
}
/// Creates a new completed trial with intermediate values and user attributes.
pub fn with_intermediate_values(
id: u64,
params: HashMap<ParamId, ParamValue>,
distributions: HashMap<ParamId, Distribution>,
param_labels: HashMap<ParamId, String>,
value: V,
intermediate_values: Vec<(u64, f64)>,
user_attrs: HashMap<String, AttrValue>,
) -> Self {
Self {
id,
params,
distributions,
param_labels,
value,
intermediate_values,
state: TrialState::Complete,
user_attrs,
constraints: Vec::new(),
}
}
/// Returns the typed value for the given parameter.
///
/// Looks up the parameter by its unique id and casts the stored
/// [`ParamValue`] to the parameter's typed value.
///
/// Returns `None` if the parameter was not used in this trial or if
/// the stored value is incompatible with the parameter type (e.g., a
/// `Float` value stored for an `IntParam`).
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize(5, |trial: &mut optimizer::Trial| {
/// let val = x.suggest(trial)?;
/// Ok::<_, optimizer::Error>(val * val)
/// })
/// .unwrap();
///
/// let best = study.best_trial().unwrap();
/// let x_val: f64 = best.get(&x).unwrap();
/// assert!((-10.0..=10.0).contains(&x_val));
/// ```
pub fn get<P: Parameter>(&self, param: &P) -> Option<P::Value> {
self.params
.get(&param.id())
.and_then(|v| param.cast_param_value(v).ok())
}
/// Returns `true` if all constraints are satisfied (values <= 0.0).
///
/// A trial with no constraints is considered feasible.
#[must_use]
pub fn is_feasible(&self) -> bool {
self.constraints.iter().all(|&c| c <= 0.0)
}
/// Gets a user attribute by key.
#[must_use]
pub fn user_attr(&self, key: &str) -> Option<&AttrValue> {
self.user_attrs.get(key)
}
/// Returns all user attributes.
#[must_use]
pub fn user_attrs(&self) -> &HashMap<String, AttrValue> {
&self.user_attrs
}
/// Validates that all floating-point fields are finite (not NaN or
/// Infinity).
///
/// Checks distribution bounds, parameter values, constraints, and
/// intermediate values. Returns a description of the first invalid
/// field found, or `Ok(())` if everything is valid.
///
/// # Errors
///
/// Returns a `String` describing the first non-finite value found.
pub fn validate(&self) -> core::result::Result<(), String> {
for (id, dist) in &self.distributions {
if let Distribution::Float(fd) = dist {
if !fd.low.is_finite() {
return Err(format!(
"trial {}: float distribution for param {id} has non-finite low bound ({})",
self.id, fd.low
));
}
if !fd.high.is_finite() {
return Err(format!(
"trial {}: float distribution for param {id} has non-finite high bound ({})",
self.id, fd.high
));
}
if let Some(step) = fd.step
&& !step.is_finite()
{
return Err(format!(
"trial {}: float distribution for param {id} has non-finite step ({step})",
self.id
));
}
}
}
for (id, pv) in &self.params {
if let ParamValue::Float(v) = pv
&& !v.is_finite()
{
return Err(format!(
"trial {}: param {id} has non-finite float value ({v})",
self.id
));
}
}
for (i, &c) in self.constraints.iter().enumerate() {
if !c.is_finite() {
return Err(format!(
"trial {}: constraint[{i}] is non-finite ({c})",
self.id
));
}
}
for &(step, v) in &self.intermediate_values {
if !v.is_finite() {
return Err(format!(
"trial {}: intermediate value at step {step} is non-finite ({v})",
self.id
));
}
}
Ok(())
}
}
/// A pending (running) trial with its parameters and distributions, but no objective value yet.
///
/// This struct represents a trial that has been started and has sampled parameters,
/// but is still running and hasn't returned an objective value. It is used with the
/// constant liar strategy for parallel optimization.
#[derive(Clone, Debug)]
pub struct PendingTrial {
/// The unique identifier for this trial.
pub id: u64,
/// The sampled parameter values, keyed by parameter id.
pub params: HashMap<ParamId, ParamValue>,
/// The parameter distributions used, keyed by parameter id.
pub distributions: HashMap<ParamId, Distribution>,
/// Human-readable labels for parameters, keyed by parameter id.
pub param_labels: HashMap<ParamId, String>,
}
impl PendingTrial {
/// Creates a new pending trial.
#[must_use]
pub fn new(
id: u64,
params: HashMap<ParamId, ParamValue>,
distributions: HashMap<ParamId, Distribution>,
param_labels: HashMap<ParamId, String>,
) -> Self {
Self {
id,
params,
distributions,
param_labels,
}
}
}
@@ -49,27 +456,49 @@ impl<V> CompletedTrial<V> {
/// the distribution and historical trial data. The trait requires
/// `Send + Sync` to support concurrent and async optimization.
///
/// # Examples
/// # Implementing a custom sampler
///
/// Implementing a custom sampler:
/// ```
/// use optimizer::sampler::{Sampler, CompletedTrial};
/// use optimizer::distribution::Distribution;
/// use optimizer::param::ParamValue;
///
/// ```ignore
/// use optimizer::{Sampler, ParamValue, Distribution, CompletedTrial};
/// struct NoisySampler {
/// noise_scale: f64,
/// seed: u64,
/// }
///
/// struct MySampler;
///
/// impl Sampler for MySampler {
/// impl Sampler for NoisySampler {
/// fn sample(
/// &self,
/// distribution: &Distribution,
/// trial_id: u64,
/// history: &[CompletedTrial],
/// ) -> ParamValue {
/// // Custom sampling logic here
/// todo!()
/// // 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.
///
+672
View File
@@ -0,0 +1,672 @@
//! MOEA/D (Multi-Objective Evolutionary Algorithm based on Decomposition) sampler.
//!
//! MOEA/D takes a fundamentally different approach from Pareto-based
//! algorithms like NSGA-II/III. It **decomposes** the multi-objective
//! problem into a set of scalar subproblems using evenly distributed
//! weight vectors (Das-Dennis points), then solves them collaboratively
//! through **neighborhood-based mating and replacement**.
//!
//! # Algorithm
//!
//! 1. **Decompose** — generate weight vectors on the unit simplex and
//! assign one scalar subproblem per weight vector.
//! 2. **Build neighborhoods** — for each subproblem, find its T nearest
//! neighbors by Euclidean distance between weight vectors.
//! 3. **Mate from neighborhood** — select parents from the neighborhood
//! of each subproblem and produce offspring via SBX crossover +
//! polynomial mutation.
//! 4. **Scalarize and update** — evaluate offspring using a scalarization
//! function and update neighboring subproblems if the offspring improves
//! their scalar value.
//! 5. **Update ideal point** — track the best value seen per objective.
//!
//! # Scalarization methods
//!
//! | Method | Formula | Best for |
//! |--------|---------|----------|
//! | [`Tchebycheff`](Decomposition::Tchebycheff) (default) | `max(wᵢ * \|fᵢ - zᵢ*\|)` | General purpose, handles non-convex fronts |
//! | [`WeightedSum`](Decomposition::WeightedSum) | `Σ(wᵢ * fᵢ)` | Convex Pareto fronts only |
//! | [`Pbi`](Decomposition::Pbi) | `d₁ + θ * d₂` | Fine-grained convergence/diversity control |
//!
//! # When to use
//!
//! - Problems where you want **evenly distributed** solutions along the
//! Pareto front (one solution per weight direction).
//! - Many-objective optimization (3+ objectives) — scales well because
//! each subproblem is a simple scalar optimization.
//! - Problems with **non-convex** Pareto fronts (use Tchebycheff or PBI).
//! - When you need explicit control over the trade-off distribution via
//! weight vectors.
//!
//! For Pareto-based approaches, see
//! [`Nsga2Sampler`](super::nsga2::Nsga2Sampler) (crowding distance) or
//! [`Nsga3Sampler`](super::nsga3::Nsga3Sampler) (reference-point niching).
//!
//! # Configuration
//!
//! | Parameter | Builder method | Default |
//! |-----------|---------------|---------|
//! | Population size | [`population_size`](MoeadSamplerBuilder::population_size) | Number of Das-Dennis weight vectors |
//! | Neighborhood size (T) | [`neighborhood_size`](MoeadSamplerBuilder::neighborhood_size) | `min(20, pop_size)` |
//! | Decomposition method | [`decomposition`](MoeadSamplerBuilder::decomposition) | Tchebycheff |
//! | Crossover probability | [`crossover_prob`](MoeadSamplerBuilder::crossover_prob) | 1.0 |
//! | SBX distribution index | [`crossover_eta`](MoeadSamplerBuilder::crossover_eta) | 20.0 |
//! | Mutation distribution index | [`mutation_eta`](MoeadSamplerBuilder::mutation_eta) | 20.0 |
//! | Random seed | [`seed`](MoeadSamplerBuilder::seed) | random |
//!
//! # Examples
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::multi_objective::MultiObjectiveStudy;
//! use optimizer::parameter::{FloatParam, Parameter};
//! use optimizer::sampler::moead::MoeadSampler;
//!
//! let sampler = MoeadSampler::with_seed(42);
//! let study =
//! MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
//!
//! let x = FloatParam::new(0.0, 1.0);
//! study
//! .optimize(100, |trial: &mut optimizer::Trial| {
//! let xv = x.suggest(trial)?;
//! Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
//! })
//! .unwrap();
//! ```
use parking_lot::Mutex;
use super::genetic::{
self, Candidate, EvolutionaryState, Phase, advance_generation, auto_divisions,
collect_evaluated_generation, crossover, das_dennis, extract_trial_params,
generate_random_candidates, mutate, sample_from_candidate, sample_random,
};
use crate::distribution::Distribution;
use crate::multi_objective::MultiObjectiveTrial;
use crate::param::ParamValue;
use crate::types::Direction;
/// Decomposition (scalarization) method for [`MoeadSampler`].
///
/// Control how multi-objective values are reduced to a single scalar
/// for each subproblem. The default is [`Tchebycheff`](Self::Tchebycheff),
/// which handles both convex and non-convex Pareto fronts.
#[derive(Debug, Clone, Default)]
pub enum Decomposition {
/// Weighted sum: `Σ(wᵢ * fᵢ)`.
///
/// Simplest method but can only find solutions on convex regions
/// of the Pareto front.
WeightedSum,
/// Tchebycheff: `max(wᵢ * |fᵢ - zᵢ*|)`.
///
/// Handles non-convex Pareto fronts. The most commonly used
/// decomposition method (default).
#[default]
Tchebycheff,
/// Penalty-based Boundary Intersection: `d₁ + θ * d₂`.
///
/// Provides fine-grained control over the convergence/diversity
/// balance via the penalty parameter `theta`. Higher `theta`
/// favors solutions closer to the weight direction.
Pbi {
/// Penalty parameter controlling the balance between convergence
/// and diversity. Default: 5.0.
theta: f64,
},
}
/// MOEA/D sampler for multi-objective optimization.
///
/// Decompose a multi-objective problem into scalar subproblems using
/// weight vectors and solve them collaboratively via neighborhood-based
/// mating. Supports [`Tchebycheff`](Decomposition::Tchebycheff),
/// [`WeightedSum`](Decomposition::WeightedSum), and
/// [`Pbi`](Decomposition::Pbi) scalarization.
///
/// Create with [`MoeadSampler::new`], [`MoeadSampler::with_seed`], or
/// [`MoeadSampler::builder`] for full configuration.
pub struct MoeadSampler {
state: Mutex<MoeadState>,
}
impl MoeadSampler {
/// Creates a new MOEA/D sampler with a random seed.
#[must_use]
pub fn new() -> Self {
Self {
state: Mutex::new(MoeadState::new(MoeadConfig::default(), None)),
}
}
/// Creates a new MOEA/D sampler with a fixed seed.
#[must_use]
pub fn with_seed(seed: u64) -> Self {
Self {
state: Mutex::new(MoeadState::new(MoeadConfig::default(), Some(seed))),
}
}
/// Creates a builder for configuring a `MoeadSampler`.
#[must_use]
pub fn builder() -> MoeadSamplerBuilder {
MoeadSamplerBuilder::default()
}
}
impl Default for MoeadSampler {
fn default() -> Self {
Self::new()
}
}
/// Builder for [`MoeadSampler`].
#[derive(Debug, Clone, Default)]
pub struct MoeadSamplerBuilder {
population_size: Option<usize>,
neighborhood_size: Option<usize>,
decomposition: Decomposition,
crossover_prob: Option<f64>,
crossover_eta: Option<f64>,
mutation_eta: Option<f64>,
seed: Option<u64>,
}
impl MoeadSamplerBuilder {
/// Sets the population size. If unset, equals the number of
/// Das-Dennis weight vectors.
#[must_use]
pub fn population_size(mut self, size: usize) -> Self {
self.population_size = Some(size);
self
}
/// Sets the neighborhood size (T). Default: `min(20, pop_size)`.
#[must_use]
pub fn neighborhood_size(mut self, size: usize) -> Self {
self.neighborhood_size = Some(size);
self
}
/// Sets the decomposition method. Default: Tchebycheff.
#[must_use]
pub fn decomposition(mut self, decomp: Decomposition) -> Self {
self.decomposition = decomp;
self
}
/// Sets the crossover probability. Default: 1.0.
#[must_use]
pub fn crossover_prob(mut self, prob: f64) -> Self {
self.crossover_prob = Some(prob);
self
}
/// Sets the SBX distribution index. Default: 20.0.
#[must_use]
pub fn crossover_eta(mut self, eta: f64) -> Self {
self.crossover_eta = Some(eta);
self
}
/// Sets the polynomial mutation distribution index. Default: 20.0.
#[must_use]
pub fn mutation_eta(mut self, eta: f64) -> Self {
self.mutation_eta = Some(eta);
self
}
/// Sets the random seed for reproducibility.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.seed = Some(seed);
self
}
/// Builds the configured [`MoeadSampler`].
#[must_use]
pub fn build(self) -> MoeadSampler {
let config = MoeadConfig {
user_population_size: self.population_size,
neighborhood_size: self.neighborhood_size,
decomposition: self.decomposition,
crossover_prob: self.crossover_prob.unwrap_or(1.0),
crossover_eta: self.crossover_eta.unwrap_or(20.0),
mutation_eta: self.mutation_eta.unwrap_or(20.0),
};
MoeadSampler {
state: Mutex::new(MoeadState::new(config, self.seed)),
}
}
}
// ---------------------------------------------------------------------------
// Internal types
// ---------------------------------------------------------------------------
#[derive(Debug, Clone)]
struct MoeadConfig {
user_population_size: Option<usize>,
neighborhood_size: Option<usize>,
decomposition: Decomposition,
crossover_prob: f64,
crossover_eta: f64,
mutation_eta: f64,
}
impl Default for MoeadConfig {
fn default() -> Self {
Self {
user_population_size: None,
neighborhood_size: None,
decomposition: Decomposition::default(),
crossover_prob: 1.0,
crossover_eta: 20.0,
mutation_eta: 20.0,
}
}
}
struct MoeadState {
evo: EvolutionaryState,
config: MoeadConfig,
/// Weight vectors (Das-Dennis), one per subproblem.
weight_vectors: Vec<Vec<f64>>,
/// Neighborhoods: for each subproblem, indices of T nearest weight vectors.
neighborhoods: Vec<Vec<usize>>,
/// Ideal point z* (best per-objective in minimize-space).
ideal_point: Vec<f64>,
/// Current population's objective values in minimize-space (one per subproblem).
population_values: Vec<Vec<f64>>,
/// Current population's parameter vectors (one per subproblem).
population_params: Vec<Vec<ParamValue>>,
/// Whether the MOEA/D state has been initialized.
initialized: bool,
}
impl MoeadState {
fn new(config: MoeadConfig, seed: Option<u64>) -> Self {
Self {
evo: EvolutionaryState::new(seed),
config,
weight_vectors: Vec::new(),
neighborhoods: Vec::new(),
ideal_point: Vec::new(),
population_values: Vec::new(),
population_params: Vec::new(),
initialized: false,
}
}
}
// ---------------------------------------------------------------------------
// MultiObjectiveSampler implementation
// ---------------------------------------------------------------------------
impl crate::multi_objective::MultiObjectiveSampler for MoeadSampler {
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[MultiObjectiveTrial],
directions: &[Direction],
) -> ParamValue {
let mut state = self.state.lock();
match &state.evo.phase {
Phase::Discovery => {
if let Some(value) =
genetic::sample_discovery(&mut state.evo, distribution, trial_id)
{
return value;
}
// Transitioned to active phase
initialize_moead(&mut state, directions);
generate_random_candidates(&mut state.evo);
sample_from_candidate(&mut state.evo, trial_id)
}
Phase::Active => {
maybe_generate_new_generation(&mut state, history, directions);
sample_from_candidate(&mut state.evo, trial_id)
}
}
}
}
/// Initialize MOEA/D: weight vectors, neighborhoods, ideal point.
fn initialize_moead(state: &mut MoeadState, directions: &[Direction]) {
let n_obj = directions.len();
// Generate weight vectors
let divisions = auto_divisions(n_obj, state.config.user_population_size.unwrap_or(100));
state.weight_vectors = das_dennis(n_obj, divisions);
let pop_size = state
.config
.user_population_size
.unwrap_or(state.weight_vectors.len())
.max(4);
// Trim or pad weight vectors to match population size
state.weight_vectors.truncate(pop_size);
while state.weight_vectors.len() < pop_size {
// Duplicate random existing weight vectors
let idx = state.evo.rng.usize(0..state.weight_vectors.len());
let w = state.weight_vectors[idx].clone();
state.weight_vectors.push(w);
}
// Compute neighborhoods
let t = state
.config
.neighborhood_size
.unwrap_or_else(|| 20.min(pop_size));
let t = t.min(pop_size);
state.neighborhoods = compute_neighborhoods(&state.weight_vectors, t);
state.evo.population_size = pop_size;
state.evo.phase = Phase::Active;
state.ideal_point = vec![f64::INFINITY; n_obj];
state.initialized = true;
}
/// Compute T-nearest neighborhoods by Euclidean distance between weight vectors.
fn compute_neighborhoods(weights: &[Vec<f64>], t: usize) -> Vec<Vec<usize>> {
let n = weights.len();
weights
.iter()
.map(|wi| {
let mut distances: Vec<(usize, f64)> = (0..n)
.map(|j| {
let d: f64 = wi
.iter()
.zip(&weights[j])
.map(|(&a, &b)| (a - b).powi(2))
.sum::<f64>()
.sqrt();
(j, d)
})
.collect();
distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal));
distances.into_iter().take(t).map(|(idx, _)| idx).collect()
})
.collect()
}
/// Convert values to minimize-space.
fn to_minimize_space(values: &[f64], directions: &[Direction]) -> Vec<f64> {
values
.iter()
.zip(directions)
.map(|(&v, d)| match d {
Direction::Minimize => v,
Direction::Maximize => -v,
})
.collect()
}
fn maybe_generate_new_generation(
state: &mut MoeadState,
history: &[MultiObjectiveTrial],
directions: &[Direction],
) {
if state.evo.candidates.is_empty() {
generate_random_candidates(&mut state.evo);
return;
}
if let Some(evaluated) = collect_evaluated_generation(&state.evo, history) {
let offspring = moead_generate_offspring(state, &evaluated, directions);
advance_generation(&mut state.evo, offspring);
}
}
// ---------------------------------------------------------------------------
// Scalarization functions
// ---------------------------------------------------------------------------
/// Weighted sum scalarization: `sum(w_i * f_i)`.
fn scalarize_weighted_sum(values: &[f64], weight: &[f64]) -> f64 {
values.iter().zip(weight).map(|(&v, &w)| w * v).sum()
}
/// Tchebycheff scalarization: `max(w_i * |f_i - z_i*|)`.
fn scalarize_tchebycheff(values: &[f64], weight: &[f64], ideal: &[f64]) -> f64 {
values
.iter()
.zip(weight)
.zip(ideal)
.map(|((&v, &w), &z)| {
let w = if w < 1e-6 { 1e-6 } else { w };
w * (v - z).abs()
})
.fold(f64::NEG_INFINITY, f64::max)
}
/// PBI scalarization: `d1 + theta * d2`.
///
/// d1 = projection onto weight direction, d2 = perpendicular distance.
fn scalarize_pbi(values: &[f64], weight: &[f64], ideal: &[f64], theta: f64) -> f64 {
let n = values.len();
// Direction from ideal to the point
let diff: Vec<f64> = values.iter().zip(ideal).map(|(&v, &z)| v - z).collect();
// Normalize weight vector
let w_norm: f64 = weight.iter().map(|&w| w * w).sum::<f64>().sqrt();
if w_norm < 1e-30 {
return f64::INFINITY;
}
let w_unit: Vec<f64> = weight.iter().map(|&w| w / w_norm).collect();
// d1 = projection of diff onto weight direction
let d1: f64 = diff.iter().zip(&w_unit).map(|(&d, &w)| d * w).sum();
// d2 = perpendicular distance
let d2_sq: f64 = (0..n)
.map(|i| {
let proj = d1 * w_unit[i];
(diff[i] - proj).powi(2)
})
.sum::<f64>();
d1 + theta * d2_sq.sqrt()
}
/// Evaluate scalarization for a given decomposition method.
fn scalarize(values: &[f64], weight: &[f64], ideal: &[f64], decomposition: &Decomposition) -> f64 {
match decomposition {
Decomposition::WeightedSum => scalarize_weighted_sum(values, weight),
Decomposition::Tchebycheff => scalarize_tchebycheff(values, weight, ideal),
Decomposition::Pbi { theta } => scalarize_pbi(values, weight, ideal, *theta),
}
}
// ---------------------------------------------------------------------------
// MOEA/D generation algorithm
// ---------------------------------------------------------------------------
fn moead_generate_offspring(
state: &mut MoeadState,
population: &[&MultiObjectiveTrial],
directions: &[Direction],
) -> Vec<Candidate> {
let pop_size = state.evo.population_size;
if population.len() < 2 {
return (0..pop_size)
.map(|_| {
let params = state
.evo
.dimensions
.iter()
.map(|d| sample_random(&mut state.evo.rng, &d.distribution))
.collect();
Candidate { params }
})
.collect();
}
// Extract current population parameters and objective values
let current_params: Vec<Vec<ParamValue>> = population
.iter()
.map(|t| extract_trial_params(t, &state.evo.dimensions, &mut state.evo.rng))
.collect();
let current_values: Vec<Vec<f64>> = population
.iter()
.map(|t| to_minimize_space(&t.values, directions))
.collect();
// Update ideal point
for vals in &current_values {
for (i, &v) in vals.iter().enumerate() {
if i < state.ideal_point.len() && v < state.ideal_point[i] {
state.ideal_point[i] = v;
}
}
}
// Assign each solution to its best subproblem via scalarization
// and select the best solution for each subproblem as its representative
let n_weights = state.weight_vectors.len();
let mut best_for_subproblem: Vec<usize> = Vec::with_capacity(n_weights);
for j in 0..n_weights {
let mut best_idx = 0;
let mut best_val = f64::INFINITY;
for (k, vals) in current_values.iter().enumerate() {
let s = scalarize(
vals,
&state.weight_vectors[j],
&state.ideal_point,
&state.config.decomposition,
);
if s < best_val {
best_val = s;
best_idx = k;
}
}
best_for_subproblem.push(best_idx);
}
// Store current population state
state.population_values = current_values;
state.population_params = current_params;
// Generate offspring: for each subproblem, mate from neighborhood
let mut offspring = Vec::with_capacity(pop_size);
for i in 0..pop_size.min(state.neighborhoods.len()) {
let neighborhood = &state.neighborhoods[i];
// Pick two parents from the neighborhood using subproblem assignments
let n1 = neighborhood[state.evo.rng.usize(0..neighborhood.len())];
let n2 = neighborhood[state.evo.rng.usize(0..neighborhood.len())];
let p1_idx = best_for_subproblem[n1 % best_for_subproblem.len()];
let p2_idx = best_for_subproblem[n2 % best_for_subproblem.len()];
let p1 = &state.population_params[p1_idx];
let p2 = &state.population_params[p2_idx];
let (mut child1, _child2) = crossover(
&mut state.evo.rng,
p1,
p2,
&state.evo.dimensions,
state.config.crossover_prob,
state.config.crossover_eta,
);
mutate(
&mut state.evo.rng,
&mut child1,
&state.evo.dimensions,
state.config.mutation_eta,
);
offspring.push(Candidate { params: child1 });
}
// If pop_size > neighborhoods, fill remaining with random neighborhood crossover
while offspring.len() < pop_size {
let i = state.evo.rng.usize(0..state.neighborhoods.len());
let neighborhood = &state.neighborhoods[i];
let n1 = neighborhood[state.evo.rng.usize(0..neighborhood.len())];
let n2 = neighborhood[state.evo.rng.usize(0..neighborhood.len())];
let p1_idx = best_for_subproblem[n1 % best_for_subproblem.len()];
let p2_idx = best_for_subproblem[n2 % best_for_subproblem.len()];
let (mut child1, _) = crossover(
&mut state.evo.rng,
&state.population_params[p1_idx],
&state.population_params[p2_idx],
&state.evo.dimensions,
state.config.crossover_prob,
state.config.crossover_eta,
);
mutate(
&mut state.evo.rng,
&mut child1,
&state.evo.dimensions,
state.config.mutation_eta,
);
offspring.push(Candidate { params: child1 });
}
offspring
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_scalarize_weighted_sum() {
let values = [1.0, 2.0, 3.0];
let weight = [0.5, 0.3, 0.2];
let result = scalarize_weighted_sum(&values, &weight);
assert!((result - (0.5 + 0.6 + 0.6)).abs() < 1e-10);
}
#[test]
fn test_scalarize_tchebycheff() {
let values = [3.0, 2.0];
let weight = [0.5, 0.5];
let ideal = [1.0, 1.0];
let result = scalarize_tchebycheff(&values, &weight, &ideal);
// max(0.5 * |3-1|, 0.5 * |2-1|) = max(1.0, 0.5) = 1.0
assert!((result - 1.0).abs() < 1e-10);
}
#[test]
fn test_scalarize_pbi() {
let values = [2.0, 2.0];
let weight = [1.0, 1.0];
let ideal = [0.0, 0.0];
let result = scalarize_pbi(&values, &weight, &ideal, 5.0);
// d1 = projection of (2,2) onto (1/√2, 1/√2) = 2*√2
// d2 = 0 (point is on the weight direction)
let expected_d1 = 2.0 * (2.0_f64).sqrt();
assert!((result - expected_d1).abs() < 1e-10);
}
#[test]
fn test_compute_neighborhoods() {
let weights = vec![vec![1.0, 0.0], vec![0.5, 0.5], vec![0.0, 1.0]];
let neighborhoods = compute_neighborhoods(&weights, 2);
assert_eq!(neighborhoods.len(), 3);
// Each neighborhood should have 2 entries
for n in &neighborhoods {
assert_eq!(n.len(), 2);
}
// First weight [1,0] should be closest to itself and [0.5,0.5]
assert_eq!(neighborhoods[0][0], 0); // itself
assert_eq!(neighborhoods[0][1], 1); // nearest neighbor
}
}
+831
View File
@@ -0,0 +1,831 @@
//! Multi-Objective Tree-Parzen Estimator (MOTPE) sampler.
//!
//! MOTPE extends TPE to multi-objective optimization by replacing the gamma-based
//! split with Pareto non-dominated sorting. This lets the sampler propose
//! parameters that push the Pareto front forward across all objectives
//! simultaneously.
//!
//! # Algorithm
//!
//! In single-objective TPE, trials are sorted by value and split at a gamma
//! percentile into good/bad groups. MOTPE replaces this with:
//!
//! 1. Compute non-dominated sorting on all completed trials.
//! 2. Use the Pareto front (rank 0) as "good" trials.
//! 3. Use dominated trials (rank 1+) as "bad" trials.
//! 4. Build KDE l(x) from good, g(x) from bad.
//! 5. Sample candidates and score by l(x)/g(x).
//!
//! # When to use
//!
//! - You have 2+ objectives and want model-guided search (not pure evolutionary).
//! - Your objectives are relatively smooth and continuous.
//! - You want a Pareto-aware version of TPE without the overhead of full
//! population-based algorithms like NSGA-II or NSGA-III.
//!
//! For single-objective problems, use [`TpeSampler`](super::tpe::TpeSampler) instead.
//! For many-objective (3+) problems with reference-point decomposition, consider
//! NSGA-III or MOEA/D.
//!
//! # Configuration
//!
//! - `n_startup_trials` — number of random trials before MOTPE kicks in (default: 11)
//! - `n_ei_candidates` — candidates evaluated per sample (default: 24)
//! - `kde_bandwidth` — optional fixed KDE bandwidth; `None` uses Scott's rule
//! - `seed` — optional seed for reproducibility
//!
//! # Examples
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::multi_objective::MultiObjectiveStudy;
//! use optimizer::parameter::{FloatParam, Parameter};
//! use optimizer::sampler::motpe::MotpeSampler;
//!
//! let sampler = MotpeSampler::builder().seed(42).build();
//! let study =
//! MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
//!
//! let x = FloatParam::new(0.0, 1.0);
//! study
//! .optimize(30, |trial: &mut optimizer::Trial| {
//! let xv = x.suggest(trial)?;
//! Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
//! })
//! .unwrap();
//!
//! let front = study.pareto_front();
//! assert!(!front.is_empty());
//! ```
use core::sync::atomic::{AtomicU64, Ordering};
use crate::distribution::Distribution;
use crate::multi_objective::{MultiObjectiveSampler, MultiObjectiveTrial};
use crate::param::ParamValue;
use crate::sampler::common;
use crate::sampler::tpe::common as tpe_common;
use crate::types::{Direction, TrialState};
use crate::{pareto, rng_util};
/// Multi-Objective TPE (MOTPE) sampler for multi-objective Bayesian optimization.
///
/// Use Pareto non-dominated sorting to split completed trials into "good"
/// (non-dominated, rank 0) and "bad" (dominated) groups, then fit kernel
/// density estimators to each group and sample new points that maximize
/// l(x)/g(x).
///
/// During the startup phase (fewer than `n_startup_trials` completed),
/// MOTPE falls back to random sampling.
///
/// # When to use
///
/// Use `MotpeSampler` when optimizing 2+ objectives and you want a
/// model-guided sampler that adapts proposals based on the current
/// Pareto front. For single-objective problems, use
/// [`TpeSampler`](super::tpe::TpeSampler) instead.
///
/// # Examples
///
/// ```
/// use optimizer::Direction;
/// use optimizer::multi_objective::MultiObjectiveStudy;
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::motpe::MotpeSampler;
///
/// let sampler = MotpeSampler::builder()
/// .n_startup_trials(10)
/// .n_ei_candidates(24)
/// .seed(42)
/// .build();
///
/// let study =
/// MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
///
/// let x = FloatParam::new(0.0, 1.0);
/// study
/// .optimize(30, |trial: &mut optimizer::Trial| {
/// let xv = x.suggest(trial)?;
/// Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
/// })
/// .unwrap();
///
/// let front = study.pareto_front();
/// assert!(!front.is_empty());
/// ```
pub struct MotpeSampler {
/// Number of trials before MOTPE kicks in (uses random sampling before this).
n_startup_trials: usize,
/// Number of candidate samples to evaluate when selecting the next point.
n_ei_candidates: usize,
/// Optional fixed bandwidth for KDE. If None, uses Scott's rule.
kde_bandwidth: Option<f64>,
/// Base seed for deterministic per-call RNG derivation (no mutex needed).
seed: u64,
/// Monotonic counter to disambiguate calls with identical (`trial_id`, distribution).
call_seq: AtomicU64,
}
impl MotpeSampler {
/// Creates a new MOTPE sampler with default settings.
///
/// Defaults:
/// - `n_startup_trials`: 11
/// - `n_ei_candidates`: 24
/// - `kde_bandwidth`: None (Scott's rule)
#[must_use]
pub fn new() -> Self {
Self {
n_startup_trials: 11,
n_ei_candidates: 24,
kde_bandwidth: None,
seed: fastrand::u64(..),
call_seq: AtomicU64::new(0),
}
}
/// Creates a new MOTPE sampler with a fixed seed.
#[must_use]
pub fn with_seed(seed: u64) -> Self {
Self {
n_startup_trials: 11,
n_ei_candidates: 24,
kde_bandwidth: None,
seed,
call_seq: AtomicU64::new(0),
}
}
/// Creates a builder for configuring a MOTPE sampler.
#[must_use]
pub fn builder() -> MotpeSamplerBuilder {
MotpeSamplerBuilder::new()
}
/// Splits trials into good (non-dominated) and bad (dominated) groups
/// using Pareto non-dominated sorting.
fn split_trials<'a>(
history: &'a [MultiObjectiveTrial],
directions: &[Direction],
) -> (Vec<&'a MultiObjectiveTrial>, Vec<&'a MultiObjectiveTrial>) {
let complete: Vec<(usize, &MultiObjectiveTrial)> = history
.iter()
.enumerate()
.filter(|(_, t)| t.state == TrialState::Complete)
.collect();
if complete.is_empty() {
return (vec![], vec![]);
}
let values: Vec<Vec<f64>> = complete.iter().map(|(_, t)| t.values.clone()).collect();
let constraints: Vec<Vec<f64>> = complete
.iter()
.map(|(_, t)| t.constraints.clone())
.collect();
let has_constraints = constraints.iter().any(|c| !c.is_empty());
let fronts = if has_constraints {
pareto::fast_non_dominated_sort_constrained(&values, directions, &constraints)
} else {
pareto::fast_non_dominated_sort(&values, directions)
};
if fronts.is_empty() {
return (vec![], vec![]);
}
// Front 0 = good (non-dominated), everything else = bad
let good: Vec<&MultiObjectiveTrial> = fronts[0].iter().map(|&i| complete[i].1).collect();
let bad: Vec<&MultiObjectiveTrial> = fronts[1..]
.iter()
.flatten()
.map(|&i| complete[i].1)
.collect();
(good, bad)
}
}
impl Default for MotpeSampler {
fn default() -> Self {
Self::new()
}
}
impl MotpeSampler {
fn sample_float(
&self,
d: &crate::distribution::FloatDistribution,
good_trials: &[&MultiObjectiveTrial],
bad_trials: &[&MultiObjectiveTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let target_dist = Distribution::Float(d.clone());
let good_values: Vec<f64> = good_trials
.iter()
.filter_map(|t| {
t.distributions.iter().find_map(|(id, dist)| {
if *dist == target_dist {
t.params.get(id).and_then(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
} else {
None
}
})
})
.collect();
let bad_values: Vec<f64> = bad_trials
.iter()
.filter_map(|t| {
t.distributions.iter().find_map(|(id, dist)| {
if *dist == target_dist {
t.params.get(id).and_then(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
} else {
None
}
})
})
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return ParamValue::Float(rng_util::f64_range(rng, d.low, d.high));
}
let value = tpe_common::sample_tpe_float(
d,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
rng,
);
ParamValue::Float(value)
}
fn sample_int(
&self,
d: &crate::distribution::IntDistribution,
good_trials: &[&MultiObjectiveTrial],
bad_trials: &[&MultiObjectiveTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let target_dist = Distribution::Int(d.clone());
let good_values: Vec<i64> = good_trials
.iter()
.filter_map(|t| {
t.distributions.iter().find_map(|(id, dist)| {
if *dist == target_dist {
t.params.get(id).and_then(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
} else {
None
}
})
})
.collect();
let bad_values: Vec<i64> = bad_trials
.iter()
.filter_map(|t| {
t.distributions.iter().find_map(|(id, dist)| {
if *dist == target_dist {
t.params.get(id).and_then(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
} else {
None
}
})
})
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return common::sample_random(rng, &Distribution::Int(d.clone()));
}
let value = tpe_common::sample_tpe_int(
d,
good_values,
bad_values,
self.n_ei_candidates,
self.kde_bandwidth,
rng,
);
ParamValue::Int(value)
}
#[allow(clippy::unused_self)]
fn sample_categorical(
&self,
d: &crate::distribution::CategoricalDistribution,
good_trials: &[&MultiObjectiveTrial],
bad_trials: &[&MultiObjectiveTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
let target_dist = Distribution::Categorical(d.clone());
let good_indices: Vec<usize> = good_trials
.iter()
.filter_map(|t| {
t.distributions.iter().find_map(|(id, dist)| {
if *dist == target_dist {
t.params.get(id).and_then(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
} else {
None
}
})
})
.collect();
let bad_indices: Vec<usize> = bad_trials
.iter()
.filter_map(|t| {
t.distributions.iter().find_map(|(id, dist)| {
if *dist == target_dist {
t.params.get(id).and_then(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
} else {
None
}
})
})
.collect();
if good_indices.is_empty() || bad_indices.is_empty() {
return common::sample_random(rng, &Distribution::Categorical(d.clone()));
}
let index =
tpe_common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng);
ParamValue::Categorical(index)
}
}
impl MultiObjectiveSampler for MotpeSampler {
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[MultiObjectiveTrial],
directions: &[Direction],
) -> ParamValue {
let seq = self.call_seq.fetch_add(1, Ordering::Relaxed);
let mut rng = fastrand::Rng::with_seed(rng_util::mix_seed(
self.seed,
trial_id,
rng_util::distribution_fingerprint(distribution).wrapping_add(seq),
));
// Fall back to random sampling during startup phase
let n_complete = history
.iter()
.filter(|t| t.state == TrialState::Complete)
.count();
if n_complete < self.n_startup_trials {
return common::sample_random(&mut rng, distribution);
}
// Split trials into good (Pareto front) and bad (dominated)
let (good_trials, bad_trials) = Self::split_trials(history, directions);
if good_trials.is_empty() || bad_trials.is_empty() {
return common::sample_random(&mut rng, distribution);
}
match distribution {
Distribution::Float(d) => self.sample_float(d, &good_trials, &bad_trials, &mut rng),
Distribution::Int(d) => self.sample_int(d, &good_trials, &bad_trials, &mut rng),
Distribution::Categorical(d) => {
self.sample_categorical(d, &good_trials, &bad_trials, &mut rng)
}
}
}
}
/// Builder for configuring a [`MotpeSampler`].
///
/// # Defaults
///
/// - `n_startup_trials`: 11
/// - `n_ei_candidates`: 24
/// - `kde_bandwidth`: None (Scott's rule)
/// - `seed`: None (OS entropy)
///
/// # Examples
///
/// ```
/// use optimizer::sampler::motpe::MotpeSamplerBuilder;
///
/// let sampler = MotpeSamplerBuilder::new()
/// .n_startup_trials(15)
/// .n_ei_candidates(32)
/// .seed(42)
/// .build();
/// ```
#[derive(Debug, Clone)]
pub struct MotpeSamplerBuilder {
n_startup_trials: usize,
n_ei_candidates: usize,
kde_bandwidth: Option<f64>,
seed: Option<u64>,
}
impl MotpeSamplerBuilder {
/// Creates a new builder with default settings.
#[must_use]
pub fn new() -> Self {
Self {
n_startup_trials: 11,
n_ei_candidates: 24,
kde_bandwidth: None,
seed: None,
}
}
/// Sets the number of startup trials before MOTPE sampling begins.
#[must_use]
pub fn n_startup_trials(mut self, n: usize) -> Self {
self.n_startup_trials = n;
self
}
/// Sets the number of EI candidates to evaluate per sample.
#[must_use]
pub fn n_ei_candidates(mut self, n: usize) -> Self {
self.n_ei_candidates = n;
self
}
/// Sets a fixed bandwidth for the kernel density estimator.
///
/// By default, Scott's rule is used for automatic bandwidth selection.
#[must_use]
pub fn kde_bandwidth(mut self, bandwidth: f64) -> Self {
self.kde_bandwidth = Some(bandwidth);
self
}
/// Sets a seed for reproducible sampling.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.seed = Some(seed);
self
}
/// Builds the configured [`MotpeSampler`].
#[must_use]
pub fn build(self) -> MotpeSampler {
MotpeSampler {
n_startup_trials: self.n_startup_trials,
n_ei_candidates: self.n_ei_candidates,
kde_bandwidth: self.kde_bandwidth,
seed: self.seed.unwrap_or_else(|| fastrand::u64(..)),
call_seq: AtomicU64::new(0),
}
}
}
impl Default for MotpeSamplerBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
#[allow(
clippy::similar_names,
clippy::cast_sign_loss,
clippy::cast_precision_loss
)]
mod tests {
use std::collections::HashMap;
use super::*;
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
use crate::parameter::ParamId;
use crate::trial::Trial;
fn create_mo_trial(
id: u64,
values: Vec<f64>,
params: Vec<(ParamId, ParamValue, Distribution)>,
) -> MultiObjectiveTrial {
let mut param_map = HashMap::new();
let mut dist_map = HashMap::new();
let label_map = HashMap::new();
for (param_id, pv, dist) in params {
param_map.insert(param_id, pv);
dist_map.insert(param_id, dist);
}
MultiObjectiveTrial {
id,
params: param_map,
distributions: dist_map,
param_labels: label_map,
values,
state: TrialState::Complete,
user_attrs: HashMap::new(),
constraints: Vec::new(),
}
}
#[test]
fn test_motpe_startup_random_sampling() {
let sampler = MotpeSampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let directions = [Direction::Minimize, Direction::Minimize];
// With no history, should use random sampling
let history: Vec<MultiObjectiveTrial> = vec![];
for i in 0..50 {
let value = sampler.sample(&dist, i, &history, &directions);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v));
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn test_motpe_split_pareto() {
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let directions = [Direction::Minimize, Direction::Minimize];
let x_id = ParamId::new();
// Create trials: Pareto front = {(0.1, 0.9), (0.5, 0.5), (0.9, 0.1)}
// Dominated = {(0.6, 0.8), (0.8, 0.7)}
let history = vec![
create_mo_trial(
0,
vec![0.1, 0.9],
vec![(x_id, ParamValue::Float(0.1), dist.clone())],
),
create_mo_trial(
1,
vec![0.5, 0.5],
vec![(x_id, ParamValue::Float(0.5), dist.clone())],
),
create_mo_trial(
2,
vec![0.9, 0.1],
vec![(x_id, ParamValue::Float(0.9), dist.clone())],
),
create_mo_trial(
3,
vec![0.6, 0.8],
vec![(x_id, ParamValue::Float(0.6), dist.clone())],
),
create_mo_trial(
4,
vec![0.8, 0.7],
vec![(x_id, ParamValue::Float(0.8), dist.clone())],
),
];
let (good, bad) = MotpeSampler::split_trials(&history, &directions);
assert_eq!(good.len(), 3, "Pareto front should have 3 members");
assert_eq!(bad.len(), 2, "2 dominated trials");
}
#[test]
fn test_motpe_samples_float() {
let sampler = MotpeSampler::builder()
.n_startup_trials(5)
.n_ei_candidates(24)
.seed(42)
.build();
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let directions = [Direction::Minimize, Direction::Minimize];
let x_id = ParamId::new();
// Build history where values near 0.3 are on the Pareto front
let mut history = Vec::new();
for i in 0..20 {
let x = f64::from(i) / 20.0;
// Pareto front: f1 = (x - 0.3)^2, f2 = (x - 0.3)^2 + 0.1
// Best solutions cluster around x = 0.3
let f1 = (x - 0.3).powi(2);
let f2 = (x - 0.7).powi(2);
history.push(create_mo_trial(
i as u64,
vec![f1, f2],
vec![(x_id, ParamValue::Float(x), dist.clone())],
));
}
// MOTPE should produce values within [0, 1]
for i in 0..50 {
let value = sampler.sample(&dist, 100 + i, &history, &directions);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v), "Value {v} out of range");
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn test_motpe_int_sampling() {
let sampler = MotpeSampler::builder().n_startup_trials(5).seed(42).build();
let dist = Distribution::Int(IntDistribution {
low: 0,
high: 100,
log_scale: false,
step: None,
});
let directions = [Direction::Minimize, Direction::Minimize];
let x_id = ParamId::new();
let mut history = Vec::new();
for i in 0..20 {
let x = i * 5;
let f1 = ((x as f64) - 30.0).powi(2);
let f2 = ((x as f64) - 70.0).powi(2);
history.push(create_mo_trial(
i as u64,
vec![f1, f2],
vec![(x_id, ParamValue::Int(x), dist.clone())],
));
}
for i in 0..50 {
let value = sampler.sample(&dist, 100 + i, &history, &directions);
if let ParamValue::Int(v) = value {
assert!((0..=100).contains(&v), "Value {v} out of range");
} else {
panic!("Expected Int value");
}
}
}
#[test]
fn test_motpe_categorical_sampling() {
let sampler = MotpeSampler::builder().n_startup_trials(5).seed(42).build();
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 3 });
let directions = [Direction::Minimize, Direction::Minimize];
let cat_id = ParamId::new();
// Category 1 is on the Pareto front, others are dominated
let mut history = Vec::new();
for i in 0..15 {
let category = i % 3;
let (f1, f2) = match category {
0 => (0.8, 0.8), // dominated
1 => (0.1, 0.9), // Pareto front
2 => (0.9, 0.1), // Pareto front
_ => unreachable!(),
};
history.push(create_mo_trial(
i as u64,
vec![f1, f2],
vec![(
cat_id,
ParamValue::Categorical(category as usize),
dist.clone(),
)],
));
}
let mut counts = vec![0usize; 3];
for i in 0..200 {
let value = sampler.sample(&dist, 100 + i, &history, &directions);
if let ParamValue::Categorical(idx) = value {
assert!(idx < 3, "Category {idx} out of range");
counts[idx] += 1;
} else {
panic!("Expected Categorical value");
}
}
// Categories 1 and 2 (on Pareto front) should dominate category 0
assert!(
counts[1] + counts[2] > counts[0],
"Pareto-front categories should be sampled more: {counts:?}"
);
}
#[test]
fn test_motpe_reproducibility() {
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let directions = [Direction::Minimize, Direction::Minimize];
let x_id = ParamId::new();
let history: Vec<MultiObjectiveTrial> = (0..20)
.map(|i| {
let x = f64::from(i) / 20.0;
create_mo_trial(
i as u64,
vec![x, 1.0 - x],
vec![(x_id, ParamValue::Float(x), dist.clone())],
)
})
.collect();
let sampler1 = MotpeSampler::builder()
.seed(12345)
.n_startup_trials(5)
.build();
let sampler2 = MotpeSampler::builder()
.seed(12345)
.n_startup_trials(5)
.build();
for i in 0..10 {
let v1 = sampler1.sample(&dist, i, &history, &directions);
let v2 = sampler2.sample(&dist, i, &history, &directions);
assert_eq!(v1, v2, "Samples should be identical with same seed");
}
}
#[test]
fn test_motpe_with_study() {
use crate::multi_objective::MultiObjectiveStudy;
use crate::parameter::{FloatParam, Parameter};
let sampler = MotpeSampler::builder().seed(42).build();
let study = MultiObjectiveStudy::with_sampler(
vec![Direction::Minimize, Direction::Minimize],
sampler,
);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(30, |trial: &mut Trial| {
let xv = x.suggest(trial)?;
Ok::<_, crate::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty(), "Should have Pareto-optimal solutions");
// All front solutions should have values summing to ~1.0
for trial in &front {
let sum: f64 = trial.values.iter().sum();
assert!(
(sum - 1.0).abs() < 0.01,
"Pareto front values should sum to ~1.0, got {sum}"
);
}
}
#[test]
fn test_motpe_builder_defaults() {
let sampler = MotpeSamplerBuilder::new().build();
assert_eq!(sampler.n_startup_trials, 11);
assert_eq!(sampler.n_ei_candidates, 24);
assert!(sampler.kde_bandwidth.is_none());
}
#[test]
fn test_motpe_builder_custom() {
let sampler = MotpeSamplerBuilder::new()
.n_startup_trials(20)
.n_ei_candidates(48)
.kde_bandwidth(0.5)
.seed(99)
.build();
assert_eq!(sampler.n_startup_trials, 20);
assert_eq!(sampler.n_ei_candidates, 48);
assert_eq!(sampler.kde_bandwidth, Some(0.5));
}
}
+420
View File
@@ -0,0 +1,420 @@
//! NSGA-II (Non-dominated Sorting Genetic Algorithm II) sampler.
//!
//! NSGA-II is one of the most widely used evolutionary multi-objective
//! optimization algorithms. It ranks the population using **non-dominated
//! sorting** (fast O(MN²) algorithm) and breaks ties within the same
//! Pareto front using **crowding distance**, which favors solutions in
//! less-crowded regions of the objective space.
//!
//! # Algorithm
//!
//! Each generation proceeds as follows:
//!
//! 1. **Non-dominated sorting** — partition the combined parent+offspring
//! population into Pareto fronts F₁, F₂, …
//! 2. **Crowding distance** — for each front, compute per-solution crowding
//! distance (sum of normalized neighbor gaps in each objective).
//! 3. **Selection** — fill the next population front-by-front. When a front
//! only partially fits, prefer solutions with higher crowding distance.
//! 4. **Binary tournament** — select parents using (rank, crowding distance)
//! comparisons.
//! 5. **SBX crossover + polynomial mutation** — generate offspring.
//!
//! # When to use
//!
//! - Two-objective problems where you want a well-spread Pareto front.
//! - General-purpose multi-objective optimization with moderate population
//! sizes.
//! - Problems that benefit from diversity preservation via crowding distance.
//!
//! For problems with **three or more objectives**, consider
//! [`Nsga3Sampler`](super::nsga3::Nsga3Sampler) (reference-point niching)
//! or [`MoeadSampler`](super::moead::MoeadSampler) (decomposition).
//!
//! # Configuration
//!
//! | Parameter | Builder method | Default |
//! |-----------|---------------|---------|
//! | Population size | [`population_size`](Nsga2SamplerBuilder::population_size) | `4 + floor(3 * ln(n_params))`, min 4 |
//! | Crossover probability | [`crossover_prob`](Nsga2SamplerBuilder::crossover_prob) | 0.9 |
//! | SBX distribution index | [`crossover_eta`](Nsga2SamplerBuilder::crossover_eta) | 20.0 |
//! | Mutation distribution index | [`mutation_eta`](Nsga2SamplerBuilder::mutation_eta) | 20.0 |
//! | Random seed | [`seed`](Nsga2SamplerBuilder::seed) | random |
//!
//! # Examples
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::multi_objective::MultiObjectiveStudy;
//! use optimizer::parameter::{FloatParam, Parameter};
//! use optimizer::sampler::nsga2::Nsga2Sampler;
//!
//! let sampler = Nsga2Sampler::with_seed(42);
//! let study =
//! MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
//!
//! let x = FloatParam::new(0.0, 1.0);
//! study
//! .optimize(50, |trial: &mut optimizer::Trial| {
//! let xv = x.suggest(trial)?;
//! Ok::<_, optimizer::Error>(vec![xv * xv, (xv - 1.0).powi(2)])
//! })
//! .unwrap();
//! ```
use parking_lot::Mutex;
use super::genetic::{
self, Candidate, EvolutionaryState, Phase, advance_generation, collect_evaluated_generation,
crossover, extract_trial_params, finalize_discovery, generate_random_candidates, mutate,
sample_from_candidate, sample_random,
};
use crate::distribution::Distribution;
use crate::multi_objective::MultiObjectiveTrial;
use crate::param::ParamValue;
use crate::pareto;
use crate::types::Direction;
/// NSGA-II sampler for multi-objective optimization.
///
/// Use non-dominated sorting with crowding-distance tie-breaking to
/// evolve a well-spread Pareto front. Best suited for bi-objective
/// problems; for 3+ objectives prefer [`Nsga3Sampler`](super::nsga3::Nsga3Sampler).
///
/// Create with [`Nsga2Sampler::new`], [`Nsga2Sampler::with_seed`], or
/// [`Nsga2Sampler::builder`] for full configuration.
pub struct Nsga2Sampler {
state: Mutex<Nsga2State>,
}
impl Nsga2Sampler {
/// Creates a new NSGA-II sampler with a random seed.
#[must_use]
pub fn new() -> Self {
Self {
state: Mutex::new(Nsga2State::new(Nsga2Config::default(), None)),
}
}
/// Creates a new NSGA-II sampler with a fixed seed.
#[must_use]
pub fn with_seed(seed: u64) -> Self {
Self {
state: Mutex::new(Nsga2State::new(Nsga2Config::default(), Some(seed))),
}
}
/// Creates a builder for configuring an `Nsga2Sampler`.
#[must_use]
pub fn builder() -> Nsga2SamplerBuilder {
Nsga2SamplerBuilder::default()
}
}
impl Default for Nsga2Sampler {
fn default() -> Self {
Self::new()
}
}
/// Builder for [`Nsga2Sampler`].
#[derive(Debug, Clone, Default)]
pub struct Nsga2SamplerBuilder {
population_size: Option<usize>,
crossover_prob: Option<f64>,
crossover_eta: Option<f64>,
mutation_eta: Option<f64>,
seed: Option<u64>,
}
impl Nsga2SamplerBuilder {
/// Sets the population size. Default: `4 + floor(3 * ln(n_params))`, minimum 4.
#[must_use]
pub fn population_size(mut self, size: usize) -> Self {
self.population_size = Some(size);
self
}
/// Sets the crossover probability. Default: 0.9.
#[must_use]
pub fn crossover_prob(mut self, prob: f64) -> Self {
self.crossover_prob = Some(prob);
self
}
/// Sets the SBX distribution index. Default: 20.0.
#[must_use]
pub fn crossover_eta(mut self, eta: f64) -> Self {
self.crossover_eta = Some(eta);
self
}
/// Sets the polynomial mutation distribution index. Default: 20.0.
#[must_use]
pub fn mutation_eta(mut self, eta: f64) -> Self {
self.mutation_eta = Some(eta);
self
}
/// Sets the random seed for reproducibility.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.seed = Some(seed);
self
}
/// Builds the configured [`Nsga2Sampler`].
#[must_use]
pub fn build(self) -> Nsga2Sampler {
let config = Nsga2Config {
user_population_size: self.population_size,
crossover_prob: self.crossover_prob.unwrap_or(0.9),
crossover_eta: self.crossover_eta.unwrap_or(20.0),
mutation_eta: self.mutation_eta.unwrap_or(20.0),
};
Nsga2Sampler {
state: Mutex::new(Nsga2State::new(config, self.seed)),
}
}
}
// ---------------------------------------------------------------------------
// Internal types
// ---------------------------------------------------------------------------
#[derive(Clone, Debug)]
struct Nsga2Config {
user_population_size: Option<usize>,
crossover_prob: f64,
crossover_eta: f64,
mutation_eta: f64,
}
impl Default for Nsga2Config {
fn default() -> Self {
Self {
user_population_size: None,
crossover_prob: 0.9,
crossover_eta: 20.0,
mutation_eta: 20.0,
}
}
}
struct Nsga2State {
evo: EvolutionaryState,
config: Nsga2Config,
}
impl Nsga2State {
fn new(config: Nsga2Config, seed: Option<u64>) -> Self {
Self {
evo: EvolutionaryState::new(seed),
config,
}
}
}
// ---------------------------------------------------------------------------
// MultiObjectiveSampler implementation
// ---------------------------------------------------------------------------
impl crate::multi_objective::MultiObjectiveSampler for Nsga2Sampler {
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[MultiObjectiveTrial],
directions: &[Direction],
) -> ParamValue {
let mut state = self.state.lock();
match &state.evo.phase {
Phase::Discovery => {
if let Some(value) =
genetic::sample_discovery(&mut state.evo, distribution, trial_id)
{
return value;
}
// Transitioned to active phase
let user_pop = state.config.user_population_size;
finalize_discovery(&mut state.evo, user_pop);
generate_random_candidates(&mut state.evo);
sample_from_candidate(&mut state.evo, trial_id)
}
Phase::Active => {
maybe_generate_new_generation(&mut state, history, directions);
sample_from_candidate(&mut state.evo, trial_id)
}
}
}
}
/// Check if all candidates in the current generation have been evaluated;
/// if so, run NSGA-II selection and generate offspring.
fn maybe_generate_new_generation(
state: &mut Nsga2State,
history: &[MultiObjectiveTrial],
directions: &[Direction],
) {
if state.evo.candidates.is_empty() {
generate_random_candidates(&mut state.evo);
return;
}
if let Some(evaluated) = collect_evaluated_generation(&state.evo, history) {
let offspring = nsga2_generate_offspring(state, &evaluated, directions);
advance_generation(&mut state.evo, offspring);
}
}
// ---------------------------------------------------------------------------
// NSGA-II generation algorithm
// ---------------------------------------------------------------------------
/// Performs NSGA-II selection: non-dominated sort + crowding distance,
/// then selects `pop_size` parents from the population.
fn nsga2_select(
state: &mut Nsga2State,
population: &[&MultiObjectiveTrial],
directions: &[Direction],
) -> (Vec<Vec<ParamValue>>, Vec<usize>, Vec<f64>) {
let pop_size = state.evo.population_size;
let values: Vec<Vec<f64>> = population.iter().map(|t| t.values.clone()).collect();
let constraints: Vec<Vec<f64>> = population.iter().map(|t| t.constraints.clone()).collect();
let has_constraints = constraints.iter().any(|c| !c.is_empty());
let fronts = if has_constraints {
pareto::fast_non_dominated_sort_constrained(&values, directions, &constraints)
} else {
pareto::fast_non_dominated_sort(&values, directions)
};
let n = population.len();
let mut rank = vec![0_usize; n];
let mut crowding = vec![0.0_f64; n];
for (front_rank, front) in fronts.iter().enumerate() {
let cd = pareto::crowding_distance_indexed(front, &values);
for (i, &idx) in front.iter().enumerate() {
rank[idx] = front_rank;
crowding[idx] = cd[i];
}
}
let mut selected: Vec<usize> = Vec::with_capacity(pop_size);
for front in &fronts {
if selected.len() + front.len() <= pop_size {
selected.extend_from_slice(front);
} else {
let remaining = pop_size - selected.len();
let mut front_sorted: Vec<usize> = front.clone();
front_sorted.sort_by(|&a, &b| {
crowding[b]
.partial_cmp(&crowding[a])
.unwrap_or(core::cmp::Ordering::Equal)
});
selected.extend_from_slice(&front_sorted[..remaining]);
break;
}
}
while selected.len() < pop_size {
selected.push(state.evo.rng.usize(0..n));
}
let parents: Vec<Vec<ParamValue>> = selected
.iter()
.map(|&idx| {
extract_trial_params(population[idx], &state.evo.dimensions, &mut state.evo.rng)
})
.collect();
let sel_rank: Vec<usize> = selected.iter().map(|&i| rank[i]).collect();
let sel_crowding: Vec<f64> = selected.iter().map(|&i| crowding[i]).collect();
(parents, sel_rank, sel_crowding)
}
/// Runs NSGA-II selection and generates offspring candidates.
fn nsga2_generate_offspring(
state: &mut Nsga2State,
population: &[&MultiObjectiveTrial],
directions: &[Direction],
) -> Vec<Candidate> {
let pop_size = state.evo.population_size;
if population.len() < 2 {
return (0..pop_size)
.map(|_| {
let params = state
.evo
.dimensions
.iter()
.map(|d| sample_random(&mut state.evo.rng, &d.distribution))
.collect();
Candidate { params }
})
.collect();
}
let (parents, sel_rank, sel_crowding) = nsga2_select(state, population, directions);
let mut offspring = Vec::with_capacity(pop_size);
while offspring.len() < pop_size {
let p1 = tournament_select(&mut state.evo.rng, &sel_rank, &sel_crowding, parents.len());
let p2 = tournament_select(&mut state.evo.rng, &sel_rank, &sel_crowding, parents.len());
let (mut child1, mut child2) = crossover(
&mut state.evo.rng,
&parents[p1],
&parents[p2],
&state.evo.dimensions,
state.config.crossover_prob,
state.config.crossover_eta,
);
mutate(
&mut state.evo.rng,
&mut child1,
&state.evo.dimensions,
state.config.mutation_eta,
);
mutate(
&mut state.evo.rng,
&mut child2,
&state.evo.dimensions,
state.config.mutation_eta,
);
offspring.push(Candidate { params: child1 });
if offspring.len() < pop_size {
offspring.push(Candidate { params: child2 });
}
}
offspring
}
/// Tournament selection: pick 2 random individuals, return index of winner.
/// Winner has lower rank; ties broken by higher crowding distance.
fn tournament_select(
rng: &mut fastrand::Rng,
ranks: &[usize],
crowding: &[f64],
n: usize,
) -> usize {
let a = rng.usize(0..n);
let b = rng.usize(0..n);
if ranks[a] < ranks[b] {
a
} else if ranks[b] < ranks[a] {
b
} else if crowding[a] >= crowding[b] {
a
} else {
b
}
}
+773
View File
@@ -0,0 +1,773 @@
//! NSGA-III (Non-dominated Sorting Genetic Algorithm III) sampler.
//!
//! NSGA-III extends NSGA-II to handle **many-objective** (3+) problems
//! where crowding distance loses effectiveness. Instead of crowding
//! distance, it uses **reference-point-based niching** with structured
//! Das-Dennis reference points distributed on the unit simplex to guide
//! the population toward a well-diversified Pareto front.
//!
//! # Algorithm
//!
//! Each generation proceeds as follows:
//!
//! 1. **Non-dominated sorting** — same as NSGA-II, partition the
//! combined population into Pareto fronts F₁, F₂, …
//! 2. **Normalize objectives** — translate by ideal point and scale by
//! intercepts so all objectives lie in roughly \[0, 1\].
//! 3. **Associate with reference points** — assign each solution to the
//! closest Das-Dennis reference direction by perpendicular distance.
//! 4. **Niching selection** — when the last front only partially fits,
//! prefer solutions associated with under-represented reference points
//! (lowest niche count first, closest distance second).
//! 5. **SBX crossover + polynomial mutation** — generate offspring via
//! rank-based tournament selection.
//!
//! # When to use
//!
//! - **Three or more objectives** — NSGA-III maintains diversity far
//! better than NSGA-II as the number of objectives grows.
//! - Problems where you want a **uniformly distributed** Pareto front
//! guided by structured reference points.
//! - Scales well up to ~10 objectives with appropriate division settings.
//!
//! For bi-objective problems, [`Nsga2Sampler`](super::nsga2::Nsga2Sampler)
//! is simpler and equally effective. For decomposition-based optimization,
//! see [`MoeadSampler`](super::moead::MoeadSampler).
//!
//! # Configuration
//!
//! | Parameter | Builder method | Default |
//! |-----------|---------------|---------|
//! | Population size | [`population_size`](Nsga3SamplerBuilder::population_size) | Number of Das-Dennis reference points |
//! | Das-Dennis divisions (H) | [`n_divisions`](Nsga3SamplerBuilder::n_divisions) | Auto-chosen from population size and objectives |
//! | Crossover probability | [`crossover_prob`](Nsga3SamplerBuilder::crossover_prob) | 1.0 |
//! | SBX distribution index | [`crossover_eta`](Nsga3SamplerBuilder::crossover_eta) | 30.0 |
//! | Mutation distribution index | [`mutation_eta`](Nsga3SamplerBuilder::mutation_eta) | 20.0 |
//! | Random seed | [`seed`](Nsga3SamplerBuilder::seed) | random |
//!
//! # Examples
//!
//! ```
//! use optimizer::Direction;
//! use optimizer::multi_objective::MultiObjectiveStudy;
//! use optimizer::parameter::{FloatParam, Parameter};
//! use optimizer::sampler::nsga3::Nsga3Sampler;
//!
//! let sampler = Nsga3Sampler::with_seed(42);
//! let study = MultiObjectiveStudy::with_sampler(
//! vec![
//! Direction::Minimize,
//! Direction::Minimize,
//! Direction::Minimize,
//! ],
//! sampler,
//! );
//!
//! let x = FloatParam::new(0.0, 1.0);
//! let y = FloatParam::new(0.0, 1.0);
//! study
//! .optimize(100, |trial: &mut optimizer::Trial| {
//! let xv = x.suggest(trial)?;
//! let yv = y.suggest(trial)?;
//! Ok::<_, optimizer::Error>(vec![xv, yv, (1.0 - xv - yv).abs()])
//! })
//! .unwrap();
//! ```
use parking_lot::Mutex;
use super::genetic::{
self, Candidate, EvolutionaryState, Phase, advance_generation, auto_divisions,
collect_evaluated_generation, crossover, das_dennis, extract_trial_params,
generate_random_candidates, mutate, sample_from_candidate, sample_random,
};
use crate::distribution::Distribution;
use crate::multi_objective::MultiObjectiveTrial;
use crate::param::ParamValue;
use crate::pareto;
use crate::types::Direction;
/// NSGA-III sampler for multi-objective optimization.
///
/// Use reference-point niching with Das-Dennis structured points to
/// maintain diversity in many-objective (3+) problems. For bi-objective
/// problems, [`Nsga2Sampler`](super::nsga2::Nsga2Sampler) is simpler.
///
/// Create with [`Nsga3Sampler::new`], [`Nsga3Sampler::with_seed`], or
/// [`Nsga3Sampler::builder`] for full configuration.
pub struct Nsga3Sampler {
state: Mutex<Nsga3State>,
}
impl Nsga3Sampler {
/// Creates a new NSGA-III sampler with a random seed.
#[must_use]
pub fn new() -> Self {
Self {
state: Mutex::new(Nsga3State::new(Nsga3Config::default(), None)),
}
}
/// Creates a new NSGA-III sampler with a fixed seed.
#[must_use]
pub fn with_seed(seed: u64) -> Self {
Self {
state: Mutex::new(Nsga3State::new(Nsga3Config::default(), Some(seed))),
}
}
/// Creates a builder for configuring an `Nsga3Sampler`.
#[must_use]
pub fn builder() -> Nsga3SamplerBuilder {
Nsga3SamplerBuilder::default()
}
}
impl Default for Nsga3Sampler {
fn default() -> Self {
Self::new()
}
}
/// Builder for [`Nsga3Sampler`].
#[derive(Debug, Clone, Default)]
pub struct Nsga3SamplerBuilder {
population_size: Option<usize>,
n_divisions: Option<usize>,
crossover_prob: Option<f64>,
crossover_eta: Option<f64>,
mutation_eta: Option<f64>,
seed: Option<u64>,
}
impl Nsga3SamplerBuilder {
/// Sets the population size. If unset, equals the number of
/// Das-Dennis reference points.
#[must_use]
pub fn population_size(mut self, size: usize) -> Self {
self.population_size = Some(size);
self
}
/// Sets the number of divisions (H) for Das-Dennis reference points.
/// If unset, automatically chosen based on population size and number
/// of objectives.
#[must_use]
pub fn n_divisions(mut self, h: usize) -> Self {
self.n_divisions = Some(h);
self
}
/// Sets the crossover probability. Default: 1.0.
#[must_use]
pub fn crossover_prob(mut self, prob: f64) -> Self {
self.crossover_prob = Some(prob);
self
}
/// Sets the SBX distribution index. Default: 30.0.
#[must_use]
pub fn crossover_eta(mut self, eta: f64) -> Self {
self.crossover_eta = Some(eta);
self
}
/// Sets the polynomial mutation distribution index. Default: 20.0.
#[must_use]
pub fn mutation_eta(mut self, eta: f64) -> Self {
self.mutation_eta = Some(eta);
self
}
/// Sets the random seed for reproducibility.
#[must_use]
pub fn seed(mut self, seed: u64) -> Self {
self.seed = Some(seed);
self
}
/// Builds the configured [`Nsga3Sampler`].
#[must_use]
pub fn build(self) -> Nsga3Sampler {
let config = Nsga3Config {
user_population_size: self.population_size,
n_divisions: self.n_divisions,
crossover_prob: self.crossover_prob.unwrap_or(1.0),
crossover_eta: self.crossover_eta.unwrap_or(30.0),
mutation_eta: self.mutation_eta.unwrap_or(20.0),
};
Nsga3Sampler {
state: Mutex::new(Nsga3State::new(config, self.seed)),
}
}
}
// ---------------------------------------------------------------------------
// Internal types
// ---------------------------------------------------------------------------
#[derive(Clone, Debug)]
struct Nsga3Config {
user_population_size: Option<usize>,
n_divisions: Option<usize>,
crossover_prob: f64,
crossover_eta: f64,
mutation_eta: f64,
}
impl Default for Nsga3Config {
fn default() -> Self {
Self {
user_population_size: None,
n_divisions: None,
crossover_prob: 1.0,
crossover_eta: 30.0,
mutation_eta: 20.0,
}
}
}
struct Nsga3State {
evo: EvolutionaryState,
config: Nsga3Config,
/// Das-Dennis reference points (lazily generated once objectives are known).
reference_points: Vec<Vec<f64>>,
/// Best value seen per objective (minimize-space).
ideal_point: Vec<f64>,
/// Whether reference points have been initialized.
initialized: bool,
}
impl Nsga3State {
fn new(config: Nsga3Config, seed: Option<u64>) -> Self {
Self {
evo: EvolutionaryState::new(seed),
config,
reference_points: Vec::new(),
ideal_point: Vec::new(),
initialized: false,
}
}
}
// ---------------------------------------------------------------------------
// MultiObjectiveSampler implementation
// ---------------------------------------------------------------------------
impl crate::multi_objective::MultiObjectiveSampler for Nsga3Sampler {
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
history: &[MultiObjectiveTrial],
directions: &[Direction],
) -> ParamValue {
let mut state = self.state.lock();
match &state.evo.phase {
Phase::Discovery => {
if let Some(value) =
genetic::sample_discovery(&mut state.evo, distribution, trial_id)
{
return value;
}
// Transitioned to active phase
initialize_nsga3(&mut state, directions);
generate_random_candidates(&mut state.evo);
sample_from_candidate(&mut state.evo, trial_id)
}
Phase::Active => {
maybe_generate_new_generation(&mut state, history, directions);
sample_from_candidate(&mut state.evo, trial_id)
}
}
}
}
/// Initialize NSGA-III: generate reference points and set population size.
fn initialize_nsga3(state: &mut Nsga3State, directions: &[Direction]) {
let n_obj = directions.len();
// Determine divisions
let divisions = state
.config
.n_divisions
.unwrap_or_else(|| auto_divisions(n_obj, state.config.user_population_size.unwrap_or(100)));
state.reference_points = das_dennis(n_obj, divisions);
let n_ref = state.reference_points.len();
// Population size = number of reference points (or user override, at least n_ref)
let pop_size = state.config.user_population_size.unwrap_or(n_ref).max(4);
state.evo.population_size = pop_size;
state.evo.phase = Phase::Active;
state.ideal_point = vec![f64::INFINITY; n_obj];
state.initialized = true;
}
fn maybe_generate_new_generation(
state: &mut Nsga3State,
history: &[MultiObjectiveTrial],
directions: &[Direction],
) {
if state.evo.candidates.is_empty() {
generate_random_candidates(&mut state.evo);
return;
}
if let Some(evaluated) = collect_evaluated_generation(&state.evo, history) {
let offspring = nsga3_generate_offspring(state, &evaluated, directions);
advance_generation(&mut state.evo, offspring);
}
}
// ---------------------------------------------------------------------------
// NSGA-III selection algorithm
// ---------------------------------------------------------------------------
/// Normalize objectives to minimize-space.
fn to_minimize_space(values: &[f64], directions: &[Direction]) -> Vec<f64> {
values
.iter()
.zip(directions)
.map(|(&v, d)| match d {
Direction::Minimize => v,
Direction::Maximize => -v,
})
.collect()
}
/// Update ideal point with new observations.
fn update_ideal_point(ideal: &mut [f64], normalized_values: &[Vec<f64>]) {
for vals in normalized_values {
for (i, &v) in vals.iter().enumerate() {
if v < ideal[i] {
ideal[i] = v;
}
}
}
}
/// Compute Achievement Scalarizing Function (ASF) for extreme point finding.
fn asf(point: &[f64], weight: &[f64], ideal: &[f64]) -> f64 {
point
.iter()
.zip(weight)
.zip(ideal)
.map(|((&p, &w), &z)| {
let w = if w < 1e-6 { 1e-6 } else { w };
(p - z) / w
})
.fold(f64::NEG_INFINITY, f64::max)
}
/// Find intercepts for normalization via extreme points.
///
/// For each objective, find the point with best ASF (using a weight vector
/// that emphasizes that objective). The intercepts are where the hyperplane
/// through the extreme points crosses each axis.
fn find_intercepts(normalized_values: &[Vec<f64>], ideal: &[f64]) -> Vec<f64> {
let n_obj = ideal.len();
let n = normalized_values.len();
if n == 0 || n_obj == 0 {
return vec![1.0; n_obj];
}
// Find extreme points (one per objective)
let mut extreme_indices = Vec::with_capacity(n_obj);
for obj in 0..n_obj {
let mut weight = vec![1e-6; n_obj];
weight[obj] = 1.0;
let mut best_idx = 0;
let mut best_asf = f64::INFINITY;
for (i, vals) in normalized_values.iter().enumerate() {
let a = asf(vals, &weight, ideal);
if a < best_asf {
best_asf = a;
best_idx = i;
}
}
extreme_indices.push(best_idx);
}
// Try to compute hyperplane intercepts
// For stability, if the extreme points are degenerate, fall back to
// max - ideal per objective
let mut intercepts = Vec::with_capacity(n_obj);
for obj in 0..n_obj {
let max_val = normalized_values
.iter()
.map(|v| v[obj])
.fold(f64::NEG_INFINITY, f64::max);
let intercept = max_val - ideal[obj];
intercepts.push(if intercept > 1e-10 { intercept } else { 1.0 });
}
intercepts
}
/// Normalize objective values: subtract ideal, divide by intercepts.
fn normalize_objectives(values: &[Vec<f64>], ideal: &[f64], intercepts: &[f64]) -> Vec<Vec<f64>> {
values
.iter()
.map(|v| {
v.iter()
.zip(ideal)
.zip(intercepts)
.map(|((&val, &z), &a)| {
let norm = if a > 1e-10 { a } else { 1.0 };
(val - z) / norm
})
.collect()
})
.collect()
}
/// Perpendicular distance from a point to a reference line (direction vector).
fn perpendicular_distance(point: &[f64], reference: &[f64]) -> f64 {
let dot: f64 = point.iter().zip(reference).map(|(&p, &r)| p * r).sum();
let ref_norm_sq: f64 = reference.iter().map(|&r| r * r).sum();
if ref_norm_sq < 1e-30 {
return f64::INFINITY;
}
let proj_scalar = dot / ref_norm_sq;
let dist_sq: f64 = point
.iter()
.zip(reference)
.map(|(&p, &r)| {
let proj = proj_scalar * r;
(p - proj).powi(2)
})
.sum();
dist_sq.sqrt()
}
/// Associate each solution with its nearest reference point.
/// Returns (`closest_ref_idx`, distance) for each solution.
fn associate_to_reference_points(
normalized: &[Vec<f64>],
reference_points: &[Vec<f64>],
) -> Vec<(usize, f64)> {
normalized
.iter()
.map(|point| {
let mut best_ref = 0;
let mut best_dist = f64::INFINITY;
for (j, rp) in reference_points.iter().enumerate() {
let d = perpendicular_distance(point, rp);
if d < best_dist {
best_dist = d;
best_ref = j;
}
}
(best_ref, best_dist)
})
.collect()
}
/// NSGA-III niching-based selection from the last front.
///
/// `already_selected` are indices into the combined population that are
/// already accepted (from fronts 0..L-1). `last_front` contains indices
/// from front L. We need to pick `remaining` more from `last_front`.
fn niching_select(
rng: &mut fastrand::Rng,
associations: &[(usize, f64)],
already_selected: &[usize],
last_front: &[usize],
n_reference_points: usize,
remaining: usize,
) -> Vec<usize> {
// Count niche per reference point for already selected
let mut niche_count = vec![0_usize; n_reference_points];
for &idx in already_selected {
niche_count[associations[idx].0] += 1;
}
// Build per-reference-point candidate lists from the last front
let mut ref_candidates: Vec<Vec<(usize, f64)>> = vec![Vec::new(); n_reference_points];
for &idx in last_front {
let (ref_idx, dist) = associations[idx];
ref_candidates[ref_idx].push((idx, dist));
}
let mut selected = Vec::with_capacity(remaining);
let mut excluded = vec![false; associations.len()];
for _ in 0..remaining {
// Find minimum niche count among reference points that still have candidates
let min_count = (0..n_reference_points)
.filter(|&j| ref_candidates[j].iter().any(|&(idx, _)| !excluded[idx]))
.map(|j| niche_count[j])
.min();
let Some(min_count) = min_count else {
break;
};
// Collect reference points with this minimum count that have candidates
let min_refs: Vec<usize> = (0..n_reference_points)
.filter(|&j| {
niche_count[j] == min_count
&& ref_candidates[j].iter().any(|&(idx, _)| !excluded[idx])
})
.collect();
if min_refs.is_empty() {
break;
}
// Pick a random reference point from the minimum set
let chosen_ref = min_refs[rng.usize(0..min_refs.len())];
// Available candidates for this reference point
let available: Vec<(usize, f64)> = ref_candidates[chosen_ref]
.iter()
.filter(|&&(idx, _)| !excluded[idx])
.copied()
.collect();
if available.is_empty() {
continue;
}
let chosen_idx = if min_count == 0 {
// Pick closest to reference line
available
.iter()
.min_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(core::cmp::Ordering::Equal))
.unwrap()
.0
} else {
// Pick random
available[rng.usize(0..available.len())].0
};
selected.push(chosen_idx);
excluded[chosen_idx] = true;
niche_count[chosen_ref] += 1;
}
selected
}
/// Perform NSGA-III selection: non-dominated sort + reference-point niching.
fn nsga3_select(
state: &mut Nsga3State,
population: &[&MultiObjectiveTrial],
directions: &[Direction],
) -> (Vec<Vec<ParamValue>>, Vec<usize>) {
let pop_size = state.evo.population_size;
let n_obj = directions.len();
// Convert to minimize-space
let min_values: Vec<Vec<f64>> = population
.iter()
.map(|t| to_minimize_space(&t.values, directions))
.collect();
// Non-dominated sort
let constraints: Vec<Vec<f64>> = population.iter().map(|t| t.constraints.clone()).collect();
let has_constraints = constraints.iter().any(|c| !c.is_empty());
let fronts = if has_constraints {
pareto::fast_non_dominated_sort_constrained(
&min_values,
&vec![Direction::Minimize; n_obj],
&constraints,
)
} else {
pareto::fast_non_dominated_sort(&min_values, &vec![Direction::Minimize; n_obj])
};
// Fill front-by-front
let mut selected: Vec<usize> = Vec::with_capacity(pop_size);
let mut last_front_idx = None;
for (fi, front) in fronts.iter().enumerate() {
if selected.len() + front.len() <= pop_size {
selected.extend_from_slice(front);
} else {
last_front_idx = Some(fi);
break;
}
}
// If we filled exactly or all fronts fit, done
if selected.len() < pop_size
&& let Some(lf_idx) = last_front_idx
{
// Need niching from the last partial front
let remaining = pop_size - selected.len();
// Update ideal point
update_ideal_point(&mut state.ideal_point, &min_values);
// Find intercepts and normalize
let intercepts = find_intercepts(&min_values, &state.ideal_point);
let normalized = normalize_objectives(&min_values, &state.ideal_point, &intercepts);
// Associate all solutions with reference points
let associations = associate_to_reference_points(&normalized, &state.reference_points);
// Select from last front using niching
let last_front = &fronts[lf_idx];
let additional = niching_select(
&mut state.evo.rng,
&associations,
&selected,
last_front,
state.reference_points.len(),
remaining,
);
selected.extend(additional);
}
// Pad if needed
let n = population.len();
while selected.len() < pop_size {
selected.push(state.evo.rng.usize(0..n));
}
let params = selected
.iter()
.map(|&idx| {
extract_trial_params(population[idx], &state.evo.dimensions, &mut state.evo.rng)
})
.collect();
(params, selected)
}
/// Tournament selection based on rank only (no crowding distance in NSGA-III).
fn tournament_select_rank(rng: &mut fastrand::Rng, ranks: &[usize], n: usize) -> usize {
let a = rng.usize(0..n);
let b = rng.usize(0..n);
if ranks[a] <= ranks[b] { a } else { b }
}
fn nsga3_generate_offspring(
state: &mut Nsga3State,
population: &[&MultiObjectiveTrial],
directions: &[Direction],
) -> Vec<Candidate> {
let pop_size = state.evo.population_size;
if population.len() < 2 {
return (0..pop_size)
.map(|_| {
let params = state
.evo
.dimensions
.iter()
.map(|d| sample_random(&mut state.evo.rng, &d.distribution))
.collect();
Candidate { params }
})
.collect();
}
// Initialize reference points and ideal on first generation
if !state.initialized {
initialize_nsga3(state, directions);
}
let (parents, selected_indices) = nsga3_select(state, population, directions);
// Assign Pareto front ranks for tournament selection
let n_obj = directions.len();
let min_values: Vec<Vec<f64>> = population
.iter()
.map(|t| to_minimize_space(&t.values, directions))
.collect();
let fronts = pareto::fast_non_dominated_sort(&min_values, &vec![Direction::Minimize; n_obj]);
// Build rank lookup for population indices
let mut pop_rank = vec![0_usize; population.len()];
for (front_rank, front) in fronts.iter().enumerate() {
for &idx in front {
if idx < pop_rank.len() {
pop_rank[idx] = front_rank;
}
}
}
// Map population ranks to selected parent indices
let parent_ranks: Vec<usize> = selected_indices
.iter()
.map(|&idx| {
if idx < pop_rank.len() {
pop_rank[idx]
} else {
0
}
})
.collect();
let mut offspring = Vec::with_capacity(pop_size);
while offspring.len() < pop_size {
let p1 = tournament_select_rank(&mut state.evo.rng, &parent_ranks, parents.len());
let p2 = tournament_select_rank(&mut state.evo.rng, &parent_ranks, parents.len());
let (mut child1, mut child2) = crossover(
&mut state.evo.rng,
&parents[p1],
&parents[p2],
&state.evo.dimensions,
state.config.crossover_prob,
state.config.crossover_eta,
);
mutate(
&mut state.evo.rng,
&mut child1,
&state.evo.dimensions,
state.config.mutation_eta,
);
mutate(
&mut state.evo.rng,
&mut child2,
&state.evo.dimensions,
state.config.mutation_eta,
);
offspring.push(Candidate { params: child1 });
if offspring.len() < pop_size {
offspring.push(Candidate { params: child2 });
}
}
offspring
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_perpendicular_distance() {
// Point (1, 0) to reference line (1, 1) (45-degree line)
let d = perpendicular_distance(&[1.0, 0.0], &[1.0, 1.0]);
// Projection is (0.5, 0.5), distance = sqrt(0.25 + 0.25) = sqrt(0.5)
assert!((d - (0.5_f64).sqrt()).abs() < 1e-10);
}
#[test]
fn test_perpendicular_distance_on_line() {
// Point on the reference line
let d = perpendicular_distance(&[2.0, 2.0], &[1.0, 1.0]);
assert!(d < 1e-10);
}
#[test]
fn test_normalize_objectives() {
let values = vec![vec![2.0, 4.0], vec![4.0, 2.0]];
let ideal = vec![1.0, 1.0];
let intercepts = vec![3.0, 3.0];
let normalized = normalize_objectives(&values, &ideal, &intercepts);
assert!((normalized[0][0] - 1.0 / 3.0).abs() < 1e-10);
assert!((normalized[0][1] - 1.0).abs() < 1e-10);
}
}
+93 -74
View File
@@ -1,18 +1,50 @@
//! Random sampler implementation.
//! Random sampler — uniform independent sampling.
//!
//! [`RandomSampler`] draws each parameter value independently and uniformly
//! at random, ignoring trial history entirely. It respects log-scale and
//! step-size constraints defined by the parameter distribution.
//!
//! # When to use
//!
//! - **Baseline comparison** — run Random alongside smarter samplers to
//! quantify their benefit.
//! - **Startup phase** — many model-based samplers (TPE, GP, CMA-ES) use
//! random sampling for their first *n* trials before fitting a surrogate.
//! - **Very high dimensions** — when the search space is too large for
//! structured exploration, random search with enough budget can be
//! surprisingly competitive.
//!
//! For better uniform coverage without model fitting, consider
//! `SobolSampler` (requires the `sobol`
//! feature flag).
//!
//! # Example
//!
//! ```
//! use optimizer::prelude::*;
//! use optimizer::sampler::random::RandomSampler;
//!
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
//! ```
use parking_lot::Mutex;
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use core::sync::atomic::{AtomicU64, Ordering};
use crate::distribution::Distribution;
use crate::multi_objective::{MultiObjectiveSampler, MultiObjectiveTrial};
use crate::param::ParamValue;
use crate::rng_util;
use crate::sampler::{CompletedTrial, Sampler};
use crate::types::Direction;
/// A simple random sampler that samples uniformly from distributions.
/// Uniform independent random sampler.
///
/// 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.
/// Sample each parameter value uniformly at random, respecting log-scale and
/// step-size constraints. Trial history is ignored — every sample is drawn
/// independently.
///
/// This is the default sampler used by [`Study::new`](crate::Study::new)
/// and during the startup phase of model-based samplers such as
/// [`TpeSampler`](super::tpe::TpeSampler).
///
/// # Examples
///
@@ -26,7 +58,9 @@ use crate::sampler::{CompletedTrial, Sampler};
/// let sampler = RandomSampler::with_seed(42);
/// ```
pub struct RandomSampler {
rng: Mutex<StdRng>,
seed: u64,
/// Monotonic counter to disambiguate calls with identical (`trial_id`, distribution).
call_seq: AtomicU64,
}
impl RandomSampler {
@@ -34,7 +68,8 @@ impl RandomSampler {
#[must_use]
pub fn new() -> Self {
Self {
rng: Mutex::new(StdRng::from_os_rng()),
seed: fastrand::u64(..),
call_seq: AtomicU64::new(0),
}
}
@@ -44,11 +79,33 @@ impl RandomSampler {
#[must_use]
pub fn with_seed(seed: u64) -> Self {
Self {
rng: Mutex::new(StdRng::seed_from_u64(seed)),
seed,
call_seq: AtomicU64::new(0),
}
}
}
/// Default multi-objective sampler that delegates to [`RandomSampler`].
pub(crate) struct RandomMultiObjectiveSampler(RandomSampler);
impl RandomMultiObjectiveSampler {
pub(crate) fn new() -> Self {
Self(RandomSampler::new())
}
}
impl MultiObjectiveSampler for RandomMultiObjectiveSampler {
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
_history: &[MultiObjectiveTrial],
_directions: &[Direction],
) -> ParamValue {
self.0.sample(distribution, trial_id, &[])
}
}
impl Default for RandomSampler {
fn default() -> Self {
Self::new()
@@ -60,55 +117,17 @@ impl Sampler for RandomSampler {
fn sample(
&self,
distribution: &Distribution,
_trial_id: u64,
trial_id: u64,
_history: &[CompletedTrial],
) -> ParamValue {
let mut rng = self.rng.lock();
let seq = self.call_seq.fetch_add(1, Ordering::Relaxed);
let mut rng = fastrand::Rng::with_seed(rng_util::mix_seed(
self.seed,
trial_id,
rng_util::distribution_fingerprint(distribution).wrapping_add(seq),
));
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)
}
}
super::common::sample_random(&mut rng, distribution)
}
}
@@ -128,8 +147,8 @@ mod tests {
step: None,
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v));
} else {
@@ -148,8 +167,8 @@ mod tests {
step: None,
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value {
assert!((1e-5..=1.0).contains(&v));
} else {
@@ -168,8 +187,8 @@ mod tests {
step: Some(0.25),
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v));
// Check it's on the step grid
@@ -192,8 +211,8 @@ mod tests {
step: None,
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value {
assert!((0..=10).contains(&v));
} else {
@@ -212,8 +231,8 @@ mod tests {
step: None,
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value {
assert!((1..=1000).contains(&v));
} else {
@@ -232,8 +251,8 @@ mod tests {
step: Some(2),
});
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value {
assert!((0..=10).contains(&v));
// Check it's on the step grid: 0, 2, 4, 6, 8, 10
@@ -249,8 +268,8 @@ mod tests {
let sampler = RandomSampler::with_seed(42);
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 });
for _ in 0..100 {
let value = sampler.sample(&dist, 0, &[]);
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Categorical(idx) = value {
assert!(idx < 5);
} else {
@@ -270,9 +289,9 @@ mod tests {
step: None,
});
for _ in 0..10 {
let v1 = sampler1.sample(&dist, 0, &[]);
let v2 = sampler2.sample(&dist, 0, &[]);
for i in 0..10 {
let v1 = sampler1.sample(&dist, i, &[]);
let v2 = sampler2.sample(&dist, i, &[]);
assert_eq!(v1, v2);
}
}
+456
View File
@@ -0,0 +1,456 @@
//! Quasi-random sampler using Sobol low-discrepancy sequences.
//!
//! [`SobolSampler`] generates points from a Sobol sequence (scrambled via the
//! Burley 2020 algorithm) to fill the parameter space more uniformly than
//! pure random sampling. Where [`RandomSampler`](super::random::RandomSampler)
//! may cluster points in some regions by chance, Sobol sequences are
//! constructed to spread points evenly across all dimensions.
//!
//! # When to use
//!
//! - **Better-than-random baseline** — when you want uniform coverage
//! without the cost of model fitting (TPE, GP, etc.).
//! - **Startup phase replacement** — use Sobol instead of random for the
//! initial exploration phase of adaptive samplers.
//! - **Moderate dimensionality** — Sobol uniformity is strongest up to
//! ~20 dimensions; beyond that the advantage over random sampling
//! diminishes.
//! - **Deterministic exploration** — Sobol sequences are fully deterministic
//! for a given seed, making experiments reproducible.
//!
//! # How it works
//!
//! Each trial maps to a Sobol sequence index, and each parameter within a
//! trial maps to a separate Sobol dimension. The resulting quasi-random
//! point in \[0, 1) is then scaled to the parameter's distribution (linear,
//! log-scale, or step grid).
//!
//! **Important:** parameters must be suggested in the same order across
//! trials for consistent dimension assignment.
//!
//! Requires the **`sobol`** feature flag:
//!
//! ```toml
//! [dependencies]
//! optimizer = { version = "...", features = ["sobol"] }
//! ```
//!
//! # Example
//!
//! ```
//! use optimizer::prelude::*;
//! use optimizer::sampler::sobol::SobolSampler;
//!
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, SobolSampler::with_seed(42));
//! ```
use std::collections::HashMap;
use parking_lot::Mutex;
use sobol_burley::sample;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::sampler::{CompletedTrial, Sampler};
/// Internal state for tracking per-trial dimension counters.
struct SobolState {
/// Next Sobol dimension for each in-flight trial.
dimensions: HashMap<u64, u32>,
}
/// Quasi-random sampler using Sobol low-discrepancy sequences.
///
/// Produce better uniform coverage of the parameter space than
/// [`RandomSampler`](super::random::RandomSampler) by using a
/// scrambled Sobol sequence (Burley 2020). Useful as a standalone
/// baseline or as a drop-in replacement for the random startup
/// phase of model-based samplers.
///
/// Each trial uses a different Sobol sequence index, and each parameter
/// within a trial maps to a different Sobol dimension. Parameters must be
/// suggested in the same order across trials for consistent dimension
/// assignment.
///
/// Sobol sequences are most effective in moderate dimensions (up to ~20).
/// For very high dimensions, the uniformity advantage diminishes.
///
/// Requires the **`sobol`** feature flag.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::sobol::SobolSampler;
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, SobolSampler::new());
/// ```
pub struct SobolSampler {
seed: u32,
state: Mutex<SobolState>,
}
impl SobolSampler {
/// Creates a new Sobol sampler with a default seed of 0.
#[must_use]
pub fn new() -> Self {
Self::with_seed(0)
}
/// Creates a new Sobol sampler with the given seed.
///
/// Different seeds produce statistically independent Sobol sequences.
/// Using the same seed will produce the same sequence of sampled values.
#[must_use]
#[allow(clippy::cast_possible_truncation)]
pub fn with_seed(seed: u64) -> Self {
Self {
seed: seed as u32,
state: Mutex::new(SobolState {
dimensions: HashMap::new(),
}),
}
}
}
impl Default for SobolSampler {
fn default() -> Self {
Self::new()
}
}
impl Sampler for SobolSampler {
#[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
fn sample(
&self,
distribution: &Distribution,
trial_id: u64,
_history: &[CompletedTrial],
) -> ParamValue {
let mut state = self.state.lock();
let dimension = state.dimensions.entry(trial_id).or_insert(0);
let dim = *dimension;
*dimension = dim + 1;
// Use trial_id as the Sobol sequence index.
let index = trial_id as u32;
// Generate a quasi-random point in [0, 1).
let point = f64::from(sample(index, dim, self.seed));
map_point_to_distribution(point, distribution)
}
}
/// Maps a uniform [0, 1) point to a value within the given distribution.
#[allow(
clippy::cast_possible_truncation,
clippy::cast_precision_loss,
clippy::cast_sign_loss
)]
fn map_point_to_distribution(point: f64, distribution: &Distribution) -> ParamValue {
match distribution {
Distribution::Float(d) => {
let value = if d.log_scale {
let log_low = d.low.ln();
let log_high = d.high.ln();
(log_low + point * (log_high - log_low)).exp()
} else if let Some(step) = d.step {
let n_steps = ((d.high - d.low) / step).floor() as i64;
let k = (point * (n_steps + 1) as f64).floor() as i64;
let k = k.min(n_steps);
d.low + (k as f64) * step
} else {
d.low + point * (d.high - d.low)
};
ParamValue::Float(value)
}
Distribution::Int(d) => {
let value = if d.log_scale {
let log_low = (d.low as f64).ln();
let log_high = (d.high as f64).ln();
let raw = (log_low + point * (log_high - log_low)).exp().round() as i64;
raw.clamp(d.low, d.high)
} else if let Some(step) = d.step {
let n_steps = (d.high - d.low) / step;
let k = (point * (n_steps + 1) as f64).floor() as i64;
let k = k.min(n_steps);
d.low + k * step
} else {
let range = d.high - d.low + 1;
let k = (point * range as f64).floor() as i64;
(d.low + k).min(d.high)
};
ParamValue::Int(value)
}
Distribution::Categorical(d) => {
let index = (point * d.n_choices as f64).floor() as usize;
let index = index.min(d.n_choices - 1);
ParamValue::Categorical(index)
}
}
}
#[cfg(test)]
#[allow(
clippy::cast_possible_truncation,
clippy::cast_precision_loss,
clippy::cast_sign_loss
)]
mod tests {
use super::*;
use crate::distribution::{CategoricalDistribution, FloatDistribution, IntDistribution};
#[test]
fn float_within_bounds() {
let sampler = SobolSampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: -5.0,
high: 5.0,
log_scale: false,
step: None,
});
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value {
assert!(
(-5.0..=5.0).contains(&v),
"value {v} out of bounds at trial {i}"
);
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn float_log_scale_within_bounds() {
let sampler = SobolSampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: 1e-5,
high: 1.0,
log_scale: true,
step: None,
});
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value {
assert!(
(1e-5..=1.0).contains(&v),
"value {v} out of bounds at trial {i}"
);
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn float_step_respects_grid() {
let sampler = SobolSampler::with_seed(42);
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: Some(0.25),
});
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value {
assert!((0.0..=1.0).contains(&v), "value {v} out of bounds");
let k = (v / 0.25).round() as i64;
let expected = k as f64 * 0.25;
assert!((v - expected).abs() < 1e-10, "value {v} not on step grid");
} else {
panic!("Expected Float value");
}
}
}
#[test]
fn int_within_bounds() {
let sampler = SobolSampler::with_seed(42);
let dist = Distribution::Int(IntDistribution {
low: 0,
high: 10,
log_scale: false,
step: None,
});
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value {
assert!(
(0..=10).contains(&v),
"value {v} out of bounds at trial {i}"
);
} else {
panic!("Expected Int value");
}
}
}
#[test]
fn int_log_scale_within_bounds() {
let sampler = SobolSampler::with_seed(42);
let dist = Distribution::Int(IntDistribution {
low: 1,
high: 1000,
log_scale: true,
step: None,
});
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value {
assert!(
(1..=1000).contains(&v),
"value {v} out of bounds at trial {i}"
);
} else {
panic!("Expected Int value");
}
}
}
#[test]
fn int_step_respects_grid() {
let sampler = SobolSampler::with_seed(42);
let dist = Distribution::Int(IntDistribution {
low: 0,
high: 10,
log_scale: false,
step: Some(2),
});
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Int(v) = value {
assert!((0..=10).contains(&v), "value {v} out of bounds");
assert!(v % 2 == 0, "value {v} not on step grid");
} else {
panic!("Expected Int value");
}
}
}
#[test]
fn categorical_within_bounds() {
let sampler = SobolSampler::with_seed(42);
let dist = Distribution::Categorical(CategoricalDistribution { n_choices: 5 });
for i in 0..100 {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Categorical(idx) = value {
assert!(idx < 5, "index {idx} out of bounds at trial {i}");
} else {
panic!("Expected Categorical value");
}
}
}
#[test]
fn deterministic_with_same_seed() {
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let sampler1 = SobolSampler::with_seed(42);
let sampler2 = SobolSampler::with_seed(42);
for i in 0..20 {
let v1 = sampler1.sample(&dist, i, &[]);
let v2 = sampler2.sample(&dist, i, &[]);
assert_eq!(v1, v2, "mismatch at trial {i}");
}
}
#[test]
fn different_seeds_produce_different_sequences() {
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let sampler1 = SobolSampler::with_seed(0);
let sampler2 = SobolSampler::with_seed(12345);
let mut any_different = false;
for i in 0..20 {
let v1 = sampler1.sample(&dist, i, &[]);
let v2 = sampler2.sample(&dist, i, &[]);
if v1 != v2 {
any_different = true;
break;
}
}
assert!(
any_different,
"different seeds should produce different sequences"
);
}
#[test]
fn better_coverage_than_random() {
// Sobol sequence should cover [0,1] more uniformly than random.
// We measure this by checking that the Sobol samples fill all
// 10 equal-width bins with only 20 samples.
let sampler = SobolSampler::with_seed(0);
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
let n_bins = 10;
let n_samples = 20;
let mut bins = vec![0u32; n_bins];
for i in 0..n_samples {
let value = sampler.sample(&dist, i, &[]);
if let ParamValue::Float(v) = value {
let bin = ((v * n_bins as f64).floor() as usize).min(n_bins - 1);
bins[bin] += 1;
}
}
let filled_bins = bins.iter().filter(|&&c| c > 0).count();
assert!(
filled_bins >= 8,
"Expected at least 8/10 bins filled, got {filled_bins}: {bins:?}"
);
}
#[test]
fn multi_parameter_uses_different_dimensions() {
// When sampling multiple parameters per trial, each parameter
// should get a different Sobol dimension, producing different values.
let sampler = SobolSampler::with_seed(0);
let dist = Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
});
// Sample two parameters for trial 0.
let v1 = sampler.sample(&dist, 0, &[]);
let v2 = sampler.sample(&dist, 0, &[]);
// They should differ (different Sobol dimensions for the same index).
assert_ne!(
v1, v2,
"multi-parameter samples should use different dimensions"
);
}
}
+216
View File
@@ -0,0 +1,216 @@
//! Shared TPE sampling functions used by both `TpeSampler` and `MotpeSampler`.
use crate::distribution::{FloatDistribution, IntDistribution};
use crate::kde::KernelDensityEstimator;
use crate::rng_util;
/// Samples using TPE for float distributions.
pub(crate) fn sample_tpe_float(
dist: &FloatDistribution,
good_values: Vec<f64>,
bad_values: Vec<f64>,
n_ei_candidates: usize,
kde_bandwidth: Option<f64>,
rng: &mut fastrand::Rng,
) -> f64 {
let (low, high, log_scale, step) = (dist.low, dist.high, dist.log_scale, dist.step);
// Transform to internal space (log space if needed)
let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
let i_low = low.ln();
let i_high = high.ln();
let g = {
let mut v = good_values;
for x in &mut v {
*x = x.ln();
}
v
};
let b = {
let mut v = bad_values;
for x in &mut v {
*x = x.ln();
}
v
};
(i_low, i_high, g, b)
} else {
(low, high, good_values, bad_values)
};
// Fit KDEs to good and bad groups
let l_kde = match kde_bandwidth {
Some(bw) => KernelDensityEstimator::with_bandwidth(good_internal, bw),
None => KernelDensityEstimator::new(good_internal),
};
let g_kde = match kde_bandwidth {
Some(bw) => KernelDensityEstimator::with_bandwidth(bad_internal, bw),
None => KernelDensityEstimator::new(bad_internal),
};
// If KDE construction fails, fall back to uniform sampling
let (Ok(l_kde), Ok(g_kde)) = (l_kde, g_kde) else {
return rng_util::f64_range(rng, low, high);
};
// Generate candidates from l(x) and select the one with best l(x)/g(x) ratio
let mut best_candidate = internal_low;
let mut best_ratio = f64::NEG_INFINITY;
for _ in 0..n_ei_candidates {
let candidate = l_kde.sample(rng).clamp(internal_low, internal_high);
let l_density = l_kde.pdf(candidate);
let g_density = g_kde.pdf(candidate);
// Compute l(x)/g(x) ratio, handling zero density
let ratio = if g_density < f64::EPSILON {
if l_density > f64::EPSILON {
f64::INFINITY
} else {
0.0
}
} else {
l_density / g_density
};
if ratio > best_ratio {
best_ratio = ratio;
best_candidate = candidate;
}
}
// Transform back from internal space
let mut value = if log_scale {
best_candidate.exp()
} else {
best_candidate
};
// Apply step constraint if present
if let Some(step) = step {
let k = ((value - low) / step).round();
value = low + k * step;
}
// Ensure value is within bounds
value.clamp(low, high)
}
/// Samples using TPE for integer distributions.
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
pub(crate) fn sample_tpe_int(
dist: &IntDistribution,
good_values: Vec<i64>,
bad_values: Vec<i64>,
n_ei_candidates: usize,
kde_bandwidth: Option<f64>,
rng: &mut fastrand::Rng,
) -> i64 {
let (low, high, log_scale, step) = (dist.low, dist.high, dist.log_scale, dist.step);
// Convert to floats for KDE
let good_floats: Vec<f64> = good_values.into_iter().map(|v| v as f64).collect();
let bad_floats: Vec<f64> = bad_values.into_iter().map(|v| v as f64).collect();
let float_dist = FloatDistribution {
low: low as f64,
high: high as f64,
log_scale,
step: step.map(|s| s as f64),
};
// Use float TPE sampling
let float_value = sample_tpe_float(
&float_dist,
good_floats,
bad_floats,
n_ei_candidates,
kde_bandwidth,
rng,
);
// Round to nearest integer
let int_value = float_value.round() as i64;
// Apply step constraint if present
let int_value = if let Some(step) = step {
let k = ((int_value - low) as f64 / step as f64).round() as i64;
low + k * step
} else {
int_value
};
// Ensure value is within bounds
int_value.clamp(low, high)
}
/// Samples using TPE for categorical distributions.
#[allow(clippy::cast_precision_loss)]
pub(crate) fn sample_tpe_categorical(
n_choices: usize,
good_indices: &[usize],
bad_indices: &[usize],
rng: &mut fastrand::Rng,
) -> usize {
// Stack-allocate for the common case (<=32 choices), heap for rare large cases
let mut good_buf = [0usize; 32];
let mut bad_buf = [0usize; 32];
let mut weight_buf = [0.0f64; 32];
let mut good_vec;
let mut bad_vec;
let mut weight_vec;
let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) =
if n_choices <= 32 {
(
&mut good_buf[..n_choices],
&mut bad_buf[..n_choices],
&mut weight_buf[..n_choices],
)
} else {
good_vec = vec![0usize; n_choices];
bad_vec = vec![0usize; n_choices];
weight_vec = vec![0.0f64; n_choices];
(&mut good_vec, &mut bad_vec, &mut weight_vec)
};
// Count occurrences in good and bad groups
for &idx in good_indices {
if idx < n_choices {
good_counts[idx] += 1;
}
}
for &idx in bad_indices {
if idx < n_choices {
bad_counts[idx] += 1;
}
}
// Add smoothing (Laplace smoothing) to avoid zero probabilities
let good_total = good_indices.len() as f64 + n_choices as f64;
let bad_total = bad_indices.len() as f64 + n_choices as f64;
// Calculate l(x)/g(x) ratio for each category
for i in 0..n_choices {
let l_prob = (good_counts[i] as f64 + 1.0) / good_total;
let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total;
weights[i] = l_prob / g_prob;
}
// Sample proportionally to weights
let total_weight: f64 = weights.iter().sum();
let threshold = rng.f64() * total_weight;
let mut cumulative = 0.0;
for (i, &w) in weights.iter().enumerate() {
cumulative += w;
if cumulative >= threshold {
return i;
}
}
// Fallback to last index (shouldn't happen)
n_choices - 1
}
+670
View File
@@ -0,0 +1,670 @@
use core::fmt::Debug;
use crate::Error;
/// A strategy for computing the gamma quantile in TPE.
///
/// The gamma value determines what fraction of trials are considered "good"
/// when splitting the trial history. Different strategies can adapt this
/// fraction based on the number of completed trials.
///
/// # Implementation Notes
///
/// - The returned gamma must be in the range (0.0, 1.0)
/// - Implementations should be deterministic for reproducibility
/// - The `clone_box` method enables trait object cloning
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::GammaStrategy;
///
/// #[derive(Debug, Clone)]
/// struct ConstantGamma(f64);
///
/// impl GammaStrategy for ConstantGamma {
/// fn gamma(&self, _n_trials: usize) -> f64 {
/// self.0
/// }
///
/// fn clone_box(&self) -> Box<dyn GammaStrategy> {
/// Box::new(self.clone())
/// }
/// }
/// ```
pub trait GammaStrategy: Send + Sync + Debug {
/// Computes the gamma quantile based on the number of completed trials.
///
/// # Arguments
///
/// * `n_trials` - The number of completed trials in the history.
///
/// # Returns
///
/// A gamma value in the range (0.0, 1.0). Values outside this range
/// will be clamped by the sampler.
fn gamma(&self, n_trials: usize) -> f64;
/// Creates a boxed clone of this strategy.
///
/// This method enables cloning of trait objects, which is necessary
/// for the builder pattern and sampler configuration.
fn clone_box(&self) -> Box<dyn GammaStrategy>;
}
impl Clone for Box<dyn GammaStrategy> {
fn clone(&self) -> Self {
self.clone_box()
}
}
/// A fixed gamma strategy that returns a constant value.
///
/// This is the simplest strategy and the default behavior of TPE.
/// The gamma value remains constant regardless of the number of trials.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{FixedGamma, TpeSampler};
///
/// // Use 15% of trials as "good"
/// let sampler = TpeSampler::builder()
/// .gamma_strategy(FixedGamma::new(0.15).unwrap())
/// .build()
/// .unwrap();
/// ```
#[derive(Debug, Clone, Copy)]
pub struct FixedGamma {
gamma: f64,
}
impl FixedGamma {
/// Creates a new fixed gamma strategy.
///
/// # Arguments
///
/// * `gamma` - The constant gamma value to use.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if gamma is not in (0.0, 1.0).
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::FixedGamma;
///
/// let strategy = FixedGamma::new(0.25).unwrap();
/// assert!((strategy.value() - 0.25).abs() < f64::EPSILON);
/// ```
pub fn new(gamma: f64) -> crate::Result<Self> {
if gamma <= 0.0 || gamma >= 1.0 {
return Err(Error::InvalidGamma(gamma));
}
Ok(Self { gamma })
}
/// Returns the fixed gamma value.
#[must_use]
pub fn value(&self) -> f64 {
self.gamma
}
}
impl Default for FixedGamma {
/// Creates a fixed gamma strategy with the default value of 0.25.
fn default() -> Self {
Self { gamma: 0.25 }
}
}
impl GammaStrategy for FixedGamma {
fn gamma(&self, _n_trials: usize) -> f64 {
self.gamma
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A linear gamma strategy that interpolates between min and max values.
///
/// The gamma value increases linearly from `gamma_min` to `gamma_max` as the
/// number of trials grows from 0 to `n_trials_max`. Beyond `n_trials_max`,
/// gamma remains at `gamma_max`.
///
/// This strategy is useful when you want to be more explorative early on
/// (smaller gamma = fewer "good" trials) and more exploitative later
/// (larger gamma = more "good" trials).
///
/// # Formula
///
/// ```text
/// gamma = gamma_min + (gamma_max - gamma_min) * min(n_trials / n_trials_max, 1.0)
/// ```
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, LinearGamma, TpeSampler};
///
/// let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap();
///
/// // At 0 trials: gamma = 0.1
/// assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON);
///
/// // At 50 trials: gamma = 0.25 (midpoint)
/// assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON);
///
/// // At 100+ trials: gamma = 0.4
/// assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON);
/// assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON);
/// ```
#[derive(Debug, Clone, Copy)]
pub struct LinearGamma {
gamma_min: f64,
gamma_max: f64,
n_trials_max: usize,
}
impl LinearGamma {
/// Creates a new linear gamma strategy.
///
/// # Arguments
///
/// * `gamma_min` - The minimum gamma value (at 0 trials).
/// * `gamma_max` - The maximum gamma value (at `n_trials_max` trials).
/// * `n_trials_max` - The number of trials at which gamma reaches its maximum.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_min` is not in (0.0, 1.0)
/// - `gamma_max` is not in (0.0, 1.0)
/// - `gamma_min > gamma_max`
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::LinearGamma;
///
/// // Gamma goes from 0.1 to 0.3 over 50 trials
/// let strategy = LinearGamma::new(0.1, 0.3, 50).unwrap();
/// ```
pub fn new(gamma_min: f64, gamma_max: f64, n_trials_max: usize) -> crate::Result<Self> {
if gamma_min <= 0.0 || gamma_min >= 1.0 {
return Err(Error::InvalidGamma(gamma_min));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
if gamma_min > gamma_max {
return Err(Error::InvalidGamma(gamma_min));
}
Ok(Self {
gamma_min,
gamma_max,
n_trials_max,
})
}
/// Returns the minimum gamma value.
#[must_use]
pub fn gamma_min(&self) -> f64 {
self.gamma_min
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
/// Returns the number of trials at which gamma reaches its maximum.
#[must_use]
pub fn n_trials_max(&self) -> usize {
self.n_trials_max
}
}
impl Default for LinearGamma {
/// Creates a linear gamma strategy with default values:
/// - `gamma_min`: 0.10
/// - `gamma_max`: 0.25
/// - `n_trials_max`: 100
fn default() -> Self {
Self {
gamma_min: 0.10,
gamma_max: 0.25,
n_trials_max: 100,
}
}
}
impl GammaStrategy for LinearGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if self.n_trials_max == 0 {
return self.gamma_max;
}
let t = (n_trials as f64 / self.n_trials_max as f64).min(1.0);
self.gamma_min + (self.gamma_max - self.gamma_min) * t
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A square root gamma strategy inspired by Optuna's default behavior.
///
/// The gamma value is computed based on the inverse square root of the number
/// of trials, providing a balance between exploration and exploitation that
/// naturally adapts as more data becomes available.
///
/// # Formula
///
/// ```text
/// n_good = max(1, floor(gamma_factor / sqrt(n_trials)))
/// gamma = min(gamma_max, n_good / n_trials)
/// ```
///
/// When `n_trials` is 0, returns `gamma_max`.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, SqrtGamma, TpeSampler};
///
/// let strategy = SqrtGamma::default();
///
/// // Gamma decreases as trials increase
/// let g10 = strategy.gamma(10);
/// let g100 = strategy.gamma(100);
/// assert!(g10 > g100, "Gamma should decrease with more trials");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct SqrtGamma {
gamma_factor: f64,
gamma_max: f64,
}
impl SqrtGamma {
/// Creates a new square root gamma strategy.
///
/// # Arguments
///
/// * `gamma_factor` - The factor controlling how quickly gamma decreases.
/// Higher values mean more "good" trials at any given point.
/// * `gamma_max` - The maximum gamma value (used when `n_trials` is small).
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_factor` is not positive
/// - `gamma_max` is not in (0.0, 1.0)
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::SqrtGamma;
///
/// let strategy = SqrtGamma::new(1.0, 0.25).unwrap();
/// ```
pub fn new(gamma_factor: f64, gamma_max: f64) -> crate::Result<Self> {
if gamma_factor <= 0.0 {
return Err(Error::InvalidGamma(gamma_factor));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
Ok(Self {
gamma_factor,
gamma_max,
})
}
/// Returns the gamma factor.
#[must_use]
pub fn gamma_factor(&self) -> f64 {
self.gamma_factor
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
}
impl Default for SqrtGamma {
/// Creates a square root gamma strategy with default values:
/// - `gamma_factor`: 1.0
/// - `gamma_max`: 0.25
fn default() -> Self {
Self {
gamma_factor: 1.0,
gamma_max: 0.25,
}
}
}
impl GammaStrategy for SqrtGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if n_trials == 0 {
return self.gamma_max;
}
let n_good = (self.gamma_factor / (n_trials as f64).sqrt()).max(1.0);
(n_good / n_trials as f64).min(self.gamma_max)
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
/// A Hyperopt-style gamma strategy.
///
/// This strategy computes gamma as `min(gamma_max, (gamma_base + 1) / n_trials)`,
/// which is inspired by the original Hyperopt TPE implementation.
///
/// # Formula
///
/// ```text
/// gamma = min(gamma_max, (gamma_base + 1) / n_trials)
/// ```
///
/// When `n_trials` is 0, returns `gamma_max`.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::{GammaStrategy, HyperoptGamma};
///
/// // With gamma_base=24 and gamma_max=0.5:
/// // - At n=25: gamma = min(0.5, 25/25) = 0.5 (capped)
/// // - At n=100: gamma = min(0.5, 25/100) = 0.25
/// let strategy = HyperoptGamma::new(24.0, 0.5).unwrap();
///
/// // Early trials have higher gamma
/// let g50 = strategy.gamma(50);
/// let g200 = strategy.gamma(200);
/// assert!(g50 > g200, "Gamma should decrease with more trials");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct HyperoptGamma {
gamma_base: f64,
gamma_max: f64,
}
impl HyperoptGamma {
/// Creates a new Hyperopt-style gamma strategy.
///
/// # Arguments
///
/// * `gamma_base` - The base value added to 1 in the numerator.
/// * `gamma_max` - The maximum gamma value.
///
/// # Errors
///
/// Returns `Error::InvalidGamma` if:
/// - `gamma_base` is negative
/// - `gamma_max` is not in (0.0, 1.0)
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::HyperoptGamma;
///
/// let strategy = HyperoptGamma::new(24.0, 0.25).unwrap();
/// ```
pub fn new(gamma_base: f64, gamma_max: f64) -> crate::Result<Self> {
if gamma_base < 0.0 {
return Err(Error::InvalidGamma(gamma_base));
}
if gamma_max <= 0.0 || gamma_max >= 1.0 {
return Err(Error::InvalidGamma(gamma_max));
}
Ok(Self {
gamma_base,
gamma_max,
})
}
/// Returns the gamma base value.
#[must_use]
pub fn gamma_base(&self) -> f64 {
self.gamma_base
}
/// Returns the maximum gamma value.
#[must_use]
pub fn gamma_max(&self) -> f64 {
self.gamma_max
}
}
impl Default for HyperoptGamma {
/// Creates a Hyperopt-style gamma strategy with default values:
/// - `gamma_base`: 24.0
/// - `gamma_max`: 0.25
fn default() -> Self {
Self {
gamma_base: 24.0,
gamma_max: 0.25,
}
}
}
impl GammaStrategy for HyperoptGamma {
#[allow(clippy::cast_precision_loss)]
fn gamma(&self, n_trials: usize) -> f64 {
if n_trials == 0 {
return self.gamma_max;
}
((self.gamma_base + 1.0) / n_trials as f64).min(self.gamma_max)
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(*self)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sampler::tpe::TpeSampler;
#[test]
fn test_fixed_gamma_default() {
let strategy = FixedGamma::default();
assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);
assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);
assert!((strategy.value() - 0.25).abs() < f64::EPSILON);
}
#[test]
fn test_fixed_gamma_custom() {
let strategy = FixedGamma::new(0.15).unwrap();
assert!((strategy.gamma(0) - 0.15).abs() < f64::EPSILON);
assert!((strategy.gamma(50) - 0.15).abs() < f64::EPSILON);
assert!((strategy.gamma(1000) - 0.15).abs() < f64::EPSILON);
}
#[test]
fn test_fixed_gamma_invalid() {
assert!(FixedGamma::new(0.0).is_err());
assert!(FixedGamma::new(1.0).is_err());
assert!(FixedGamma::new(-0.1).is_err());
assert!(FixedGamma::new(1.5).is_err());
}
#[test]
fn test_linear_gamma_default() {
let strategy = LinearGamma::default();
assert!((strategy.gamma(0) - 0.10).abs() < f64::EPSILON);
assert!((strategy.gamma(50) - 0.175).abs() < f64::EPSILON); // midpoint
assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);
assert!((strategy.gamma(200) - 0.25).abs() < f64::EPSILON); // capped
}
#[test]
fn test_linear_gamma_custom() {
let strategy = LinearGamma::new(0.1, 0.4, 100).unwrap();
assert!((strategy.gamma(0) - 0.1).abs() < f64::EPSILON);
assert!((strategy.gamma(50) - 0.25).abs() < f64::EPSILON);
assert!((strategy.gamma(100) - 0.4).abs() < f64::EPSILON);
assert!((strategy.gamma(200) - 0.4).abs() < f64::EPSILON);
}
#[test]
fn test_linear_gamma_invalid() {
assert!(LinearGamma::new(0.0, 0.5, 100).is_err());
assert!(LinearGamma::new(0.1, 1.0, 100).is_err());
assert!(LinearGamma::new(0.5, 0.2, 100).is_err()); // min > max
}
#[test]
fn test_sqrt_gamma_default() {
let strategy = SqrtGamma::default();
// At n=0, returns gamma_max
assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);
// gamma decreases with more trials
let g10 = strategy.gamma(10);
let g100 = strategy.gamma(100);
assert!(g10 > g100);
}
#[test]
fn test_sqrt_gamma_custom() {
let strategy = SqrtGamma::new(2.0, 0.5).unwrap();
assert!((strategy.gamma(0) - 0.5).abs() < f64::EPSILON);
// At n=4: n_good = max(1, 2/2) = 1, gamma = 1/4 = 0.25
let g4 = strategy.gamma(4);
assert!((g4 - 0.25).abs() < f64::EPSILON);
}
#[test]
fn test_sqrt_gamma_invalid() {
assert!(SqrtGamma::new(0.0, 0.25).is_err()); // factor must be positive
assert!(SqrtGamma::new(-1.0, 0.25).is_err());
assert!(SqrtGamma::new(1.0, 0.0).is_err());
assert!(SqrtGamma::new(1.0, 1.0).is_err());
}
#[test]
fn test_hyperopt_gamma_default() {
let strategy = HyperoptGamma::default();
// At n=0, returns gamma_max
assert!((strategy.gamma(0) - 0.25).abs() < f64::EPSILON);
// At n=100: (24+1)/100 = 0.25, so capped to 0.25
assert!((strategy.gamma(100) - 0.25).abs() < f64::EPSILON);
// At n=200: (24+1)/200 = 0.125
assert!((strategy.gamma(200) - 0.125).abs() < f64::EPSILON);
}
#[test]
fn test_hyperopt_gamma_custom() {
let strategy = HyperoptGamma::new(9.0, 0.5).unwrap();
// At n=20: (9+1)/20 = 0.5, capped to 0.5
assert!((strategy.gamma(20) - 0.5).abs() < f64::EPSILON);
// At n=100: (9+1)/100 = 0.1
assert!((strategy.gamma(100) - 0.1).abs() < f64::EPSILON);
}
#[test]
fn test_hyperopt_gamma_invalid() {
assert!(HyperoptGamma::new(-1.0, 0.25).is_err());
assert!(HyperoptGamma::new(24.0, 0.0).is_err());
assert!(HyperoptGamma::new(24.0, 1.0).is_err());
}
#[test]
fn test_gamma_strategy_clone_box() {
let fixed: Box<dyn GammaStrategy> = Box::new(FixedGamma::new(0.3).unwrap());
let cloned = fixed.clone();
assert!((cloned.gamma(0) - 0.3).abs() < f64::EPSILON);
let linear: Box<dyn GammaStrategy> = Box::new(LinearGamma::default());
let cloned = linear.clone();
assert!((cloned.gamma(0) - 0.10).abs() < f64::EPSILON);
}
#[test]
fn test_tpe_with_linear_gamma_strategy() {
let sampler = TpeSampler::builder()
.gamma_strategy(LinearGamma::new(0.1, 0.3, 50).unwrap())
.n_startup_trials(5)
.seed(42)
.build()
.unwrap();
// Verify the strategy is applied
let g = sampler.gamma_strategy().gamma(25);
assert!((g - 0.2).abs() < f64::EPSILON); // midpoint of 0.1 to 0.3
}
#[test]
fn test_gamma_overrides_gamma_strategy() {
// When gamma() is called after gamma_strategy(), it should take precedence
let sampler = TpeSampler::builder()
.gamma_strategy(SqrtGamma::default())
.gamma(0.15) // This should override
.build()
.unwrap();
// Should use fixed gamma of 0.15
assert!((sampler.gamma_strategy().gamma(0) - 0.15).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(100) - 0.15).abs() < f64::EPSILON);
}
#[test]
fn test_gamma_strategy_overrides_gamma() {
// When gamma_strategy() is called after gamma(), it should take precedence
let sampler = TpeSampler::builder()
.gamma(0.15)
.gamma_strategy(SqrtGamma::default()) // This should override
.build()
.unwrap();
// Should use SqrtGamma - gamma decreases with trials
let g10 = sampler.gamma_strategy().gamma(10);
let g100 = sampler.gamma_strategy().gamma(100);
assert!(g10 > g100, "SqrtGamma should decrease with more trials");
}
#[test]
fn test_custom_gamma_strategy() {
#[derive(Debug, Clone)]
struct DoubleGamma;
impl GammaStrategy for DoubleGamma {
fn gamma(&self, n_trials: usize) -> f64 {
// Double the trial count-based calculation, capped at 0.5
#[allow(clippy::cast_precision_loss)]
(0.01 * n_trials as f64).min(0.5)
}
fn clone_box(&self) -> Box<dyn GammaStrategy> {
Box::new(self.clone())
}
}
let sampler = TpeSampler::builder()
.gamma_strategy(DoubleGamma)
.build()
.unwrap();
assert!((sampler.gamma_strategy().gamma(10) - 0.1).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(50) - 0.5).abs() < f64::EPSILON);
assert!((sampler.gamma_strategy().gamma(100) - 0.5).abs() < f64::EPSILON);
}
}
+73
View File
@@ -0,0 +1,73 @@
//! Tree-Parzen Estimator (TPE) sampler family for Bayesian optimization.
//!
//! TPE is a sequential model-based optimization algorithm that models P(x|y) instead
//! of P(y|x). It splits completed trials into "good" (below the gamma quantile) and
//! "bad" groups, fits a kernel density estimator (KDE) to each, and proposes new
//! points by maximizing the l(x)/g(x) ratio — an approximation of Expected Improvement.
//!
//! # Samplers
//!
//! | Sampler | Models parameters | Best for |
//! |---------|-------------------|----------|
//! | [`TpeSampler`] | Independently | General-purpose single-objective optimization |
//! | [`MultivariateTpeSampler`] | Jointly | Problems with correlated parameters |
//!
//! # Gamma strategies
//!
//! The gamma quantile controls how many trials are considered "good". This module
//! provides four built-in strategies via the [`GammaStrategy`] trait:
//!
//! | Strategy | Formula | Default |
//! |----------|---------|---------|
//! | [`FixedGamma`] | Constant value | gamma = 0.25 |
//! | [`LinearGamma`] | Linear ramp from min to max | 0.10 → 0.25 over 100 trials |
//! | [`SqrtGamma`] | 1/√n decay (Optuna-style) | factor = 1.0, max = 0.25 |
//! | [`HyperoptGamma`] | (base+1)/n (Hyperopt-style) | base = 24, max = 0.25 |
//!
//! You can also implement [`GammaStrategy`] for a custom splitting rule.
//!
//! # Search-space utilities
//!
//! The [`search_space`] submodule provides [`IntersectionSearchSpace`] for computing
//! the common parameter set across trials, and [`GroupDecomposedSearchSpace`] for
//! splitting parameters into independent groups based on co-occurrence.
//!
//! # Examples
//!
//! Basic TPE with default settings:
//!
//! ```
//! use optimizer::sampler::tpe::TpeSampler;
//! use optimizer::{Direction, Study};
//!
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, TpeSampler::new());
//! ```
//!
//! Multivariate TPE for correlated parameters:
//!
//! ```
//! use optimizer::sampler::tpe::MultivariateTpeSampler;
//! use optimizer::{Direction, Study};
//!
//! let sampler = MultivariateTpeSampler::builder()
//! .gamma(0.15)
//! .n_startup_trials(20)
//! .group(true)
//! .seed(42)
//! .build()
//! .unwrap();
//! let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
//! ```
pub(crate) mod common;
mod gamma;
mod multivariate;
mod sampler;
pub mod search_space;
pub use gamma::{FixedGamma, GammaStrategy, HyperoptGamma, LinearGamma, SqrtGamma};
pub use multivariate::{
ConstantLiarStrategy, MultivariateTpeSampler, MultivariateTpeSamplerBuilder,
};
pub use sampler::{TpeSampler, TpeSamplerBuilder};
pub use search_space::{GroupDecomposedSearchSpace, IntersectionSearchSpace};
+623
View File
@@ -0,0 +1,623 @@
//! Core multivariate TPE sampling logic.
//!
//! Contains the main sampling engine: group decomposition, single-group multivariate
//! TPE, candidate selection, independent fallbacks, and value conversion.
use std::collections::HashMap;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::parameter::ParamId;
use crate::sampler::CompletedTrial;
use super::MultivariateTpeSampler;
impl MultivariateTpeSampler {
/// Samples parameters by decomposing the search space into independent groups.
///
/// When `group=true`, this method analyzes the trial history to identify groups of
/// parameters that always appear together, then samples each group independently
/// using multivariate TPE. This is more efficient when parameters naturally partition
/// into independent subsets (e.g., due to conditional search spaces).
///
/// # Arguments
///
/// * `search_space` - The full search space containing all parameters to sample.
/// * `history` - Completed trials from the optimization history.
///
/// # Returns
///
/// A `HashMap` mapping parameter names to their sampled values.
pub(crate) fn sample_with_groups(
&self,
search_space: &HashMap<ParamId, Distribution>,
history: &[CompletedTrial],
) -> HashMap<ParamId, ParamValue> {
use std::collections::HashSet;
use crate::sampler::tpe::GroupDecomposedSearchSpace;
// Decompose the search space into independent parameter groups
let groups = GroupDecomposedSearchSpace::calculate(history);
let mut result: HashMap<ParamId, ParamValue> = HashMap::new();
// Sample each group independently
for group in &groups {
// Build a sub-search space for this group
let group_search_space: HashMap<ParamId, Distribution> = search_space
.iter()
.filter(|(id, _)| group.contains(id))
.map(|(id, dist)| (*id, dist.clone()))
.collect();
if group_search_space.is_empty() {
continue;
}
// Filter history to trials that have at least one parameter in this group
let group_history: Vec<&CompletedTrial> = history
.iter()
.filter(|trial| {
trial
.distributions
.keys()
.any(|param_id| group.contains(param_id))
})
.collect();
// Build completed trials from references for the group
// We need to create a temporary slice for sample_group_internal
let group_history_owned: Vec<CompletedTrial> =
group_history.iter().map(|t| (*t).clone()).collect();
// Sample this group using multivariate TPE
let mut rng = self.rng.lock();
let group_result =
self.sample_single_group(&group_search_space, &group_history_owned, &mut rng);
drop(rng);
// Merge group results into the main result
for (id, value) in group_result {
result.insert(id, value);
}
}
// Handle parameters not in any group (sample independently)
let grouped_params: HashSet<ParamId> = groups.iter().flatten().copied().collect();
let ungrouped_params: HashMap<ParamId, Distribution> = search_space
.iter()
.filter(|(id, _)| !grouped_params.contains(id) && !result.contains_key(id))
.map(|(id, dist)| (*id, dist.clone()))
.collect();
if !ungrouped_params.is_empty() {
// Sample ungrouped parameters uniformly (no history for them)
let mut rng = self.rng.lock();
for (id, dist) in &ungrouped_params {
let value = crate::sampler::common::sample_random(&mut rng, dist);
result.insert(*id, value);
}
}
result
}
/// Validates observations and fits multivariate KDEs for good and bad groups.
///
/// Returns `None` if observations are invalid or KDE construction fails.
fn try_fit_kdes(
good_obs: Vec<Vec<f64>>,
bad_obs: Vec<Vec<f64>>,
expected_dims: usize,
) -> Option<(crate::kde::MultivariateKDE, crate::kde::MultivariateKDE)> {
use crate::kde::MultivariateKDE;
let valid = !good_obs.is_empty()
&& !bad_obs.is_empty()
&& good_obs.iter().all(|obs| obs.len() == expected_dims)
&& bad_obs.iter().all(|obs| obs.len() == expected_dims);
if !valid {
return None;
}
let good_kde = MultivariateKDE::new(good_obs).ok()?;
let bad_kde = MultivariateKDE::new(bad_obs).ok()?;
Some((good_kde, bad_kde))
}
/// Samples parameters as a single group using multivariate TPE.
///
/// This is the core multivariate TPE sampling logic, used both in non-grouped mode
/// and for sampling individual groups in grouped mode.
///
/// # Arguments
///
/// * `search_space` - The search space for this group of parameters.
/// * `history` - Completed trials to use for model fitting.
/// * `rng` - Random number generator (caller must hold lock).
///
/// # Returns
///
/// A `HashMap` mapping parameter names to their sampled values.
pub(crate) fn sample_single_group(
&self,
search_space: &HashMap<ParamId, Distribution>,
history: &[CompletedTrial],
rng: &mut fastrand::Rng,
) -> HashMap<ParamId, ParamValue> {
use crate::sampler::tpe::IntersectionSearchSpace;
use crate::sampler::tpe::common;
// Early returns for cases requiring random sampling
if history.len() < self.n_startup_trials {
return self.sample_all_uniform(search_space, rng);
}
let intersection = IntersectionSearchSpace::calculate(history);
if intersection.is_empty() {
return self.sample_all_independent_with_rng(search_space, history, rng);
}
let filtered = self.filter_trials(history, &intersection);
if filtered.len() < 2 {
return self.sample_all_independent_with_rng(search_space, history, rng);
}
let (good, bad) = self.split_trials(&filtered);
// Sample categorical parameters using TPE with l(x)/g(x) ratio
let mut result: HashMap<ParamId, ParamValue> = HashMap::new();
for (param_id, dist) in &intersection {
if let Distribution::Categorical(d) = dist {
let good_indices = Self::extract_categorical_indices(&good, *param_id);
let bad_indices = Self::extract_categorical_indices(&bad, *param_id);
let idx =
common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng);
result.insert(*param_id, ParamValue::Categorical(idx));
}
}
// Collect continuous parameters
let mut param_order: Vec<ParamId> = intersection
.iter()
.filter(|(_, dist)| !matches!(dist, Distribution::Categorical(_)))
.map(|(id, _)| *id)
.collect();
if param_order.is_empty() {
self.fill_remaining_independent_with_rng(
search_space,
&intersection,
history,
&mut result,
rng,
);
return result;
}
param_order.sort();
// Extract observations, validate, and fit KDEs
let good_obs = self.extract_observations(&good, &param_order);
let bad_obs = self.extract_observations(&bad, &param_order);
let Some((good_kde, bad_kde)) = Self::try_fit_kdes(good_obs, bad_obs, param_order.len())
else {
self.fill_remaining_independent_with_rng(
search_space,
&intersection,
history,
&mut result,
rng,
);
return result;
};
// Compute parameter bounds for each dimension so candidates are clamped
let bounds: Vec<(f64, f64)> = param_order
.iter()
.filter_map(|id| {
intersection.get(id).and_then(|dist| match dist {
Distribution::Float(d) => Some((d.low, d.high)),
#[allow(clippy::cast_precision_loss)]
Distribution::Int(d) => Some((d.low as f64, d.high as f64)),
Distribution::Categorical(_) => None,
})
})
.collect();
let selected = self.select_candidate_with_rng(&good_kde, &bad_kde, &bounds, rng);
// Map selected values to parameter ids
for (idx, param_id) in param_order.iter().enumerate() {
if let Some(dist) = intersection.get(param_id) {
let value = selected[idx];
let param_value = self.convert_to_param_value(value, dist);
if let Some(pv) = param_value {
result.insert(*param_id, pv);
}
}
}
// Fill remaining parameters using independent TPE sampling
self.fill_remaining_independent_with_rng(
search_space,
&intersection,
history,
&mut result,
rng,
);
result
}
/// Converts a raw f64 value to a `ParamValue` based on the distribution.
#[allow(clippy::unused_self)]
pub(crate) fn convert_to_param_value(
&self,
value: f64,
dist: &Distribution,
) -> Option<ParamValue> {
match dist {
Distribution::Float(d) => {
let clamped = value.clamp(d.low, d.high);
let stepped = if let Some(step) = d.step {
let steps = ((clamped - d.low) / step).round();
(d.low + steps * step).clamp(d.low, d.high)
} else {
clamped
};
Some(ParamValue::Float(stepped))
}
Distribution::Int(d) => {
#[allow(clippy::cast_possible_truncation)]
let int_value = value.round() as i64;
let clamped = int_value.clamp(d.low, d.high);
let stepped = if let Some(step) = d.step {
let steps = (clamped - d.low) / step;
(d.low + steps * step).clamp(d.low, d.high)
} else {
clamped
};
Some(ParamValue::Int(stepped))
}
Distribution::Categorical(_) => None,
}
}
/// Clamps each dimension of a candidate to the corresponding parameter bounds.
fn clamp_candidate(candidate: &mut [f64], bounds: &[(f64, f64)]) {
for (val, &(lo, hi)) in candidate.iter_mut().zip(bounds.iter()) {
*val = val.clamp(lo, hi);
}
}
/// Selects the best candidate from a set of samples using the joint acquisition function.
///
/// This method implements the core TPE selection criterion: it generates candidates
/// from the "good" KDE (l(x)) and selects the one that maximizes the ratio l(x)/g(x),
/// which is equivalent to maximizing `log(l(x)) - log(g(x))`.
///
/// Candidates are clamped to parameter bounds before evaluation so the acquisition
/// function scores the values that will actually be used.
#[must_use]
#[cfg(test)]
pub(crate) fn select_candidate(
&self,
good_kde: &crate::kde::MultivariateKDE,
bad_kde: &crate::kde::MultivariateKDE,
bounds: &[(f64, f64)],
) -> Vec<f64> {
let mut rng = self.rng.lock();
// Generate candidates from the good distribution, clamped to bounds
let candidates: Vec<Vec<f64>> = (0..self.n_ei_candidates)
.map(|_| {
let mut c = good_kde.sample(&mut rng);
Self::clamp_candidate(&mut c, bounds);
c
})
.collect();
// Compute log(l(x)) - log(g(x)) for each candidate
// This is equivalent to log(l(x)/g(x)) which we want to maximize
let log_ratios: Vec<f64> = candidates
.iter()
.map(|candidate| {
let log_l = good_kde.log_pdf(candidate);
let log_g = bad_kde.log_pdf(candidate);
log_l - log_g
})
.collect();
// Find the candidate with the maximum log ratio
let mut best_idx = 0;
let mut best_ratio = f64::NEG_INFINITY;
for (idx, &ratio) in log_ratios.iter().enumerate() {
// Handle NaN by treating it as worse than any finite value
if ratio > best_ratio || (best_ratio.is_nan() && !ratio.is_nan()) {
best_ratio = ratio;
best_idx = idx;
}
}
candidates.into_iter().nth(best_idx).unwrap_or_default()
}
/// Selects the best candidate using an external RNG.
///
/// This variant accepts an external RNG, used when the caller already holds the lock.
/// Candidates are clamped to parameter bounds before evaluation.
pub(crate) fn select_candidate_with_rng(
&self,
good_kde: &crate::kde::MultivariateKDE,
bad_kde: &crate::kde::MultivariateKDE,
bounds: &[(f64, f64)],
rng: &mut fastrand::Rng,
) -> Vec<f64> {
// Generate candidates from the good distribution, clamped to bounds
let candidates: Vec<Vec<f64>> = (0..self.n_ei_candidates)
.map(|_| {
let mut c = good_kde.sample(rng);
Self::clamp_candidate(&mut c, bounds);
c
})
.collect();
// Compute log(l(x)) - log(g(x)) for each candidate
let log_ratios: Vec<f64> = candidates
.iter()
.map(|candidate| {
let log_l = good_kde.log_pdf(candidate);
let log_g = bad_kde.log_pdf(candidate);
log_l - log_g
})
.collect();
// Find the candidate with the maximum log ratio
let mut best_idx = 0;
let mut best_ratio = f64::NEG_INFINITY;
for (idx, &ratio) in log_ratios.iter().enumerate() {
if ratio > best_ratio || (best_ratio.is_nan() && !ratio.is_nan()) {
best_ratio = ratio;
best_idx = idx;
}
}
candidates.into_iter().nth(best_idx).unwrap_or_default()
}
/// Fills remaining parameters using independent TPE sampling with an external RNG.
///
/// This variant accepts an external RNG, used when the caller already holds the lock.
pub(crate) fn fill_remaining_independent_with_rng(
&self,
search_space: &HashMap<ParamId, Distribution>,
_intersection: &HashMap<ParamId, Distribution>,
history: &[CompletedTrial],
result: &mut HashMap<ParamId, ParamValue>,
rng: &mut fastrand::Rng,
) {
// Identify parameters not in result, sorted for deterministic RNG consumption
let mut missing_params: Vec<(&ParamId, &Distribution)> = search_space
.iter()
.filter(|(id, _)| !result.contains_key(id))
.collect();
if missing_params.is_empty() {
return;
}
missing_params.sort_by_key(|(id, _)| *id);
// Split trials for independent sampling
let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::<Vec<_>>());
for (param_id, dist) in missing_params {
let value =
self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng);
result.insert(*param_id, value);
}
}
/// Samples all parameters using independent TPE sampling.
///
/// This is used as a complete fallback when no intersection search space exists.
/// Parameters are sorted by `ParamId` for deterministic RNG consumption order.
#[cfg(test)]
pub(crate) fn sample_all_independent(
&self,
search_space: &HashMap<ParamId, Distribution>,
history: &[CompletedTrial],
) -> HashMap<ParamId, ParamValue> {
// Split trials for independent sampling
let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::<Vec<_>>());
let mut rng = self.rng.lock();
let mut result = HashMap::new();
let mut sorted: Vec<_> = search_space.iter().collect();
sorted.sort_by_key(|(id, _)| *id);
for (param_id, dist) in sorted {
let value =
self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, &mut rng);
result.insert(*param_id, value);
}
result
}
/// Samples all parameters using independent TPE sampling with an external RNG.
///
/// This variant accepts an external RNG, used when the caller already holds the lock.
/// Parameters are sorted by `ParamId` for deterministic RNG consumption order.
pub(crate) fn sample_all_independent_with_rng(
&self,
search_space: &HashMap<ParamId, Distribution>,
history: &[CompletedTrial],
rng: &mut fastrand::Rng,
) -> HashMap<ParamId, ParamValue> {
// Split trials for independent sampling
let (good_trials, bad_trials) = self.split_trials(&history.iter().collect::<Vec<_>>());
let mut result = HashMap::new();
let mut sorted: Vec<_> = search_space.iter().collect();
sorted.sort_by_key(|(id, _)| *id);
for (param_id, dist) in sorted {
let value =
self.sample_independent_tpe(*param_id, dist, &good_trials, &bad_trials, rng);
result.insert(*param_id, value);
}
result
}
/// Samples a single parameter using independent TPE.
///
/// This method extracts values for the given parameter from good and bad trials,
/// fits univariate KDEs, and samples using the TPE acquisition function.
pub(crate) fn sample_independent_tpe(
&self,
param_id: ParamId,
distribution: &Distribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
match distribution {
Distribution::Float(d) => {
self.sample_independent_float(param_id, d, good_trials, bad_trials, rng)
}
Distribution::Int(d) => {
self.sample_independent_int(param_id, d, good_trials, bad_trials, rng)
}
Distribution::Categorical(d) => {
self.sample_independent_categorical(param_id, d, good_trials, bad_trials, rng)
}
}
}
fn sample_independent_float(
&self,
param_id: ParamId,
d: &crate::distribution::FloatDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
use crate::sampler::tpe::common;
let good_values: Vec<f64> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_id))
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<f64> = bad_trials
.iter()
.filter_map(|t| t.params.get(&param_id))
.filter_map(|v| match v {
ParamValue::Float(f) => Some(*f),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return crate::sampler::common::sample_random(rng, &Distribution::Float(d.clone()));
}
let value =
common::sample_tpe_float(d, good_values, bad_values, self.n_ei_candidates, None, rng);
ParamValue::Float(value)
}
fn sample_independent_int(
&self,
param_id: ParamId,
d: &crate::distribution::IntDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
use crate::sampler::tpe::common;
let good_values: Vec<i64> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_id))
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
let bad_values: Vec<i64> = bad_trials
.iter()
.filter_map(|t| t.params.get(&param_id))
.filter_map(|v| match v {
ParamValue::Int(i) => Some(*i),
_ => None,
})
.filter(|&v| v >= d.low && v <= d.high)
.collect();
if good_values.is_empty() || bad_values.is_empty() {
return crate::sampler::common::sample_random(rng, &Distribution::Int(d.clone()));
}
let value =
common::sample_tpe_int(d, good_values, bad_values, self.n_ei_candidates, None, rng);
ParamValue::Int(value)
}
#[allow(clippy::unused_self)]
fn sample_independent_categorical(
&self,
param_id: ParamId,
d: &crate::distribution::CategoricalDistribution,
good_trials: &[&CompletedTrial],
bad_trials: &[&CompletedTrial],
rng: &mut fastrand::Rng,
) -> ParamValue {
use crate::sampler::tpe::common;
let good_indices: Vec<usize> = good_trials
.iter()
.filter_map(|t| t.params.get(&param_id))
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
let bad_indices: Vec<usize> = bad_trials
.iter()
.filter_map(|t| t.params.get(&param_id))
.filter_map(|v| match v {
ParamValue::Categorical(i) => Some(*i),
_ => None,
})
.filter(|&i| i < d.n_choices)
.collect();
if good_indices.is_empty() || bad_indices.is_empty() {
return crate::sampler::common::sample_random(
rng,
&Distribution::Categorical(d.clone()),
);
}
let idx = common::sample_tpe_categorical(d.n_choices, &good_indices, &bad_indices, rng);
ParamValue::Categorical(idx)
}
}
File diff suppressed because it is too large Load Diff
+220
View File
@@ -0,0 +1,220 @@
//! Trial processing for the multivariate TPE sampler.
//!
//! Contains constant-liar imputation, trial filtering, good/bad splitting,
//! observation extraction, and categorical index extraction.
use std::collections::HashMap;
use crate::distribution::Distribution;
use crate::param::ParamValue;
use crate::parameter::ParamId;
use crate::sampler::{CompletedTrial, PendingTrial};
use super::{ConstantLiarStrategy, MultivariateTpeSampler};
impl MultivariateTpeSampler {
/// Imputes objective values for pending trials based on the constant liar strategy.
///
/// In parallel optimization, multiple trials may be running simultaneously. This method
/// assigns "lie" values to pending trials so they can be included in the model fitting,
/// which helps avoid redundant exploration of the same region.
///
/// # Arguments
///
/// * `pending_trials` - Trials that are currently running and have no objective value yet.
/// * `completed_trials` - Trials that have completed and have objective values.
///
/// # Returns
///
/// A vector of `CompletedTrial` objects containing both the original completed trials
/// and the pending trials with imputed values. If the strategy is `None`, returns
/// only the completed trials (pending trials are ignored).
#[must_use]
pub fn impute_pending_trials(
&self,
pending_trials: &[PendingTrial],
completed_trials: &[CompletedTrial],
) -> Vec<CompletedTrial> {
// Start with a copy of completed trials
let mut result: Vec<CompletedTrial> = completed_trials.to_vec();
// If strategy is None or no pending trials, just return completed trials
if matches!(self.constant_liar, ConstantLiarStrategy::None) || pending_trials.is_empty() {
return result;
}
// Compute the imputation value based on strategy
let imputed_value = self.compute_imputation_value(completed_trials);
// Convert pending trials to completed trials with imputed values
for pending in pending_trials {
result.push(CompletedTrial::new(
pending.id,
pending.params.clone(),
pending.distributions.clone(),
HashMap::new(),
imputed_value,
));
}
result
}
/// Computes the imputation value based on the constant liar strategy.
///
/// This is a helper method used by [`impute_pending_trials`](Self::impute_pending_trials).
#[allow(clippy::cast_precision_loss)]
pub(crate) fn compute_imputation_value(&self, completed_trials: &[CompletedTrial]) -> f64 {
match self.constant_liar {
ConstantLiarStrategy::None => 0.0, // This case is handled before calling this method
ConstantLiarStrategy::Mean => {
if completed_trials.is_empty() {
0.0
} else {
let sum: f64 = completed_trials.iter().map(|t| t.value).sum();
sum / completed_trials.len() as f64
}
}
ConstantLiarStrategy::Best => {
// Best means minimum for minimization problems
completed_trials
.iter()
.map(|t| t.value)
.fold(f64::INFINITY, f64::min)
}
ConstantLiarStrategy::Worst => {
// Worst means maximum for minimization problems
completed_trials
.iter()
.map(|t| t.value)
.fold(f64::NEG_INFINITY, f64::max)
}
ConstantLiarStrategy::Custom(v) => v,
}
}
/// Filters trials to those containing all parameters in the search space.
///
/// Only trials that contain ALL parameters in the search space are included,
/// ensuring we can model the joint distribution over all parameters.
#[must_use]
pub fn filter_trials<'a>(
&self,
history: &'a [CompletedTrial],
search_space: &HashMap<ParamId, Distribution>,
) -> Vec<&'a CompletedTrial> {
history
.iter()
.filter(|trial| {
// Include trial only if it has ALL parameters in the search space
search_space
.keys()
.all(|param_id| trial.params.contains_key(param_id))
})
.collect()
}
/// Splits filtered trials into good and bad groups based on the gamma quantile.
///
/// The gamma value is computed dynamically using the configured [`super::GammaStrategy`].
/// Trials are sorted by objective value (ascending for minimization), and the
/// gamma quantile determines the split point.
#[allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss
)]
#[must_use]
pub fn split_trials<'a>(
&self,
trials: &[&'a CompletedTrial],
) -> (Vec<&'a CompletedTrial>, Vec<&'a CompletedTrial>) {
if trials.is_empty() {
return (vec![], vec![]);
}
// Sort trials by objective value (ascending for minimization)
let mut sorted_indices: Vec<usize> = (0..trials.len()).collect();
sorted_indices.sort_by(|&a, &b| {
trials[a]
.value
.partial_cmp(&trials[b].value)
.unwrap_or(core::cmp::Ordering::Equal)
});
// Compute gamma using the strategy and clamp to valid range
let gamma = self
.gamma_strategy
.gamma(trials.len())
.clamp(f64::EPSILON, 1.0 - f64::EPSILON);
// Calculate the split point (gamma quantile)
// Ensure at least 1 trial in each group if possible
let n_good = ((trials.len() as f64 * gamma).ceil() as usize)
.max(1)
.min(trials.len().saturating_sub(1));
// Handle edge case: if we have only 1 trial, put it in good
if trials.len() == 1 {
return (vec![trials[0]], vec![]);
}
let good: Vec<_> = sorted_indices[..n_good]
.iter()
.map(|&i| trials[i])
.collect();
let bad: Vec<_> = sorted_indices[n_good..]
.iter()
.map(|&i| trials[i])
.collect();
(good, bad)
}
/// Extracts parameter values from trials as a numeric observation matrix.
///
/// Each row in the output represents one trial's parameter values in the specified order.
/// Categorical parameters are skipped.
#[must_use]
#[allow(clippy::cast_precision_loss)]
pub fn extract_observations(
&self,
trials: &[&CompletedTrial],
param_order: &[ParamId],
) -> Vec<Vec<f64>> {
trials
.iter()
.map(|trial| {
param_order
.iter()
.filter_map(|param_id| {
trial.params.get(param_id).and_then(|value| match value {
crate::param::ParamValue::Float(f) => Some(*f),
crate::param::ParamValue::Int(i) => Some(*i as f64),
crate::param::ParamValue::Categorical(_) => None, // Skip categorical
})
})
.collect()
})
.collect()
}
/// Extracts categorical indices from trials for a specific parameter.
pub(crate) fn extract_categorical_indices(
trials: &[&CompletedTrial],
param_id: ParamId,
) -> Vec<usize> {
trials
.iter()
.filter_map(|trial| {
trial.params.get(&param_id).and_then(|value| {
if let ParamValue::Categorical(idx) = value {
Some(*idx)
} else {
None
}
})
})
.collect()
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+362
View File
@@ -0,0 +1,362 @@
//! JSONL-based journal storage backend.
//!
//! [`JournalStorage`] persists completed trials as one JSON object per
//! line ([JSONL / JSON Lines](https://jsonlines.org/)) while keeping a
//! full copy in memory for fast read access.
//!
//! # File format
//!
//! Each line is a self-contained JSON serialization of a
//! [`CompletedTrial<V>`](crate::sampler::CompletedTrial). The file
//! is append-only — no existing lines are ever modified or deleted.
//!
//! ```text
//! {"id":0,"params":{...},"value":1.23,"state":"Completed",...}
//! {"id":1,"params":{...},"value":0.87,"state":"Completed",...}
//! ```
//!
//! # File locking
//!
//! Concurrent access is coordinated with `fs2` file locks:
//!
//! - **Writes** acquire an *exclusive* lock so only one process
//! appends at a time.
//! - **Reads** ([`refresh`](super::Storage::refresh)) acquire a
//! *shared* lock so readers never see a partially written line.
//!
//! This makes it safe for multiple processes to share the same JSONL
//! file — for example, distributed workers each running their own
//! [`Study`](crate::Study) with a `JournalStorage` pointing to a
//! shared path.
//!
//! # Resuming a study
//!
//! Use [`JournalStorage::open`] to reload previously persisted trials
//! and continue optimization from where you left off:
//!
//! ```no_run
//! use optimizer::prelude::*;
//! use optimizer::storage::JournalStorage;
//!
//! // First run — creates the file.
//! let storage = JournalStorage::<f64>::new("trials.jsonl");
//! let mut study = Study::builder().minimize().storage(storage).build();
//! study
//! .optimize(50, |trial: &mut optimizer::Trial| {
//! let x = FloatParam::new(-5.0, 5.0).suggest(trial)?;
//! Ok::<_, optimizer::Error>(x * x)
//! })
//! .unwrap();
//!
//! // Later run — reloads previous 50 trials, then adds 50 more.
//! let storage = JournalStorage::<f64>::open("trials.jsonl").unwrap();
//! let mut study = Study::builder().minimize().storage(storage).build();
//! study
//! .optimize(50, |trial: &mut optimizer::Trial| {
//! let x = FloatParam::new(-5.0, 5.0).suggest(trial)?;
//! Ok::<_, optimizer::Error>(x * x)
//! })
//! .unwrap();
//! ```
//!
//! # When to use
//!
//! - **Persistence** — survive process crashes or intentional restarts.
//! - **Multi-process** — several workers collaborating on a single study.
//! - **Inspection** — `cat trials.jsonl | jq .` for quick debugging.
//!
//! For pure in-memory usage without disk I/O, use
//! [`MemoryStorage`](super::MemoryStorage) instead (the default).
//!
//! # Security considerations
//!
//! Both [`JournalStorage::open`] and [`refresh`](super::Storage::refresh)
//! read the entire JSONL file into memory. A very large file will
//! consume memory proportional to its size, which could lead to
//! out-of-memory conditions.
//!
//! If your application accepts externally-provided file paths, consider:
//!
//! - Checking the file size before calling `open`.
//! - Validating or sanitizing untrusted JSONL content.
//! - Imposing an upper bound on the number of trials or file size you
//! are willing to load.
use core::marker::PhantomData;
use core::sync::atomic::{AtomicU64, Ordering};
use std::fs::{File, OpenOptions};
use std::io::{BufRead, BufReader, Read as _, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use fs2::FileExt;
use parking_lot::{Mutex, RwLock};
use serde::Serialize;
use serde::de::DeserializeOwned;
use super::{MemoryStorage, Storage};
use crate::sampler::CompletedTrial;
/// Append-only JSONL storage backend with file locking.
///
/// Trials are kept in memory (via an inner [`MemoryStorage`]) for fast
/// read access and simultaneously appended to a JSONL file on disk.
/// Multiple processes can safely share the same file thanks to
/// `fs2` file locks — writes use an exclusive lock, reads use a
/// shared lock.
///
/// The type parameter `V` is the objective value type (typically
/// `f64`). It must implement [`Serialize`](serde::Serialize) and
/// [`DeserializeOwned`](serde::de::DeserializeOwned) so trials can be
/// written to and read from disk.
///
/// See the [`storage`](super) module docs for file format details
/// and a resumption example.
///
/// # Example
///
/// ```no_run
/// use optimizer::prelude::*;
/// use optimizer::storage::JournalStorage;
///
/// let storage = JournalStorage::<f64>::new("trials.jsonl");
/// let mut study = Study::builder().minimize().storage(storage).build();
/// ```
///
/// # Security considerations
///
/// File contents are loaded into memory in full; see the
/// module-level docs for details and mitigations.
pub struct JournalStorage<V = f64> {
memory: MemoryStorage<V>,
path: PathBuf,
/// Serialise in-process writes and refreshes so they don't race.
io_lock: Mutex<()>,
/// Byte offset of last-read position for incremental refresh.
file_offset: AtomicU64,
_marker: PhantomData<V>,
}
impl<V: Serialize + DeserializeOwned + Send + Sync> JournalStorage<V> {
/// Create a new journal storage that writes to the given path.
///
/// The file does not need to exist yet — it will be created on the
/// first write. Existing trials in the file are **not** loaded
/// until [`refresh`](Storage::refresh) is called (which happens
/// automatically at the start of each trial via the [`Study`](crate::Study)).
///
/// To pre-load existing trials at construction time, use
/// [`JournalStorage::open`] instead.
#[must_use]
pub fn new(path: impl AsRef<Path>) -> Self {
let path = path
.as_ref()
.canonicalize()
.unwrap_or_else(|_| path.as_ref().to_path_buf());
Self {
memory: MemoryStorage::new(),
path,
io_lock: Mutex::new(()),
file_offset: AtomicU64::new(0),
_marker: PhantomData,
}
}
/// Open an existing journal file and load all stored trials.
///
/// If the file does not exist, return an empty storage (no error).
/// This is the primary way to **resume** a study after a restart.
///
/// # Errors
///
/// Return a [`Storage`](crate::Error::Storage) error if the file
/// exists but cannot be read or parsed.
pub fn open(path: impl AsRef<Path>) -> crate::Result<Self> {
let path = path
.as_ref()
.canonicalize()
.unwrap_or_else(|_| path.as_ref().to_path_buf());
let (trials, offset) = load_trials_from_file(&path)?;
Ok(Self {
memory: MemoryStorage::with_trials(trials),
path,
io_lock: Mutex::new(()),
file_offset: AtomicU64::new(offset),
_marker: PhantomData,
})
}
/// Append a single trial to the JSONL file (best-effort).
///
/// Does **not** advance `file_offset` — that is left to `refresh`
/// so that externally-written data between the old offset and our
/// write is never skipped.
fn write_to_file(&self, trial: &CompletedTrial<V>) -> crate::Result<()> {
let _guard = self.io_lock.lock();
let mut file = OpenOptions::new()
.create(true)
.truncate(false)
.read(true)
.write(true)
.open(&self.path)
.map_err(|e| crate::Error::Storage(e.to_string()))?;
file.lock_exclusive()
.map_err(|e| crate::Error::Storage(e.to_string()))?;
file.seek(SeekFrom::End(0))
.map_err(|e| crate::Error::Storage(e.to_string()))?;
let line =
serde_json::to_string(trial).map_err(|e| crate::Error::Storage(e.to_string()))?;
writeln!(file, "{line}").map_err(|e| crate::Error::Storage(e.to_string()))?;
file.sync_data()
.map_err(|e| crate::Error::Storage(e.to_string()))?;
file.unlock()
.map_err(|e| crate::Error::Storage(e.to_string()))?;
Ok(())
}
}
impl<V: Serialize + DeserializeOwned + Send + Sync> Storage<V> for JournalStorage<V> {
fn push(&self, trial: CompletedTrial<V>) {
// Best-effort persist; the trial stays in memory regardless.
let _ = self.write_to_file(&trial);
self.memory.push(trial);
}
fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>> {
self.memory.trials_arc()
}
fn next_trial_id(&self) -> u64 {
self.memory.next_trial_id()
}
fn peek_next_trial_id(&self) -> u64 {
self.memory.peek_next_trial_id()
}
fn refresh(&self) -> bool {
let _guard = self.io_lock.lock();
let Ok(file) = File::open(&self.path) else {
return false;
};
if file.lock_shared().is_err() {
return false;
}
let offset = self.file_offset.load(Ordering::SeqCst);
let file_size = if let Ok(m) = file.metadata() {
m.len()
} else {
let _ = file.unlock();
return false;
};
if file_size <= offset {
let _ = file.unlock();
return false;
}
let mut buf = String::new();
let mut handle = &file;
if handle.seek(SeekFrom::Start(offset)).is_err() {
let _ = file.unlock();
return false;
}
if handle.read_to_string(&mut buf).is_err() {
let _ = file.unlock();
return false;
}
let _ = file.unlock();
let bytes_read = buf.len() as u64;
let new_offset = offset + bytes_read;
let mut new_trials = Vec::new();
for line in buf.lines() {
let line = line.trim();
if line.is_empty() {
continue;
}
let trial: CompletedTrial<V> = match serde_json::from_str(line) {
Ok(t) => t,
Err(_) => return false,
};
if trial.validate().is_err() {
return false;
}
new_trials.push(trial);
}
if new_trials.is_empty() {
self.file_offset.fetch_max(new_offset, Ordering::SeqCst);
return false;
}
let mut mem_guard = self.memory.trials_arc().write();
// Deduplicate: only add trials whose IDs are not already in memory.
let existing_ids: std::collections::HashSet<u64> = mem_guard.iter().map(|t| t.id).collect();
new_trials.retain(|t| !existing_ids.contains(&t.id));
if let Some(max_id) = new_trials.iter().map(|t| t.id).max() {
self.memory.bump_next_id(max_id + 1);
}
let added = !new_trials.is_empty();
mem_guard.extend(new_trials);
self.file_offset.fetch_max(new_offset, Ordering::SeqCst);
added
}
}
/// Read all trials from a JSONL file. Returns an empty vec (and
/// offset 0) if the file does not exist. The returned `u64` is the
/// file size at the time of reading, suitable for initialising the
/// incremental-refresh offset.
fn load_trials_from_file<V: DeserializeOwned>(
path: &Path,
) -> crate::Result<(Vec<CompletedTrial<V>>, u64)> {
let file = match File::open(path) {
Ok(f) => f,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok((Vec::new(), 0)),
Err(e) => return Err(crate::Error::Storage(e.to_string())),
};
file.lock_shared()
.map_err(|e| crate::Error::Storage(e.to_string()))?;
let file_size = file
.metadata()
.map_err(|e| crate::Error::Storage(e.to_string()))?
.len();
let reader = BufReader::new(&file);
let mut trials = Vec::new();
for line in reader.lines() {
let line = line.map_err(|e| crate::Error::Storage(e.to_string()))?;
let line = line.trim();
if line.is_empty() {
continue;
}
let trial: CompletedTrial<V> =
serde_json::from_str(line).map_err(|e| crate::Error::Storage(e.to_string()))?;
trial.validate().map_err(crate::Error::Storage)?;
trials.push(trial);
}
file.unlock()
.map_err(|e| crate::Error::Storage(e.to_string()))?;
Ok((trials, file_size))
}
+111
View File
@@ -0,0 +1,111 @@
//! In-memory storage backend.
//!
//! [`MemoryStorage`] is the default backend used by every
//! [`Study`](crate::Study). Trials are stored in a
//! `Vec<CompletedTrial<V>>` behind a [`parking_lot::RwLock`] for
//! thread-safe access.
//!
//! # When to use
//!
//! - **Single-process** studies where persistence is not needed.
//! - **Testing** or **prototyping** — zero configuration required.
//! - When you want the **fastest** possible read/write performance
//! (no disk I/O).
//!
//! For persistent storage that survives process restarts, see
//! `JournalStorage` (requires the `journal` feature).
//!
//! # Example
//!
//! ```
//! use optimizer::prelude::*;
//! use optimizer::storage::MemoryStorage;
//!
//! // Explicit memory storage (equivalent to the default)
//! let storage = MemoryStorage::<f64>::new();
//! let study = Study::builder().minimize().storage(storage).build();
//! ```
use core::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use parking_lot::RwLock;
use super::Storage;
use crate::sampler::CompletedTrial;
/// In-memory trial storage (the default).
///
/// Wrap a `Vec<CompletedTrial<V>>` behind a read-write lock so that
/// trials can be appended from any thread. This is the backend that
/// [`Study`](crate::Study) uses when no explicit storage is provided.
///
/// Use [`with_trials`](Self::with_trials) to seed a study with
/// previously collected data.
///
/// # Example
///
/// ```
/// use optimizer::storage::{MemoryStorage, Storage};
///
/// let storage = MemoryStorage::<f64>::new();
/// assert_eq!(storage.trials_arc().read().len(), 0);
/// ```
pub struct MemoryStorage<V> {
trials: Arc<RwLock<Vec<CompletedTrial<V>>>>,
next_id: AtomicU64,
}
impl<V> MemoryStorage<V> {
/// Create a new, empty in-memory store.
#[must_use]
pub fn new() -> Self {
Self {
trials: Arc::new(RwLock::new(Vec::new())),
next_id: AtomicU64::new(0),
}
}
/// Create an in-memory store pre-populated with `trials`.
///
/// The internal ID counter is set to one past the highest trial ID
/// so that subsequent trials receive unique IDs.
#[must_use]
pub fn with_trials(trials: Vec<CompletedTrial<V>>) -> Self {
let next_id = trials.iter().map(|t| t.id).max().map_or(0, |id| id + 1);
Self {
trials: Arc::new(RwLock::new(trials)),
next_id: AtomicU64::new(next_id),
}
}
/// Ensure the ID counter is at least `min_value`.
#[cfg(feature = "journal")]
pub(crate) fn bump_next_id(&self, min_value: u64) {
self.next_id.fetch_max(min_value, Ordering::SeqCst);
}
}
impl<V> Default for MemoryStorage<V> {
fn default() -> Self {
Self::new()
}
}
impl<V: Send + Sync> Storage<V> for MemoryStorage<V> {
fn push(&self, trial: CompletedTrial<V>) {
self.trials.write().push(trial);
}
fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>> {
&self.trials
}
fn next_trial_id(&self) -> u64 {
self.next_id.fetch_add(1, Ordering::SeqCst)
}
fn peek_next_trial_id(&self) -> u64 {
self.next_id.load(Ordering::SeqCst)
}
}
+94
View File
@@ -0,0 +1,94 @@
//! Trial storage backends.
//!
//! The [`Storage`] trait defines how completed trials are persisted and
//! retrieved. Every [`Study`](crate::Study) owns an `Arc<dyn Storage<V>>`
//! so storage is transparently shared across threads.
//!
//! # Available backends
//!
//! | Backend | Description | Feature flag |
//! |---------|-------------|-------------|
//! | [`MemoryStorage`] | In-memory `Vec` behind a read-write lock (the default) | — |
//! | `JournalStorage` | JSONL file with `fs2` file locking for multi-process sharing | `journal` |
//!
//! # When to swap backends
//!
//! The default [`MemoryStorage`] is sufficient for single-process studies
//! where persistence is not needed. Switch to `JournalStorage` when you
//! want to:
//!
//! - **Resume** a study after a process restart.
//! - **Share state** across multiple processes writing to the same file.
//! - **Inspect** trial history in a human-readable JSONL file.
//!
//! # Implementing a custom backend
//!
//! Implement the [`Storage`] trait to plug in your own backend (e.g. a
//! database). The trait requires four methods: [`push`](Storage::push),
//! [`trials_arc`](Storage::trials_arc), [`next_trial_id`](Storage::next_trial_id),
//! and optionally [`refresh`](Storage::refresh) for external data sources.
//!
//! Inject your storage into a study via the builder:
//!
//! ```
//! use optimizer::prelude::*;
//! use optimizer::storage::MemoryStorage;
//!
//! let storage = MemoryStorage::<f64>::new();
//! let study = Study::builder().minimize().storage(storage).build();
//! ```
#[cfg(feature = "journal")]
mod journal;
use std::sync::Arc;
#[cfg(feature = "journal")]
pub use journal::JournalStorage;
use parking_lot::RwLock;
mod memory;
pub use memory::MemoryStorage;
use crate::sampler::CompletedTrial;
/// Trait for storing and retrieving completed trials.
///
/// Every [`Study`](crate::Study) owns an `Arc<dyn Storage<V>>`. The
/// default implementation is [`MemoryStorage`], which keeps trials in
/// a plain `Vec` behind a read-write lock.
///
/// Implementations must be `Send + Sync` because a study may be shared
/// across threads (e.g. via `Study::optimize_parallel`).
pub trait Storage<V>: Send + Sync {
/// Append a completed trial to the store.
fn push(&self, trial: CompletedTrial<V>);
/// Return a reference to the in-memory trial buffer.
///
/// All implementations must maintain an `Arc<RwLock<Vec<…>>>` that
/// reflects the current set of trials. Callers may acquire a read
/// lock for efficient, allocation-free access.
fn trials_arc(&self) -> &Arc<RwLock<Vec<CompletedTrial<V>>>>;
/// Atomically return the next unique trial ID.
///
/// Each call increments an internal counter so that consecutive
/// calls always produce distinct IDs.
fn next_trial_id(&self) -> u64;
/// Return the current value of the next-trial-ID counter without incrementing.
///
/// This is used for persistence (e.g. `Study::save`) to capture the
/// counter's exact position, including IDs assigned to failed trials
/// that are not stored.
fn peek_next_trial_id(&self) -> u64;
/// Reload from an external source (e.g. a file written by another
/// process). Return `true` if the in-memory buffer was updated.
///
/// The default implementation is a no-op that returns `false`.
fn refresh(&self) -> bool {
false
}
}
-1125
View File
File diff suppressed because it is too large Load Diff
+379
View File
@@ -0,0 +1,379 @@
use crate::sampler::CompletedTrial;
use crate::types::TrialState;
use super::Study;
impl<V> Study<V>
where
V: PartialOrd,
{
/// Return the trial with the best objective value.
///
/// The "best" trial depends on the optimization direction:
/// - `Direction::Minimize`: Returns the trial with the lowest objective value.
/// - `Direction::Maximize`: Returns the trial with the highest objective value.
///
/// When constraints are present, feasible trials always rank above infeasible
/// trials. Among infeasible trials, those with lower total constraint violation
/// are preferred.
///
/// # Errors
///
/// Returns `Error::NoCompletedTrials` if no trials have been completed.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
///
/// // Error when no trials completed
/// assert!(study.best_trial().is_err());
///
/// let x_param = FloatParam::new(0.0, 1.0);
///
/// let mut trial1 = study.create_trial();
/// let _ = x_param.suggest(&mut trial1);
/// study.complete_trial(trial1, 0.8);
///
/// let mut trial2 = study.create_trial();
/// let _ = x_param.suggest(&mut trial2);
/// study.complete_trial(trial2, 0.3);
///
/// let best = study.best_trial().unwrap();
/// assert_eq!(best.value, 0.3); // Minimize: lower is better
/// ```
pub fn best_trial(&self) -> crate::Result<CompletedTrial<V>>
where
V: Clone,
{
let trials = self.storage.trials_arc().read();
let direction = self.direction;
let best = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.max_by(|a, b| Self::compare_trials(a, b, direction))
.ok_or(crate::Error::NoCompletedTrials)?;
Ok(best.clone())
}
/// Return the best objective value found so far.
///
/// The "best" value depends on the optimization direction:
/// - `Direction::Minimize`: Returns the lowest objective value.
/// - `Direction::Maximize`: Returns the highest objective value.
///
/// # Errors
///
/// Returns `Error::NoCompletedTrials` if no trials have been completed.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Maximize);
///
/// // Error when no trials completed
/// assert!(study.best_value().is_err());
///
/// let x_param = FloatParam::new(0.0, 1.0);
///
/// let mut trial1 = study.create_trial();
/// let _ = x_param.suggest(&mut trial1);
/// study.complete_trial(trial1, 0.3);
///
/// let mut trial2 = study.create_trial();
/// let _ = x_param.suggest(&mut trial2);
/// study.complete_trial(trial2, 0.8);
///
/// let best = study.best_value().unwrap();
/// assert_eq!(best, 0.8); // Maximize: higher is better
/// ```
pub fn best_value(&self) -> crate::Result<V>
where
V: Clone,
{
self.best_trial().map(|trial| trial.value)
}
/// Return the top `n` trials sorted by objective value.
///
/// For `Direction::Minimize`, returns trials with the lowest values.
/// For `Direction::Maximize`, returns trials with the highest values.
/// Only includes completed trials (not failed or pruned).
///
/// If fewer than `n` completed trials exist, returns all of them.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0);
///
/// for val in [5.0, 1.0, 3.0] {
/// let mut t = study.create_trial();
/// let _ = x.suggest(&mut t);
/// study.complete_trial(t, val);
/// }
///
/// let top2 = study.top_trials(2);
/// assert_eq!(top2.len(), 2);
/// assert!(top2[0].value <= top2[1].value);
/// ```
#[must_use]
pub fn top_trials(&self, n: usize) -> Vec<CompletedTrial<V>>
where
V: Clone,
{
let trials = self.storage.trials_arc().read();
let direction = self.direction;
// Sort indices instead of cloning all trials, then clone only the top N.
let mut indices: Vec<usize> = trials
.iter()
.enumerate()
.filter(|(_, t)| t.state == TrialState::Complete)
.map(|(i, _)| i)
.collect();
// Sort best-first: reverse the compare_trials ordering (which is designed for max_by)
indices.sort_by(|&a, &b| Self::compare_trials(&trials[b], &trials[a], direction));
indices.truncate(n);
indices.iter().map(|&i| trials[i].clone()).collect()
}
}
impl<V> Study<V>
where
V: PartialOrd + Clone + Into<f64>,
{
/// Compute parameter importance scores using Spearman rank correlation.
///
/// For each parameter, the absolute Spearman correlation between its values
/// and the objective values is computed across all completed trials. Scores
/// are normalized so they sum to 1.0 and sorted in descending order.
///
/// Parameters that appear in fewer than 2 trials are omitted.
/// Returns an empty `Vec` if the study has fewer than 2 completed trials.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0).name("x");
///
/// study
/// .optimize(20, |trial: &mut optimizer::Trial| {
/// let xv = x.suggest(trial)?;
/// Ok::<_, optimizer::Error>(xv * xv)
/// })
/// .unwrap();
///
/// let importance = study.param_importance();
/// assert_eq!(importance.len(), 1);
/// assert_eq!(importance[0].0, "x");
/// ```
#[must_use]
#[allow(clippy::cast_precision_loss)]
pub fn param_importance(&self) -> Vec<(String, f64)> {
use std::collections::BTreeSet;
use crate::importance::spearman;
use crate::param::ParamValue;
use crate::types::TrialState;
let trials = self.storage.trials_arc().read();
let complete: Vec<_> = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.collect();
if complete.len() < 2 {
return Vec::new();
}
// Collect all parameter IDs across trials.
let all_param_ids: BTreeSet<_> = complete.iter().flat_map(|t| t.params.keys()).collect();
let mut scores: Vec<(String, f64)> = Vec::with_capacity(all_param_ids.len());
for &param_id in &all_param_ids {
// Collect (param_value_f64, objective_f64) for trials that have this param.
let mut param_vals = Vec::with_capacity(complete.len());
let mut obj_vals = Vec::with_capacity(complete.len());
for trial in &complete {
if let Some(pv) = trial.params.get(param_id) {
let f = match *pv {
ParamValue::Float(v) => v,
ParamValue::Int(v) => v as f64,
ParamValue::Categorical(v) => v as f64,
};
param_vals.push(f);
obj_vals.push(trial.value.clone().into());
}
}
if param_vals.len() < 2 {
continue;
}
let corr = spearman(&param_vals, &obj_vals).abs();
// Determine label: use param_labels if available, else "param_{id}".
let label = complete
.iter()
.find_map(|t| t.param_labels.get(param_id))
.map_or_else(|| param_id.to_string(), Clone::clone);
scores.push((label, corr));
}
// Normalize so scores sum to 1.0.
let sum: f64 = scores.iter().map(|(_, s)| *s).sum();
if sum > 0.0 {
for entry in &mut scores {
entry.1 /= sum;
}
}
// Sort descending by score.
scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(core::cmp::Ordering::Equal));
scores
}
/// Compute parameter importance using fANOVA (functional ANOVA) with
/// default configuration.
///
/// Fits a random forest to the trial data and decomposes variance into
/// per-parameter main effects and pairwise interaction effects. This is
/// more accurate than correlation-based importance ([`Self::param_importance`])
/// and can detect non-linear relationships and parameter interactions.
///
/// # Errors
///
/// Returns [`crate::Error::NoCompletedTrials`] if fewer than 2 trials have completed.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0).name("x");
/// let y = FloatParam::new(0.0, 10.0).name("y");
///
/// study
/// .optimize(30, |trial: &mut optimizer::Trial| {
/// let xv = x.suggest(trial)?;
/// let yv = y.suggest(trial)?;
/// Ok::<_, optimizer::Error>(xv * xv + 0.1 * yv)
/// })
/// .unwrap();
///
/// let result = study.fanova().unwrap();
/// assert!(!result.main_effects.is_empty());
/// ```
pub fn fanova(&self) -> crate::Result<crate::fanova::FanovaResult> {
self.fanova_with_config(&crate::fanova::FanovaConfig::default())
}
/// Compute parameter importance using fANOVA with custom configuration.
///
/// See [`Self::fanova`] for details. The [`FanovaConfig`](crate::fanova::FanovaConfig)
/// allows tuning the number of trees, tree depth, and random seed.
///
/// # Errors
///
/// Returns [`crate::Error::NoCompletedTrials`] if fewer than 2 trials have completed.
#[allow(clippy::cast_precision_loss)]
pub fn fanova_with_config(
&self,
config: &crate::fanova::FanovaConfig,
) -> crate::Result<crate::fanova::FanovaResult> {
use std::collections::BTreeSet;
use crate::fanova::compute_fanova;
use crate::param::ParamValue;
use crate::types::TrialState;
let trials = self.storage.trials_arc().read();
let complete: Vec<_> = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.collect();
if complete.len() < 2 {
return Err(crate::Error::NoCompletedTrials);
}
// Collect all parameter IDs in a stable order.
let all_param_ids: Vec<_> = {
let set: BTreeSet<_> = complete.iter().flat_map(|t| t.params.keys()).collect();
set.into_iter().collect()
};
if all_param_ids.is_empty() {
return Ok(crate::fanova::FanovaResult {
main_effects: Vec::new(),
interactions: Vec::new(),
});
}
// Build feature matrix (only trials that have all parameters).
let mut data = Vec::with_capacity(complete.len());
let mut targets = Vec::with_capacity(complete.len());
for trial in &complete {
let mut row = Vec::with_capacity(all_param_ids.len());
let mut has_all = true;
for &pid in &all_param_ids {
if let Some(pv) = trial.params.get(pid) {
row.push(match *pv {
ParamValue::Float(v) => v,
ParamValue::Int(v) => v as f64,
ParamValue::Categorical(v) => v as f64,
});
} else {
has_all = false;
break;
}
}
if has_all {
data.push(row);
targets.push(trial.value.clone().into());
}
}
if data.len() < 2 {
return Err(crate::Error::NoCompletedTrials);
}
// Build feature names from parameter labels.
let feature_names: Vec<String> = all_param_ids
.iter()
.map(|&pid| {
complete
.iter()
.find_map(|t| t.param_labels.get(pid))
.map_or_else(|| pid.to_string(), Clone::clone)
})
.collect();
Ok(compute_fanova(&data, &targets, &feature_names, config))
}
}
+285
View File
@@ -0,0 +1,285 @@
use core::ops::ControlFlow;
use std::sync::Arc;
use crate::trial::Trial;
use crate::types::TrialState;
use super::{Study, is_trial_pruned};
impl<V> Study<V>
where
V: PartialOrd,
{
/// Run async optimization with an objective.
///
/// Like [`optimize`](Self::optimize), but each evaluation is wrapped in
/// [`spawn_blocking`](tokio::task::spawn_blocking), keeping the async
/// runtime responsive for CPU-bound objectives. Trials run sequentially.
///
/// Accepts any [`Objective`](crate::Objective) implementation, including
/// plain closures. Struct-based objectives can override lifecycle hooks.
///
/// # Errors
///
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
/// Returns `Error::TaskError` if a spawned blocking task panics.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
/// # #[cfg(feature = "async")]
/// # async fn example() -> optimizer::Result<()> {
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_async(10, move |trial: &mut optimizer::Trial| {
/// let x = x_param.suggest(trial)?;
/// Ok::<_, optimizer::Error>(x * x)
/// })
/// .await?;
///
/// assert!(study.n_trials() > 0);
/// # Ok(())
/// # }
/// ```
pub async fn optimize_async<O>(&self, n_trials: usize, objective: O) -> crate::Result<()>
where
O: crate::objective::Objective<V> + Send + Sync + 'static,
O::Error: Send,
V: Clone + Default + Send + 'static,
{
#[cfg(feature = "tracing")]
let _span =
tracing::info_span!("optimize_async", n_trials, direction = ?self.direction).entered();
let objective = Arc::new(objective);
for _ in 0..n_trials {
if let ControlFlow::Break(()) = objective.before_trial(self) {
break;
}
let obj = Arc::clone(&objective);
let mut trial = self.create_trial();
let result = tokio::task::spawn_blocking(move || {
let res = obj.evaluate(&mut trial);
(trial, res)
})
.await
.map_err(|e| crate::Error::TaskError(e.to_string()))?;
match result {
(t, Ok(value)) => {
#[cfg(feature = "tracing")]
let trial_id = t.id();
let completed = t.into_completed(value, TrialState::Complete);
let flow = objective.after_trial(self, &completed);
self.storage.push(completed);
trace_info!(trial_id, "trial completed");
if let ControlFlow::Break(()) = flow {
return Ok(());
}
}
(t, Err(e)) if is_trial_pruned(&e) => {
#[cfg(feature = "tracing")]
let trial_id = t.id();
self.prune_trial(t);
trace_info!(trial_id, "trial pruned");
}
(t, Err(e)) => {
#[cfg(feature = "tracing")]
let trial_id = t.id();
self.fail_trial(t, e.to_string());
trace_debug!(trial_id, "trial failed");
}
}
}
let has_complete = self
.storage
.trials_arc()
.read()
.iter()
.any(|t| t.state == TrialState::Complete);
if !has_complete {
return Err(crate::Error::NoCompletedTrials);
}
Ok(())
}
/// Run parallel optimization with an objective.
///
/// Spawns up to `concurrency` evaluations concurrently using
/// [`spawn_blocking`](tokio::task::spawn_blocking). Results are
/// collected via a [`JoinSet`](tokio::task::JoinSet).
///
/// Accepts any [`Objective`](crate::Objective) implementation, including
/// plain closures. The [`after_trial`](crate::Objective::after_trial)
/// hook fires as each result arrives — returning `Break` stops spawning
/// new trials while in-flight tasks drain.
///
/// # Errors
///
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
/// Returns `Error::TaskError` if the semaphore is closed or a spawned task panics.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
/// # #[cfg(feature = "async")]
/// # async fn example() -> optimizer::Result<()> {
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize_parallel(10, 4, move |trial: &mut optimizer::Trial| {
/// let x = x_param.suggest(trial)?;
/// Ok::<_, optimizer::Error>(x * x)
/// })
/// .await?;
///
/// assert_eq!(study.n_trials(), 10);
/// # Ok(())
/// # }
/// ```
#[allow(clippy::missing_panics_doc, clippy::too_many_lines)]
pub async fn optimize_parallel<O>(
&self,
n_trials: usize,
concurrency: usize,
objective: O,
) -> crate::Result<()>
where
O: crate::objective::Objective<V> + Send + Sync + 'static,
O::Error: Send,
V: Clone + Default + Send + 'static,
{
use tokio::sync::Semaphore;
use tokio::task::JoinSet;
assert!(concurrency > 0, "concurrency must be at least 1");
#[cfg(feature = "tracing")]
let _span = tracing::info_span!("optimize_parallel", n_trials, concurrency, direction = ?self.direction).entered();
let objective = Arc::new(objective);
let semaphore = Arc::new(Semaphore::new(concurrency));
let mut join_set: JoinSet<(Trial, Result<V, O::Error>)> = JoinSet::new();
let mut spawned = 0;
'spawn: while spawned < n_trials {
if let ControlFlow::Break(()) = objective.before_trial(self) {
break;
}
// If the join set is full, drain one result to free a slot.
while join_set.len() >= concurrency {
let result = join_set
.join_next()
.await
.expect("join_set should not be empty")
.map_err(|e| crate::Error::TaskError(e.to_string()))?;
match result {
(t, Ok(value)) => {
#[cfg(feature = "tracing")]
let trial_id = t.id();
let completed = t.into_completed(value, TrialState::Complete);
let flow = objective.after_trial(self, &completed);
self.storage.push(completed);
trace_info!(trial_id, "trial completed");
if let ControlFlow::Break(()) = flow {
break 'spawn;
}
}
(t, Err(e)) => {
#[cfg(feature = "tracing")]
let trial_id = t.id();
if is_trial_pruned(&e) {
self.prune_trial(t);
trace_info!(trial_id, "trial pruned");
} else {
self.fail_trial(t, e.to_string());
trace_debug!(trial_id, "trial failed");
}
}
}
}
let permit = semaphore
.clone()
.acquire_owned()
.await
.map_err(|e| crate::Error::TaskError(e.to_string()))?;
let mut trial = self.create_trial();
let obj = Arc::clone(&objective);
join_set.spawn(async move {
let result = tokio::task::spawn_blocking(move || {
let res = obj.evaluate(&mut trial);
(trial, res)
})
.await
.expect("spawn_blocking should not panic");
drop(permit);
result
});
spawned += 1;
}
// Drain remaining in-flight tasks.
while let Some(result) = join_set.join_next().await {
let result = result.map_err(|e| crate::Error::TaskError(e.to_string()))?;
match result {
(t, Ok(value)) => {
#[cfg(feature = "tracing")]
let trial_id = t.id();
let completed = t.into_completed(value, TrialState::Complete);
// Still fire after_trial for bookkeeping, but don't break — we're draining.
let _ = objective.after_trial(self, &completed);
self.storage.push(completed);
trace_info!(trial_id, "trial completed");
}
(t, Err(e)) => {
#[cfg(feature = "tracing")]
let trial_id = t.id();
if is_trial_pruned(&e) {
self.prune_trial(t);
trace_info!(trial_id, "trial pruned");
} else {
self.fail_trial(t, e.to_string());
trace_debug!(trial_id, "trial failed");
}
}
}
}
let has_complete = self
.storage
.trials_arc()
.read()
.iter()
.any(|t| t.state == TrialState::Complete);
if !has_complete {
return Err(crate::Error::NoCompletedTrials);
}
Ok(())
}
}
+135
View File
@@ -0,0 +1,135 @@
use core::marker::PhantomData;
use std::collections::VecDeque;
use std::sync::Arc;
use parking_lot::Mutex;
use crate::pruner::{NopPruner, Pruner};
use crate::sampler::Sampler;
use crate::sampler::random::RandomSampler;
use crate::types::Direction;
use super::Study;
/// A builder for constructing [`Study`] instances with a fluent API.
///
/// Created via [`Study::builder()`]. Collects sampler, pruner, direction,
/// and storage options before constructing the study.
///
/// # Defaults
///
/// - Direction: [`Minimize`](Direction::Minimize)
/// - Sampler: [`RandomSampler`]
/// - Pruner: [`NopPruner`]
/// - Storage: [`MemoryStorage`](crate::storage::MemoryStorage)
///
/// # Examples
///
/// ```
/// use optimizer::prelude::*;
///
/// let study: Study<f64> = Study::builder()
/// .maximize()
/// .sampler(TpeSampler::new())
/// .pruner(MedianPruner::new(Direction::Maximize).n_warmup_steps(5))
/// .build();
///
/// assert_eq!(study.direction(), Direction::Maximize);
/// ```
pub struct StudyBuilder<V: PartialOrd = f64> {
direction: Direction,
sampler: Option<Box<dyn Sampler>>,
pruner: Option<Box<dyn Pruner>>,
storage: Option<Box<dyn crate::storage::Storage<V>>>,
_marker: PhantomData<V>,
}
impl<V: PartialOrd> StudyBuilder<V> {
/// Create a new builder with default settings.
pub(super) fn new() -> Self {
Self {
direction: Direction::Minimize,
sampler: None,
pruner: None,
storage: None,
_marker: PhantomData,
}
}
/// Set the optimization direction to minimize (the default).
#[must_use]
pub fn minimize(mut self) -> Self {
self.direction = Direction::Minimize;
self
}
/// Set the optimization direction to maximize.
#[must_use]
pub fn maximize(mut self) -> Self {
self.direction = Direction::Maximize;
self
}
/// Set the optimization direction explicitly.
#[must_use]
pub fn direction(mut self, direction: Direction) -> Self {
self.direction = direction;
self
}
/// Set the sampler used for parameter suggestions.
///
/// Defaults to [`RandomSampler`] if not specified.
#[must_use]
pub fn sampler(mut self, sampler: impl Sampler + 'static) -> Self {
self.sampler = Some(Box::new(sampler));
self
}
/// Set the pruner used for early stopping of trials.
///
/// Defaults to [`NopPruner`] (no pruning) if not specified.
#[must_use]
pub fn pruner(mut self, pruner: impl Pruner + 'static) -> Self {
self.pruner = Some(Box::new(pruner));
self
}
/// Set a custom storage backend.
///
/// Defaults to [`MemoryStorage`](crate::storage::MemoryStorage) if not specified.
#[must_use]
pub fn storage(mut self, storage: impl crate::storage::Storage<V> + 'static) -> Self {
self.storage = Some(Box::new(storage));
self
}
/// Build the [`Study`] with the configured options.
#[must_use]
pub fn build(self) -> Study<V>
where
V: Send + Sync + 'static,
{
let sampler = self
.sampler
.unwrap_or_else(|| Box::new(RandomSampler::new()));
let pruner = self.pruner.unwrap_or_else(|| Box::new(NopPruner));
let storage = self
.storage
.unwrap_or_else(|| Box::new(crate::storage::MemoryStorage::<V>::new()));
let sampler: Arc<dyn Sampler> = Arc::from(sampler);
let pruner: Arc<dyn Pruner> = Arc::from(pruner);
let storage: Arc<dyn crate::storage::Storage<V>> = Arc::from(storage);
let trial_factory = Study::make_trial_factory(&sampler, &storage, &pruner);
Study {
direction: self.direction,
sampler,
pruner,
storage,
trial_factory,
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
}
}
}
+274
View File
@@ -0,0 +1,274 @@
use core::fmt;
use crate::parameter::ParamId;
use crate::types::{Direction, TrialState};
use super::Study;
impl<V> Study<V>
where
V: PartialOrd + Clone + fmt::Display,
{
/// Write completed trials to a writer in CSV format.
///
/// Columns: `trial_id`, `value`, `state`, then one column per unique
/// parameter label, then one column per unique user-attribute key.
///
/// Parameters without labels use a generated name (`param_<id>`).
/// Pruned trials have an empty `value` cell.
///
/// # Errors
///
/// Returns an I/O error if writing fails.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0).name("x");
///
/// let mut trial = study.create_trial();
/// let _ = x.suggest(&mut trial);
/// study.complete_trial(trial, 0.42);
///
/// let mut buf = Vec::new();
/// study.to_csv(&mut buf).unwrap();
/// let csv = String::from_utf8(buf).unwrap();
/// assert!(csv.contains("trial_id"));
/// ```
pub fn to_csv(&self, mut writer: impl std::io::Write) -> std::io::Result<()> {
use std::collections::BTreeMap;
let trials = self.storage.trials_arc().read();
// Collect all unique parameter labels (sorted for deterministic column order).
let mut param_columns: BTreeMap<ParamId, String> = BTreeMap::new();
for trial in trials.iter() {
for &id in trial.params.keys() {
param_columns.entry(id).or_insert_with(|| {
trial
.param_labels
.get(&id)
.cloned()
.unwrap_or_else(|| id.to_string())
});
}
}
// Fill in labels from other trials that might have better labels.
for trial in trials.iter() {
for (&id, label) in &trial.param_labels {
param_columns.entry(id).or_insert_with(|| label.clone());
}
}
// Collect all unique attribute keys (sorted).
let mut attr_keys: Vec<String> = Vec::new();
for trial in trials.iter() {
for key in trial.user_attrs.keys() {
if !attr_keys.contains(key) {
attr_keys.push(key.clone());
}
}
}
attr_keys.sort();
let param_ids: Vec<ParamId> = param_columns.keys().copied().collect();
// Write header.
write!(writer, "trial_id,value,state")?;
for id in &param_ids {
write!(writer, ",{}", csv_escape(&param_columns[id]))?;
}
for key in &attr_keys {
write!(writer, ",{}", csv_escape(key))?;
}
writeln!(writer)?;
// Write one row per trial.
for trial in trials.iter() {
write!(writer, "{}", trial.id)?;
// Value: empty for non-complete trials.
if trial.state == TrialState::Complete {
write!(writer, ",{}", trial.value)?;
} else {
write!(writer, ",")?;
}
write!(
writer,
",{}",
match trial.state {
TrialState::Complete => "Complete",
TrialState::Pruned => "Pruned",
TrialState::Failed => "Failed",
TrialState::Running => "Running",
}
)?;
for id in &param_ids {
if let Some(pv) = trial.params.get(id) {
write!(writer, ",{pv}")?;
} else {
write!(writer, ",")?;
}
}
for key in &attr_keys {
if let Some(attr) = trial.user_attrs.get(key) {
write!(writer, ",{}", csv_escape(&format_attr(attr)))?;
} else {
write!(writer, ",")?;
}
}
writeln!(writer)?;
}
Ok(())
}
/// Export completed trials to a CSV file at the given path.
///
/// Convenience wrapper around [`to_csv`](Self::to_csv) that creates a
/// buffered file writer.
///
/// # Errors
///
/// Returns an I/O error if the file cannot be created or written.
pub fn export_csv(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
let file = std::fs::File::create(path)?;
self.to_csv(std::io::BufWriter::new(file))
}
/// Return a human-readable summary of the study.
///
/// The summary includes:
/// - Optimization direction and total trial count
/// - Breakdown by state (complete, pruned) when applicable
/// - Best trial value and parameters (if any completed trials exist)
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0).name("x");
///
/// let mut trial = study.create_trial();
/// let _ = x.suggest(&mut trial).unwrap();
/// study.complete_trial(trial, 0.42);
///
/// let summary = study.summary();
/// assert!(summary.contains("Minimize"));
/// assert!(summary.contains("0.42"));
/// ```
#[must_use]
pub fn summary(&self) -> String {
use fmt::Write;
let trials = self.storage.trials_arc().read();
let n_complete = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.count();
let n_pruned = trials
.iter()
.filter(|t| t.state == TrialState::Pruned)
.count();
let direction_str = match self.direction {
Direction::Minimize => "Minimize",
Direction::Maximize => "Maximize",
};
let mut s = format!("Study: {direction_str} | {n} trials", n = trials.len());
if n_pruned > 0 {
let _ = write!(s, " ({n_complete} complete, {n_pruned} pruned)");
}
drop(trials);
if let Ok(best) = self.best_trial() {
let _ = write!(s, "\nBest value: {} (trial #{})", best.value, best.id);
if !best.params.is_empty() {
s.push_str("\nBest parameters:");
let mut params: Vec<_> = best.params.iter().collect();
params.sort_by_key(|(id, _)| *id);
for (id, value) in params {
let label = best.param_labels.get(id).map_or("?", String::as_str);
let _ = write!(s, "\n {label} = {value}");
}
}
}
s
}
}
impl<V> fmt::Display for Study<V>
where
V: PartialOrd + Clone + fmt::Display,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.summary())
}
}
#[cfg(feature = "serde")]
impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
/// Export trials as a pretty-printed JSON array to a file.
///
/// Each element in the array is a serialized [`CompletedTrial`](crate::sampler::CompletedTrial).
/// Requires the `serde` feature.
///
/// # Errors
///
/// Returns an I/O error if the file cannot be created or written.
pub fn export_json(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
let file = std::fs::File::create(path)?;
let trials = self.trials();
serde_json::to_writer_pretty(file, &trials).map_err(std::io::Error::other)
}
}
impl Study<f64> {
/// Generate an HTML report with interactive Plotly.js charts.
///
/// Create a self-contained HTML file that can be opened in any browser.
/// See [`generate_html_report`](crate::visualization::generate_html_report)
/// for details on the included charts.
///
/// # Errors
///
/// Returns an I/O error if the file cannot be created or written.
pub fn export_html(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
crate::visualization::generate_html_report(self, path)
}
}
/// Escape a string for CSV output. If the value contains a comma, quote, or
/// newline, wrap it in double-quotes and double any embedded quotes.
fn csv_escape(s: &str) -> String {
if s.contains(',') || s.contains('"') || s.contains('\n') || s.contains('\r') {
format!("\"{}\"", s.replace('"', "\"\""))
} else {
s.to_string()
}
}
/// Format an `AttrValue` as a string for CSV cells.
fn format_attr(attr: &crate::trial::AttrValue) -> String {
use crate::trial::AttrValue;
match attr {
AttrValue::Float(v) => v.to_string(),
AttrValue::Int(v) => v.to_string(),
AttrValue::String(v) => v.clone(),
AttrValue::Bool(v) => v.to_string(),
}
}
+43
View File
@@ -0,0 +1,43 @@
use crate::sampler::CompletedTrial;
use super::Study;
impl<V> Study<V>
where
V: PartialOrd + Clone,
{
/// Return an iterator over all completed trials.
///
/// This clones the internal trial list, so it is suitable for
/// analysis and iteration but not for hot paths.
///
/// # Examples
///
/// ```
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let trial = study.create_trial();
/// study.complete_trial(trial, 1.0);
///
/// for t in study.iter() {
/// println!("Trial {} → {}", t.id, t.value);
/// }
/// ```
#[must_use]
pub fn iter(&self) -> std::vec::IntoIter<CompletedTrial<V>> {
self.trials().into_iter()
}
}
impl<V> IntoIterator for &Study<V>
where
V: PartialOrd + Clone,
{
type Item = CompletedTrial<V>;
type IntoIter = std::vec::IntoIter<CompletedTrial<V>>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
+772
View File
@@ -0,0 +1,772 @@
//! Study implementation for managing optimization trials.
use core::any::Any;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use parking_lot::{Mutex, RwLock};
use crate::param::ParamValue;
use crate::parameter::ParamId;
use crate::pruner::{NopPruner, Pruner};
use crate::sampler::random::RandomSampler;
use crate::sampler::{CompletedTrial, Sampler};
use crate::trial::Trial;
use crate::types::{Direction, TrialState};
mod analysis;
mod builder;
mod export;
mod iter;
mod optimize;
mod persistence;
#[cfg(feature = "async")]
mod async_impl;
pub use builder::StudyBuilder;
#[cfg(feature = "serde")]
pub use persistence::StudySnapshot;
/// A study manages the optimization process, tracking trials and their results.
///
/// The study is parameterized by the objective value type `V`, which defaults to `f64`.
/// The only constraint on `V` is `PartialOrd`, allowing comparison of objective values
/// to determine which trial is best.
///
/// When `V = f64`, the study passes trial history to the sampler for informed
/// parameter suggestions (e.g., TPE sampler uses history to guide sampling).
///
/// # Examples
///
/// ```
/// use optimizer::{Direction, Study};
///
/// // Create a study to minimize an objective function
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// assert_eq!(study.direction(), Direction::Minimize);
/// ```
pub struct Study<V = f64>
where
V: PartialOrd,
{
/// The optimization direction.
pub(crate) direction: Direction,
/// The sampler used to generate parameter values.
pub(crate) sampler: Arc<dyn Sampler>,
/// The pruner used to decide whether to stop trials early.
pub(crate) pruner: Arc<dyn Pruner>,
/// Trial storage backend (default: [`MemoryStorage`](crate::storage::MemoryStorage)).
pub(crate) storage: Arc<dyn crate::storage::Storage<V>>,
/// Optional factory for creating sampler-aware trials.
/// Set automatically for `Study<f64>` so that `create_trial()` and all
/// optimization methods use the sampler without requiring `_with_sampler` suffixes.
pub(crate) trial_factory: Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>,
/// Queue of parameter configurations to evaluate next.
pub(crate) enqueued_params: Arc<Mutex<VecDeque<HashMap<ParamId, ParamValue>>>>,
}
impl<V> Study<V>
where
V: PartialOrd,
{
/// Create a new study with the given optimization direction.
///
/// Uses the default `RandomSampler` for parameter sampling.
///
/// # Arguments
///
/// * `direction` - Whether to minimize or maximize the objective function.
///
/// # Examples
///
/// ```
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// assert_eq!(study.direction(), Direction::Minimize);
/// ```
#[must_use]
pub fn new(direction: Direction) -> Self
where
V: Send + Sync + 'static,
{
Self::with_sampler(direction, RandomSampler::new())
}
/// Return a [`StudyBuilder`] for constructing a study with a fluent API.
///
/// # Examples
///
/// ```
/// use optimizer::prelude::*;
///
/// let study: Study<f64> = Study::builder()
/// .minimize()
/// .sampler(TpeSampler::new())
/// .pruner(NopPruner)
/// .build();
/// ```
#[must_use]
pub fn builder() -> StudyBuilder<V> {
StudyBuilder::new()
}
/// Create a study that minimizes the objective value.
///
/// This is a shorthand for `Study::with_sampler(Direction::Minimize, sampler)`.
///
/// # Arguments
///
/// * `sampler` - The sampler to use for parameter sampling.
///
/// # Examples
///
/// ```
/// use optimizer::Study;
/// use optimizer::sampler::tpe::TpeSampler;
///
/// let study: Study<f64> = Study::minimize(TpeSampler::new());
/// assert_eq!(study.direction(), optimizer::Direction::Minimize);
/// ```
#[must_use]
pub fn minimize(sampler: impl Sampler + 'static) -> Self
where
V: Send + Sync + 'static,
{
Self::with_sampler(Direction::Minimize, sampler)
}
/// Create a study that maximizes the objective value.
///
/// This is a shorthand for `Study::with_sampler(Direction::Maximize, sampler)`.
///
/// # Arguments
///
/// * `sampler` - The sampler to use for parameter sampling.
///
/// # Examples
///
/// ```
/// use optimizer::Study;
/// use optimizer::sampler::tpe::TpeSampler;
///
/// let study: Study<f64> = Study::maximize(TpeSampler::new());
/// assert_eq!(study.direction(), optimizer::Direction::Maximize);
/// ```
#[must_use]
pub fn maximize(sampler: impl Sampler + 'static) -> Self
where
V: Send + Sync + 'static,
{
Self::with_sampler(Direction::Maximize, sampler)
}
/// Create a new study with a custom sampler.
///
/// # Arguments
///
/// * `direction` - Whether to minimize or maximize the objective function.
/// * `sampler` - The sampler to use for parameter sampling.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Maximize, sampler);
/// assert_eq!(study.direction(), Direction::Maximize);
/// ```
pub fn with_sampler(direction: Direction, sampler: impl Sampler + 'static) -> Self
where
V: Send + Sync + 'static,
{
Self::with_sampler_and_storage(
direction,
sampler,
crate::storage::MemoryStorage::<V>::new(),
)
}
/// Build a trial factory for sampler integration when `V = f64`.
pub(crate) fn make_trial_factory(
sampler: &Arc<dyn Sampler>,
storage: &Arc<dyn crate::storage::Storage<V>>,
pruner: &Arc<dyn Pruner>,
) -> Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>
where
V: 'static,
{
// Try to downcast the storage's trial buffer to the f64 specialization.
// This succeeds only when V = f64, enabling automatic sampler integration.
let trials_arc = storage.trials_arc();
let any_ref: &dyn Any = trials_arc;
let f64_trials: Option<&Arc<RwLock<Vec<CompletedTrial<f64>>>>> = any_ref.downcast_ref();
f64_trials.map(|trials| {
let sampler = Arc::clone(sampler);
let trials = Arc::clone(trials);
let pruner = Arc::clone(pruner);
let factory: Arc<dyn Fn(u64) -> Trial + Send + Sync> = Arc::new(move |id| {
Trial::with_sampler(
id,
Arc::clone(&sampler),
Arc::clone(&trials),
Arc::clone(&pruner),
)
});
factory
})
}
/// Create a study with a custom sampler and storage backend.
///
/// This is the most general constructor — all other constructors
/// delegate to this one. Use it when you need a non-default storage
/// backend (e.g., `JournalStorage`).
///
/// # Arguments
///
/// * `direction` - Whether to minimize or maximize the objective function.
/// * `sampler` - The sampler to use for parameter sampling.
/// * `storage` - The storage backend for completed trials.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::storage::MemoryStorage;
/// use optimizer::{Direction, Study};
///
/// let storage = MemoryStorage::<f64>::new();
/// let study = Study::with_sampler_and_storage(Direction::Minimize, RandomSampler::new(), storage);
/// ```
pub fn with_sampler_and_storage(
direction: Direction,
sampler: impl Sampler + 'static,
storage: impl crate::storage::Storage<V> + 'static,
) -> Self
where
V: 'static,
{
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
let pruner: Arc<dyn Pruner> = Arc::new(NopPruner);
let storage: Arc<dyn crate::storage::Storage<V>> = Arc::new(storage);
let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner);
Self {
direction,
sampler,
pruner,
storage,
trial_factory,
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
}
}
/// Return the optimization direction.
#[must_use]
pub fn direction(&self) -> Direction {
self.direction
}
/// Creates a study with a custom sampler and pruner.
///
/// Uses the default [`MemoryStorage`](crate::storage::MemoryStorage) backend.
///
/// # Arguments
///
/// * `direction` - Whether to minimize or maximize the objective function.
/// * `sampler` - The sampler to use for parameter sampling.
/// * `pruner` - The pruner to use for trial pruning.
///
/// # Examples
///
/// ```
/// use optimizer::pruner::NopPruner;
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, sampler, NopPruner);
/// ```
pub fn with_sampler_and_pruner(
direction: Direction,
sampler: impl Sampler + 'static,
pruner: impl Pruner + 'static,
) -> Self
where
V: Send + Sync + 'static,
{
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
let pruner: Arc<dyn Pruner> = Arc::new(pruner);
let storage: Arc<dyn crate::storage::Storage<V>> =
Arc::new(crate::storage::MemoryStorage::<V>::new());
let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner);
Self {
direction,
sampler,
pruner,
storage,
trial_factory,
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
}
}
/// Replace the sampler used for future parameter suggestions.
///
/// The new sampler takes effect for all subsequent calls to
/// [`create_trial`](Self::create_trial), [`ask`](Self::ask), and the
/// `optimize*` family. Already-completed trials are unaffected.
///
/// # Examples
///
/// ```
/// use optimizer::sampler::tpe::TpeSampler;
/// use optimizer::{Direction, Study};
///
/// let mut study: Study<f64> = Study::new(Direction::Minimize);
/// study.set_sampler(TpeSampler::new());
/// ```
pub fn set_sampler(&mut self, sampler: impl Sampler + 'static)
where
V: 'static,
{
self.sampler = Arc::new(sampler);
self.trial_factory = Self::make_trial_factory(&self.sampler, &self.storage, &self.pruner);
}
/// Replace the pruner used for future trials.
///
/// The new pruner takes effect for all trials created after this call.
///
/// # Examples
///
/// ```
/// use optimizer::prelude::*;
///
/// let mut study: Study<f64> = Study::new(Direction::Minimize);
/// study.set_pruner(MedianPruner::new(Direction::Minimize));
/// ```
pub fn set_pruner(&mut self, pruner: impl Pruner + 'static)
where
V: 'static,
{
self.pruner = Arc::new(pruner);
self.trial_factory = Self::make_trial_factory(&self.sampler, &self.storage, &self.pruner);
}
/// Return a reference to the study's current pruner.
#[must_use]
pub fn pruner(&self) -> &dyn Pruner {
&*self.pruner
}
/// Enqueue a specific parameter configuration to be evaluated next.
///
/// The next call to [`ask()`](Self::ask) or the next trial in [`optimize()`](Self::optimize)
/// will use these exact parameters instead of sampling from the sampler.
///
/// Multiple configurations can be enqueued; they are evaluated in FIFO order.
/// If an enqueued configuration is missing a parameter that the objective calls
/// `suggest()` on, that parameter falls back to normal sampling.
///
/// # Arguments
///
/// * `params` - A map from parameter IDs to the values to use.
///
/// # Examples
///
/// ```
/// use std::collections::HashMap;
///
/// use optimizer::parameter::{FloatParam, IntParam, ParamValue, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0);
/// let y = IntParam::new(1, 100);
///
/// // Evaluate these specific configurations first
/// study.enqueue(HashMap::from([
/// (x.id(), ParamValue::Float(0.001)),
/// (y.id(), ParamValue::Int(3)),
/// ]));
///
/// // Next trial will use x=0.001, y=3
/// let mut trial = study.ask();
/// assert_eq!(x.suggest(&mut trial).unwrap(), 0.001);
/// assert_eq!(y.suggest(&mut trial).unwrap(), 3);
/// ```
pub fn enqueue(&self, params: HashMap<ParamId, ParamValue>) {
self.enqueued_params.lock().push_back(params);
}
/// Return the trial ID of the current best trial from the given slice.
#[cfg(feature = "tracing")]
pub(crate) fn best_id(&self, trials: &[CompletedTrial<V>]) -> Option<u64> {
let direction = self.direction;
trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.max_by(|a, b| Self::compare_trials(a, b, direction))
.map(|t| t.id)
}
/// Return the number of enqueued parameter configurations.
///
/// See [`enqueue`](Self::enqueue) for how to add configurations.
#[must_use]
pub fn n_enqueued(&self) -> usize {
self.enqueued_params.lock().len()
}
/// Generate the next unique trial ID.
pub(crate) fn next_trial_id(&self) -> u64 {
self.storage.next_trial_id()
}
/// Create a new trial with a unique ID.
///
/// The trial starts in the `Running` state and can be used to suggest
/// parameter values. After the objective function is evaluated, call
/// `complete_trial` or `fail_trial` to record the result.
///
/// For `Study<f64>`, this method automatically integrates with the study's
/// sampler and trial history, so there is no need to call a separate
/// `create_trial_with_sampler()` method.
///
/// # Examples
///
/// ```
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let trial = study.create_trial();
/// assert_eq!(trial.id(), 0);
///
/// let trial2 = study.create_trial();
/// assert_eq!(trial2.id(), 1);
/// ```
#[must_use]
pub fn create_trial(&self) -> Trial {
self.storage.refresh();
let id = self.next_trial_id();
let mut trial = if let Some(factory) = &self.trial_factory {
factory(id)
} else {
Trial::new(id)
};
// If there are enqueued params, inject them into this trial
if let Some(fixed_params) = self.enqueued_params.lock().pop_front() {
trial.set_fixed_params(fixed_params);
}
trial
}
/// Record a completed trial with its objective value.
///
/// This method stores the trial's parameters, distributions, and objective
/// value in the study's history. The stored data is used by samplers to
/// inform future parameter suggestions.
///
/// # Arguments
///
/// * `trial` - The trial that was evaluated.
/// * `value` - The objective value returned by the objective function.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x_param = FloatParam::new(0.0, 1.0);
/// let mut trial = study.create_trial();
/// let x = x_param.suggest(&mut trial).unwrap();
/// let objective_value = x * x;
/// study.complete_trial(trial, objective_value);
///
/// assert_eq!(study.n_trials(), 1);
/// ```
pub fn complete_trial(&self, trial: Trial, value: V) {
let completed = trial.into_completed(value, TrialState::Complete);
self.storage.push(completed);
}
/// Record a failed trial with an error message.
///
/// Failed trials are not stored in the study's history and do not
/// contribute to future sampling decisions. This method is useful
/// when the objective function raises an error that should not stop
/// the optimization process.
///
/// # Arguments
///
/// * `trial` - The trial that failed.
/// * `_error` - An error message describing why the trial failed.
///
/// # Examples
///
/// ```
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let trial = study.create_trial();
/// study.fail_trial(trial, "objective function raised an exception");
///
/// // Failed trials are not counted
/// assert_eq!(study.n_trials(), 0);
/// ```
pub fn fail_trial(&self, mut trial: Trial, _error: impl ToString) {
trial.set_failed();
// Failed trials are not stored in completed_trials
// They could be stored in a separate list for debugging if needed
}
/// Request a new trial with suggested parameters.
///
/// This is the first half of the ask-and-tell interface. After calling
/// `ask()`, use parameter types to suggest values on the returned trial,
/// evaluate your objective externally, then pass the trial back to
/// [`tell()`](Self::tell) with the result.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x = FloatParam::new(0.0, 10.0);
///
/// let mut trial = study.ask();
/// let x_val = x.suggest(&mut trial).unwrap();
/// let value = x_val * x_val;
/// study.tell(trial, Ok::<_, &str>(value));
/// ```
#[must_use]
pub fn ask(&self) -> Trial {
self.create_trial()
}
/// Report the result of a trial obtained from [`ask()`](Self::ask).
///
/// Pass `Ok(value)` for a successful evaluation or `Err(reason)` for a
/// failure. Failed trials are not stored in the study's history.
///
/// # Examples
///
/// ```
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
///
/// let trial = study.ask();
/// study.tell(trial, Ok::<_, &str>(42.0));
/// assert_eq!(study.n_trials(), 1);
///
/// let trial = study.ask();
/// study.tell(trial, Err::<f64, _>("evaluation failed"));
/// assert_eq!(study.n_trials(), 1); // failed trials not counted
/// ```
pub fn tell(&self, trial: Trial, value: core::result::Result<V, impl ToString>) {
match value {
Ok(v) => self.complete_trial(trial, v),
Err(e) => self.fail_trial(trial, e),
}
}
/// Record a pruned trial, preserving its intermediate values.
///
/// Pruned trials are stored alongside completed trials so that samplers
/// can optionally learn from partial evaluations. The trial's state is
/// set to [`Pruned`](crate::TrialState::Pruned).
///
/// In practice you rarely call this directly — returning
/// `Err(TrialPruned)` from an objective function handles pruning
/// automatically.
///
/// # Arguments
///
/// * `trial` - The trial that was pruned.
pub fn prune_trial(&self, trial: Trial)
where
V: Default,
{
let completed = trial.into_completed(V::default(), TrialState::Pruned);
self.storage.push(completed);
}
/// Return all completed trials as a `Vec`.
///
/// The returned vector contains clones of `CompletedTrial` values, which contain
/// the trial's parameters, distributions, and objective value.
///
/// Note: This method acquires a read lock on the completed trials, so the
/// returned vector is a clone of the internal storage.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// let x_param = FloatParam::new(0.0, 1.0);
/// let mut trial = study.create_trial();
/// let _ = x_param.suggest(&mut trial);
/// study.complete_trial(trial, 0.5);
///
/// for completed in study.trials() {
/// println!("Trial {} has value {:?}", completed.id, completed.value);
/// }
/// ```
#[must_use]
pub fn trials(&self) -> Vec<CompletedTrial<V>>
where
V: Clone,
{
self.storage.trials_arc().read().clone()
}
/// Return the number of completed trials.
///
/// Pruned and failed trials are not counted. Use
/// [`n_pruned_trials()`](Self::n_pruned_trials) for the pruned count.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> = Study::new(Direction::Minimize);
/// assert_eq!(study.n_trials(), 0);
///
/// let x_param = FloatParam::new(0.0, 1.0);
/// let mut trial = study.create_trial();
/// let _ = x_param.suggest(&mut trial);
/// study.complete_trial(trial, 0.5);
/// assert_eq!(study.n_trials(), 1);
/// ```
#[must_use]
pub fn n_trials(&self) -> usize {
self.storage
.trials_arc()
.read()
.iter()
.filter(|t| t.state == TrialState::Complete)
.count()
}
/// Return the number of pruned trials.
///
/// Pruned trials are those that were stopped early by the pruner.
#[must_use]
pub fn n_pruned_trials(&self) -> usize {
self.storage
.trials_arc()
.read()
.iter()
.filter(|t| t.state == TrialState::Pruned)
.count()
}
/// Compare two completed trials using constraint-aware ranking.
///
/// 1. Feasible trials always rank above infeasible trials.
/// 2. Among feasible trials, rank by objective value (respecting direction).
/// 3. Among infeasible trials, rank by total constraint violation (lower is better).
pub(crate) fn compare_trials(
a: &CompletedTrial<V>,
b: &CompletedTrial<V>,
direction: Direction,
) -> core::cmp::Ordering {
match (a.is_feasible(), b.is_feasible()) {
(true, false) => core::cmp::Ordering::Greater,
(false, true) => core::cmp::Ordering::Less,
(false, false) => {
let va: f64 = a.constraints.iter().map(|c| c.max(0.0)).sum();
let vb: f64 = b.constraints.iter().map(|c| c.max(0.0)).sum();
vb.partial_cmp(&va).unwrap_or(core::cmp::Ordering::Equal)
}
(true, true) => {
let ordering = a.value.partial_cmp(&b.value);
match direction {
Direction::Minimize => {
ordering.map_or(core::cmp::Ordering::Equal, core::cmp::Ordering::reverse)
}
Direction::Maximize => ordering.unwrap_or(core::cmp::Ordering::Equal),
}
}
}
}
}
impl<V: PartialOrd + Send + Sync + 'static> Study<V> {
/// Create a study with a custom sampler, pruner, and storage backend.
///
/// The most flexible constructor, allowing full control over all components.
///
/// # Arguments
///
/// * `direction` - Whether to minimize or maximize the objective function.
/// * `sampler` - The sampler to use for parameter sampling.
/// * `pruner` - The pruner to use for trial pruning.
/// * `storage` - The storage backend for completed trials.
///
/// # Examples
///
/// ```
/// use optimizer::prelude::*;
/// use optimizer::storage::MemoryStorage;
///
/// let study = Study::with_sampler_pruner_and_storage(
/// Direction::Minimize,
/// TpeSampler::new(),
/// MedianPruner::new(Direction::Minimize),
/// MemoryStorage::<f64>::new(),
/// );
/// ```
pub fn with_sampler_pruner_and_storage(
direction: Direction,
sampler: impl Sampler + 'static,
pruner: impl Pruner + 'static,
storage: impl crate::storage::Storage<V> + 'static,
) -> Self {
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
let pruner: Arc<dyn Pruner> = Arc::new(pruner);
let storage: Arc<dyn crate::storage::Storage<V>> = Arc::new(storage);
let trial_factory = Self::make_trial_factory(&sampler, &storage, &pruner);
Self {
direction,
sampler,
pruner,
storage,
trial_factory,
enqueued_params: Arc::new(Mutex::new(VecDeque::new())),
}
}
}
/// Returns `true` if the error represents a pruned trial.
///
/// Checks via `Any` downcasting whether `e` is `Error::TrialPruned` or
/// the standalone `TrialPruned` struct.
pub(super) fn is_trial_pruned<E: 'static>(e: &E) -> bool {
let any: &dyn Any = e;
if let Some(err) = any.downcast_ref::<crate::Error>() {
matches!(err, crate::Error::TrialPruned)
} else {
any.downcast_ref::<crate::error::TrialPruned>().is_some()
}
}
+123
View File
@@ -0,0 +1,123 @@
use core::ops::ControlFlow;
use crate::types::TrialState;
use super::{Study, is_trial_pruned};
impl<V> Study<V>
where
V: PartialOrd,
{
/// Run optimization with an objective.
///
/// Accepts any [`Objective`](crate::Objective) implementation, including
/// plain closures (`Fn(&mut Trial) -> Result<V, E>`) thanks to the
/// blanket impl. Struct-based objectives can override
/// [`before_trial`](crate::Objective::before_trial) and
/// [`after_trial`](crate::Objective::after_trial) for early stopping.
///
/// Runs up to `n_trials` evaluations sequentially.
///
/// # Errors
///
/// Returns `Error::NoCompletedTrials` if no trials completed successfully.
///
/// # Examples
///
/// ```
/// use optimizer::parameter::{FloatParam, Parameter};
/// use optimizer::sampler::random::RandomSampler;
/// use optimizer::{Direction, Study};
///
/// let sampler = RandomSampler::with_seed(42);
/// let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
/// let x_param = FloatParam::new(-10.0, 10.0);
///
/// study
/// .optimize(10, |trial: &mut optimizer::Trial| {
/// let x = x_param.suggest(trial)?;
/// Ok::<_, optimizer::Error>(x * x)
/// })
/// .unwrap();
///
/// assert!(study.n_trials() > 0);
/// assert!(study.best_value().unwrap() >= 0.0);
/// ```
#[allow(clippy::needless_pass_by_value)]
pub fn optimize(
&self,
n_trials: usize,
objective: impl crate::objective::Objective<V>,
) -> crate::Result<()>
where
V: Clone + Default,
{
#[cfg(feature = "tracing")]
let _span =
tracing::info_span!("optimize", n_trials, direction = ?self.direction).entered();
for _ in 0..n_trials {
if let ControlFlow::Break(()) = objective.before_trial(self) {
break;
}
let mut trial = self.create_trial();
match objective.evaluate(&mut trial) {
Ok(value) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
let completed = trial.into_completed(value, TrialState::Complete);
// Fire after_trial hook before pushing to storage
let flow = objective.after_trial(self, &completed);
self.storage.push(completed);
#[cfg(feature = "tracing")]
{
tracing::info!(trial_id, "trial completed");
let trials = self.storage.trials_arc().read();
if trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.count()
== 1
|| trials.last().map(|t| t.id) == self.best_id(&trials)
{
tracing::info!(trial_id, "new best value found");
}
}
if let ControlFlow::Break(()) = flow {
return Ok(());
}
}
Err(e) if is_trial_pruned(&e) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.prune_trial(trial);
trace_info!(trial_id, "trial pruned");
}
Err(e) => {
#[cfg(feature = "tracing")]
let trial_id = trial.id();
self.fail_trial(trial, e.to_string());
trace_debug!(trial_id, "trial failed");
}
}
}
// Return error if no trials completed successfully
let has_complete = self
.storage
.trials_arc()
.read()
.iter()
.any(|t| t.state == TrialState::Complete);
if !has_complete {
return Err(crate::Error::NoCompletedTrials);
}
Ok(())
}
}
+138
View File
@@ -0,0 +1,138 @@
#[cfg(feature = "serde")]
use std::collections::HashMap;
#[cfg(feature = "serde")]
use crate::sampler::CompletedTrial;
#[cfg(feature = "serde")]
use crate::types::Direction;
#[cfg(feature = "serde")]
use super::Study;
/// A serializable snapshot of a study's state.
///
/// Since [`Study`] contains non-serializable fields (samplers, atomics, etc.),
/// this struct captures the essential state needed to save and restore a study.
///
/// # Schema versioning
///
/// The `version` field enables future schema evolution without breaking existing files.
/// The current version is `1`.
///
/// # Sampler state
///
/// Sampler state is **not** included in the snapshot. After loading, the study
/// uses a default `RandomSampler`. Call [`Study::set_sampler`] to restore
/// the desired sampler configuration.
#[cfg(feature = "serde")]
#[derive(serde::Serialize, serde::Deserialize)]
pub struct StudySnapshot<V> {
/// Schema version for forward compatibility.
pub version: u32,
/// The optimization direction.
pub direction: Direction,
/// All completed (and pruned) trials.
pub trials: Vec<CompletedTrial<V>>,
/// The next trial ID to assign.
pub next_trial_id: u64,
/// Optional metadata (creation timestamp, sampler description, etc.).
pub metadata: HashMap<String, String>,
}
#[cfg(feature = "serde")]
impl<V: PartialOrd + Clone + serde::Serialize> Study<V> {
/// Save the study state to a JSON file.
///
/// # Errors
///
/// Returns an I/O error if the file cannot be created or written.
pub fn save(&self, path: impl AsRef<std::path::Path>) -> std::io::Result<()> {
let path = path.as_ref();
let trials = self.trials();
let next_trial_id = self.storage.peek_next_trial_id();
let snapshot = StudySnapshot {
version: 1,
direction: self.direction,
trials,
next_trial_id,
metadata: HashMap::new(),
};
// Atomic write: write to a temp file in the same directory, then rename.
// This prevents corrupt files if the process crashes mid-write.
let parent = path.parent().unwrap_or(std::path::Path::new("."));
let tmp_path = parent.join(format!(
".{}.tmp",
path.file_name().unwrap_or_default().to_string_lossy()
));
let file = std::fs::File::create(&tmp_path)?;
serde_json::to_writer_pretty(file, &snapshot).map_err(std::io::Error::other)?;
std::fs::rename(&tmp_path, path)
}
}
#[cfg(feature = "serde")]
impl<V: PartialOrd + Send + Sync + Clone + serde::de::DeserializeOwned + 'static> Study<V> {
/// Load a study from a JSON file.
///
/// The loaded study uses a `RandomSampler` by default. Call
/// [`set_sampler()`](Self::set_sampler) to restore the original sampler
/// configuration.
///
/// # Errors
///
/// Returns an I/O error if the file cannot be read or parsed.
pub fn load(path: impl AsRef<std::path::Path>) -> std::io::Result<Self> {
use crate::sampler::random::RandomSampler;
let file = std::fs::File::open(path)?;
let snapshot: StudySnapshot<V> = serde_json::from_reader(file)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
let storage = crate::storage::MemoryStorage::with_trials(snapshot.trials);
Ok(Self::with_sampler_and_storage(
snapshot.direction,
RandomSampler::new(),
storage,
))
}
}
#[cfg(feature = "journal")]
impl<V> Study<V>
where
V: PartialOrd + Send + Sync + serde::Serialize + serde::de::DeserializeOwned + 'static,
{
/// Create a study backed by a JSONL journal file.
///
/// Any existing trials in the file are loaded into memory and the
/// trial ID counter is set to one past the highest stored ID. New
/// trials are written through to the file on completion.
///
/// # Arguments
///
/// * `direction` - Whether to minimize or maximize the objective function.
/// * `sampler` - The sampler to use for parameter sampling.
/// * `path` - Path to the JSONL journal file (created if absent).
///
/// # Errors
///
/// Returns a [`Storage`](crate::Error::Storage) error if loading fails.
///
/// # Examples
///
/// ```no_run
/// use optimizer::sampler::tpe::TpeSampler;
/// use optimizer::{Direction, Study};
///
/// let study: Study<f64> =
/// Study::with_journal(Direction::Minimize, TpeSampler::new(), "trials.jsonl").unwrap();
/// ```
pub fn with_journal(
direction: Direction,
sampler: impl crate::sampler::Sampler + 'static,
path: impl AsRef<std::path::Path>,
) -> crate::Result<Self> {
let storage = crate::storage::JournalStorage::<V>::open(path)?;
Ok(Self::with_sampler_and_storage(direction, sampler, storage))
}
}
+506 -630
View File
File diff suppressed because it is too large Load Diff
+4
View File
@@ -2,6 +2,7 @@
/// The direction of optimization.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Direction {
/// Minimize the objective value.
Minimize,
@@ -11,6 +12,7 @@ pub enum Direction {
/// The state of a trial in its lifecycle.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum TrialState {
/// The trial is currently running.
Running,
@@ -18,4 +20,6 @@ pub enum TrialState {
Complete,
/// The trial failed with an error.
Failed,
/// The trial was pruned (stopped early).
Pruned,
}
+500
View File
@@ -0,0 +1,500 @@
//! HTML report generation for optimization visualization.
//!
//! Generate self-contained HTML files with embedded
//! [Plotly.js](https://plotly.com/javascript/) charts for offline
//! visualization of optimization results. No feature flag is required —
//! this module is always available.
//!
//! # Charts included
//!
//! | Chart | Description |
//! |---|---|
//! | **Optimization history** | Objective value vs trial number with best-so-far line |
//! | **Slice plots** | Objective value vs each parameter (1D scatter per param) |
//! | **Parallel coordinates** | Multi-parameter relationship view (color = objective) |
//! | **Parameter importance** | Horizontal bar chart of Spearman-based importance |
//! | **Trial timeline** | Duration/index of each trial, color-coded by state |
//! | **Intermediate values** | Per-trial learning curves (if pruning data available) |
//!
//! # Usage
//!
//! Call [`Study::export_html()`](crate::Study::export_html) or
//! [`generate_html_report()`] directly:
//!
//! ```no_run
//! use optimizer::prelude::*;
//!
//! let study: Study<f64> = Study::new(Direction::Minimize);
//! # let x = FloatParam::new(0.0, 1.0);
//! # study.optimize(10, |trial: &mut optimizer::Trial| {
//! # let v = x.suggest(trial)?;
//! # Ok::<_, optimizer::Error>(v * v)
//! # }).unwrap();
//! study.export_html("report.html").unwrap();
//! ```
//!
//! The output is a single HTML file that can be opened in any browser.
//! An internet connection is needed on first load to fetch `Plotly.js`
//! from a CDN.
use core::fmt::Write as _;
use std::collections::BTreeMap;
use std::path::Path;
use crate::param::ParamValue;
use crate::parameter::ParamId;
use crate::sampler::CompletedTrial;
use crate::types::{Direction, TrialState};
/// Generate an HTML report with interactive Plotly.js charts.
///
/// Create a self-contained HTML file at `path` containing up to six
/// interactive charts. Charts that require data not present in the study
/// (e.g., intermediate values) are automatically omitted.
///
/// The report includes: optimization history, slice plots, parallel
/// coordinates, parameter importance, trial timeline, and intermediate
/// values (when available).
///
/// This is also available as [`Study::export_html()`](crate::Study::export_html).
///
/// # Errors
///
/// Return an I/O error if the file cannot be created or written.
pub fn generate_html_report(
study: &crate::Study<f64>,
path: impl AsRef<Path>,
) -> std::io::Result<()> {
let trials = study.trials();
let direction = study.direction();
let importance = study.param_importance();
let html = build_html(&trials, direction, &importance);
std::fs::write(path, html)
}
fn build_html(
trials: &[CompletedTrial<f64>],
direction: Direction,
importance: &[(String, f64)],
) -> String {
let mut html = String::with_capacity(8192);
let dir_label = match direction {
Direction::Minimize => "Minimize",
Direction::Maximize => "Maximize",
};
// Collect parameter metadata.
let param_info = collect_param_info(trials);
let has_intermediate = trials.iter().any(|t| !t.intermediate_values.is_empty());
let _ = write!(
html,
r#"<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Optimization Report</title>
<script src="https://cdn.plot.ly/plotly-2.35.2.min.js"></script>
<style>
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
background: #f5f6fa; color: #2c3e50; padding: 24px; }}
h1 {{ text-align: center; margin-bottom: 8px; font-size: 1.8em; }}
.subtitle {{ text-align: center; color: #7f8c8d; margin-bottom: 24px; }}
.chart {{ background: #fff; border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.08);
margin-bottom: 24px; padding: 16px; }}
.chart-title {{ font-size: 1.1em; font-weight: 600; margin-bottom: 8px; }}
</style>
</head>
<body>
<h1>Optimization Report</h1>
<p class="subtitle">{dir_label} &middot; {n} trials</p>
"#,
n = trials.len(),
);
// Optimization history chart.
html.push_str("<div class=\"chart\"><div class=\"chart-title\">Optimization History</div><div id=\"history\"></div></div>\n");
write_history_chart(&mut html, trials, direction);
// Slice plots.
if !param_info.is_empty() {
html.push_str("<div class=\"chart\"><div class=\"chart-title\">Slice Plots</div><div id=\"slices\"></div></div>\n");
write_slice_charts(&mut html, trials, &param_info);
}
// Parallel coordinates.
if param_info.len() >= 2 {
html.push_str("<div class=\"chart\"><div class=\"chart-title\">Parallel Coordinates</div><div id=\"parcoords\"></div></div>\n");
write_parallel_coordinates(&mut html, trials, &param_info, direction);
}
// Parameter importance.
if !importance.is_empty() {
html.push_str("<div class=\"chart\"><div class=\"chart-title\">Parameter Importance</div><div id=\"importance\"></div></div>\n");
write_importance_chart(&mut html, importance);
}
// Trial timeline.
html.push_str("<div class=\"chart\"><div class=\"chart-title\">Trial Timeline</div><div id=\"timeline\"></div></div>\n");
write_timeline_chart(&mut html, trials);
// Intermediate values.
if has_intermediate {
html.push_str("<div class=\"chart\"><div class=\"chart-title\">Intermediate Values</div><div id=\"intermediate\"></div></div>\n");
write_intermediate_chart(&mut html, trials);
}
html.push_str("</body>\n</html>\n");
html
}
/// Metadata about each parameter seen across trials.
struct ParamMeta {
label: String,
}
/// Collect parameter labels and distributions across all trials.
fn collect_param_info(trials: &[CompletedTrial<f64>]) -> BTreeMap<ParamId, ParamMeta> {
let mut info: BTreeMap<ParamId, ParamMeta> = BTreeMap::new();
for trial in trials {
for &id in trial.params.keys() {
info.entry(id).or_insert_with(|| {
let label = trial
.param_labels
.get(&id)
.cloned()
.unwrap_or_else(|| id.to_string());
ParamMeta { label }
});
}
}
info
}
// ---------------------------------------------------------------------------
// Chart generators
// ---------------------------------------------------------------------------
fn write_history_chart(html: &mut String, trials: &[CompletedTrial<f64>], direction: Direction) {
let complete: Vec<_> = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.collect();
if complete.is_empty() {
return;
}
let mut ids = Vec::with_capacity(complete.len());
let mut vals = Vec::with_capacity(complete.len());
let mut best_vals = Vec::with_capacity(complete.len());
let mut best = complete[0].value;
for t in &complete {
ids.push(t.id);
vals.push(t.value);
best = match direction {
Direction::Minimize => best.min(t.value),
Direction::Maximize => best.max(t.value),
};
best_vals.push(best);
}
let _ = write!(
html,
r##"<script>
Plotly.newPlot("history", [
{{ x: {ids:?}, y: {vals:?}, mode: "markers", name: "Objective", type: "scatter",
marker: {{ color: "#3498db", size: 6 }} }},
{{ x: {ids:?}, y: {best_vals:?}, mode: "lines", name: "Best so far", type: "scatter",
line: {{ color: "#e74c3c", width: 2 }} }}
], {{ xaxis: {{ title: "Trial" }}, yaxis: {{ title: "Objective Value" }},
margin: {{ t: 10 }}, legend: {{ x: 1, xanchor: "right", y: 1 }} }},
{{ responsive: true }});
</script>
"##,
);
}
fn write_slice_charts(
html: &mut String,
trials: &[CompletedTrial<f64>],
param_info: &BTreeMap<ParamId, ParamMeta>,
) {
let complete: Vec<_> = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.collect();
if complete.is_empty() {
return;
}
let n_params = param_info.len();
let cols = if n_params <= 2 { n_params } else { 2 };
let rows = n_params.div_ceil(cols);
// Build subplot titles and data.
let mut subplot_titles = Vec::new();
let mut traces = String::new();
for (i, (id, meta)) in param_info.iter().enumerate() {
subplot_titles.push(format!("\"{}\"", escape_js(&meta.label)));
let mut x_vals = Vec::new();
let mut y_vals = Vec::new();
for t in &complete {
if let Some(pv) = t.params.get(id) {
x_vals.push(param_value_to_f64(pv));
y_vals.push(t.value);
}
}
let subplot_idx = i + 1;
let xa = if subplot_idx == 1 {
"x".to_string()
} else {
format!("x{subplot_idx}")
};
let ya = if subplot_idx == 1 {
"y".to_string()
} else {
format!("y{subplot_idx}")
};
let _ = write!(
traces,
r##"{{ x: {x_vals:?}, y: {y_vals:?}, mode: "markers", type: "scatter",
xaxis: "{xa}", yaxis: "{ya}",
marker: {{ color: "#3498db", size: 5 }}, showlegend: false }},"##,
);
}
let _ = write!(
html,
r#"<script>
Plotly.newPlot("slices", [{traces}],
{{ grid: {{ rows: {rows}, columns: {cols}, pattern: "independent" }},
annotations: [{annotations}],
margin: {{ t: 30 }}, showlegend: false }},
{{ responsive: true }});
</script>
"#,
annotations = build_subplot_annotations(&subplot_titles, rows, cols),
);
}
fn write_parallel_coordinates(
html: &mut String,
trials: &[CompletedTrial<f64>],
param_info: &BTreeMap<ParamId, ParamMeta>,
direction: Direction,
) {
let complete: Vec<_> = trials
.iter()
.filter(|t| t.state == TrialState::Complete)
.collect();
if complete.is_empty() {
return;
}
let mut dimensions = String::new();
// Add objective value as the first dimension.
let obj_vals: Vec<f64> = complete.iter().map(|t| t.value).collect();
let _ = write!(
dimensions,
r#"{{ label: "Objective", values: {obj_vals:?} }},"#,
);
// Add each parameter as a dimension.
for (id, meta) in param_info {
let vals: Vec<f64> = complete
.iter()
.map(|t| t.params.get(id).map_or(f64::NAN, param_value_to_f64))
.collect();
let _ = write!(
dimensions,
r#"{{ label: "{label}", values: {vals:?} }},"#,
label = escape_js(&meta.label),
);
}
// Color by objective value: green = good, red = bad.
let (cmin, cmax) = min_max(&obj_vals);
let colorscale = match direction {
Direction::Minimize => r##"[[0,"#2ecc71"],[1,"#e74c3c"]]"##,
Direction::Maximize => r##"[[0,"#e74c3c"],[1,"#2ecc71"]]"##,
};
let _ = write!(
html,
r#"<script>
Plotly.newPlot("parcoords", [{{
type: "parcoords",
line: {{ color: {obj_vals:?}, colorscale: {colorscale},
cmin: {cmin}, cmax: {cmax}, showscale: true }},
dimensions: [{dimensions}]
}}], {{ margin: {{ t: 10 }} }}, {{ responsive: true }});
</script>
"#,
);
}
fn write_importance_chart(html: &mut String, importance: &[(String, f64)]) {
let names: Vec<_> = importance
.iter()
.map(|(n, _)| format!("\"{}\"", escape_js(n)))
.collect();
let values: Vec<f64> = importance.iter().map(|(_, v)| *v).collect();
let _ = write!(
html,
r##"<script>
Plotly.newPlot("importance", [{{
x: {values:?}, y: [{names}], type: "bar", orientation: "h",
marker: {{ color: "#9b59b6" }}
}}], {{ xaxis: {{ title: "Importance (|Spearman correlation|)" }},
yaxis: {{ automargin: true }}, margin: {{ t: 10, l: 120 }} }},
{{ responsive: true }});
</script>
"##,
names = names.join(","),
);
}
fn write_timeline_chart(html: &mut String, trials: &[CompletedTrial<f64>]) {
let mut ids = Vec::with_capacity(trials.len());
let mut colors = Vec::with_capacity(trials.len());
let mut labels = Vec::with_capacity(trials.len());
for t in trials {
ids.push(format!("\"Trial {}\"", t.id));
let (color, label) = match t.state {
TrialState::Complete => ("#2ecc71", "Complete"),
TrialState::Pruned => ("#f39c12", "Pruned"),
TrialState::Failed => ("#e74c3c", "Failed"),
TrialState::Running => ("#3498db", "Running"),
};
colors.push(format!("\"{color}\""));
labels.push(format!("\"{label}\""));
}
// Use trial index as a proxy for duration (no wallclock data available).
let indices: Vec<usize> = (0..trials.len()).collect();
let _ = write!(
html,
r#"<script>
Plotly.newPlot("timeline", [{{
y: [{ids}], x: {indices:?}, type: "bar", orientation: "h",
text: [{labels}], textposition: "auto",
marker: {{ color: [{colors}] }}
}}], {{ xaxis: {{ title: "Trial Index" }}, yaxis: {{ automargin: true, autorange: "reversed" }},
margin: {{ t: 10, l: 80 }}, showlegend: false }},
{{ responsive: true }});
</script>
"#,
ids = ids.join(","),
colors = colors.join(","),
labels = labels.join(","),
);
}
fn write_intermediate_chart(html: &mut String, trials: &[CompletedTrial<f64>]) {
let trials_with_iv: Vec<_> = trials
.iter()
.filter(|t| !t.intermediate_values.is_empty())
.collect();
if trials_with_iv.is_empty() {
return;
}
let mut traces = String::new();
for t in &trials_with_iv {
let steps: Vec<u64> = t.intermediate_values.iter().map(|(s, _)| *s).collect();
let values: Vec<f64> = t.intermediate_values.iter().map(|(_, v)| *v).collect();
let color = match t.state {
TrialState::Pruned => "#f39c12",
_ => "#3498db",
};
let _ = write!(
traces,
r#"{{ x: {steps:?}, y: {values:?}, mode: "lines+markers", name: "Trial {id}",
line: {{ color: "{color}", width: 1 }}, marker: {{ size: 3 }} }},"#,
id = t.id,
);
}
let _ = write!(
html,
r#"<script>
Plotly.newPlot("intermediate", [{traces}],
{{ xaxis: {{ title: "Step" }}, yaxis: {{ title: "Intermediate Value" }},
margin: {{ t: 10 }}, showlegend: true }},
{{ responsive: true }});
</script>
"#,
);
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
#[allow(clippy::cast_precision_loss)]
fn param_value_to_f64(pv: &ParamValue) -> f64 {
match *pv {
ParamValue::Float(v) => v,
ParamValue::Int(v) => v as f64,
ParamValue::Categorical(v) => v as f64,
}
}
fn escape_js(s: &str) -> String {
s.replace('\\', "\\\\")
.replace('"', "\\\"")
.replace('\n', "\\n")
}
fn min_max(vals: &[f64]) -> (f64, f64) {
let mut mn = f64::INFINITY;
let mut mx = f64::NEG_INFINITY;
for &v in vals {
if v.is_nan() {
continue;
}
if v < mn {
mn = v;
}
if v > mx {
mx = v;
}
}
// If all values were NaN, return 0.0..1.0 as a safe fallback.
if mn > mx {
return (0.0, 1.0);
}
(mn, mx)
}
/// Build Plotly annotation objects to act as subplot titles.
#[allow(clippy::cast_precision_loss)]
fn build_subplot_annotations(titles: &[String], rows: usize, cols: usize) -> String {
let mut anns = Vec::new();
for (i, title) in titles.iter().enumerate() {
let row = i / cols;
let col = i % cols;
// Compute x/y anchor in paper coordinates.
let x = if cols == 1 {
0.5
} else {
(f64::from(u32::try_from(col).unwrap_or(0))) / (cols as f64 - 1.0)
};
let y = 1.0 - (f64::from(u32::try_from(row).unwrap_or(0))) / (rows as f64).max(1.0) + 0.02;
anns.push(format!(
r#"{{ text: {title}, x: {x:.3}, y: {y:.3}, xref: "paper", yref: "paper",
showarrow: false, font: {{ size: 12 }} }}"#,
));
}
anns.join(",")
}
+176 -69
View File
@@ -4,6 +4,9 @@
#![cfg(feature = "async")]
use std::sync::Arc;
use optimizer::parameter::{FloatParam, Parameter};
use optimizer::sampler::random::RandomSampler;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, Error, Study};
@@ -13,10 +16,12 @@ async fn test_optimize_async_basic() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-10.0, 10.0);
study
.optimize_async(10, |mut trial| async move {
let x = trial.suggest_float("x", -10.0, 10.0)?;
Ok::<_, Error>((trial, x * x))
.optimize_async(10, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x * x)
})
.await
.expect("async optimization should succeed");
@@ -27,7 +32,7 @@ async fn test_optimize_async_basic() {
}
#[tokio::test]
async fn test_optimize_async_with_sampler() {
async fn test_optimize_async_with_tpe() {
let sampler = TpeSampler::builder()
.seed(42)
.n_startup_trials(5)
@@ -36,10 +41,12 @@ async fn test_optimize_async_with_sampler() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
study
.optimize_async_with_sampler(15, |mut trial| async move {
let x = trial.suggest_float("x", -5.0, 5.0)?;
Ok::<_, Error>((trial, x * x))
.optimize_async(15, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x * x)
})
.await
.expect("async optimization with sampler should succeed");
@@ -54,10 +61,12 @@ async fn test_optimize_parallel() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-10.0, 10.0);
study
.optimize_parallel(20, 4, |mut trial| async move {
let x = trial.suggest_float("x", -10.0, 10.0)?;
Ok::<_, Error>((trial, x * x))
.optimize_parallel(20, 4, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x * x)
})
.await
.expect("parallel optimization should succeed");
@@ -66,7 +75,7 @@ async fn test_optimize_parallel() {
}
#[tokio::test]
async fn test_optimize_parallel_with_sampler() {
async fn test_optimize_parallel_with_tpe() {
let sampler = TpeSampler::builder()
.seed(42)
.n_startup_trials(5)
@@ -75,11 +84,14 @@ async fn test_optimize_parallel_with_sampler() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(-5.0, 5.0);
let y_param = FloatParam::new(-5.0, 5.0);
study
.optimize_parallel_with_sampler(15, 3, |mut trial| async move {
let x = trial.suggest_float("x", -5.0, 5.0)?;
let y = trial.suggest_float("y", -5.0, 5.0)?;
Ok::<_, Error>((trial, x * x + y * y))
.optimize_parallel(15, 3, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
let y = y_param.suggest(trial)?;
Ok::<_, Error>(x * x + y * y)
})
.await
.expect("parallel optimization with sampler should succeed");
@@ -92,26 +104,8 @@ async fn test_optimize_async_all_failures() {
let study: Study<f64> = Study::new(Direction::Minimize);
let result = study
.optimize_async(5, |trial| async move {
let _ = trial;
Err::<(_, f64), &str>("always fails")
})
.await;
assert!(
matches!(result, Err(Error::NoCompletedTrials)),
"should return NoCompletedTrials when all trials fail"
);
}
#[tokio::test]
async fn test_optimize_async_with_sampler_all_failures() {
let study: Study<f64> = Study::new(Direction::Minimize);
let result = study
.optimize_async_with_sampler(5, |trial| async move {
let _ = trial;
Err::<(_, f64), &str>("always fails")
.optimize_async(5, |_trial: &mut optimizer::Trial| {
Err::<f64, &str>("always fails")
})
.await;
@@ -126,26 +120,8 @@ async fn test_optimize_parallel_all_failures() {
let study: Study<f64> = Study::new(Direction::Minimize);
let result = study
.optimize_parallel(5, 2, |trial| async move {
let _ = trial;
Err::<(_, f64), &str>("always fails")
})
.await;
assert!(
matches!(result, Err(Error::NoCompletedTrials)),
"should return NoCompletedTrials when all trials fail"
);
}
#[tokio::test]
async fn test_optimize_parallel_with_sampler_all_failures() {
let study: Study<f64> = Study::new(Direction::Minimize);
let result = study
.optimize_parallel_with_sampler(5, 2, |trial| async move {
let _ = trial;
Err::<(_, f64), &str>("always fails")
.optimize_parallel(5, 2, |_trial: &mut optimizer::Trial| {
Err::<f64, &str>("always fails")
})
.await;
@@ -162,16 +138,16 @@ async fn test_optimize_async_partial_failures() {
let counter = std::sync::atomic::AtomicUsize::new(0);
let x_param = FloatParam::new(0.0, 10.0);
study
.optimize_async(10, |mut trial| {
.optimize_async(10, move |trial: &mut optimizer::Trial| {
let count = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
async move {
if count.is_multiple_of(2) {
let x = trial.suggest_float("x", 0.0, 10.0)?;
Ok::<_, Error>((trial, x))
} else {
Err(Error::NoCompletedTrials) // Use as error type
}
if count.is_multiple_of(2) {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x)
} else {
Err(Error::NoCompletedTrials) // Use as error type
}
})
.await
@@ -186,11 +162,13 @@ async fn test_optimize_parallel_high_concurrency() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(0.0, 10.0);
// Run with concurrency higher than n_trials
study
.optimize_parallel(5, 10, |mut trial| async move {
let x = trial.suggest_float("x", 0.0, 10.0)?;
Ok::<_, Error>((trial, x))
.optimize_parallel(5, 10, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x)
})
.await
.expect("should handle high concurrency");
@@ -203,14 +181,143 @@ async fn test_optimize_parallel_single_concurrency() {
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(0.0, 10.0);
// Run with concurrency of 1 (sequential)
study
.optimize_parallel(10, 1, |mut trial| async move {
let x = trial.suggest_float("x", 0.0, 10.0)?;
Ok::<_, Error>((trial, x))
.optimize_parallel(10, 1, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x)
})
.await
.expect("should work with single concurrency");
assert_eq!(study.n_trials(), 10);
}
#[tokio::test]
async fn test_parallel_executes_concurrently() {
use std::sync::atomic::{AtomicUsize, Ordering};
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(0.0, 10.0);
let active = Arc::new(AtomicUsize::new(0));
let max_active = Arc::new(AtomicUsize::new(0));
let active_c = Arc::clone(&active);
let max_active_c = Arc::clone(&max_active);
study
.optimize_parallel(4, 4, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
let current = active_c.fetch_add(1, Ordering::SeqCst) + 1;
max_active_c.fetch_max(current, Ordering::SeqCst);
std::thread::sleep(std::time::Duration::from_millis(50));
active_c.fetch_sub(1, Ordering::SeqCst);
Ok::<_, Error>(x)
})
.await
.expect("parallel optimization should succeed");
assert_eq!(study.n_trials(), 4);
let max = max_active.load(Ordering::SeqCst);
// With 4 trials and concurrency=4, all should run concurrently
assert!(
max >= 2,
"expected at least 2 concurrent workers, but max was {max}"
);
}
#[tokio::test]
async fn test_parallel_max_concurrency_reached() {
use std::sync::atomic::{AtomicUsize, Ordering};
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(0.0, 10.0);
let active = Arc::new(AtomicUsize::new(0));
let max_active = Arc::new(AtomicUsize::new(0));
let active_c = Arc::clone(&active);
let max_active_c = Arc::clone(&max_active);
study
.optimize_parallel(8, 4, move |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
let current = active_c.fetch_add(1, Ordering::SeqCst) + 1;
// Update max_active if this is the highest seen so far
max_active_c.fetch_max(current, Ordering::SeqCst);
std::thread::sleep(std::time::Duration::from_millis(50));
active_c.fetch_sub(1, Ordering::SeqCst);
Ok::<_, Error>(x)
})
.await
.expect("parallel optimization should succeed");
assert_eq!(study.n_trials(), 8);
let max = max_active.load(Ordering::SeqCst);
assert!(
max >= 2,
"expected at least 2 concurrent workers, but max was {max}"
);
}
#[tokio::test]
async fn test_parallel_panic_returns_task_error() {
let study: Study<f64> = Study::new(Direction::Minimize);
let result = study
.optimize_parallel(3, 2, |_trial: &mut optimizer::Trial| {
panic!("boom");
#[allow(unreachable_code)]
Ok::<_, Error>(0.0)
})
.await;
match result {
Err(Error::TaskError(msg)) => {
assert!(
msg.contains("panic"),
"expected error message to contain 'panic', got: {msg}"
);
}
other => panic!("expected TaskError, got {other:?}"),
}
}
#[tokio::test]
async fn test_parallel_partial_failures_trial_count() {
use std::sync::atomic::{AtomicUsize, Ordering};
let sampler = RandomSampler::with_seed(42);
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
let x_param = FloatParam::new(0.0, 10.0);
let counter = Arc::new(AtomicUsize::new(0));
let counter_c = Arc::clone(&counter);
study
.optimize_parallel(10, 3, move |trial: &mut optimizer::Trial| {
let idx = counter_c.fetch_add(1, Ordering::SeqCst);
if idx % 2 == 1 {
return Err(Error::NoCompletedTrials);
}
let x = x_param.suggest(trial)?;
Ok::<_, Error>(x)
})
.await
.expect("should succeed with partial failures");
// Even indices succeed (0, 2, 4, 6, 8), odd indices fail
assert_eq!(study.n_trials(), 5);
}
+61
View File
@@ -0,0 +1,61 @@
use optimizer::Trial;
use optimizer::parameter::{EnumParam, Parameter};
use optimizer_derive::Categorical;
#[derive(Clone, Debug, PartialEq, Categorical)]
enum Color {
Red,
Green,
Blue,
}
#[derive(Clone, Debug, PartialEq, Categorical)]
enum SingleVariant {
Only,
}
#[test]
fn derive_categorical_n_choices() {
use optimizer::parameter::Categorical;
assert_eq!(Color::N_CHOICES, 3);
assert_eq!(SingleVariant::N_CHOICES, 1);
}
#[test]
fn derive_categorical_roundtrip() {
use optimizer::parameter::Categorical;
for i in 0..Color::N_CHOICES {
let val = Color::from_index(i);
assert_eq!(val.to_index(), i);
}
}
#[test]
fn derive_categorical_values() {
use optimizer::parameter::Categorical;
assert_eq!(Color::from_index(0), Color::Red);
assert_eq!(Color::from_index(1), Color::Green);
assert_eq!(Color::from_index(2), Color::Blue);
assert_eq!(Color::Red.to_index(), 0);
assert_eq!(Color::Green.to_index(), 1);
assert_eq!(Color::Blue.to_index(), 2);
}
#[test]
fn derive_categorical_with_enum_param() {
let mut trial = Trial::new(0);
let param = EnumParam::<Color>::new();
let color = param.suggest(&mut trial).unwrap();
assert!([Color::Red, Color::Green, Color::Blue].contains(&color));
// Cached (same param id)
let color2 = param.suggest(&mut trial).unwrap();
assert_eq!(color, color2);
}
#[test]
fn derive_categorical_suggest_via_trial() {
let mut trial = Trial::new(0);
let color = trial.suggest_param(&EnumParam::<Color>::new()).unwrap();
assert!([Color::Red, Color::Green, Color::Blue].contains(&color));
}
+230
View File
@@ -0,0 +1,230 @@
use optimizer::parameter::{FloatParam, IntParam, Parameter};
use optimizer::sampler::random::RandomSampler;
use optimizer::{Direction, Study};
#[test]
fn csv_empty_study_produces_header_only() {
let study: Study<f64> = Study::new(Direction::Minimize);
let mut buf = Vec::new();
study.to_csv(&mut buf).unwrap();
let csv = String::from_utf8(buf).unwrap();
assert_eq!(csv, "trial_id,value,state\n");
}
#[test]
fn csv_includes_all_trial_data() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 10.0).name("x");
let y = IntParam::new(1, 5).name("y");
study
.optimize(3, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let yv = y.suggest(trial)?;
Ok::<_, optimizer::Error>(xv + yv as f64)
})
.unwrap();
let mut buf = Vec::new();
study.to_csv(&mut buf).unwrap();
let csv = String::from_utf8(buf).unwrap();
let lines: Vec<&str> = csv.lines().collect();
// Header + 3 data rows.
assert_eq!(lines.len(), 4);
// Header should contain our parameter names.
let header = lines[0];
assert!(header.starts_with("trial_id,value,state"));
assert!(header.contains("x"));
assert!(header.contains("y"));
// Each data row should have the right number of columns.
let n_cols = header.split(',').count();
for line in &lines[1..] {
assert_eq!(line.split(',').count(), n_cols);
}
// All rows should have "Complete" state.
for line in &lines[1..] {
assert!(line.contains("Complete"));
}
}
#[test]
fn csv_handles_pruned_trials() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 10.0).name("x");
// First trial: complete
let mut trial = study.create_trial();
let _ = x.suggest(&mut trial).unwrap();
study.complete_trial(trial, 1.0);
// Second trial: pruned
let mut trial = study.create_trial();
let _ = x.suggest(&mut trial).unwrap();
study.prune_trial(trial);
let mut buf = Vec::new();
study.to_csv(&mut buf).unwrap();
let csv = String::from_utf8(buf).unwrap();
let lines: Vec<&str> = csv.lines().collect();
assert_eq!(lines.len(), 3); // header + 2 data rows
// Pruned trial should have empty value.
let pruned_line = lines[2];
assert!(pruned_line.contains("Pruned"));
// The value field (second column) should be empty.
let cols: Vec<&str> = pruned_line.split(',').collect();
assert_eq!(cols[2], "Pruned");
assert_eq!(cols[1], ""); // empty value for pruned
}
#[test]
fn csv_handles_different_parameter_sets() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 10.0).name("x");
let y = FloatParam::new(0.0, 10.0).name("y");
// First trial: only x
let mut trial = study.create_trial();
let xv = x.suggest(&mut trial).unwrap();
study.complete_trial(trial, xv);
// Second trial: only y
let mut trial = study.create_trial();
let yv = y.suggest(&mut trial).unwrap();
study.complete_trial(trial, yv);
let mut buf = Vec::new();
study.to_csv(&mut buf).unwrap();
let csv = String::from_utf8(buf).unwrap();
let lines: Vec<&str> = csv.lines().collect();
assert_eq!(lines.len(), 3);
// Both x and y columns should exist.
let header = lines[0];
assert!(header.contains("x"));
assert!(header.contains("y"));
// Each row has the right column count (missing params are empty).
let n_cols = header.split(',').count();
for line in &lines[1..] {
assert_eq!(line.split(',').count(), n_cols);
}
}
#[test]
fn csv_output_is_parseable() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let lr = FloatParam::new(0.001, 0.1).name("learning_rate");
let layers = IntParam::new(1, 5).name("n_layers");
study
.optimize(5, |trial: &mut optimizer::Trial| {
let l = lr.suggest(trial)?;
let n = layers.suggest(trial)?;
Ok::<_, optimizer::Error>(l * n as f64)
})
.unwrap();
let mut buf = Vec::new();
study.to_csv(&mut buf).unwrap();
let csv = String::from_utf8(buf).unwrap();
// Parse each row: every value field should be a valid f64 for complete trials.
let lines: Vec<&str> = csv.lines().collect();
for line in &lines[1..] {
let cols: Vec<&str> = line.split(',').collect();
// trial_id should be a number
cols[0].parse::<u64>().unwrap();
// value should be parseable as f64
cols[1].parse::<f64>().unwrap();
// state should be a known value
assert!(["Complete", "Pruned", "Failed", "Running"].contains(&cols[2]));
}
}
#[test]
fn export_csv_writes_file() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 10.0).name("x");
study
.optimize(3, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(xv * xv)
})
.unwrap();
let dir = std::env::temp_dir().join("optimizer_export_test");
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("test_export.csv");
study.export_csv(&path).unwrap();
let contents = std::fs::read_to_string(&path).unwrap();
assert!(contents.starts_with("trial_id,value,state"));
assert!(contents.lines().count() == 4); // header + 3 rows
// Clean up.
let _ = std::fs::remove_dir_all(&dir);
}
#[cfg(feature = "serde")]
#[test]
fn export_json_writes_file() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 10.0).name("x");
study
.optimize(3, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(xv * xv)
})
.unwrap();
let dir = std::env::temp_dir().join("optimizer_json_export_test");
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("test_export.json");
study.export_json(&path).unwrap();
let contents = std::fs::read_to_string(&path).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&contents).unwrap();
let arr = parsed.as_array().unwrap();
assert_eq!(arr.len(), 3);
// Each entry should have the expected fields.
for entry in arr {
assert!(entry.get("id").is_some());
assert!(entry.get("value").is_some());
assert!(entry.get("state").is_some());
assert!(entry.get("params").is_some());
}
// Clean up.
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn csv_includes_user_attributes() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 10.0).name("x");
study
.optimize(2, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
trial.set_user_attr("training_time_secs", 45.2);
Ok::<_, optimizer::Error>(xv * xv)
})
.unwrap();
let mut buf = Vec::new();
study.to_csv(&mut buf).unwrap();
let csv = String::from_utf8(buf).unwrap();
let header = csv.lines().next().unwrap();
assert!(header.contains("training_time_secs"));
}
+86
View File
@@ -0,0 +1,86 @@
//! Integration tests for fANOVA parameter importance.
use optimizer::prelude::*;
#[test]
fn fanova_dominant_parameter() {
// f(x, y) = x^2 — x should dominate
let x = FloatParam::new(0.0, 10.0).name("x");
let y = FloatParam::new(0.0, 10.0).name("y");
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
study
.optimize(50, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let _yv = y.suggest(trial)?;
Ok::<_, Error>(xv * xv)
})
.unwrap();
let result = study.fanova().unwrap();
assert_eq!(result.main_effects[0].0, "x");
assert!(
result.main_effects[0].1 > 0.7,
"x importance = {}",
result.main_effects[0].1
);
}
#[test]
fn fanova_interaction() {
// f(x, y) = x * y — both matter and interact
let x = FloatParam::new(0.0, 10.0).name("x");
let y = FloatParam::new(0.0, 10.0).name("y");
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(7));
study
.optimize(100, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let yv = y.suggest(trial)?;
Ok::<_, Error>(xv * yv)
})
.unwrap();
let config = FanovaConfig {
n_trees: 128,
..FanovaConfig::default()
};
let result = study.fanova_with_config(&config).unwrap();
// Should detect interaction
assert!(
!result.interactions.is_empty(),
"should detect x*y interaction"
);
}
#[test]
fn fanova_consistent_with_correlation() {
// f(x, y) = 3*x + 0.5*y — x should rank higher in both methods
let x = FloatParam::new(0.0, 10.0).name("x");
let y = FloatParam::new(0.0, 10.0).name("y");
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(99));
study
.optimize(80, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let yv = y.suggest(trial)?;
Ok::<_, Error>(3.0 * xv + 0.5 * yv)
})
.unwrap();
let corr = study.param_importance();
let fanova = study.fanova().unwrap();
// Both methods should rank x above y
assert_eq!(corr[0].0, "x", "correlation should rank x first");
assert_eq!(fanova.main_effects[0].0, "x", "fanova should rank x first");
}
#[test]
fn fanova_too_few_trials() {
let study: Study<f64> = Study::new(Direction::Minimize);
let result = study.fanova();
assert!(result.is_err(), "should error with no trials");
}
+167
View File
@@ -0,0 +1,167 @@
use optimizer::parameter::{CategoricalParam, FloatParam, IntParam, Parameter};
use optimizer::sampler::random::RandomSampler;
use optimizer::{Direction, Study};
#[test]
fn known_perfect_correlation() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 100.0).name("x");
// Objective = x, so perfect correlation.
for _ in 0..30 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 1);
assert_eq!(importance[0].0, "x");
assert!(
(importance[0].1 - 1.0).abs() < 1e-10,
"single param should be 1.0"
);
}
#[test]
fn no_effect_parameter() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 100.0).name("x");
let noise = FloatParam::new(0.0, 100.0).name("noise");
// Objective depends only on x; noise is unused in objective.
for _ in 0..50 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
let _nv = noise.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 2);
// x should have much higher importance than noise.
let x_score = importance.iter().find(|(l, _)| l == "x").unwrap().1;
let noise_score = importance.iter().find(|(l, _)| l == "noise").unwrap().1;
assert!(
x_score > noise_score,
"x ({x_score}) should outrank noise ({noise_score})"
);
// x should dominate
assert!(x_score > 0.7, "x importance {x_score} should be dominant");
}
#[test]
fn multiple_parameters_varying_importance() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 10.0).name("x");
let y = FloatParam::new(0.0, 10.0).name("y");
// Objective = 10*x + 0.01*y → x should be far more important.
for _ in 0..50 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
let yv = y.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(10.0 * xv + 0.01 * yv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 2);
assert_eq!(importance[0].0, "x", "x should rank first");
assert!(importance[0].1 > importance[1].1);
}
#[test]
fn fewer_than_two_trials_returns_empty() {
let study: Study<f64> = Study::new(Direction::Minimize);
// 0 trials
assert!(study.param_importance().is_empty());
// 1 trial
let x = FloatParam::new(0.0, 1.0).name("x");
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
assert!(study.param_importance().is_empty());
}
#[test]
fn int_parameter() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let n = IntParam::new(1, 100).name("n");
for _ in 0..30 {
let mut trial = study.ask();
let nv = n.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(nv as f64));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 1);
assert_eq!(importance[0].0, "n");
assert!(importance[0].1 > 0.9);
}
#[test]
fn categorical_parameter() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let cat = CategoricalParam::new(vec!["a", "b", "c"]).name("cat");
let x = FloatParam::new(0.0, 100.0).name("x");
// Objective depends only on x; categorical is random noise.
for _ in 0..50 {
let mut trial = study.ask();
let _c = cat.suggest(&mut trial).unwrap();
let xv = x.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 2);
let x_score = importance.iter().find(|(l, _)| l == "x").unwrap().1;
assert!(x_score > 0.5, "x should dominate over categorical noise");
}
#[test]
fn normalization_sums_to_one() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
let x = FloatParam::new(0.0, 10.0).name("x");
let y = FloatParam::new(0.0, 10.0).name("y");
let z = FloatParam::new(0.0, 10.0).name("z");
for _ in 0..50 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
let yv = y.suggest(&mut trial).unwrap();
let zv = z.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv + 0.5 * yv + 0.1 * zv));
}
let importance = study.param_importance();
let sum: f64 = importance.iter().map(|(_, s)| *s).sum();
assert!(
(sum - 1.0).abs() < 1e-10,
"scores should sum to 1.0, got {sum}"
);
}
#[test]
fn label_when_unnamed_uses_debug() {
let study: Study<f64> = Study::with_sampler(Direction::Minimize, RandomSampler::with_seed(42));
// No .name() call → label defaults to Debug representation.
let x = FloatParam::new(0.0, 10.0);
for _ in 0..10 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
study.tell(trial, Ok::<_, &str>(xv));
}
let importance = study.param_importance();
assert_eq!(importance.len(), 1);
assert!(
importance[0].0.starts_with("FloatParam"),
"expected Debug label, got {:?}",
importance[0].0
);
}
-1188
View File
File diff suppressed because it is too large Load Diff
+540
View File
@@ -0,0 +1,540 @@
//! Integration tests for the journal storage backend.
use std::collections::HashMap;
use std::sync::Arc;
use std::io::Write;
use optimizer::parameter::{FloatParam, Parameter};
use optimizer::sampler::CompletedTrial;
use optimizer::sampler::random::RandomSampler;
use optimizer::storage::{JournalStorage, Storage};
use optimizer::{Direction, Study};
fn temp_path() -> std::path::PathBuf {
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(0);
let mut path = std::env::temp_dir();
path.push(format!(
"optimizer_journal_test_{}_{}.jsonl",
std::process::id(),
COUNTER.fetch_add(1, Ordering::Relaxed)
));
path
}
fn sample_trial(id: u64, value: f64) -> CompletedTrial<f64> {
CompletedTrial::new(id, HashMap::new(), HashMap::new(), HashMap::new(), value)
}
#[test]
fn roundtrip_single_trial() {
let path = temp_path();
let storage = JournalStorage::new(&path);
storage.push(sample_trial(0, 42.0));
let loaded = storage.trials_arc().read().clone();
assert_eq!(loaded.len(), 1);
assert_eq!(loaded[0].id, 0);
assert_eq!(loaded[0].value, 42.0);
// Also verify via a fresh open from disk
let storage2 = JournalStorage::<f64>::open(&path).unwrap();
let loaded2 = storage2.trials_arc().read().clone();
assert_eq!(loaded2.len(), 1);
assert_eq!(loaded2[0].value, 42.0);
std::fs::remove_file(&path).ok();
}
#[test]
fn append_multiple_trials() {
let path = temp_path();
let storage = JournalStorage::new(&path);
for i in 0..5 {
storage.push(sample_trial(i, i as f64));
}
// Reload from disk
let storage2 = JournalStorage::<f64>::open(&path).unwrap();
let loaded = storage2.trials_arc().read().clone();
assert_eq!(loaded.len(), 5);
for (i, trial) in loaded.iter().enumerate() {
assert_eq!(trial.id, i as u64);
assert_eq!(trial.value, i as f64);
}
std::fs::remove_file(&path).ok();
}
#[test]
fn missing_file_returns_empty() {
let path = temp_path();
let storage = JournalStorage::<f64>::open(&path).unwrap();
let loaded = storage.trials_arc().read().clone();
assert!(loaded.is_empty());
}
#[test]
fn concurrent_writes() {
let path = temp_path();
let storage = Arc::new(JournalStorage::new(&path));
let mut handles = Vec::new();
for thread_id in 0..4u64 {
let s = Arc::clone(&storage);
handles.push(std::thread::spawn(move || {
for i in 0..25u64 {
let id = thread_id * 25 + i;
s.push(sample_trial(id, id as f64));
}
}));
}
for h in handles {
h.join().unwrap();
}
// Reload from disk to verify persistence
let storage2 = JournalStorage::<f64>::open(&path).unwrap();
let loaded = storage2.trials_arc().read().clone();
assert_eq!(loaded.len(), 100);
// Verify all IDs are present (order may vary)
let mut ids: Vec<u64> = loaded.iter().map(|t| t.id).collect();
ids.sort();
assert_eq!(ids, (0..100).collect::<Vec<_>>());
std::fs::remove_file(&path).ok();
}
#[test]
fn study_with_journal_integration() {
let path = temp_path();
let x = FloatParam::new(-10.0, 10.0);
// First "process": run some trials
{
let study =
Study::with_journal(Direction::Minimize, RandomSampler::with_seed(1), &path).unwrap();
study
.optimize(5, |trial: &mut optimizer::Trial| {
let val = x.suggest(trial)?;
Ok::<_, optimizer::Error>(val * val)
})
.unwrap();
assert_eq!(study.n_trials(), 5);
}
// Second "process": loads the same file, sees existing trials
let study2 =
Study::with_journal(Direction::Minimize, RandomSampler::with_seed(2), &path).unwrap();
assert_eq!(study2.n_trials(), 5);
// Continue optimizing
study2
.optimize(5, |trial: &mut optimizer::Trial| {
let val = x.suggest(trial)?;
Ok::<_, optimizer::Error>(val * val)
})
.unwrap();
assert_eq!(study2.n_trials(), 10);
// Verify all 10 written to disk
let storage3 = JournalStorage::<f64>::open(&path).unwrap();
let loaded = storage3.trials_arc().read().clone();
assert_eq!(loaded.len(), 10);
std::fs::remove_file(&path).ok();
}
#[test]
fn ids_are_unique_after_reload() {
let path = temp_path();
// First batch
{
let study =
Study::with_journal(Direction::Minimize, RandomSampler::with_seed(1), &path).unwrap();
study
.optimize(3, |trial: &mut optimizer::Trial| {
let _ = FloatParam::new(0.0, 1.0).suggest(trial)?;
Ok::<_, optimizer::Error>(1.0)
})
.unwrap();
}
// Second batch — IDs should continue from 3
let study =
Study::with_journal(Direction::Minimize, RandomSampler::with_seed(2), &path).unwrap();
study
.optimize(3, |trial: &mut optimizer::Trial| {
let _ = FloatParam::new(0.0, 1.0).suggest(trial)?;
Ok::<_, optimizer::Error>(1.0)
})
.unwrap();
let all = study.trials();
let mut ids: Vec<u64> = all.iter().map(|t| t.id).collect();
ids.sort();
// All 6 IDs should be unique
ids.dedup();
assert_eq!(ids.len(), 6);
std::fs::remove_file(&path).ok();
}
#[test]
fn pruned_trials_are_stored() {
let path = temp_path();
let study =
Study::with_journal(Direction::Minimize, RandomSampler::with_seed(1), &path).unwrap();
// Complete one, prune one
let x = FloatParam::new(0.0, 1.0);
study
.optimize(3, |trial: &mut optimizer::Trial| {
let _ = x.suggest(trial)?;
if trial.id() == 1 {
Err(optimizer::TrialPruned)?;
}
Ok::<_, optimizer::Error>(1.0)
})
.unwrap();
let storage2 = JournalStorage::<f64>::open(&path).unwrap();
let loaded = storage2.trials_arc().read().clone();
assert_eq!(loaded.len(), 3);
assert!(
loaded
.iter()
.any(|t| t.state == optimizer::TrialState::Pruned)
);
std::fs::remove_file(&path).ok();
}
#[test]
fn rejects_non_finite_values_in_journal() {
// serde_json rejects 1e999 ("number out of range"), so non-finite
// floats cannot sneak in through standard JSON. Verify the overall
// loading path catches the error regardless of which layer rejects it.
let path = temp_path();
std::fs::write(
&path,
r#"{"id":0,"params":{},"distributions":{"0":{"Float":{"low":0.0,"high":1e999,"log_scale":false,"step":null}}},"param_labels":{},"value":1.0,"intermediate_values":[],"state":"Complete","user_attrs":{},"constraints":[]}"#,
)
.unwrap();
assert!(JournalStorage::<f64>::open(&path).is_err());
std::fs::remove_file(&path).ok();
}
#[test]
fn validate_rejects_non_finite_distribution_bound() {
use optimizer::distribution::{Distribution, FloatDistribution};
let pid = FloatParam::new(0.0, 1.0).id();
let mut trial = sample_trial(0, 1.0);
trial.distributions.insert(
pid,
Distribution::Float(FloatDistribution {
low: 0.0,
high: f64::INFINITY,
log_scale: false,
step: None,
}),
);
let err = trial.validate().unwrap_err();
assert!(err.contains("non-finite"), "unexpected: {err}");
}
#[test]
fn validate_rejects_nan_constraint() {
let mut trial = sample_trial(0, 1.0);
trial.constraints.push(f64::NAN);
let err = trial.validate().unwrap_err();
assert!(err.contains("non-finite"), "unexpected: {err}");
}
#[test]
fn validate_rejects_non_finite_param_value() {
use optimizer::param::ParamValue;
let pid = FloatParam::new(0.0, 1.0).id();
let mut trial = sample_trial(0, 1.0);
trial
.params
.insert(pid, ParamValue::Float(f64::NEG_INFINITY));
let err = trial.validate().unwrap_err();
assert!(err.contains("non-finite"), "unexpected: {err}");
}
#[test]
fn validate_rejects_nan_intermediate_value() {
let mut trial = sample_trial(0, 1.0);
trial.intermediate_values.push((0, f64::NAN));
let err = trial.validate().unwrap_err();
assert!(err.contains("non-finite"), "unexpected: {err}");
}
#[test]
fn validate_accepts_valid_trial() {
use optimizer::distribution::{Distribution, FloatDistribution};
use optimizer::param::ParamValue;
let pid = FloatParam::new(0.0, 1.0).id();
let mut trial = sample_trial(0, 1.0);
trial.params.insert(pid, ParamValue::Float(0.5));
trial.distributions.insert(
pid,
Distribution::Float(FloatDistribution {
low: 0.0,
high: 1.0,
log_scale: false,
step: None,
}),
);
trial.constraints.push(-1.0);
trial.intermediate_values.push((0, 0.5));
assert!(trial.validate().is_ok());
}
#[test]
fn accepts_valid_journal_with_distributions() {
let path = temp_path();
std::fs::write(
&path,
r#"{"id":0,"params":{"0":{"Float":0.5}},"distributions":{"0":{"Float":{"low":0.0,"high":1.0,"log_scale":false,"step":null}}},"param_labels":{},"value":0.25,"intermediate_values":[],"state":"Complete","user_attrs":{},"constraints":[-1.0]}"#,
)
.unwrap();
let storage = JournalStorage::<f64>::open(&path).unwrap();
let loaded = storage.trials_arc().read().clone();
assert_eq!(loaded.len(), 1);
assert_eq!(loaded[0].value, 0.25);
std::fs::remove_file(&path).ok();
}
#[test]
fn refresh_skips_own_writes() {
let path = temp_path();
let storage = JournalStorage::new(&path);
for i in 0..5 {
storage.push(sample_trial(i, i as f64));
// Our own push advanced the offset, so refresh should find nothing new.
assert!(!storage.refresh(), "refresh returned true after push {i}");
}
assert_eq!(storage.trials_arc().read().len(), 5);
std::fs::remove_file(&path).ok();
}
#[test]
fn refresh_picks_up_external_writes() {
let path = temp_path();
let storage = JournalStorage::new(&path);
// Push 3 trials through the storage (advances offset).
for i in 0..3 {
storage.push(sample_trial(i, i as f64));
}
assert_eq!(storage.trials_arc().read().len(), 3);
// Simulate an external process appending 2 more lines directly.
{
let mut file = std::fs::OpenOptions::new()
.append(true)
.open(&path)
.unwrap();
for i in 3..5u64 {
let trial = sample_trial(i, i as f64);
let line = serde_json::to_string(&trial).unwrap();
writeln!(file, "{line}").unwrap();
}
file.sync_all().unwrap();
}
// refresh() should pick up the 2 external trials.
assert!(storage.refresh(), "refresh should detect external writes");
assert_eq!(storage.trials_arc().read().len(), 5);
// A second refresh should be a no-op.
assert!(
!storage.refresh(),
"second refresh should return false (no new data)"
);
assert_eq!(storage.trials_arc().read().len(), 5);
std::fs::remove_file(&path).ok();
}
// ── Corrupted / malicious journal file tests ────────────────────────
fn valid_trial_line_with_id(id: u64) -> String {
format!(
r#"{{"id":{id},"params":{{}},"distributions":{{}},"param_labels":{{}},"value":1.0,"intermediate_values":[],"state":"Complete","user_attrs":{{}},"constraints":[]}}"#
)
}
#[test]
fn empty_file_loads_as_empty_storage() {
let path = temp_path();
std::fs::write(&path, "").unwrap();
let storage = JournalStorage::<f64>::open(&path).unwrap();
assert_eq!(storage.trials_arc().read().len(), 0);
std::fs::remove_file(&path).ok();
}
#[test]
fn whitespace_only_lines_are_skipped() {
let path = temp_path();
std::fs::write(&path, " \n\t\n\n").unwrap();
let storage = JournalStorage::<f64>::open(&path).unwrap();
assert_eq!(storage.trials_arc().read().len(), 0);
std::fs::remove_file(&path).ok();
}
#[test]
fn truncated_json_line_returns_error() {
let path = temp_path();
std::fs::write(&path, r#"{"id":0,"params":{"#).unwrap();
assert!(JournalStorage::<f64>::open(&path).is_err());
std::fs::remove_file(&path).ok();
}
#[test]
fn invalid_json_syntax_returns_error() {
let path = temp_path();
std::fs::write(&path, "not valid json\n").unwrap();
assert!(JournalStorage::<f64>::open(&path).is_err());
std::fs::remove_file(&path).ok();
}
#[test]
fn missing_required_field_returns_error() {
let path = temp_path();
// Missing params, distributions, param_labels, etc.
std::fs::write(&path, r#"{"id":0,"value":1.0}"#).unwrap();
assert!(JournalStorage::<f64>::open(&path).is_err());
std::fs::remove_file(&path).ok();
}
#[test]
fn extra_fields_are_ignored() {
let path = temp_path();
let line = r#"{"id":0,"params":{},"distributions":{},"param_labels":{},"value":0.5,"intermediate_values":[],"state":"Complete","user_attrs":{},"constraints":[],"foo":"bar","extra_number":42}"#;
std::fs::write(&path, format!("{line}\n")).unwrap();
let storage = JournalStorage::<f64>::open(&path).unwrap();
let loaded = storage.trials_arc().read().clone();
assert_eq!(loaded.len(), 1);
assert_eq!(loaded[0].id, 0);
assert_eq!(loaded[0].value, 0.5);
std::fs::remove_file(&path).ok();
}
#[test]
fn out_of_bounds_categorical_index_loads() {
let path = temp_path();
// Categorical param with index 999, but distribution only has 3 choices.
// validate() does not check categorical bounds, so this should load.
let line = r#"{"id":0,"params":{"0":{"Categorical":999}},"distributions":{"0":{"Categorical":{"n_choices":3}}},"param_labels":{},"value":1.0,"intermediate_values":[],"state":"Complete","user_attrs":{},"constraints":[]}"#;
std::fs::write(&path, format!("{line}\n")).unwrap();
let storage = JournalStorage::<f64>::open(&path).unwrap();
let loaded = storage.trials_arc().read().clone();
assert_eq!(loaded.len(), 1);
std::fs::remove_file(&path).ok();
}
#[test]
fn valid_lines_before_corruption_are_not_loaded() {
let path = temp_path();
let content = format!(
"{}\n{}\n{}\n",
valid_trial_line_with_id(0),
valid_trial_line_with_id(1),
"CORRUPTED LINE"
);
std::fs::write(&path, content).unwrap();
// load_trials_from_file is all-or-nothing: the corrupted third line
// makes the entire open() fail.
assert!(JournalStorage::<f64>::open(&path).is_err());
std::fs::remove_file(&path).ok();
}
#[test]
fn refresh_rejects_corrupted_external_append() {
let path = temp_path();
let storage = JournalStorage::new(&path);
// Push 2 valid trials through the storage API.
storage.push(sample_trial(0, 1.0));
storage.push(sample_trial(1, 2.0));
assert_eq!(storage.trials_arc().read().len(), 2);
// Simulate an external process appending corrupted JSON.
{
let mut file = std::fs::OpenOptions::new()
.append(true)
.open(&path)
.unwrap();
writeln!(file, "CORRUPTED LINE").unwrap();
file.sync_all().unwrap();
}
// refresh() should reject the corrupted data and return false.
assert!(!storage.refresh());
// Memory still has only the original 2 trials.
assert_eq!(storage.trials_arc().read().len(), 2);
std::fs::remove_file(&path).ok();
}
#[test]
fn refresh_rejects_truncated_external_append() {
let path = temp_path();
let storage = JournalStorage::new(&path);
// Push 2 valid trials through the storage API.
storage.push(sample_trial(0, 1.0));
storage.push(sample_trial(1, 2.0));
assert_eq!(storage.trials_arc().read().len(), 2);
// Simulate an external process appending truncated JSON.
{
let mut file = std::fs::OpenOptions::new()
.append(true)
.open(&path)
.unwrap();
writeln!(file, r#"{{"id":2,"params":{{"#).unwrap();
file.sync_all().unwrap();
}
// refresh() should reject the truncated data and return false.
assert!(!storage.refresh());
// Memory still has only the original 2 trials.
assert_eq!(storage.trials_arc().read().len(), 2);
std::fs::remove_file(&path).ok();
}
+739
View File
@@ -0,0 +1,739 @@
//! Integration tests for multi-objective optimization.
use optimizer::Direction;
use optimizer::multi_objective::MultiObjectiveStudy;
use optimizer::parameter::{CategoricalParam, FloatParam, Parameter};
use optimizer::sampler::Decomposition;
use optimizer::sampler::moead::MoeadSampler;
use optimizer::sampler::nsga2::Nsga2Sampler;
use optimizer::sampler::nsga3::Nsga3Sampler;
// ---------------------------------------------------------------------------
// Pareto utility tests (via public MultiObjectiveStudy)
// ---------------------------------------------------------------------------
#[test]
fn test_basic_two_objective_random() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty(), "Pareto front should be non-empty");
// Verify no solution in the front dominates another
for a in &front {
for b in &front {
if core::ptr::eq(a, b) {
continue;
}
let a_dom_b = a.values[0] <= b.values[0]
&& a.values[1] <= b.values[1]
&& (a.values[0] < b.values[0] || a.values[1] < b.values[1]);
assert!(
!a_dom_b,
"Front solution {:?} dominates {:?}",
a.values, b.values
);
}
}
}
#[test]
fn test_dimension_mismatch_error() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let x = FloatParam::new(0.0, 1.0);
let result = study.optimize(1, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
// Return wrong number of values
Ok::<_, optimizer::Error>(vec![xv])
});
assert!(result.is_err());
let err = result.unwrap_err();
assert!(
matches!(
err,
optimizer::Error::ObjectiveDimensionMismatch {
expected: 2,
got: 1
}
),
"Expected ObjectiveDimensionMismatch, got: {err}"
);
}
#[test]
fn test_ask_tell() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Maximize]);
let x = FloatParam::new(0.0, 10.0);
for _ in 0..10 {
let mut trial = study.ask();
let xv = x.suggest(&mut trial).unwrap();
study
.tell(trial, Ok::<_, &str>(vec![xv, 10.0 - xv]))
.unwrap();
}
assert_eq!(study.n_trials(), 10);
let front = study.pareto_front();
assert!(!front.is_empty());
}
#[test]
fn test_ask_tell_dimension_mismatch() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let trial = study.ask();
let result = study.tell(trial, Ok::<_, &str>(vec![1.0, 2.0, 3.0]));
assert!(result.is_err());
}
#[test]
fn test_n_trials_counting() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
assert_eq!(study.n_trials(), 0);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(5, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
assert_eq!(study.n_trials(), 5);
}
#[test]
fn test_three_objectives() {
let study = MultiObjectiveStudy::new(vec![
Direction::Minimize,
Direction::Minimize,
Direction::Maximize,
]);
let x = FloatParam::new(0.0, 1.0);
let y = FloatParam::new(0.0, 1.0);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let yv = y.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, yv, 1.0 - xv - yv])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty());
assert_eq!(study.n_objectives(), 3);
}
#[test]
fn test_directions_accessor() {
let dirs = vec![Direction::Minimize, Direction::Maximize];
let study = MultiObjectiveStudy::new(dirs.clone());
assert_eq!(study.directions(), &dirs);
assert_eq!(study.n_objectives(), 2);
}
#[test]
fn test_trials_accessor() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(3, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
let trials = study.trials();
assert_eq!(trials.len(), 3);
for t in &trials {
assert_eq!(t.values.len(), 2);
}
}
// ---------------------------------------------------------------------------
// NSGA-II sampler tests
// ---------------------------------------------------------------------------
#[test]
fn test_nsga2_zdt1() {
// ZDT1 benchmark: minimize both objectives
let n_vars = 5;
let params: Vec<FloatParam> = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect();
let sampler = Nsga2Sampler::builder().population_size(20).seed(42).build();
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
study
.optimize(200, |trial: &mut optimizer::Trial| {
let xs: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()?;
let f1 = xs[0];
let g = 1.0 + 9.0 * xs[1..].iter().sum::<f64>() / (n_vars - 1) as f64;
let f2 = g * (1.0 - (f1 / g).sqrt());
Ok::<_, optimizer::Error>(vec![f1, f2])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty(), "Pareto front should be non-empty");
// Verify no dominated solutions in the front
for a in &front {
for b in &front {
if core::ptr::eq(a, b) {
continue;
}
let a_dom_b = a.values[0] <= b.values[0]
&& a.values[1] <= b.values[1]
&& (a.values[0] < b.values[0] || a.values[1] < b.values[1]);
assert!(
!a_dom_b,
"Front solution {:?} dominates {:?}",
a.values, b.values
);
}
}
}
#[test]
fn test_nsga2_with_seed_reproducible() {
let x = FloatParam::new(0.0, 1.0);
let y = FloatParam::new(0.0, 1.0);
let run = |seed: u64| -> Vec<Vec<f64>> {
let sampler = Nsga2Sampler::with_seed(seed);
let study = MultiObjectiveStudy::with_sampler(
vec![Direction::Minimize, Direction::Minimize],
sampler,
);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let yv = y.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, yv])
})
.unwrap();
study.trials().iter().map(|t| t.values.clone()).collect()
};
let r1 = run(123);
let r2 = run(123);
assert_eq!(r1, r2, "Same seed should produce same results");
let r3 = run(456);
assert_ne!(r1, r3, "Different seeds should produce different results");
}
#[test]
fn test_nsga2_builder() {
let sampler = Nsga2Sampler::builder()
.population_size(10)
.crossover_prob(0.8)
.crossover_eta(15.0)
.mutation_eta(25.0)
.seed(42)
.build();
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
assert_eq!(study.n_trials(), 30);
}
#[test]
fn test_nsga2_categorical_params() {
let sampler = Nsga2Sampler::with_seed(42);
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
let x = FloatParam::new(0.0, 1.0);
let cat = CategoricalParam::new(vec!["a", "b", "c"]);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let cv = cat.suggest(trial)?;
let bonus = match cv {
"a" => 0.0,
"b" => 0.5,
_ => 1.0,
};
Ok::<_, optimizer::Error>(vec![xv + bonus, 1.0 - xv])
})
.unwrap();
assert_eq!(study.n_trials(), 30);
let front = study.pareto_front();
assert!(!front.is_empty());
}
#[test]
fn test_nsga2_constraints() {
let sampler = Nsga2Sampler::with_seed(42);
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(50, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
// Constraint: x >= 0.3 (i.e. 0.3 - x <= 0)
trial.set_constraints(vec![0.3 - xv]);
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty());
// Check that feasible solutions exist on the front
let feasible_count = front.iter().filter(|t| t.is_feasible()).count();
assert!(
feasible_count > 0,
"Should have feasible solutions on front"
);
}
#[test]
fn test_multi_objective_trial_get() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let x = FloatParam::new(0.0, 10.0).name("x");
study
.optimize(5, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, 10.0 - xv])
})
.unwrap();
let front = study.pareto_front();
for t in &front {
let xv: f64 = t.get(&x).unwrap();
assert!((0.0..=10.0).contains(&xv));
}
}
#[test]
fn test_multi_objective_trial_is_feasible() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(10, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
trial.set_constraints(vec![0.5 - xv]); // feasible if x >= 0.5
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
let trials = study.trials();
for t in &trials {
let xv = t.values[0];
if xv >= 0.5 {
assert!(t.is_feasible());
} else {
assert!(!t.is_feasible());
}
}
}
#[test]
fn test_multi_objective_trial_user_attrs() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(3, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
trial.set_user_attr("iteration", 42_i64);
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
let trials = study.trials();
for t in &trials {
assert!(t.user_attr("iteration").is_some());
}
}
#[test]
fn test_tell_with_failure() {
let study = MultiObjectiveStudy::new(vec![Direction::Minimize, Direction::Minimize]);
let trial = study.ask();
study
.tell(trial, Err::<Vec<f64>, _>("evaluation failed"))
.unwrap();
// Failed trial not counted
assert_eq!(study.n_trials(), 0);
}
// ---------------------------------------------------------------------------
// NSGA-III sampler tests
// ---------------------------------------------------------------------------
#[test]
fn test_nsga3_zdt1() {
let n_vars = 5;
let params: Vec<FloatParam> = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect();
let sampler = Nsga3Sampler::builder().population_size(20).seed(42).build();
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
study
.optimize(200, |trial: &mut optimizer::Trial| {
let xs: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()?;
let f1 = xs[0];
let g = 1.0 + 9.0 * xs[1..].iter().sum::<f64>() / (n_vars - 1) as f64;
let f2 = g * (1.0 - (f1 / g).sqrt());
Ok::<_, optimizer::Error>(vec![f1, f2])
})
.unwrap();
let front = study.pareto_front();
assert!(
!front.is_empty(),
"NSGA-III Pareto front should be non-empty"
);
// Verify no dominated solutions in the front
for a in &front {
for b in &front {
if core::ptr::eq(a, b) {
continue;
}
let a_dom_b = a.values[0] <= b.values[0]
&& a.values[1] <= b.values[1]
&& (a.values[0] < b.values[0] || a.values[1] < b.values[1]);
assert!(
!a_dom_b,
"Front solution {:?} dominates {:?}",
a.values, b.values
);
}
}
}
#[test]
fn test_nsga3_four_objectives() {
// DTLZ2 with 4 objectives
let n_obj = 4;
let n_vars = n_obj + 4; // k = 5 decision variables beyond the first (n_obj-1)
let params: Vec<FloatParam> = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect();
let sampler = Nsga3Sampler::builder().population_size(50).seed(42).build();
let directions = vec![Direction::Minimize; n_obj];
let study = MultiObjectiveStudy::with_sampler(directions, sampler);
study
.optimize(500, |trial: &mut optimizer::Trial| {
let xs: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()?;
// DTLZ2 formulation
let g: f64 = xs[n_obj - 1..]
.iter()
.map(|&xi| (xi - 0.5).powi(2))
.sum::<f64>();
let mut objectives = vec![0.0_f64; n_obj];
for i in 0..n_obj {
let mut f = 1.0 + g;
for xj in &xs[..(n_obj - 1 - i)] {
f *= (xj * core::f64::consts::FRAC_PI_2).cos();
}
if i > 0 {
f *= (xs[n_obj - 1 - i] * core::f64::consts::FRAC_PI_2).sin();
}
objectives[i] = f;
}
Ok::<_, optimizer::Error>(objectives)
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty(), "4-objective front should be non-empty");
// All front solutions should have 4 objectives
for t in &front {
assert_eq!(t.values.len(), 4);
}
}
#[test]
fn test_nsga3_reproducible() {
let x = FloatParam::new(0.0, 1.0);
let y = FloatParam::new(0.0, 1.0);
let run = |seed: u64| -> Vec<Vec<f64>> {
let sampler = Nsga3Sampler::with_seed(seed);
let study = MultiObjectiveStudy::with_sampler(
vec![Direction::Minimize, Direction::Minimize],
sampler,
);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let yv = y.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, yv])
})
.unwrap();
study.trials().iter().map(|t| t.values.clone()).collect()
};
let r1 = run(123);
let r2 = run(123);
assert_eq!(r1, r2, "Same seed should produce same results");
let r3 = run(456);
assert_ne!(r1, r3, "Different seeds should produce different results");
}
#[test]
fn test_nsga3_builder() {
let sampler = Nsga3Sampler::builder()
.population_size(12)
.n_divisions(4)
.crossover_prob(0.9)
.crossover_eta(20.0)
.mutation_eta(20.0)
.seed(42)
.build();
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
assert_eq!(study.n_trials(), 30);
}
#[test]
fn test_nsga3_constraints() {
let sampler = Nsga3Sampler::with_seed(42);
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(50, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
trial.set_constraints(vec![0.3 - xv]);
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty());
let feasible_count = front.iter().filter(|t| t.is_feasible()).count();
assert!(
feasible_count > 0,
"Should have feasible solutions on front"
);
}
// ---------------------------------------------------------------------------
// MOEA/D sampler tests
// ---------------------------------------------------------------------------
#[test]
fn test_moead_zdt1_tchebycheff() {
let n_vars = 5;
let params: Vec<FloatParam> = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect();
let sampler = MoeadSampler::builder().population_size(20).seed(42).build();
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
study
.optimize(200, |trial: &mut optimizer::Trial| {
let xs: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()?;
let f1 = xs[0];
let g = 1.0 + 9.0 * xs[1..].iter().sum::<f64>() / (n_vars - 1) as f64;
let f2 = g * (1.0 - (f1 / g).sqrt());
Ok::<_, optimizer::Error>(vec![f1, f2])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty(), "MOEA/D Pareto front should be non-empty");
for a in &front {
for b in &front {
if core::ptr::eq(a, b) {
continue;
}
let a_dom_b = a.values[0] <= b.values[0]
&& a.values[1] <= b.values[1]
&& (a.values[0] < b.values[0] || a.values[1] < b.values[1]);
assert!(
!a_dom_b,
"Front solution {:?} dominates {:?}",
a.values, b.values
);
}
}
}
#[test]
fn test_moead_zdt1_weighted_sum() {
let n_vars = 3;
let params: Vec<FloatParam> = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect();
let sampler = MoeadSampler::builder()
.population_size(20)
.decomposition(Decomposition::WeightedSum)
.seed(42)
.build();
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
study
.optimize(200, |trial: &mut optimizer::Trial| {
let xs: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()?;
let f1 = xs[0];
let g = 1.0 + 9.0 * xs[1..].iter().sum::<f64>() / (n_vars - 1) as f64;
let f2 = g * (1.0 - (f1 / g).sqrt());
Ok::<_, optimizer::Error>(vec![f1, f2])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty());
}
#[test]
fn test_moead_zdt1_pbi() {
let n_vars = 3;
let params: Vec<FloatParam> = (0..n_vars).map(|_| FloatParam::new(0.0, 1.0)).collect();
let sampler = MoeadSampler::builder()
.population_size(20)
.decomposition(Decomposition::Pbi { theta: 5.0 })
.seed(42)
.build();
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
study
.optimize(200, |trial: &mut optimizer::Trial| {
let xs: Vec<f64> = params
.iter()
.map(|p| p.suggest(trial))
.collect::<Result<_, _>>()?;
let f1 = xs[0];
let g = 1.0 + 9.0 * xs[1..].iter().sum::<f64>() / (n_vars - 1) as f64;
let f2 = g * (1.0 - (f1 / g).sqrt());
Ok::<_, optimizer::Error>(vec![f1, f2])
})
.unwrap();
let front = study.pareto_front();
assert!(!front.is_empty());
}
#[test]
fn test_moead_reproducible() {
let x = FloatParam::new(0.0, 1.0);
let y = FloatParam::new(0.0, 1.0);
let run = |seed: u64| -> Vec<Vec<f64>> {
let sampler = MoeadSampler::with_seed(seed);
let study = MultiObjectiveStudy::with_sampler(
vec![Direction::Minimize, Direction::Minimize],
sampler,
);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
let yv = y.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, yv])
})
.unwrap();
study.trials().iter().map(|t| t.values.clone()).collect()
};
let r1 = run(123);
let r2 = run(123);
assert_eq!(r1, r2, "Same seed should produce same results");
let r3 = run(456);
assert_ne!(r1, r3, "Different seeds should produce different results");
}
#[test]
fn test_moead_builder() {
let sampler = MoeadSampler::builder()
.population_size(15)
.neighborhood_size(5)
.decomposition(Decomposition::Tchebycheff)
.crossover_prob(0.9)
.crossover_eta(20.0)
.mutation_eta(20.0)
.seed(42)
.build();
let study =
MultiObjectiveStudy::with_sampler(vec![Direction::Minimize, Direction::Minimize], sampler);
let x = FloatParam::new(0.0, 1.0);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let xv = x.suggest(trial)?;
Ok::<_, optimizer::Error>(vec![xv, 1.0 - xv])
})
.unwrap();
assert_eq!(study.n_trials(), 30);
}
+240
View File
@@ -0,0 +1,240 @@
use optimizer::parameter::{
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter,
};
use optimizer::{Direction, Study, Trial};
#[test]
fn suggest_float_param_via_trial() {
let param = FloatParam::new(0.0, 1.0);
let mut trial = Trial::new(0);
let x = trial.suggest_param(&param).unwrap();
assert!((0.0..=1.0).contains(&x));
// Cached
let x2 = trial.suggest_param(&param).unwrap();
assert_eq!(x, x2);
}
#[test]
fn suggest_float_log_param_via_trial() {
let param = FloatParam::new(1e-5, 1e-1).log_scale();
let mut trial = Trial::new(0);
let lr = trial.suggest_param(&param).unwrap();
assert!((1e-5..=1e-1).contains(&lr));
}
#[test]
fn suggest_float_step_param_via_trial() {
let param = FloatParam::new(0.0, 1.0).step(0.25);
let mut trial = Trial::new(0);
let x = trial.suggest_param(&param).unwrap();
assert!((0.0..=1.0).contains(&x));
}
#[test]
fn suggest_int_param_via_trial() {
let param = IntParam::new(1, 10);
let mut trial = Trial::new(0);
let n = trial.suggest_param(&param).unwrap();
assert!((1..=10).contains(&n));
// Cached
let n2 = trial.suggest_param(&param).unwrap();
assert_eq!(n, n2);
}
#[test]
fn suggest_int_log_param_via_trial() {
let param = IntParam::new(1, 1024).log_scale();
let mut trial = Trial::new(0);
let batch = trial.suggest_param(&param).unwrap();
assert!((1..=1024).contains(&batch));
}
#[test]
fn suggest_int_step_param_via_trial() {
let param = IntParam::new(32, 512).step(32);
let mut trial = Trial::new(0);
let units = trial.suggest_param(&param).unwrap();
assert!((32..=512).contains(&units));
assert_eq!((units - 32) % 32, 0);
}
#[test]
fn suggest_categorical_param_via_trial() {
let choices = vec!["sgd", "adam", "rmsprop"];
let param = CategoricalParam::new(choices.clone());
let mut trial = Trial::new(0);
let opt = trial.suggest_param(&param).unwrap();
assert!(choices.contains(&opt));
// Cached
let opt2 = trial.suggest_param(&param).unwrap();
assert_eq!(opt, opt2);
}
#[test]
fn suggest_bool_param_via_trial() {
let param = BoolParam::new();
let mut trial = Trial::new(0);
let val = trial.suggest_param(&param).unwrap();
let _ = val;
// Cached
let val2 = trial.suggest_param(&param).unwrap();
assert_eq!(val, val2);
}
#[derive(Clone, Debug, PartialEq)]
enum Activation {
Relu,
Sigmoid,
Tanh,
}
impl Categorical for Activation {
const N_CHOICES: usize = 3;
fn from_index(index: usize) -> Self {
match index {
0 => Activation::Relu,
1 => Activation::Sigmoid,
2 => Activation::Tanh,
_ => panic!("invalid index"),
}
}
fn to_index(&self) -> usize {
match self {
Activation::Relu => 0,
Activation::Sigmoid => 1,
Activation::Tanh => 2,
}
}
}
#[test]
fn suggest_enum_param_via_trial() {
let param = EnumParam::<Activation>::new();
let mut trial = Trial::new(0);
let act = trial.suggest_param(&param).unwrap();
assert!([Activation::Relu, Activation::Sigmoid, Activation::Tanh].contains(&act));
// Cached
let act2 = trial.suggest_param(&param).unwrap();
assert_eq!(act, act2);
}
#[test]
fn parameter_conflict_detection() {
let float_param = FloatParam::new(0.0, 1.0);
let int_param = IntParam::new(0, 10);
let mut trial = Trial::new(0);
let _ = trial.suggest_param(&float_param).unwrap();
// Different param type with different id - no conflict
let result = trial.suggest_param(&int_param);
assert!(result.is_ok());
// Different bounds for same param type but different id - no conflict
let float_param2 = FloatParam::new(0.0, 2.0);
let result = trial.suggest_param(&float_param2);
assert!(result.is_ok());
}
#[test]
fn validation_prevents_suggest() {
let mut trial = Trial::new(0);
assert!(trial.suggest_param(&FloatParam::new(1.0, 0.0)).is_err());
assert!(
trial
.suggest_param(&FloatParam::new(-1.0, 1.0).log_scale())
.is_err()
);
assert!(
trial
.suggest_param(&FloatParam::new(0.0, 1.0).step(-0.1))
.is_err()
);
assert!(trial.suggest_param(&IntParam::new(10, 1)).is_err());
assert!(
trial
.suggest_param(&IntParam::new(0, 10).log_scale())
.is_err()
);
assert!(trial.suggest_param(&IntParam::new(0, 10).step(-1)).is_err());
assert!(
trial
.suggest_param(&CategoricalParam::<&str>::new(vec![]))
.is_err()
);
}
#[test]
fn parameter_api_with_study() {
let x_param = FloatParam::new(-5.0, 5.0);
let n_param = IntParam::new(1, 10);
let dropout_param = BoolParam::new();
let opt_param = CategoricalParam::new(vec!["sgd", "adam"]);
let study: Study<f64> = Study::new(Direction::Minimize);
study
.optimize(5, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
let n = n_param.suggest(trial)?;
let dropout = dropout_param.suggest(trial)?;
let opt = opt_param.suggest(trial)?;
let _ = (n, dropout, opt);
Ok::<_, optimizer::Error>(x * x)
})
.unwrap();
let best = study.best_trial().unwrap();
assert!(best.value >= 0.0);
}
#[test]
fn parameter_suggest_method() {
let param = FloatParam::new(0.0, 1.0);
let mut trial = Trial::new(0);
let x = param.suggest(&mut trial).unwrap();
assert!((0.0..=1.0).contains(&x));
}
#[test]
fn existing_suggest_methods_still_work() {
let mut trial = Trial::new(0);
let x_param = FloatParam::new(0.0, 1.0);
let x = x_param.suggest(&mut trial).unwrap();
assert!((0.0..=1.0).contains(&x));
let lr_param = FloatParam::new(1e-5, 1e-1).log_scale();
let lr = lr_param.suggest(&mut trial).unwrap();
assert!((1e-5..=1e-1).contains(&lr));
let step_param = FloatParam::new(0.0, 1.0).step(0.25);
let step = step_param.suggest(&mut trial).unwrap();
assert!((0.0..=1.0).contains(&step));
let n_param = IntParam::new(1, 10);
let n = n_param.suggest(&mut trial).unwrap();
assert!((1..=10).contains(&n));
let batch_param = IntParam::new(1, 1024).log_scale();
let batch = batch_param.suggest(&mut trial).unwrap();
assert!((1..=1024).contains(&batch));
let units_param = IntParam::new(32, 512).step(32);
let units = units_param.suggest(&mut trial).unwrap();
assert!((32..=512).contains(&units));
let opt_param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
let opt = opt_param.suggest(&mut trial).unwrap();
assert!(["sgd", "adam", "rmsprop"].contains(&opt));
let flag_param = BoolParam::new();
let flag = flag_param.suggest(&mut trial).unwrap();
let _ = flag;
}
+2
View File
@@ -0,0 +1,2 @@
mod median;
mod threshold;
+228
View File
@@ -0,0 +1,228 @@
use std::collections::HashMap;
use optimizer::Direction;
use optimizer::pruner::{MedianPruner, Pruner};
use optimizer::sampler::CompletedTrial;
/// Helper to build a completed trial with given intermediate values.
fn trial_with_values(id: u64, intermediate_values: Vec<(u64, f64)>) -> CompletedTrial {
CompletedTrial::with_intermediate_values(
id,
HashMap::new(),
HashMap::new(),
HashMap::new(),
0.0,
intermediate_values,
HashMap::new(),
)
}
// --- Minimize direction ---
#[test]
fn prune_when_worse_than_median_minimize() {
let pruner = MedianPruner::new(Direction::Minimize);
// 3 completed trials with values at step 2: [1.0, 2.0, 3.0] => median = 2.0
let completed = vec![
trial_with_values(0, vec![(0, 0.5), (1, 0.8), (2, 1.0)]),
trial_with_values(1, vec![(0, 0.6), (1, 1.5), (2, 2.0)]),
trial_with_values(2, vec![(0, 0.7), (1, 2.0), (2, 3.0)]),
];
// Current trial value at step 2 is 2.5 > median 2.0 => prune
let current = vec![(0, 0.5), (1, 1.0), (2, 2.5)];
assert!(pruner.should_prune(3, 2, &current, &completed));
}
#[test]
fn no_prune_when_better_than_median_minimize() {
let pruner = MedianPruner::new(Direction::Minimize);
let completed = vec![
trial_with_values(0, vec![(0, 0.5), (1, 0.8), (2, 1.0)]),
trial_with_values(1, vec![(0, 0.6), (1, 1.5), (2, 2.0)]),
trial_with_values(2, vec![(0, 0.7), (1, 2.0), (2, 3.0)]),
];
// Current trial value at step 2 is 1.5 < median 2.0 => don't prune
let current = vec![(0, 0.5), (1, 1.0), (2, 1.5)];
assert!(!pruner.should_prune(3, 2, &current, &completed));
}
// --- Maximize direction ---
#[test]
fn prune_when_worse_than_median_maximize() {
let pruner = MedianPruner::new(Direction::Maximize);
// Values at step 1: [5.0, 7.0, 9.0] => median = 7.0
let completed = vec![
trial_with_values(0, vec![(0, 3.0), (1, 5.0)]),
trial_with_values(1, vec![(0, 4.0), (1, 7.0)]),
trial_with_values(2, vec![(0, 5.0), (1, 9.0)]),
];
// Current value 6.0 < median 7.0 => prune (worse for maximize)
let current = vec![(0, 4.0), (1, 6.0)];
assert!(pruner.should_prune(3, 1, &current, &completed));
}
#[test]
fn no_prune_when_better_than_median_maximize() {
let pruner = MedianPruner::new(Direction::Maximize);
let completed = vec![
trial_with_values(0, vec![(0, 3.0), (1, 5.0)]),
trial_with_values(1, vec![(0, 4.0), (1, 7.0)]),
trial_with_values(2, vec![(0, 5.0), (1, 9.0)]),
];
// Current value 8.0 > median 7.0 => don't prune
let current = vec![(0, 4.0), (1, 8.0)];
assert!(!pruner.should_prune(3, 1, &current, &completed));
}
// --- Warmup steps ---
#[test]
fn no_prune_during_warmup() {
let pruner = MedianPruner::new(Direction::Minimize).n_warmup_steps(5);
let completed = vec![trial_with_values(0, vec![(0, 1.0), (1, 1.0), (2, 1.0)])];
// Step 2 < warmup 5 => never prune, even if value is terrible
let current = vec![(0, 100.0), (1, 100.0), (2, 100.0)];
assert!(!pruner.should_prune(1, 2, &current, &completed));
}
#[test]
fn prune_after_warmup() {
let pruner = MedianPruner::new(Direction::Minimize).n_warmup_steps(2);
let completed = vec![trial_with_values(0, vec![(0, 1.0), (1, 1.0), (2, 1.0)])];
// Step 2 >= warmup 2 => pruning allowed; current 100.0 > median 1.0
let current = vec![(0, 100.0), (1, 100.0), (2, 100.0)];
assert!(pruner.should_prune(1, 2, &current, &completed));
}
// --- n_min_trials ---
#[test]
fn no_prune_when_fewer_than_n_min_trials() {
let pruner = MedianPruner::new(Direction::Minimize).n_min_trials(3);
// Only 2 completed trials — below threshold of 3
let completed = vec![
trial_with_values(0, vec![(0, 1.0)]),
trial_with_values(1, vec![(0, 2.0)]),
];
let current = vec![(0, 100.0)];
assert!(!pruner.should_prune(2, 0, &current, &completed));
}
#[test]
fn prune_when_at_least_n_min_trials() {
let pruner = MedianPruner::new(Direction::Minimize).n_min_trials(3);
// 3 completed trials with step 0: [1.0, 2.0, 3.0] => median 2.0
let completed = vec![
trial_with_values(0, vec![(0, 1.0)]),
trial_with_values(1, vec![(0, 2.0)]),
trial_with_values(2, vec![(0, 3.0)]),
];
// 5.0 > median 2.0 => prune
let current = vec![(0, 5.0)];
assert!(pruner.should_prune(3, 0, &current, &completed));
}
// --- No completed trials with values at step ---
#[test]
fn no_prune_when_no_completed_trials_at_step() {
let pruner = MedianPruner::new(Direction::Minimize);
// Completed trials only have values at step 0, not step 5
let completed = vec![
trial_with_values(0, vec![(0, 1.0)]),
trial_with_values(1, vec![(0, 2.0)]),
];
let current = vec![(0, 0.5), (5, 100.0)];
assert!(!pruner.should_prune(2, 5, &current, &completed));
}
// --- Median calculation edge cases ---
#[test]
fn correct_median_with_even_number_of_trials() {
let pruner = MedianPruner::new(Direction::Minimize);
// 4 trials at step 0: [1.0, 2.0, 3.0, 4.0] => median = 2.5
let completed = vec![
trial_with_values(0, vec![(0, 1.0)]),
trial_with_values(1, vec![(0, 2.0)]),
trial_with_values(2, vec![(0, 3.0)]),
trial_with_values(3, vec![(0, 4.0)]),
];
// 2.6 > 2.5 => prune
let current = vec![(0, 2.6)];
assert!(pruner.should_prune(4, 0, &current, &completed));
// 2.4 < 2.5 => don't prune
let current = vec![(0, 2.4)];
assert!(!pruner.should_prune(4, 0, &current, &completed));
}
#[test]
fn correct_median_with_odd_number_of_trials() {
let pruner = MedianPruner::new(Direction::Minimize);
// 5 trials at step 0: [1.0, 2.0, 3.0, 4.0, 5.0] => median = 3.0
let completed = vec![
trial_with_values(0, vec![(0, 1.0)]),
trial_with_values(1, vec![(0, 2.0)]),
trial_with_values(2, vec![(0, 3.0)]),
trial_with_values(3, vec![(0, 4.0)]),
trial_with_values(4, vec![(0, 5.0)]),
];
// 3.5 > 3.0 => prune
let current = vec![(0, 3.5)];
assert!(pruner.should_prune(5, 0, &current, &completed));
// 2.5 < 3.0 => don't prune
let current = vec![(0, 2.5)];
assert!(!pruner.should_prune(5, 0, &current, &completed));
}
// --- Non-contiguous step numbers ---
#[test]
fn works_with_non_contiguous_steps() {
let pruner = MedianPruner::new(Direction::Minimize);
// Steps are 0, 10, 100 — non-contiguous
let completed = vec![
trial_with_values(0, vec![(0, 1.0), (10, 2.0), (100, 3.0)]),
trial_with_values(1, vec![(0, 1.5), (10, 2.5), (100, 4.0)]),
trial_with_values(2, vec![(0, 2.0), (10, 3.0), (100, 5.0)]),
];
// At step 100: [3.0, 4.0, 5.0] => median = 4.0
let current = vec![(0, 1.0), (10, 2.0), (100, 4.5)];
assert!(pruner.should_prune(3, 100, &current, &completed));
let current = vec![(0, 1.0), (10, 2.0), (100, 3.5)];
assert!(!pruner.should_prune(3, 100, &current, &completed));
}
// --- No intermediate values for current trial ---
#[test]
fn no_prune_when_no_intermediate_values() {
let pruner = MedianPruner::new(Direction::Minimize);
let completed = vec![trial_with_values(0, vec![(0, 1.0)])];
assert!(!pruner.should_prune(1, 0, &[], &completed));
}
// --- Pruned trials are excluded from median calculation ---
#[test]
fn pruned_trials_excluded_from_median() {
use optimizer::TrialState;
let pruner = MedianPruner::new(Direction::Minimize);
let mut pruned = trial_with_values(0, vec![(0, 0.1)]);
pruned.state = TrialState::Pruned;
// Only the completed trial (value 5.0) counts. Pruned trial (0.1) is excluded.
let completed = vec![pruned, trial_with_values(1, vec![(0, 5.0)])];
// 3.0 < 5.0 => don't prune (only 1 completed trial with median 5.0)
let current = vec![(0, 3.0)];
assert!(!pruner.should_prune(2, 0, &current, &completed));
// 6.0 > 5.0 => prune
let current = vec![(0, 6.0)];
assert!(pruner.should_prune(2, 0, &current, &completed));
}
+64
View File
@@ -0,0 +1,64 @@
use optimizer::pruner::{Pruner, ThresholdPruner};
#[test]
fn prune_when_value_exceeds_upper_threshold() {
let pruner = ThresholdPruner::new().upper(10.0);
let values = vec![(0, 5.0), (1, 8.0), (2, 11.0)];
assert!(pruner.should_prune(0, 2, &values, &[]));
}
#[test]
fn prune_when_value_falls_below_lower_threshold() {
let pruner = ThresholdPruner::new().lower(0.0);
let values = vec![(0, 5.0), (1, 2.0), (2, -1.0)];
assert!(pruner.should_prune(0, 2, &values, &[]));
}
#[test]
fn no_prune_when_value_within_bounds() {
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
let values = vec![(0, 3.0), (1, 5.0), (2, 7.0)];
assert!(!pruner.should_prune(0, 2, &values, &[]));
}
#[test]
fn no_prune_when_no_intermediate_values() {
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
assert!(!pruner.should_prune(0, 0, &[], &[]));
}
#[test]
fn works_with_only_upper_set() {
let pruner = ThresholdPruner::new().upper(5.0);
let below = vec![(0, 3.0)];
let above = vec![(0, 6.0)];
assert!(!pruner.should_prune(0, 0, &below, &[]));
assert!(pruner.should_prune(0, 0, &above, &[]));
}
#[test]
fn works_with_only_lower_set() {
let pruner = ThresholdPruner::new().lower(2.0);
let above = vec![(0, 5.0)];
let below = vec![(0, 1.0)];
assert!(!pruner.should_prune(0, 0, &above, &[]));
assert!(pruner.should_prune(0, 0, &below, &[]));
}
#[test]
fn works_with_both_thresholds_set() {
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
// Within bounds
assert!(!pruner.should_prune(0, 0, &[(0, 5.0)], &[]));
// Exceeds upper
assert!(pruner.should_prune(0, 0, &[(0, 15.0)], &[]));
// Below lower
assert!(pruner.should_prune(0, 0, &[(0, -3.0)], &[]));
// At exact boundary (not pruned — strictly greater/less)
assert!(!pruner.should_prune(0, 0, &[(0, 10.0)], &[]));
assert!(!pruner.should_prune(0, 0, &[(0, 0.0)], &[]));
}
+133
View File
@@ -0,0 +1,133 @@
//! Integration tests for the BOHB sampler.
#![allow(
clippy::cast_sign_loss,
clippy::cast_precision_loss,
clippy::cast_possible_truncation
)]
use optimizer::parameter::{FloatParam, Parameter};
use optimizer::sampler::bohb::BohbSampler;
use optimizer::{Direction, Error, Study, TrialPruned};
#[test]
fn bohb_converges_on_quadratic() {
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(9)
.reduction_factor(3)
.min_points_in_model(5)
.seed(42)
.build()
.unwrap();
let pruner = bohb.matching_pruner(Direction::Minimize);
let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
let x_param = FloatParam::new(-10.0, 10.0);
study
.optimize(60, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
// Report intermediate values at budget steps 1, 3, 9
let obj = (x - 3.0).powi(2);
// Simulate budget-based evaluation with noise decreasing at higher budgets
trial.report(1, obj + 5.0);
trial.report(3, obj + 1.0);
trial.report(9, obj);
Ok::<_, Error>(obj)
})
.expect("optimization should succeed");
let best = study.best_trial().expect("should have trials");
assert!(
best.value < 10.0,
"BOHB should find a reasonable solution, got {}",
best.value
);
}
#[test]
fn bohb_with_pruning() {
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(27)
.reduction_factor(3)
.min_points_in_model(3)
.seed(123)
.build()
.unwrap();
let pruner = bohb.matching_pruner(Direction::Minimize);
let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
let x_param = FloatParam::new(-5.0, 5.0);
study
.optimize(40, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
let obj = x * x;
// Report at each rung step and check for pruning
for &step in &[1u64, 3, 9, 27] {
let noisy_obj = obj + 10.0 / step as f64;
trial.report(step, noisy_obj);
if trial.should_prune() {
return Err(TrialPruned.into());
}
}
Ok::<_, Error>(obj)
})
.expect("optimization should succeed");
// Verify we have completed trials
let best = study.best_trial().expect("should have at least one trial");
assert!(
best.value < 25.0,
"best value {} should be reasonable",
best.value
);
}
#[test]
fn bohb_uses_budget_conditioned_history() {
// Verify that BOHB conditions on budget level by testing that samples
// are influenced by intermediate values, not just final values.
let bohb = BohbSampler::builder()
.min_resource(1)
.max_resource(9)
.reduction_factor(3)
.min_points_in_model(3)
.seed(42)
.build()
.unwrap();
let pruner = bohb.matching_pruner(Direction::Minimize);
let study: Study<f64> = Study::with_sampler_and_pruner(Direction::Minimize, bohb, pruner);
let x_param = FloatParam::new(0.0, 10.0);
study
.optimize(30, |trial: &mut optimizer::Trial| {
let x = x_param.suggest(trial)?;
// Intermediate values that guide optimization toward x=2
trial.report(1, (x - 2.0).powi(2) + 1.0);
trial.report(3, (x - 2.0).powi(2) + 0.5);
trial.report(9, (x - 2.0).powi(2));
Ok::<_, Error>((x - 2.0).powi(2))
})
.expect("optimization should succeed");
let best = study.best_trial().unwrap();
let best_x: f64 = best.get(&x_param).unwrap();
// Should find x reasonably close to 2.0
assert!(
(best_x - 2.0).abs() < 5.0,
"BOHB should explore near x=2, got x={best_x}"
);
}

Some files were not shown because too many files have changed in this diff Show More