[RFC] Add construct_inputs class method to acquisition functions

[Balandat]

Context:
#783 introduces a new construct_inputs class method on AcquisitionFunctions that can be used to construct the arguments to the respective acquisition function constructor in a programmatic way based on a general set of inputs. This is analogous to the construct_inputs methods on the Models that we already have in place. It would be an optional method and when called on acquisition functions that do not implement custom handling would just raise a NotImplementedError.

Why do this:
The main use case for these is to simplify the instantiation of acquisition functions based on a general set of inputs with a consistent call signature in higher-level libraries using botorch as a modeling / optimization engine under the hood. Specifically, this would significantly simplify putting together custom combinations of Surrogate models, Acquisition functions, and optimizers in the modular botorch model setup in Ax: https://github.com/facebook/Ax/tree/master/ax/models/torch/botorch_modular. That avoids having to write a big old registry in Ax that is not only ugly but will also be very hard to keep in sync with the BoTorch changes.

Why not do this:
This adds a bunch of additional complexity and code that's not actually used directly within BoTorch. This means that even though this functionality is optional, the more complicated implementations of our stock acquisition functions could scare away users from using BoTorch in the first place. It also means additional maintenance overhead.

I am curious to hear what people think about this proposal.

[Balandat]

@sdaulton @danielrjiang @dme65 @qingfeng10 @bletham @eytan @bkarrer @liangshi7 @mshvartsman @saitcakmak @wjmaddox @henrymoss, any thoughts on this? Specifically about whether you think this additional complexity / overhead would make it harder (or at lest perceived to be harder) to work with BoTorch.

[Balandat]

OK how about a registry that we auto-populate by means of a decorator? We can key directly on the class, that way there shouldn't really be any big if/else trees. Sth like the following:

from typing import Type, Callable, Any

from botorch.acquisition.acquisition import AcquisitionFunction
from botorch.acquisition.analytic import ExpectedImprovement
from botorch.acquisition.monte_carlo import qUpperConfidenceBound


ACQF_INPUT_CONSTRUCTOR_REGISTRY = {}


def acqf_input_constructor(
    acqf_class: Type[AcquisitionFunction]
) -> Callable[[], AcquisitionFunction]:
    if acqf_class in ACQF_INPUT_CONSTRUCTOR_REGISTRY:
        raise ValueError(f"Cannot register duplicate arg constructor for acquisition class ({acqf_class})")

    def decorator(method):
        ACQF_INPUT_CONSTRUCTOR_REGISTRY[acqf_class] = method
        return method

    return decorator


@acqf_input_constructor(ExpectedImprovement)
def construct_inputs_expected_improvement(*args, **kwargs):
    # TODO: Real stuff
    return {"model": "model", "best_f": args[0]}


@acqf_input_constructor(qUpperConfidenceBound)
def construct_inputs_q_upper_confidence_bound(*args, **kwargs):
    # TODO: Real stuff 
    return {"model": "model", "beta": args[0]}

If we also wrote sth simple like

def get_acqf_input_constructor(acqf_class: Type[AcquisitionFunction]):
    if acqf_class not in ACQF_INPUT_CONSTRUCTOR_REGISTRY:
        raise RuntimeError(
            "Input constructor for acquisition class {acqf_class} not registered. "
            "Use the `@acqf_input_constructor` decorator to register a new method."
        ) 
    return ACQF_INPUT_CONSTRUCTOR_REGISTRY[acqf_class]

then upstream we could use this as follows

from botorch.acquisition.utils import get_acqf_input_constructor

,,,

input_arg_constructor = get_acqf_input_constructor(botorch_acqf_class)
acqf = acqf_class(**input_arg_constructor(training_data, model, ...))

Thoughts?

[lena-kashtelyan]

Love this solution, it's very creative and seems like the best of all the worlds.

[sdaulton]

Yeah I like this solution. I am definitely in favor of keeping the construct_inputs logic outside of the core acquisition classes. This is an elegant solution

[Balandat]

sg, will put up a PR for this

[Balandat]

PR is here: #788