Skip to content

Add BONSAI-style greedy pruning for NChooseK constraints#747

Draft
jduerholt wants to merge 2 commits intomainfrom
feature/bonsai
Draft

Add BONSAI-style greedy pruning for NChooseK constraints#747
jduerholt wants to merge 2 commits intomainfrom
feature/bonsai

Conversation

@jduerholt
Copy link
Copy Markdown
Contributor

Motivation

BONSAI for handling of NChooseK constraint (https://arxiv.org/abs/2602.07144).

@dlinzner-bcs: one could do the same in the DOE module.

Have you read the Contributing Guidelines on pull requests?

Yes.

Have you updated CHANGELOG.md?

Not yet.

Test Plan

Unit tests.

jduerholt and others added 2 commits March 26, 2026 17:13
Implement greedy post-processing pruning based on the BONSAI algorithm
(https://arxiv.org/abs/2602.07144) to handle NChooseK constraints more
effectively. Instead of encoding NChooseK as nonlinear constraints during
acquisition optimization, candidates are optimized without them and then
pruned in post-processing by greedily zeroing the feature with smallest
acquisition function impact until the constraint is satisfied.

- Add Domain.is_nchoosek_pruning_applicable() to check if pruning is safe
  (no NChooseK feature appears in any other constraint)
- Add exclude_nchoosek param to get_nonlinear_constraints()
- Skip NChooseK in optimize_acqf when pruning is applicable
- Override _postprocess_candidates in BotorchStrategy with Ax-style
  greedy pruning that conditions on already-pruned candidates for q>1

Co-Authored-By: Claude Opus 4.6 <[email protected]>
Relax is_nchoosek_pruning_applicable() to allow NChooseK features that
also participate in linear equality/inequality constraints (only block
for nonlinear/product/interpoint overlap). When linear overlap exists,
pruning variants use QP projection (LinearProjection) to find a feasible
warm-start with x_k=0, then local optimize_acqf to find the best
achievable AF value — giving accurate greedy decisions. The simple
zeroing path is preserved when no linear overlap exists.

- Relax is_nchoosek_pruning_applicable: linear overlap now allowed
- Add Domain.has_nchoosek_linear_overlap() to select pruning path
- Add QP + local optimize_acqf variant construction in botorch.py
- Add _build_pruning_domain to create modified domain for QP projection
- Add end-to-end test with NChooseK + mixture constraint (sum=1)

Co-Authored-By: Claude Opus 4.6 <[email protected]>
@jduerholt jduerholt marked this pull request as draft March 26, 2026 21:44
@jduerholt
Copy link
Copy Markdown
Contributor Author

@Jimbo994, if you are interested ;)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant