Multi fidelity acquisition functions and recommender logic

baybe/acquisition/__init__.py

@@ -3,6 +3,7 @@
 from baybe.acquisition.acqfs import (
     ExpectedImprovement,
     LogExpectedImprovement,
+    MultiFidelityUpperConfidenceBound,
     PosteriorMean,
     PosteriorStandardDeviation,
     ProbabilityOfImprovement,
@@ -13,6 +14,7 @@
     qLogNoisyExpectedHypervolumeImprovement,
     qLogNoisyExpectedImprovement,
     qLogNParEGO,
+    qMultiFidelityKnowledgeGradient,
     qNegIntegratedPosteriorVariance,
     qNoisyExpectedHypervolumeImprovement,
     qNoisyExpectedImprovement,
@@ -30,6 +32,7 @@
 EI = ExpectedImprovement
 qEI = qExpectedImprovement
 qKG = qKnowledgeGradient
+qMFKG = qMultiFidelityKnowledgeGradient
 LogEI = LogExpectedImprovement
 qLogEI = qLogExpectedImprovement
 qNEI = qNoisyExpectedImprovement
@@ -38,6 +41,7 @@
 PI = ProbabilityOfImprovement
 qPI = qProbabilityOfImprovement
 UCB = UpperConfidenceBound
+MFUCB = MultiFidelityUpperConfidenceBound
 qUCB = qUpperConfidenceBound
 qTS = qThompsonSampling
 qNEHVI = qNoisyExpectedHypervolumeImprovement
@@ -47,6 +51,7 @@
     ######################### Acquisition functions
     # Knowledge Gradient
     "qKnowledgeGradient",
+    "qMultiFidelityKnowledgeGradient",
     # Posterior Statistics
     "PosteriorMean",
     "PosteriorStandardDeviation",
@@ -67,6 +72,7 @@
     # Upper Confidence Bound
     "UpperConfidenceBound",
     "qUpperConfidenceBound",
+    "MultiFidelityUpperConfidenceBound",
     # Thompson Sampling
     "qThompsonSampling",
     # Hypervolume Improvement
@@ -77,6 +83,7 @@
     ######################### Abbreviations
     # Knowledge Gradient
     "qKG",
+    "qMFKG",
     # Posterior Statistics
     "PM",
     "PSTD",
@@ -97,6 +104,7 @@
     # Upper Confidence Bound
     "UCB",
     "qUCB",
+    "MFUCB",
     # Thompson Sampling
     "qTS",
     # Hypervolume Improvement

baybe/acquisition/_builder.py

@@ -23,11 +23,13 @@
     _ExpectedHypervolumeImprovement,
     qExpectedHypervolumeImprovement,
     qLogExpectedHypervolumeImprovement,
+    qMultiFidelityKnowledgeGradient,
     qNegIntegratedPosteriorVariance,
     qThompsonSampling,
 )
 from baybe.acquisition.base import AcquisitionFunction, _get_botorch_acqf_class
