Support Multi-Output, Multi-Fidelity optimization (#705)

* Create base classes; split TaskInput into Categorical and Continuous

* Create separate MFKGStrategy

* Initial sketch of fidelity kernel and pydantic checks

* Create default downsampling kernel for MFHVKG strategy

* Add MOMF benchmark, and minor fixes to data model

* Fix some typing issues

* Fix testing bugs with (Cat)TaskInputs

* Fix (Cat)TaskInput in botorch data model

* Add MFHVKG specs

* Add NotImplementedError for cat task inputs

* Create MFHVKG test case

* Fix init points function not taking `generator` arg

* Create simple test to evaluate qMFHVKG strategy

* Use deterministic surrogate as fidelity cost model

* Revert fidelity cost to Input

* Add fidelity costs to CategoricalTaskInput

* DiscreteTaskInput and related changes

* Default and validation of the cost task model

* Create some test specs

* Update test to reflect changes

* Fix sampler and objective causing bugs with output shapes

* Subset model outputs so that fantasize works

* Minor tweaks to tests

* Better handling of fixing task fidelities

* Create test for objective in fidelity output

* Improve fixing inputs and filtering outputs

* First draft of tutorial

* Remove discrete task input

* Separate cost model from surrogate

* Fix tests

* Fix example notebook, and give info about cost model

* Clean up spec tests

* Explore benchmark in notebook

* Fix more tests

* Add `d` argument for botorch v0.17

* Add WedgeKernel to aggregation (allowing for conditional+mfkvkg)

* Remove Fidelity kernel from kernel types

* Refactor `_generate_specs_surrogates`

* Rename `MultiFidelityStrategy`

* Remove unused code

* Add `current_value` docstring; replace `project` with partial

* Minor cleanups

* Fix bug in default specs

* Tighten registration of MF data models and add missing specs

- Mark `TaskInput` and `FidelityKernel` as abstract via `type: Any`
  with docstrings, so they no longer pretend to be concrete classes
  outside the `AnyFeature`/`AnyKernel` unions.
- Add a valid spec for `ContinuousTaskInput` so the serialization
  roundtrip test actually exercises it.
- Wire `MultiFidelityVarianceBasedStrategy` and
  `MultiFidelityHVKGStrategy` into the `AnyPredictive` union.
- Add a `validate_ref_point` validator to `MultiFidelityHVKGStrategy`
  matching `MoboStrategy`, so `dict[str, float]` and `None` reference
  points are normalized to `ExplicitReferencePoint` before the
  functional class consumes them.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* `MFHVKGStrategy` now inherits from `MoboStrategy`

* Remove `validate_ref_point` from `MFHVKG`

* Revert PEP585 type annotation

---------

Co-authored-by: Johannes P. Dürholt <johannespeter.duerholt@evonik.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 3 people

_quarto.yml

@@ -62,7 +62,7 @@ quartodoc:
             - BotorchStrategy
             - EntingStrategy
             - MoboStrategy
-            - MultiFidelityStrategy
+            - MultiFidelityVarianceBasedStrategy
             - MultiobjectiveStrategy
             - AdditiveSoboStrategy
             - CustomSoboStrategy
@@ -116,7 +116,7 @@ quartodoc:
             - MoboStrategy
             - BotorchStrategy
             - EntingStrategy
-            - MultiFidelityStrategy
+            - MultiFidelityVarianceBasedStrategy
             - ActiveLearningStrategy
             - ShortestPathStrategy
             - StepwiseStrategy

bofire/benchmarks/api.py

@@ -16,6 +16,7 @@
     TNK,
     ZDT1,
     CrossCoupling,
+    MOMFBraninCurrin,
     SnarBenchmark,
 )
 from bofire.benchmarks.single import (
@@ -39,6 +40,7 @@
     SnarBenchmark,
     BNH,
     TNK,
+    MOMFBraninCurrin,
 ]
 AnySingleBenchmark = Union[
     Ackley,

bofire/benchmarks/multi.py

@@ -22,6 +22,7 @@
     CategoricalDescriptorInput,
     ContinuousInput,
     ContinuousOutput,
+    ContinuousTaskInput,
     Input,
 )
 from bofire.data_models.objectives.api import (
@@ -700,3 +701,63 @@ def _get_base_cost(self, base, mmol_base):
             "TEA": 0.01,
         }
         return float(base_prices[base] * mmol_base)
