Add TopK downselection for initial batch generation.

[CompRhys]

Motivation

In order to get facebook/Ax#2938 over the line with initial candidate generation that obey the constraints we want to use the existing tooling within botorch. The hard coded logic currently in Ax uses topk to downselect the sobol samples. To make a change there that will not impact existing users we then need to implement topk downselection in botorch.

Test Plan

TODO:

• add tests for initialize_q_batch_topk

Related PRs

facebook/Ax#2938. (#2610 was initially intended to play part of this solution but then I realized that the pattern I wanted to use was conflating repeats and the batch dimension.)

[Balandat]

couple of small nits, o/w this looks great.

[facebook-github-bot]

@Balandat has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[saitcakmak]

Other initializers use Boltzmann sampling through initialize_q_batch helper to achieve a more diverse set of ICs. The motivation is that this will hopefully lead to different local-optima, making it more likely to discover the global optimum. If you have multiple raw samples that are close by, TopK can get you ICs that all lead to the same local optima. Should we follow the same logic in initialize_q_batch instead?

[Balandat]

That would make sense, I'd be interested in seeing a benchmark how that affects performance of SEBO

[CompRhys]

I would be entirely happy to not merge this if it will end up as dead code if theres a decision to swap the ic default for SEBO to use Boltzmann sampling. The motivation here was that I want to use the polytope sampling from botorch such that I can add inequality constraints to SEBO, adding a bunch of switching logic is bad in terms of maintenance there so the minimal change to keep the same behavior for how it's currently setup is to add this here.

The longer term idea for the cleanest way to do it I had was based on this #2610 where SEBO default could just pass the X_pareto points as initial batch conditions and then have botorch just fill up the rest of the requested n_restarts without needing to call anything in Ax from botorch to do the initial sampling that currently uses topk to get 1/2 of the initial conditions.

[codecov]

Codecov Report

Attention: Patch coverage is 89.65517% with 3 lines in your changes missing coverage. Please review.

Project coverage is 99.97%. Comparing base (3f2e2c7) to head (8e27422).

Files with missing lines	Patch %	Lines
botorch/optim/initializers.py	88.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2636      +/-   ##
==========================================
- Coverage   99.98%   99.97%   -0.02%     
==========================================
  Files         196      196              
  Lines       17365    17382      +17     
==========================================
+ Hits        17363    17377      +14     
- Misses          2        5       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.