Multi-fidelity surrogate models

baybe/kernels/__init__.py

@@ -5,6 +5,7 @@
 """
 
 from baybe.kernels.basic import (
+    IndexKernel,
     LinearKernel,
     MaternKernel,
     PeriodicKernel,
@@ -18,6 +19,7 @@
 
 __all__ = [
     "AdditiveKernel",
+    "IndexKernel",
     "LinearKernel",
     "MaternKernel",
     "PeriodicKernel",

baybe/parameters/fidelity.py

@@ -21,8 +21,8 @@
     validate_is_finite,
     validate_unique_values,
 )
+from baybe.settings import active_settings
 from baybe.utils.conversion import nonstring_to_tuple
-from baybe.utils.numerical import DTypeFloatNumpy
 
 
 def _convert_zeta(
@@ -107,7 +107,9 @@ def values(self) -> tuple[str | bool, ...]:
     @cached_property
     def comp_df(self) -> pd.DataFrame:
         return pd.DataFrame(
-            range(len(self.values)), dtype=DTypeFloatNumpy, columns=[self.name]
+            range(len(self.values)),
+            dtype=active_settings.DTypeFloatNumpy,
+            columns=[self.name],
         )
 
 
@@ -159,5 +161,7 @@ def values(self) -> tuple[float, ...]:
     @cached_property
     def comp_df(self) -> pd.DataFrame:
         return pd.DataFrame(
-            {self.name: self.values}, index=self.values, dtype=DTypeFloatNumpy
+            {self.name: self.values},
+            index=self.values,
+            dtype=active_settings.DTypeFloatNumpy,
         )

baybe/recommenders/pure/bayesian/base.py

@@ -44,6 +44,13 @@ def _autoreplicate(surrogate: SurrogateProtocol, /) -> SurrogateProtocol:
 class BayesianRecommender(PureRecommender, ABC):
     """An abstract class for Bayesian Recommenders."""
 
+    # TODO: Factory defaults the surrogate to a GaussianProcessesSurrogate always.
+    #   Surrogate and kernel defaults should be different for searchspaces with
+    #   CategoricalFidelityParameter or NumericalDiscreteFidelityParameter.
+    #   This can be achieved without the user having to specify the surroagte model,
+    #   e.g., by
+    #   * using a dispatcher factory which decides surrogate model on fit time
+    #   * having a "_setup_surrogate" method similar to the acquisition function logic
     _surrogate_model: SurrogateProtocol = field(
         alias="surrogate_model",
         factory=GaussianProcessSurrogate,

baybe/searchspace/core.py

@@ -15,6 +15,10 @@
 from baybe.constraints.base import Constraint
 from baybe.parameters import TaskParameter
 from baybe.parameters.base import Parameter
+from baybe.parameters.fidelity import (
+    CategoricalFidelityParameter,
+    NumericalDiscreteFidelityParameter,
+)
 from baybe.searchspace.continuous import SubspaceContinuous
 from baybe.searchspace.discrete import (
     MemorySize,
@@ -48,6 +52,29 @@ class SearchSpaceType(Enum):
     """Flag for hybrid search spaces resp. compatibility with hybrid search spaces."""
 
 
+class SearchSpaceTaskType(Enum):
+    """Enum class for different types of task and/or fidelity subspaces."""
+
+    SINGLETASK = "SINGLETASK"
+    """Flag for search spaces with no task parameters."""
+
+    CATEGORICALTASK = "CATEGORICALTASK"
+    """Flag for search spaces with a categorical task parameter."""
+
+    NUMERICALFIDELITY = "NUMERICALFIDELITY"
+    """Flag for search spaces with a discrete numerical (ordered) fidelity parameter."""
+
+    CATEGORICALFIDELITY = "CATEGORICALFIDELITY"
+    """Flag for search spaces with a categorical (unordered) fidelity parameter."""
+
+    # TODO: Distinguish between multiple task parameter and mixed task parameter types.
+    # In future versions, multiple task/fidelity parameters may be allowed. For now,
+    # they are disallowed, whether the task-like parameters are different or the same
+    # class.
+    MULTIPLETASKPARAMETER = "MULTIPLETASKPARAMETER"
+    """Flag for search spaces with mixed task and fidelity parameters."""
+
+
 @define
 class SearchSpace(SerialMixin):
     """Class for managing the overall search space.
@@ -275,6 +302,24 @@ def task_idx(self) -> int | None:
         #   --> Fix this when refactoring the data
         return cast(int, self.discrete.comp_rep.columns.get_loc(task_param.name))
 