+
+
+class MOMFBraninCurrin(Benchmark):
+    def __init__(self, **kwargs):
+        inputs = Inputs(
+            features=[
+                *(ContinuousInput(key=f"x{i}", bounds=(0.0, 1.0)) for i in range(2)),
+                ContinuousTaskInput(key="fidelity", bounds=(0.0, 1.0)),
+            ]
+        )
+
+        outputs = Outputs(
+            features=[
+                ContinuousOutput(key="branin", objective=MinimizeObjective()),
+                ContinuousOutput(key="currin", objective=MinimizeObjective()),
+            ]
+        )
+
+        self._domain = Domain(
+            inputs=inputs,
+            outputs=outputs,
+        )
+
+        super().__init__(**kwargs)
+
+    def _branin(self, X: np.ndarray) -> np.ndarray:
+        x1 = X[..., 0]
+        x2 = X[..., 1]
+        s = X[..., 2]
+
+        x11 = 15 * x1 - 5
+        x22 = 15 * x2
+        b = 5.1 / (4 * math.pi**2) - 0.01 * (1 - s)
+        c = 5 / math.pi - 0.1 * (1 - s)
+        r = 6
+        t = 1 / (8 * math.pi) + 0.05 * (1 - s)
+        y = (x22 - b * x11**2 + c * x11 - r) ** 2 + 10 * (1 - t) * np.cos(x11) + 10
+        B = 21 - y
+        return B / 22
+
+    def _currin(self, X: np.ndarray) -> np.ndarray:
+        x1 = X[..., 0]
+        x2 = X[..., 1]
+        s = X[..., 2]
+        A = 2300 * x1**3 + 1900 * x1**2 + 2092 * x1 + 60
+        B = 100 * x1**3 + 500 * x1**2 + 4 * x1 + 20
+        y = (1 - 0.1 * (1 - s) * np.exp(-1 / (2 * x2))) * A / B
+        C = -y + 14
+        return C / 15
+
+    def _f(self, candidates: pd.DataFrame) -> pd.DataFrame:
+        X = candidates.to_numpy()
+        return pd.DataFrame(
+            {
+                "branin": self._branin(X),
+                "valid_branin": 1,
+                "currin": self._currin(X),
+                "valid_currin": 1,
+            }
+        )

bofire/benchmarks/single.py

@@ -15,10 +15,10 @@
 from bofire.data_models.features.api import (
     CategoricalDescriptorInput,
     CategoricalInput,
+    CategoricalTaskInput,
     ContinuousInput,
     ContinuousOutput,
     DiscreteInput,
-    TaskInput,
 )
 from bofire.data_models.objectives.api import MaximizeObjective, MinimizeObjective
 from bofire.utils.torch_tools import tkwargs
@@ -487,7 +487,9 @@ def __init__(self, use_constraints: bool = False, **kwargs):
         self.use_constraints = use_constraints
         inputs = []
 
-        inputs.append(TaskInput(key="task_id", categories=["task_1", "task_2"]))
+        inputs.append(
+            CategoricalTaskInput(key="task_id", categories=["task_1", "task_2"])
+        )
         inputs.append(ContinuousInput(key="x_1", bounds=[-6, 6]))
         inputs.append(ContinuousInput(key="x_2", bounds=[-6, 6]))
 

bofire/data_models/acquisition_functions/acquisition_function.py

@@ -83,6 +83,14 @@ class qLogNEHVI(MultiObjectiveAcquisitionFunction):
     n_mc_samples: IntPowerOfTwo = 512
 
 
+class qMFHVKG(MultiObjectiveAcquisitionFunction):
+    type: Literal["qMFHVKG"] = "qMFHVKG"
+    alpha: Annotated[float, Field(ge=0)] = 0.0
+    num_fantasies: int = 8
+    num_pareto: int = 10
+    n_mc_samples: IntPowerOfTwo = 32
+
+
 class qNegIntPosVar(SingleObjectiveAcquisitionFunction):
     type: Literal["qNegIntPosVar"] = "qNegIntPosVar"
     n_mc_samples: IntPowerOfTwo = 512

bofire/data_models/acquisition_functions/api.py

@@ -9,6 +9,7 @@
     qLogNEHVI,
     qLogNEI,
     qLogPF,
+    qMFHVKG,
     qNegIntPosVar,
     qNEHVI,
     qNEI,
@@ -31,6 +32,7 @@
     qLogEHVI,
     qNEHVI,
     qLogNEHVI,
+    qMFHVKG,
     qNegIntPosVar,
     qLogPF,
 )
@@ -39,6 +41,8 @@
     qNEI, qEI, qSR, qUCB, qPI, qLogEI, qLogNEI, qLogPF
 )
 
-AnyMultiObjectiveAcquisitionFunction = tagged_union(qEHVI, qLogEHVI, qNEHVI, qLogNEHVI)
+AnyMultiObjectiveAcquisitionFunction = tagged_union(
+    qEHVI, qLogEHVI, qNEHVI, qLogNEHVI, qMFHVKG
+)
 
 AnyActiveLearningAcquisitionFunction = qNegIntPosVar

bofire/data_models/domain/features.py

@@ -37,13 +37,13 @@
     CategoricalInput,
     CategoricalMolecularInput,
     CategoricalOutput,
+    CategoricalTaskInput,
     ContinuousInput,
     ContinuousOutput,
     DiscreteInput,
     Feature,
     Input,
     Output,
-    TaskInput,
 )
 from bofire.data_models.features.feature import get_encoded_name
 from bofire.data_models.filters import filter_by_attribute, filter_by_class
@@ -350,12 +350,14 @@ class Inputs(_BaseFeatures[AnyInput]):
     def validate_only_one_task_input(cls, features: Sequence[AnyInput]):
         filtered = filter_by_class(
             features,
-            includes=TaskInput,
+            includes=CategoricalTaskInput,
             excludes=None,
             exact=False,
         )
         if len(filtered) > 1:
-            raise ValueError(f"Only one `TaskInput` is allowed, got {len(filtered)}.")
+            raise ValueError(
+                f"Only one `CategoricalTaskInput` is allowed, got {len(filtered)}."
+            )
         return features
 
     def get_fixed(self) -> Inputs:

bofire/data_models/features/api.py

@@ -24,7 +24,11 @@
     ContinuousMolecularInput,
 )
 from bofire.data_models.features.numerical import NumericalInput
-from bofire.data_models.features.task import TaskInput
+from bofire.data_models.features.task import (
+    CategoricalTaskInput,
+    ContinuousTaskInput,
+    TaskInput,
+)
 from bofire.data_models.unions import tagged_union
 
 
@@ -37,7 +41,8 @@
     ContinuousDescriptorInput,
     CategoricalDescriptorInput,
     CategoricalMolecularInput,
-    TaskInput,
+    CategoricalTaskInput,
+    ContinuousTaskInput,
     SumFeature,
     MeanFeature,
     WeightedMeanFeature,
@@ -57,7 +62,8 @@
     ContinuousDescriptorInput,
     CategoricalDescriptorInput,
     CategoricalMolecularInput,
-    TaskInput,
+    ContinuousTaskInput,
+    CategoricalTaskInput,
     ContinuousMolecularInput,
 )
 

bofire/data_models/features/task.py

@@ -1,15 +1,29 @@
-from typing import ClassVar, List, Literal
+from typing import Any, ClassVar, Literal
 
 import numpy as np
 from pydantic import model_validator
 
 from bofire.data_models.features.categorical import CategoricalInput
+from bofire.data_models.features.continuous import ContinuousInput
+from bofire.data_models.features.feature import Input
 
 
-class TaskInput(CategoricalInput):
+class TaskInput(Input):
+    """Abstract base class for task-encoding inputs.
+
+    This class is not directly instantiable and is not part of any
+    ``AnyFeature``/``AnyInput`` union. Use :class:`CategoricalTaskInput` or
+    :class:`ContinuousTaskInput` instead. It exists solely so that strategies
+    can use ``isinstance(feat, TaskInput)`` to detect either flavour.
+    """
+
+    type: Any
+
+
+class CategoricalTaskInput(TaskInput, CategoricalInput):
     order_id: ClassVar[int] = 8
-    type: Literal["TaskInput"] = "TaskInput"
-    fidelities: List[int] = []
+    type: Literal["CategoricalTaskInput"] = "CategoricalTaskInput"
+    fidelities: list[int] = []
 
     @model_validator(mode="after")
     def validate_fidelities(self):
@@ -19,10 +33,15 @@ def validate_fidelities(self):
                 self.fidelities.append(0)
         if len(self.fidelities) != n_tasks:
             raise ValueError(
-                "Length of fidelity lists must be equal to the number of tasks",
+                "Length of fidelity list must be equal to the number of tasks",
             )
         if list(set(self.fidelities)) != list(range(np.max(self.fidelities) + 1)):
             raise ValueError(
                 "Fidelities must be a list containing integers, starting from 0 and increasing by 1",
             )
         return self
+
+
+class ContinuousTaskInput(TaskInput, ContinuousInput):
+    order_id: ClassVar[int] = 11
+    type: Literal["ContinuousTaskInput"] = "ContinuousTaskInput"  # type: ignore

bofire/data_models/kernels/aggregation.py

@@ -6,6 +6,7 @@
     IndexKernel,
     PositiveIndexKernel,
 )
+from bofire.data_models.kernels.conditional import WedgeKernel
 from bofire.data_models.kernels.continuous import (
     InfiniteWidthBNNKernel,
     LinearKernel,
@@ -14,6 +15,7 @@
     RBFKernel,
     SphericalLinearKernel,
 )
+from bofire.data_models.kernels.fidelity import DownsamplingKernel
 from bofire.data_models.kernels.kernel import AggregationKernel
 from bofire.data_models.kernels.molecular import TanimotoKernel
 from bofire.data_models.kernels.shape import WassersteinKernel
@@ -32,6 +34,8 @@ class AdditiveKernel(AggregationKernel):
             IndexKernel,
             PositiveIndexKernel,
             TanimotoKernel,
+            DownsamplingKernel,
+            WedgeKernel,
             "AdditiveKernel",
             "MultiplicativeKernel",
             "ScaleKernel",
@@ -53,6 +57,8 @@ class MultiplicativeKernel(AggregationKernel):
             PositiveIndexKernel,
             AdditiveKernel,
             TanimotoKernel,
+            DownsamplingKernel,
+            WedgeKernel,
             "MultiplicativeKernel",
             "ScaleKernel",
         ]
@@ -72,6 +78,8 @@ class ScaleKernel(AggregationKernel):
         AdditiveKernel,
         MultiplicativeKernel,
         TanimotoKernel,
+        DownsamplingKernel,
+        WedgeKernel,
         "ScaleKernel",
         WassersteinKernel,
     ]