Unify acquisition rules and function builders

[joelberkeley]

WORK IN PROGRESS

see also #136 #137 #91

Here I propose to unify and generalise the acquisition rules and acquisition function builders, reduce boilerplate, and provide a number of utilities which cast the acquisition optimizers and acquisition function reducers as or with standard design patterns.

Let's start with where we are right now. An acquisition rule is defined by

class AcquisitionRule(ABC, Generic[S, SP_contra]):
    @abstractmethod
    def acquire(
        self,
        search_space: SP_contra,
        datasets: Mapping[str, Dataset],
        models: Mapping[str, ProbabilisticModel],
        state: S | None,
    ) -> tuple[TensorType, S]:
        ...

and an acquisition function builder by

class AcquisitionFunctionBuilder(ABC):
    @abstractmethod
    def prepare_acquisition_function(
        self, datasets: Mapping[str, Dataset], models: Mapping[str, ProbabilisticModel]
    ) -> AcquisitionFunction:
        ...

There are three differences between these two signatures

a) the AcquisitionRule accepts a search space
b) the AcquisitionRule has a notion of state
c) the AcquisitionFunctionBuilder produces a AcquisitionFunction while the AcquisitionRule produces a TensorType

The first two may need changing. For a) ... For b) we can either: note that TrustRegion is the only rule that uses state, and remove the state from the acquisition rule by making the process of generating an acquisition space a seperate step to acquiring a new point; or we can make the AcquisitionFunctionBuilder stateful, which would allow optimizations such as passing the Pareto set between steps, thus avoiding the need to recompute all of it on each step.

For the third, we can parametrize over the return type. For illustration, we'll look at the simple stateless case where nothing receives the search space. Here, we can use the interface

class Step(Generic[S, T], ABC):
    def go(self, datasets: Mapping[str, Dataset], models: Mapping[str, ProbabilisticModel]) -> T:
        ...

An AcquisitionRule is then a Callable[[SP_contra], Step[TensorType]] and a function builder is a Step[AcquisitionFunction].

We'll take a look at representation. Step is an interface with one method, and most if not all implementations also only have one method, so common design approaches suggest this should be a function

Step = Callable[[Mapping[str, Dataset], Mapping[str, ProbabilisticModel], Tuple[T]]

One benefit of this is that defining a function requires less boilerplate than implementing an ABC. The other is that it encourages composition over inheritance.

We'll now turn to combinators over Step. These utilities provide natural, yet abstract and powerful tools for working with the values produced by these objects, marked above by type T. Between them, they would comprise most of the architecture of trieste. First, we can map over Steps,

def map(step: Step[T], f: Callable[[T], U]) -> Step[U]:
    return lambda data, models: f(step(data, models))

with which we can define an acquisition rule as

AcquisitionOptimizer = Callable[[SP_contra, AcquisitionFunction], TensorType]
AcquisitionRule = Callable[[SP_contra], Step[TensorType]]
AcquisitionFunctionBuilder = Step[AcquisitionFunction]

def efficient_global_optimization(
    builder: AcquisitionFunctionBuilder, optimizer: AcquisitionOptimizer[SP_contra]
) -> AcquisitionRule[SP_contra]:
    return lambda search_space: map(builder, lambda f: optimizer(search_space, f))

We also get other things for free with map, such as modifying the acquisition function that would be produced by a builder before it's built

def negate(f: AcquisitionFunction) -> AcquisitionFunction:
    return lambda at: -f(at)

negative_lcb_builder = map(lcb_builder, negate)

Second, we can define reduce:

def product(f: AcquisitionFunction, ff: AcquisitionFunction) -> AcquisitionFunction:
    return lambda x: f(x) * ff(x)

def reduce(steps: list[Step[T]], f: Callable[[T, T], T]) -> Step[T]:
    ...

reduce(product, [EI(), POF()])

===============================

First, we can map2 (from Applicative) over Steps, which takes a Step producing a T and another producing a U, and using an appropriate function, produce a Step producing a V, as

def map2(step_t: Step[S, T], step_u: Step[S, U], f: Callable[[T, U], V]) -> Step[S, V]:
    def step_v(data: Mapping[str, Dataset], models: Mapping[str, ProbabilisticModel], state: S) -> tuple[V, S]:
        new_state1, t = step_t(data, models, state)
        new_state2, u = step_u(data, models, new_state1)
        return f(t, u), new_state2

    return step_v

with which we can define an acquisition rule in terms of the trust region algorithm, the acquisition function builder and an optimizer

AcquisitionOptimizer = Callable[[SP_contra, AcquisitionFunction], TensorType]

def efficient_global_optimization(
    builder: Step[S, AcquisitionFunction], search_space: Step[S, SearchSpace], optimizer: AcquisitionOptimizer[SP_contra]
) -> Step[SP_contra, TensorType]:
    return map2(builder, search_space, optimizer)

Note this isn't Python's built-in map, but an example of the more general concept:

def map(step: Step[S, T], f: Callable[[T], U]) -> Step[S, U]:
    def new_step(data: Mapping[str, Dataset], models: Mapping[str, ProbabilisticModel], state: S) -> tuple[U, S]:
        new_state, t = step(data, models)
        return f(t), new_state

    return new_step