-from baybe.acquisition.utils import make_partitioning
+from baybe.acquisition.custom_acqfs import MultiFidelityUpperConfidenceBound
+from baybe.acquisition.utils import make_MFUCB_dicts, make_partitioning
 from baybe.exceptions import (
     IncompatibilityError,
     IncompleteMeasurementsError,
@@ -75,16 +77,21 @@ class BotorchAcquisitionArgs:
     # Optional, depending on the specific acquisition function being used
     best_f: float | None = _OPT_FIELD
     beta: float | None = _OPT_FIELD
+    costs_dict: dict[Any, tuple[float, ...]] = _OPT_FIELD
+    current_value: Tensor | None = _OPT_FIELD
+    fidelities_dict: dict[Any, tuple[Any, ...]] = _OPT_FIELD
     maximize: bool | None = _OPT_FIELD
     mc_points: Tensor | None = _OPT_FIELD
     num_fantasies: int | None = _OPT_FIELD
     objective: MCAcquisitionObjective | None = _OPT_FIELD
     partitioning: BoxDecomposition | None = _OPT_FIELD
     posterior_transform: PosteriorTransform | None = _OPT_FIELD
+    project: Callable[[Tensor], Tensor] | None = _OPT_FIELD
     prune_baseline: bool | None = _OPT_FIELD
     ref_point: Tensor | None = _OPT_FIELD
     X_baseline: Tensor | None = _OPT_FIELD
     X_pending: Tensor | None = _OPT_FIELD
+    zetas_dict: dict[Any, tuple[float, ...]] = _OPT_FIELD
 
     def collect(self) -> dict[str, Any]:
         """Collect the assigned arguments into a dictionary."""
@@ -202,6 +209,9 @@ def build(self) -> BoAcquisitionFunction:
         self._set_mc_points()
         self._set_ref_point()
         self._set_partitioning()
+        self._set_current_value()
+        self._set_projection()
+        self._set_MFUCB_dicts()
 
         botorch_acqf = self._botorch_acqf_cls(**self._args.collect())
         self.set_default_sample_shape(botorch_acqf)
@@ -264,6 +274,81 @@ def _set_best_f(self) -> None:
             case _:
                 raise NotImplementedError("This line should be impossible to reach.")
 
+    def _set_current_value(self) -> None:
+        """Set current value maximising posterior mean in qMFKG."""
+        if not isinstance(self.acqf, qMultiFidelityKnowledgeGradient):
+            return
+
+        from botorch.optim import optimize_acqf_mixed
+
+        if isinstance(self.acqf, qMultiFidelityKnowledgeGradient):
+            from botorch.acquisition import PosteriorMean
+            from botorch.acquisition.fixed_feature import (
+                FixedFeatureAcquisitionFunction,
+            )
+
+            curr_val_acqf = FixedFeatureAcquisitionFunction(
+                acq_function=PosteriorMean(self._botorch_surrogate),
+                d=len(self.searchspace.parameters),
+                columns=[
+                    self.searchspace._fidelity_idx,
+                ],
+                values=[
+                    1.0,
+                ],
+            )
+
+            # Jordan MHS NOTE: This is fast-and-loose use of mixed space optimization.
+            # Changes will be made with the next PR which uses a notion of wrapped acqfs
+            # for setting a current value but also for defining cost aware wrappers.
+
+            candidates_comp = self.searchspace.discrete.comp_rep
+            num_comp_columns = len(candidates_comp.columns)
+            candidates_comp.columns = list(range(num_comp_columns))  # type: ignore
+            candidates_comp_dict = candidates_comp.to_dict("records")
+
+            # Possible TODO. Align num_restarts and raw_samples with that defined by the
+            # user for the main acquisition function.
+            _, current_value = optimize_acqf_mixed(
+                acq_function=curr_val_acqf,
+                bounds=torch.from_numpy(self.searchspace.comp_rep_bounds.values),
+                fixed_features_list=candidates_comp_dict,  # type: ignore[arg-type]
+                q=1,
+                num_restarts=10,
+                raw_samples=64,
+            )
+
+        self._args.current_value = current_value
+
+    def _set_projection(self) -> None:
+        """Set projection to the target fidelity for qMFKG."""
+        if not isinstance(self.acqf, (qMultiFidelityKnowledgeGradient)):
+            return
+
+        assert self.searchspace._fidelity_idx is not None  # for mypy
+
+        target_fidelities = {self.searchspace._fidelity_idx: 1.0}
+
+        num_dims = len(self.searchspace.parameters)
+
+        def target_fidelity_projection(X: Tensor) -> Tensor:
+            from botorch.acquisition.utils import project_to_target_fidelity
+
+            return project_to_target_fidelity(X, target_fidelities, num_dims)
+
+        self._args.project = target_fidelity_projection
+
+    def _set_MFUCB_dicts(self) -> None:
+        """Set value, fidelities and cost dictionaries for MFUCB."""
+        if not isinstance(self.acqf, MultiFidelityUpperConfidenceBound):
+            return
+
+        fidelities_dict, costs_dict, zetas_dict = make_MFUCB_dicts(self.searchspace)
+
+        self._args.fidelities_dict = fidelities_dict
+        self._args.costs_dict = costs_dict
+        self._args.zetas_dict = zetas_dict
+
     def set_default_sample_shape(self, acqf: BoAcquisitionFunction, /):
         """Apply temporary workaround for Thompson sampling."""
         # TODO: Needs redesign once bandits are supported more generally

baybe/acquisition/acqfs.py

@@ -13,7 +13,7 @@
 from attr.converters import optional as optional_c
 from attr.validators import optional as optional_v
 from attrs import AttrsInstance, define, field, fields
-from attrs.validators import gt, instance_of, le
+from attrs.validators import ge, gt, instance_of, le
 from typing_extensions import override
 
 from baybe.acquisition.base import AcquisitionFunction
@@ -156,6 +156,22 @@ class qKnowledgeGradient(AcquisitionFunction):
     memory footprint and wall time."""
 
 
+@define(frozen=True)
+class qMultiFidelityKnowledgeGradient(AcquisitionFunction):
+    """Monte Carlo based knowledge gradient.
+
+    This acquisition function currently only supports purely continuous spaces.
+    """
+
+    abbreviation: ClassVar[str] = "qMFKG"
+
+    num_fantasies: int = field(validator=[instance_of(int), gt(0)], default=128)
+    """Number of fantasies to draw for approximating the knowledge gradient.
+
+    More samples result in a better approximation, at the expense of both increased
+    memory footprint and wall time."""
+
+
 ########################################################################################
 ### Posterior Statistics
 @define(frozen=True)
@@ -289,6 +305,31 @@ class qUpperConfidenceBound(AcquisitionFunction):
     """See :paramref:`UpperConfidenceBound.beta`."""
 
 
+@define(frozen=True)
+class MultiFidelityUpperConfidenceBound(AcquisitionFunction):
+    """Two stage acquisition function of Kandasamy et al (2016).
+
+    Stage 1: Choose design features based on argmax_x (softmin_m (UCB_m(x) + zeta_m)).
+
+    Stage 2: Choose cheapest fidelity satisfying a cost-aware informativeness threshold.
+    """
+
+    abbreviation: ClassVar[str] = "MFUCB"
+
+    softmin_temperature: float = field(
+        converter=float, validator=[finite_float, ge(0.0)], default=1e-2
+    )
+    """Softmin smoothing parameter."""
+
+    beta: float = field(converter=float, validator=finite_float, default=0.2)
+    """See :paramref:`UpperConfidenceBound.beta`."""
+
+    @override
+    @classproperty
+    def supports_batching(cls) -> bool:
+        return False
+
+
 ########################################################################################
 ### ThompsonSampling
 @define(frozen=True)

baybe/acquisition/base.py

@@ -165,11 +165,14 @@ def _get_botorch_acqf_class(
     """Extract the BoTorch acquisition class for the given BayBE acquisition class."""
     import botorch
 
+    from baybe.acquisition import custom_acqfs
+
     for cls in baybe_acqf_cls.mro():
         if (
             acqf_cls := getattr(botorch.acquisition, cls.__name__, False)
             or getattr(botorch.acquisition.multi_objective, cls.__name__, False)
             or getattr(botorch.acquisition.multi_objective.parego, cls.__name__, False)
+            or getattr(custom_acqfs, cls.__name__, False)
         ):
             if is_abstract(acqf_cls):
                 continue

baybe/acquisition/custom_acqfs/__init__.py

@@ -0,0 +1,10 @@
+"""Custom acquisition functions."""
+
+from baybe.acquisition.custom_acqfs.two_stage import (
+    MultiFidelityUpperConfidenceBound,
+)
+
+__all__ = [
+    # Multi fidelity acquisition functions
+    "MultiFidelityUpperConfidenceBound",
+]