How to match data, models and acquisitions

[joelberkeley]

We'd like a way to match data and models with each other, and specify which data and models are to be used in which acquisitions, We'd like a design ...

• that only requires as much complexity as the problem at hand. In particular, basic setups would require little code
• where all parts of trieste are compatible with all other parts that could be used with it. At the moment, single-objective acquisitions require multi-objective data to be in separate datasets, but there's not inherent need for that to be the case.

The scenarios we want to support include:

• single-objective
• multi-objective where all objectives correspond to the same query points and can be reasonably expressed with the same data (that is, the observations can all be in the same tensor)
• multi-objective where some objectives can be contained in the same tensor, but not all can, or it doesn't make sense to do that
• multi-objective where no objectives can be expressed in the same tensor, or it doesn't make sense to do that
• all of the above with and without constraints
• ...

Open questions
Q: Is the Dataset concept helping or hindering?

@joelberkeley to ask another Secondmind dev for their input once requirements are gathered

[joelberkeley]

I'm including this from an earlier discussion for reference. I now think it might not be radical enough ...

One option is: since there is always at least one objective, we can special-case the objective data and model with their own arguments, e.g. for the acquisition rule,

def acquire(
    self,
    search_space: Box,
    objective_data: Dataset,
    objective_model: ProbabilisticModel
    additional_data: Mapping[str, Dataset],
    additional_models: Mapping[str, ProbabilisticModel]
    state: None
)

This may become confusing for multi-objective in the case where we need more than one objective dataset and model, but I'm not sure we'll ever need that: the current multi-objective work suggests we'd want users to collect all objectives into a single dataset.

[joelberkeley]

Another option is to make everything (bayesian optimizer, rules, acquisition builders) generic in the container type. This is the most coherent way of doing this that I've thought of, but it requires higher-kinded types, which Python's type hints don't support so I used a third party library returns. I think the complexity of the returns API, the knowledge required to work with higher-kinded types, and using third party tools, makes this approach infeasible. It's here for completeness. You can check out this branch if you want to see more

from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TypeVar, Generic, Mapping, Iterator, Callable
from returns.primitives.hkt import Kind1, SupportsKind1, dekind

T_co = TypeVar("T_co", covariant=True)
U = TypeVar("U")
V = TypeVar("V")
G = TypeVar("G", bound="Grouping")


class Grouping(ABC, Generic[T_co]):
    @abstractmethod
    def map(self: G, f: Callable[[T_co], U]) -> Kind1[G, U]:
        ...

    @abstractmethod
    def zip_with(self: G, other: Kind1[G, U], f: Callable[[T_co, U], V]) -> Kind1[G, V]:
        ...

    @abstractmethod
    def __iter__(self) -> Iterator[T_co]:
        ...


class One(Grouping[T_co], SupportsKind1["One", T_co]):
    def __init__(self, o: T_co):
        self._o = o

    def get(self) -> T_co:
        return self._o

    def map(self, f: Callable[[T_co], U]) -> One[U]:
        return One(f(self._o))

    def zip_with(self, other: Kind1[One, U], f: Callable[[T_co, U], V]) -> One[V]:
        return One(f(self.get(), other.get()))

    def __iter__(self) -> Iterator[T_co]:
        return iter((self._o,))


class Many(Grouping[T_co], SupportsKind1["Many", T_co]):
    def __init__(self, m: Mapping[str, T_co]):
        self._m = m

    def __getitem__(self, key: str) -> T_co:
        return self._m[key]

    def map(self, f: Callable[[T_co], U]) -> Many[U]:
        return Many({k: f(v) for k, v in self._m.items()})

    def zip_with(self, other: Kind1[Many, U], f: Callable[[T_co, U], V]) -> Many[V]:
        return Many({k: f(self[k], dekind(other)[k]) for k in self._m})

    def __iter__(self) -> Iterator[T_co]:
        return iter(self._m.values())

and an example of usage:

class AcquisitionRule(ABC, Generic[S, SP, G]):
    @abstractmethod
    def acquire(
        self,
        search_space: SP,
        datasets: Kind1[G, Dataset],
        models: Kind1[G, ModelInterface],
        state: Optional[S],
    ) -> Tuple[QueryPoints, S]:
        ...

[joelberkeley]

Another option is to make separate APIs for single and multi-model cases. I did try this out in a fair bit of detail, and it works OK. Other people weren't too keen on it, and I can kind of understand why.

One downside of this was there was a fair bit of duplication. I didn't see an obvious way to avoid it.

Here's a branch where I tried this out

[joelberkeley]

Another thought: ExpectedConstrainedImprovement is the only obvious rule/function that needs more than one dataset/model in a non-trivial way. You could ask "how would I design this if I wasn't constrained by ExpectedConstrainedImprovement?", then see if it's possible to include ExpectedConstrainedImprovement in that design in a sensible way.