+    @property
+    def fidelity_idx(self) -> int | None:
+        """The column index of the task parameter in computational representation."""
+        try:
+            # See TODO [16932] and TODO [11611]
+            fidelity_param = next(
+                p
+                for p in self.parameters
+                if isinstance(
+                    p,
+                    (CategoricalFidelityParameter, NumericalDiscreteFidelityParameter),
+                )
+            )
+        except StopIteration:
+            return None
+
+        return cast(int, self.discrete.comp_rep.columns.get_loc(fidelity_param.name))
+
     @property
     def n_tasks(self) -> int:
         """The number of tasks encoded in the search space."""
@@ -287,6 +332,105 @@ def n_tasks(self) -> int:
             return 1
         return len(task_param.values)
 
+    @property
+    def n_fidelities(self) -> int:
+        """The number of tasks encoded in the search space."""
+        # See TODO [16932]
+        try:
+            fidelity_param = next(
+                p
+                for p in self.parameters
+                if isinstance(
+                    p,
+                    (CategoricalFidelityParameter, NumericalDiscreteFidelityParameter),
+                )
+            )
+            return len(fidelity_param.values)
+
+        # When there are no fidelity parameters, we effectively have a single fidelity
+        except StopIteration:
+            return 1
+
+    @property
+    def n_task_dimensions(self) -> int:
+        """The number of task dimensions."""
+        try:
+            # See TODO [16932]
+            fidelity_param = next(
+                p for p in self.parameters if isinstance(p, (TaskParameter,))
+            )
+        except StopIteration:
+            fidelity_param = None
+
+        return 1 if fidelity_param is not None else 0
+
+    @property
+    def n_fidelity_dimensions(self) -> int:
+        """The number of fidelity dimensions."""
+        try:
+            # See TODO [16932]
+            fidelity_param = next(
+                p
+                for p in self.parameters
+                if isinstance(
+                    p,
+                    (CategoricalFidelityParameter, NumericalDiscreteFidelityParameter),
+                )
+            )
+        except StopIteration:
+            fidelity_param = None
+
+        return 1 if fidelity_param is not None else 0
+
+    @property
+    def task_type(self) -> SearchSpaceTaskType:
+        """Return the task type of the search space.
+
+        Raises:
+            ValueError: If searchspace contains more than one task/fidelity parameter.
+            ValueError: An unrecognised fidelity parameter type is in SearchSpace.
+        """
+        task_like_parameters = (
+            TaskParameter,
+            CategoricalFidelityParameter,
+            NumericalDiscreteFidelityParameter,
+        )
+
+        n_task_like_parameters = sum(
+            isinstance(p, (task_like_parameters)) for p in self.parameters
+        )
+
+        if n_task_like_parameters == 0:
+            return SearchSpaceTaskType.SINGLETASK
+        elif n_task_like_parameters > 1:
+            # TODO: commute this validation further downstream.
+            # In case of user-defined custom models which allow for multiple task
+            # parameters, this should be later in recommender logic.
+            #  * Should this be an IncompatibilityError?
+            raise ValueError(
+                "SearchSpace must not contain more than one task/fidelity parameter."
+            )
+            return SearchSpaceTaskType.MULTIPLETASKPARAMETER
+
+        if self.n_task_dimensions == 1:
+            return SearchSpaceTaskType.CATEGORICALTASK
+
+        if self.n_fidelity_dimensions == 1:
+            n_categorical_fidelity_dims = sum(
+                isinstance(p, CategoricalFidelityParameter) for p in self.parameters
+            )
+            if n_categorical_fidelity_dims == 1:
+                return SearchSpaceTaskType.CATEGORICALFIDELITY
+
+            n_numerical_disc_fidelity_dims = sum(
+                isinstance(p, NumericalDiscreteFidelityParameter)
+                for p in self.parameters
+            )
+            if n_numerical_disc_fidelity_dims == 1:
+                return SearchSpaceTaskType.NUMERICALFIDELITY
+
+        raise RuntimeError("This line should be impossible to reach.")
+
     def get_comp_rep_parameter_indices(self, name: str, /) -> tuple[int, ...]:
         """Find a parameter's column indices in the computational representation.
 

baybe/surrogates/bandit.py

@@ -32,6 +32,9 @@ class BetaBernoulliMultiArmedBanditSurrogate(Surrogate):
     supports_transfer_learning: ClassVar[bool] = False
     # See base class.
 
+    supports_multi_fidelity: ClassVar[bool] = False
+    # See base class.
+
     prior: BetaPrior = field(factory=lambda: BetaPrior(1, 1))
     """The beta prior for the win rates of the bandit arms. Uniform by default."""
 

baybe/surrogates/base.py

@@ -86,6 +86,10 @@ class Surrogate(ABC, SurrogateProtocol, SerialMixin):
     """Class variable encoding whether or not the surrogate supports transfer
     learning."""
 
+    supports_multi_fidelity: ClassVar[bool]
+    """Class variable encoding whether or not the surrogate supports multi fidelity
+    Bayesian optimization."""
+
     supports_multi_output: ClassVar[bool] = False
     """Class variable encoding whether or not the surrogate is multi-output
     compatible."""
@@ -428,6 +432,14 @@ def fit(
                 f"support transfer learning."
             )
 
+        # Check if multi fidelity capabilities are needed
+        if (searchspace.n_fidelities > 1) and (not self.supports_multi_fidelity):
+            raise ValueError(
+                f"The search space contains fidelity parameters but the selected "
+                f"surrogate model type ({self.__class__.__name__}) does not "
+                f"support multi fidelity Bayesian optimisation."
+            )
+
         # Block partial measurements
         handle_missing_values(measurements, [t.name for t in objective.targets])
 
@@ -472,6 +484,11 @@ def __str__(self) -> str:
                 self.supports_transfer_learning,
                 single_line=True,
             ),
+            to_string(
+                "Supports Multi Fidelity",
+                self.supports_multi_fidelity,
+                single_line=True,
+            ),
         ]
         return to_string(self.__class__.__name__, *fields)
 

baybe/surrogates/custom.py

@@ -70,6 +70,9 @@ class CustomONNXSurrogate(IndependentGaussianSurrogate):
     supports_transfer_learning: ClassVar[bool] = False
     # See base class.
 
+    supports_multi_fidelity: ClassVar[bool] = False
+    # See base class.
+
     onnx_input_name: str = field(validator=validators.instance_of(str))
     """The input name used for constructing the ONNX str."""
 

baybe/surrogates/gaussian_process/components/kernel.py

@@ -10,6 +10,7 @@
 from baybe.kernels.base import Kernel
 from baybe.kernels.composite import ProductKernel
 from baybe.parameters.categorical import TaskParameter
+from baybe.parameters.fidelity import CategoricalFidelityParameter
 from baybe.parameters.selector import (
     ParameterSelectorProtocol,
     TypeSelector,
@@ -79,15 +80,30 @@ def _default_base_kernel_factory(self) -> KernelFactoryProtocol:
             BayBENumericalKernelFactory,
         )
 
-        return BayBENumericalKernelFactory(TypeSelector((TaskParameter,), exclude=True))
+        return BayBENumericalKernelFactory(
+            TypeSelector(
+                (
+                    TaskParameter,
+                    CategoricalFidelityParameter,
+                ),
+                exclude=True,
+            )
+        )
 
     @task_kernel_factory.default
     def _default_task_kernel_factory(self) -> KernelFactoryProtocol:
         from baybe.surrogates.gaussian_process.presets.baybe import (
             BayBETaskKernelFactory,
         )
 
-        return BayBETaskKernelFactory(TypeSelector((TaskParameter,)))
+        return BayBETaskKernelFactory(
+            TypeSelector(
+                (
+                    TaskParameter,
+                    CategoricalFidelityParameter,
+                )
+            )
+        )
 
     @override
     def __call__(

baybe/surrogates/gaussian_process/core.py

@@ -52,6 +52,8 @@
     from torch import Tensor
 
 
+# TODO Jordan MHS: _ModelContext is used by fidelity surrogate models now so may deserve
+# its own file.
 @define
 class _ModelContext:
     """Model context for :class:`GaussianProcessSurrogate`."""
@@ -80,6 +82,27 @@ def n_tasks(self) -> int:
         """The number of tasks."""
         return self.searchspace.n_tasks
 
+    @property
+    def n_fidelity_dimensions(self) -> int:
+        """The number of fidelity dimensions."""
+        # Possible TODO: Generalize to multiple fidelity dimensions
+        return 1 if self.searchspace.fidelity_idx is not None else 0
+
+    @property
+    def is_multi_fidelity(self) -> bool:
+        """Are there any fidelity dimensions?"""
+        return self.n_fidelity_dimensions > 0
+
+    @property
+    def fidelity_idx(self) -> int | None:
+        """The computational column index of the task parameter, if available."""
+        return self.searchspace.fidelity_idx
+
+    @property
+    def n_fidelities(self) -> int:
+        """The number of fidelities."""
+        return self.searchspace.n_fidelities
+
     @property
     def parameter_bounds(self) -> Tensor:
         """Get the search space parameter bounds in BoTorch Format."""
@@ -93,7 +116,7 @@ def numerical_indices(self) -> list[int]:
         return [
             i
             for i in range(len(self.searchspace.comp_rep_columns))
-            if i != self.task_idx
+            if i not in (self.task_idx, self.fidelity_idx)
         ]