Compare commits
140 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| db192eb1cb | |||
| 80a7c27551 | |||
| 86db6361d6 | |||
| 11a8534b38 | |||
| 1cd18c16f9 | |||
| 9f2ce2a482 | |||
| 8c025b4525 | |||
| cfa8426312 | |||
| 8d06d213db | |||
| 479554ffb2 | |||
| fd5c5453e5 | |||
| 7c0cce160a | |||
| 8394a747a5 | |||
| 0e0e9ff38f | |||
| 12efbaabab | |||
| 3d3dcb4c26 | |||
| cc50e1ad4b | |||
| 38a20bbc42 | |||
| f0798ca58a | |||
| a502d7e1b8 | |||
| d80d2bb1fb | |||
| 4781107ede | |||
| 9ab94002d9 | |||
| 1b158e4d1b | |||
| ddbbe294bc | |||
| 5da63d8976 | |||
| 74cf0643eb | |||
| baebbae64c | |||
| 750bbad0d8 | |||
| 271748700f | |||
| c561caaa34 | |||
| 04d822479e | |||
| 8713a93636 | |||
| 3723f91849 | |||
| 968eb67d07 | |||
| 2d88526d8d | |||
| cc3937d4f7 | |||
| 6f1c070b7d | |||
| 12f7f25baa | |||
| e68a3788be | |||
| d20f09c66a | |||
| 47b5f9cec8 | |||
| c20a53dfba | |||
| d81d1de4ff | |||
| 964b4d5749 | |||
| ee59c9cdd0 | |||
| 6a4b27c46f | |||
| cc0dce5daf | |||
| 1f1e9d1779 | |||
| 3581187cd3 | |||
| 349ffcbed3 | |||
| 88e2ac0ff0 | |||
| d122582bf7 | |||
| 705687a42e | |||
| 03deedd775 | |||
| 047a788afd | |||
| d2e36ec304 | |||
| d14f4e0792 | |||
| ed80c59795 | |||
| ed1e9e31da | |||
| 36c6262dbc | |||
| 0c295b4bc7 | |||
| 0a6f2345a8 | |||
| 24a0bdd473 | |||
| 8239cc58a1 | |||
| 906e5296de | |||
| 947701e83f | |||
| b3bf4ccb71 | |||
| 7bf64d6cf6 | |||
| 624283c766 | |||
| 64c7aa2448 | |||
| abb1b4c229 | |||
| dff58340a4 | |||
| e359392a00 | |||
| 0e54356345 | |||
| d851aad548 | |||
| 95402dc9b6 | |||
| f873722763 | |||
| ba31df69c7 | |||
| bcc4549e66 | |||
| 7d79111d81 | |||
| a9ea097377 | |||
| 7d70220da0 | |||
| b5c0d6353b | |||
| 08fde0c707 | |||
| 147a64708d | |||
| b36137bc96 | |||
| f3941a4b3e | |||
| 0cae3b4227 | |||
| 09480a973b | |||
| aa9734587a | |||
| a1dd6d393c | |||
| a53ba4f472 | |||
| 9b4a8321b5 | |||
| db0314a1d1 | |||
| df5fcedecf | |||
| 5061df9876 | |||
| 5fe0a75f78 | |||
| dcf3403261 | |||
| bb79d98027 | |||
| f4b0631178 | |||
| 52f3c074dc | |||
| f67188e3a7 | |||
| 72da883add | |||
| bc278d83a3 | |||
| c586650df4 | |||
| cc78a92ed6 | |||
| 12c35e7cb4 | |||
| a885d10436 | |||
| 41150057d6 | |||
| a522fb2f34 | |||
| ec04757740 | |||
| 432c74b927 | |||
| 8199ba6b43 | |||
| abe15bdd7b | |||
| 4d8af3242b | |||
| 2651e61b2c | |||
| d88280c152 | |||
| 0e8fc82201 | |||
| d7ad3f60db | |||
| 5dc81fa0ab | |||
| 55e50e6afb | |||
| 2fef1ab3c3 | |||
| aa1718fed7 | |||
| f9227df96e | |||
| 02ec9498c4 | |||
| 183a78c09e | |||
| ce817c1f38 | |||
| b9f8844e60 | |||
| d1c8925fc3 | |||
| 4ce4311c68 | |||
| 293b992e75 | |||
| af8a9f7638 | |||
| b9a90ebec3 | |||
| 24e921b7b9 | |||
| 6363540269 | |||
| 9f3c0f0d3b | |||
| 747e1928b9 | |||
| 15a43d9485 | |||
| 1091429697 |
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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
@@ -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
@@ -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
|
||||
@@ -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
@@ -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
@@ -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.
|
||||
@@ -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
|
||||
@@ -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.
|
||||
|
||||
[](https://docs.rs/optimizer)
|
||||
[](https://crates.io/crates/optimizer)
|
||||
[](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
|
||||
|
||||
|
||||
@@ -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), ¶ms, |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), ¶ms, |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);
|
||||
@@ -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 ¶ms {
|
||||
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);
|
||||
@@ -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
@@ -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"
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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}");
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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(|(¢er_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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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(¶m.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(())
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
File diff suppressed because it is too large
Load Diff
+660
@@ -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 ¤t_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());
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)], &[]));
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &completed);
|
||||
let strict_prunes = strict.should_prune(1, 19, ¤t, &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, ¤t, &[]));
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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)),
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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(¶m_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
@@ -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
@@ -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(¶m_id)` returns `Option<&ParamValue>`.
|
||||
//! - **Find the best trial:**
|
||||
//! `history.iter().min_by(|a, b| a.value.partial_cmp(&b.value).unwrap())`.
|
||||
//! - **Filter by state:**
|
||||
//! `history.iter().filter(|t| t.state == TrialState::Complete)`.
|
||||
//! - **Check feasibility:**
|
||||
//! `trial.is_feasible()` returns `true` when all constraints are ≤ 0.
|
||||
//!
|
||||
//! # Thread safety
|
||||
//!
|
||||
//! The [`Sampler`] trait requires `Send + Sync`. [`Study`](crate::Study) stores
|
||||
//! the sampler as `Arc<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(¶m.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.
|
||||
///
|
||||
|
||||
@@ -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 ¤t_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
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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, ¶m_order);
|
||||
let bad_obs = self.extract_observations(&bad, ¶m_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(¶m_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(¶m_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(¶m_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(¶m_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(¶m_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(¶m_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
@@ -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(¶m_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
@@ -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))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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 ¶m_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(¶m_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))
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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())),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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 ¶m_ids {
|
||||
write!(writer, ",{}", csv_escape(¶m_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 ¶m_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(),
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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} · {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, ¶m_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, ¶m_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
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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();
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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(¶m).unwrap();
|
||||
assert!((0.0..=1.0).contains(&x));
|
||||
|
||||
// Cached
|
||||
let x2 = trial.suggest_param(¶m).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(¶m).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(¶m).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(¶m).unwrap();
|
||||
assert!((1..=10).contains(&n));
|
||||
|
||||
// Cached
|
||||
let n2 = trial.suggest_param(¶m).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(¶m).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(¶m).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(¶m).unwrap();
|
||||
assert!(choices.contains(&opt));
|
||||
|
||||
// Cached
|
||||
let opt2 = trial.suggest_param(¶m).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(¶m).unwrap();
|
||||
let _ = val;
|
||||
|
||||
// Cached
|
||||
let val2 = trial.suggest_param(¶m).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(¶m).unwrap();
|
||||
assert!([Activation::Relu, Activation::Sigmoid, Activation::Tanh].contains(&act));
|
||||
|
||||
// Cached
|
||||
let act2 = trial.suggest_param(¶m).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;
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
mod median;
|
||||
mod threshold;
|
||||
@@ -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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &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, ¤t, &completed));
|
||||
// 2.4 < 2.5 => don't prune
|
||||
let current = vec![(0, 2.4)];
|
||||
assert!(!pruner.should_prune(4, 0, ¤t, &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, ¤t, &completed));
|
||||
// 2.5 < 3.0 => don't prune
|
||||
let current = vec![(0, 2.5)];
|
||||
assert!(!pruner.should_prune(5, 0, ¤t, &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, ¤t, &completed));
|
||||
|
||||
let current = vec![(0, 1.0), (10, 2.0), (100, 3.5)];
|
||||
assert!(!pruner.should_prune(3, 100, ¤t, &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, ¤t, &completed));
|
||||
|
||||
// 6.0 > 5.0 => prune
|
||||
let current = vec![(0, 6.0)];
|
||||
assert!(pruner.should_prune(2, 0, ¤t, &completed));
|
||||
}
|
||||
@@ -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)], &[]));
|
||||
}
|
||||
@@ -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
Reference in New Issue
Block a user