diff --git a/CHANGELOG.md b/CHANGELOG.md index de3165f616..58cfa2b479 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,8 +15,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Gaussian process component factories - Support for GPyTorch objects (kernels, means, likelihood) as Gaussian process components, enabling full low-level customization +- Configurable fitting criterion for Gaussian process hyperparameter optimization - Factories for all Gaussian process components -- `CHEN`, `EDBO` and `EDBO_SMOOTHED` presets for `GaussianProcessSurrogate` +- `BOTORCH`, `CHEN`, `EDBO` and `EDBO_SMOOTHED` presets for `GaussianProcessSurrogate` - `TypeSelector` and `NameSelector` classes for parameter selection in kernel factories - `parameter_names` attribute to basic kernels for controlling the considered parameters - `ParameterKind` flag enum for classifying parameters by their role and automatic diff --git a/baybe/surrogates/gaussian_process/components/__init__.py b/baybe/surrogates/gaussian_process/components/__init__.py index a9e11f4afe..bb83f995e4 100644 --- a/baybe/surrogates/gaussian_process/components/__init__.py +++ b/baybe/surrogates/gaussian_process/components/__init__.py @@ -1,5 +1,10 @@ """Gaussian process surrogate components.""" +from baybe.surrogates.gaussian_process.components.fit_criterion import ( + FitCriterion, + FitCriterionFactoryProtocol, + PlainFitCriterionFactory, +) from baybe.surrogates.gaussian_process.components.kernel import ( KernelFactoryProtocol, PlainKernelFactory, @@ -15,6 +20,10 @@ ) __all__ = [ + # Fit Criterion + "FitCriterion", + "FitCriterionFactoryProtocol", + "PlainFitCriterionFactory", # Kernel "KernelFactoryProtocol", "PlainKernelFactory", diff --git a/baybe/surrogates/gaussian_process/components/_gpytorch.py b/baybe/surrogates/gaussian_process/components/_gpytorch.py new file mode 100644 index 0000000000..b7efc7bcbe --- /dev/null +++ b/baybe/surrogates/gaussian_process/components/_gpytorch.py @@ -0,0 +1,71 @@ +"""Custom GPyTorch components.""" + +import torch +from botorch.models.multitask import _compute_multitask_mean +from botorch.models.utils.gpytorch_modules import MIN_INFERRED_NOISE_LEVEL +from gpytorch.constraints import GreaterThan +from gpytorch.likelihoods.hadamard_gaussian_likelihood import HadamardGaussianLikelihood +from gpytorch.means import MultitaskMean +from gpytorch.means.multitask_mean import Mean +from gpytorch.priors import LogNormalPrior +from torch import Tensor +from torch.nn import Module + + +class HadamardConstantMean(Mean): + """A GPyTorch mean function implementing BoTorch's multitask mean logic. + + While GPyTorch already provides a :class:`~gpytorch.means.MultitaskMean` class, it + computes mean values for all (input, task)-pairs (where input means all parameters + except the task parameter), i.e. it intrinsically applies a Cartesian expansion. + However, for the regular transfer learning setting, we only need the means for the + pairs that are actually observed/requested. BoTorch subselects the relevant means + from the GPyTorch output in `MultiTaskGP.forward`, i.e. it uses a class-based + approach to define its special logic for the multitask case. In contrast, BayBE uses + a composition approach, which is more flexible but requires that the logic is + injected via a self-contained `Mean` object, which is what this class provides. + + Note: + Analogous to GPyTorch's + https://github.com/cornellius-gp/gpytorch/blob/main/gpytorch/likelihoods/hadamard_gaussian_likelihood.py + but where the logic is applied to the mean function, i.e. we learn a different + (constant) mean for each task. + """ + + def __init__(self, mean_module: Module, num_tasks: int, task_feature: int): + super().__init__() + self.multitask_mean = MultitaskMean(mean_module, num_tasks=num_tasks) + self.task_feature = task_feature + + def forward(self, x: Tensor) -> Tensor: + # Adapted from https://github.com/meta-pytorch/botorch/blob/e0f4f5b941b5949a4a1171bf8d4ee9f74f146f3a/botorch/models/multitask.py#L397 + + # Convert task feature to positive index + task_feature = self.task_feature % x.shape[-1] + + # Split input into task and non-task components + x_before = x[..., :task_feature] + task_idcs = x[..., task_feature : task_feature + 1] + x_after = x[..., task_feature + 1 :] + + return _compute_multitask_mean( + self.multitask_mean, x_before, task_idcs, x_after + ) + + +def make_botorch_multitask_likelihood( + num_tasks: int, task_feature: int +) -> HadamardGaussianLikelihood: + """Adapted from :class:`botorch.models.multitask.MultiTaskGP`.""" + noise_prior = LogNormalPrior(loc=-4.0, scale=1.0) + return HadamardGaussianLikelihood( + num_tasks=num_tasks, + batch_shape=torch.Size(), + noise_prior=noise_prior, + noise_constraint=GreaterThan( + MIN_INFERRED_NOISE_LEVEL, + transform=None, + initial_value=noise_prior.mode, + ), + task_feature_index=task_feature, + ) diff --git a/baybe/surrogates/gaussian_process/components/fit_criterion.py b/baybe/surrogates/gaussian_process/components/fit_criterion.py new file mode 100644 index 0000000000..c822251985 --- /dev/null +++ b/baybe/surrogates/gaussian_process/components/fit_criterion.py @@ -0,0 +1,46 @@ +"""Fitting criteria for the Gaussian process surrogate.""" + +from __future__ import annotations + +from enum import Enum +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood + from gpytorch.mlls import MarginalLogLikelihood + from gpytorch.models import GP as GPyTorchModel + + +class FitCriterion(Enum): + """Available fitting criteria for GP hyperparameter optimization.""" + + MARGINAL_LOG_LIKELIHOOD = "MARGINAL_LOG_LIKELIHOOD" + """Exact marginal log-likelihood.""" + + LEAVE_ONE_OUT_PSEUDOLIKELIHOOD = "LEAVE_ONE_OUT_PSEUDOLIKELIHOOD" + """Leave-one-out cross-validation pseudo-likelihood.""" + + def to_gpytorch( + self, likelihood: GPyTorchLikelihood, model: GPyTorchModel + ) -> MarginalLogLikelihood: + """Create the corresponding GPyTorch MLL object.""" + import gpytorch + + mll_class = { + FitCriterion.MARGINAL_LOG_LIKELIHOOD: gpytorch.ExactMarginalLogLikelihood, + FitCriterion.LEAVE_ONE_OUT_PSEUDOLIKELIHOOD: gpytorch.mlls.LeaveOneOutPseudoLikelihood, # noqa: E501 + }[self] + return mll_class(likelihood, model) + + +# Delayed import to avoid circular dependency +from baybe.surrogates.gaussian_process.components.generic import ( # noqa: E402 + GPComponentFactoryProtocol, + PlainGPComponentFactory, +) + +FitCriterionFactoryProtocol = GPComponentFactoryProtocol[FitCriterion] +"""A protocol defining the interface for fit criterion factories.""" + +PlainFitCriterionFactory = PlainGPComponentFactory[FitCriterion] +"""A trivial factory that returns a fixed fit criterion.""" diff --git a/baybe/surrogates/gaussian_process/components/generic.py b/baybe/surrogates/gaussian_process/components/generic.py index c1977146e5..52ed9a66cd 100644 --- a/baybe/surrogates/gaussian_process/components/generic.py +++ b/baybe/surrogates/gaussian_process/components/generic.py @@ -14,8 +14,9 @@ from baybe.searchspace import SearchSpace from baybe.serialization.core import block_serialization_hook, converter from baybe.serialization.mixin import SerialMixin +from baybe.surrogates.gaussian_process.components.fit_criterion import FitCriterion -BayBEGPComponent: TypeAlias = Kernel +BayBEGPComponent: TypeAlias = Kernel | FitCriterion if TYPE_CHECKING: from gpytorch.kernels import Kernel as GPyTorchKernel @@ -44,15 +45,24 @@ class GPComponentType(Enum): LIKELIHOOD = "LIKELIHOOD" """Gaussian process likelihood.""" + CRITERION = "CRITERION" + """Gaussian process fitting criterion.""" + def get_types(self) -> tuple[type, ...]: """Get the accepted BayBE and GPyTorch types for this component.""" - types = [] + types: list[type[GPComponent]] = [] # Add BayBE type if applicable if self is GPComponentType.KERNEL: from baybe.kernels.base import Kernel types.append(Kernel) + elif self is GPComponentType.CRITERION: + from baybe.surrogates.gaussian_process.components.fit_criterion import ( + FitCriterion, + ) + + types.append(FitCriterion) # Add GPyTorch type if available if sys.modules.get("gpytorch") is not None: @@ -85,7 +95,7 @@ def _is_gpytorch_component_class(obj: Any, /) -> bool: def _validate_component(instance: Any, attribute: Attribute, value: Any) -> None: """Validate that an object is a BayBE or a GPyTorch GP component.""" - if isinstance(value, Kernel) or _is_gpytorch_component_class(type(value)): + if isinstance(value, BayBEGPComponent) or _is_gpytorch_component_class(type(value)): return raise TypeError( diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 8b6ddff280..57c7e8486a 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -102,7 +102,7 @@ def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None: @override def __call__( self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor - ) -> Kernel: + ) -> Kernel | GPyTorchKernel: """Construct the kernel, validating parameter kinds before construction.""" if self.parameter_selector is not None: params = [p for p in searchspace.parameters if self.parameter_selector(p)] @@ -115,7 +115,7 @@ def __call__( @abstractmethod def _make( self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor - ) -> Kernel: + ) -> Kernel | GPyTorchKernel: """Construct the kernel.""" @@ -171,10 +171,43 @@ def _default_task_kernel_factory(self) -> KernelFactoryProtocol: def __call__( self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor ) -> Kernel: + if searchspace.task_idx is None: + raise IncompatibleSearchSpaceError( + f"'{type(self).__name__}' can only be used with a searchspace that " + f"contains a '{TaskParameter.__name__}'." + ) + base_kernel = self.base_kernel_factory(searchspace, train_x, train_y) task_kernel = self.task_kernel_factory(searchspace, train_x, train_y) if isinstance(base_kernel, Kernel): base_kernel = base_kernel.to_gpytorch(searchspace) if isinstance(task_kernel, Kernel): task_kernel = task_kernel.to_gpytorch(searchspace) + + # Ensure correct partitioning between base and task kernels active dimensions + all_idcs = set(range(len(searchspace.comp_rep_columns))) + allowed_task_idcs = {searchspace.task_idx} + allowed_base_idcs = all_idcs - allowed_task_idcs + base_idcs = ( + set(dims) + if (dims := base_kernel.active_dims.tolist()) is not None + else None + ) + task_idcs = ( + set(dims) + if (dims := task_kernel.active_dims.tolist()) is not None + else None + ) + + if base_idcs is not None and (base_idcs > allowed_base_idcs): + raise ValueError( + f"The base kernel's 'active_dims' {base_idcs} must be a subset of " + f"the non-task indices {allowed_base_idcs}." + ) + if task_idcs != allowed_task_idcs: + raise ValueError( + f"The task kernel's 'active_dims' {task_idcs} does not match " + f"the task index {allowed_task_idcs}." + ) + return base_kernel * task_kernel diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 2b1a3f361e..f8a979a1a3 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -19,6 +19,10 @@ from baybe.parameters.categorical import TaskParameter from baybe.searchspace.core import SearchSpace from baybe.surrogates.base import Surrogate +from baybe.surrogates.gaussian_process.components.fit_criterion import ( + FitCriterion, + FitCriterionFactoryProtocol, +) from baybe.surrogates.gaussian_process.components.generic import ( GPComponentType, to_component_factory, @@ -35,6 +39,7 @@ GaussianProcessPreset, ) from baybe.surrogates.gaussian_process.presets.baybe import ( + BayBEFitCriterionFactory, BayBEKernelFactory, BayBELikelihoodFactory, BayBEMeanFactory, @@ -178,6 +183,21 @@ class GaussianProcessSurrogate(Surrogate): * :class:`gpytorch.likelihoods.Likelihood` """ + criterion_factory: FitCriterionFactoryProtocol = field( + alias="criterion_or_factory", + factory=BayBEFitCriterionFactory, + converter=partial( # type: ignore[misc] + to_component_factory, component_type=GPComponentType.CRITERION + ), + validator=is_callable(), + ) + """The fitting criterion for Gaussian process hyperparameter optimization. + + Accepts: + * :class:`.components.fit_criterion.FitCriterion` + * :class:`.components.fit_criterion.FitCriterionFactoryProtocol` + """ + # TODO: type should be Optional[botorch.models.SingleTaskGP] but is currently # omitted due to: https://github.com/python-attrs/cattrs/issues/531 _model = field(init=False, default=None, eq=False) @@ -195,6 +215,7 @@ def from_preset( likelihood_or_factory: LikelihoodFactoryProtocol | GPyTorchLikelihood | None = None, + criterion_or_factory: FitCriterion | FitCriterionFactoryProtocol | None = None, ) -> Self: """Create a Gaussian process surrogate from one of the defined presets.""" preset = GaussianProcessPreset(preset) @@ -204,13 +225,18 @@ def from_preset( ) module = importlib.import_module(module_name) - kernel = kernel_or_factory or getattr(module, "PresetKernelFactory")() - mean = mean_or_factory or getattr(module, "PresetMeanFactory")() - likelihood = ( - likelihood_or_factory or getattr(module, "PresetLikelihoodFactory")() + kernel = kernel_or_factory or getattr(module, "PRESET_KERNEL_FACTORY") + mean = mean_or_factory or getattr(module, "PRESET_MEAN_FACTORY") + likelihood = likelihood_or_factory or getattr( + module, "PRESET_LIKELIHOOD_FACTORY" + ) + criterion = criterion_or_factory or getattr( + module, "PRESET_FIT_CRITERION_FACTORY" ) - return cls(kernel, mean, likelihood) + gp = cls(kernel, mean, likelihood, criterion) + gp._custom_kernel = False # preset are first-party features + return gp @override def to_botorch(self) -> GPyTorchModel: @@ -237,7 +263,6 @@ def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior: @override def _fit(self, train_x: Tensor, train_y: Tensor) -> None: import botorch - import gpytorch from botorch.models.transforms import Normalize, Standardize assert self._searchspace is not None # provided by base class @@ -281,6 +306,9 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: ### Likelihood likelihood = self.likelihood_factory(context.searchspace, train_x, train_y) + ### Criterion + criterion = self.criterion_factory(context.searchspace, train_x, train_y) + ### Model construction and fitting self._model = botorch.models.SingleTaskGP( train_x, @@ -291,18 +319,7 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: covar_module=kernel, likelihood=likelihood, ) - - # TODO: This is still a temporary workaround to avoid overfitting seen in - # low-dimensional TL cases. More robust settings are being researched. - if context.n_task_dimensions > 0: - mll = gpytorch.mlls.LeaveOneOutPseudoLikelihood( - self._model.likelihood, self._model - ) - else: - mll = gpytorch.ExactMarginalLogLikelihood( - self._model.likelihood, self._model - ) - + mll = criterion.to_gpytorch(self._model.likelihood, self._model) botorch.fit.fit_gpytorch_mll(mll) @override @@ -311,6 +328,9 @@ def __str__(self) -> str: to_string("Kernel factory", self.kernel_factory, single_line=True), to_string("Mean factory", self.mean_factory, single_line=True), to_string("Likelihood factory", self.likelihood_factory, single_line=True), + to_string( + "Fit criterion factory", self.criterion_factory, single_line=True + ), ] return to_string(super().__str__(), *fields) diff --git a/baybe/surrogates/gaussian_process/presets/__init__.py b/baybe/surrogates/gaussian_process/presets/__init__.py index 434fbf560f..189733131b 100644 --- a/baybe/surrogates/gaussian_process/presets/__init__.py +++ b/baybe/surrogates/gaussian_process/presets/__init__.py @@ -1,20 +1,35 @@ """Gaussian process surrogate presets.""" +# Criterion +from baybe.surrogates.gaussian_process.components.fit_criterion import FitCriterion + # Default preset from baybe.surrogates.gaussian_process.presets.baybe import ( + BayBEFitCriterionFactory, BayBEKernelFactory, BayBELikelihoodFactory, BayBEMeanFactory, ) +# BoTorch preset +from baybe.surrogates.gaussian_process.presets.botorch import ( + BotorchKernelFactory, + BotorchLikelihoodFactory, + BotorchMeanFactory, +) + # Chen preset -from baybe.surrogates.gaussian_process.presets.chen import CHENKernelFactory +from baybe.surrogates.gaussian_process.presets.chen import ( + CHENFitCriterionFactory, + CHENKernelFactory, +) # Core from baybe.surrogates.gaussian_process.presets.core import GaussianProcessPreset # EDBO preset from baybe.surrogates.gaussian_process.presets.edbo import ( + EDBOFitCriterionFactory, EDBOKernelFactory, EDBOLikelihoodFactory, EDBOMeanFactory, @@ -22,6 +37,7 @@ # Smoothed EDBO preset from baybe.surrogates.gaussian_process.presets.edbo_smoothed import ( + SmoothedEDBOFitCriterionFactory, SmoothedEDBOKernelFactory, SmoothedEDBOLikelihoodFactory, SmoothedEDBOMeanFactory, @@ -29,18 +45,27 @@ __all__ = [ # Core + "FitCriterion", "GaussianProcessPreset", # Default BayBE preset + "BayBEFitCriterionFactory", "BayBEKernelFactory", "BayBELikelihoodFactory", "BayBEMeanFactory", + # BoTorch preset + "BotorchKernelFactory", + "BotorchLikelihoodFactory", + "BotorchMeanFactory", # Chen preset + "CHENFitCriterionFactory", "CHENKernelFactory", # EDBO preset + "EDBOFitCriterionFactory", "EDBOKernelFactory", "EDBOLikelihoodFactory", "EDBOMeanFactory", # Smoothed EDBO preset + "SmoothedEDBOFitCriterionFactory", "SmoothedEDBOKernelFactory", "SmoothedEDBOLikelihoodFactory", "SmoothedEDBOMeanFactory", diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py index 690fc318cd..355fe766ae 100644 --- a/baybe/surrogates/gaussian_process/presets/baybe.py +++ b/baybe/surrogates/gaussian_process/presets/baybe.py @@ -17,6 +17,10 @@ to_parameter_selector, ) from baybe.searchspace.core import SearchSpace +from baybe.surrogates.gaussian_process.components.fit_criterion import ( + FitCriterion, + FitCriterionFactoryProtocol, +) from baybe.surrogates.gaussian_process.components.kernel import _PureKernelFactory from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory from baybe.surrogates.gaussian_process.presets.edbo_smoothed import ( @@ -85,7 +89,24 @@ def _make( BayBELikelihoodFactory = SmoothedEDBOLikelihoodFactory """The factory providing the default likelihood for Gaussian process surrogates.""" -# Aliases for generic preset imports -PresetKernelFactory = BayBEKernelFactory -PresetMeanFactory = BayBEMeanFactory -PresetLikelihoodFactory = BayBELikelihoodFactory + +@define +class BayBEFitCriterionFactory(FitCriterionFactoryProtocol): + """The factory providing the default fitting criterion for Gaussian process surrogates.""" # noqa: E501 + + @override + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> FitCriterion: + return ( + FitCriterion.MARGINAL_LOG_LIKELIHOOD + if searchspace.task_idx is None + else FitCriterion.LEAVE_ONE_OUT_PSEUDOLIKELIHOOD + ) + + +# Preset defaults +PRESET_KERNEL_FACTORY = BayBEKernelFactory() +PRESET_MEAN_FACTORY = BayBEMeanFactory() +PRESET_LIKELIHOOD_FACTORY = BayBELikelihoodFactory() +PRESET_FIT_CRITERION_FACTORY = BayBEFitCriterionFactory() diff --git a/baybe/surrogates/gaussian_process/presets/botorch.py b/baybe/surrogates/gaussian_process/presets/botorch.py new file mode 100644 index 0000000000..8af5db4d5d --- /dev/null +++ b/baybe/surrogates/gaussian_process/presets/botorch.py @@ -0,0 +1,149 @@ +"""BoTorch preset for Gaussian process surrogates.""" + +from __future__ import annotations + +import gc +from itertools import chain +from typing import TYPE_CHECKING, ClassVar + +from attrs import define +from typing_extensions import override + +from baybe.kernels.base import Kernel +from baybe.parameters.enum import _ParameterKind +from baybe.searchspace.core import SearchSpace +from baybe.surrogates.gaussian_process.components import LikelihoodFactoryProtocol +from baybe.surrogates.gaussian_process.components._gpytorch import ( + make_botorch_multitask_likelihood, +) +from baybe.surrogates.gaussian_process.components.kernel import ( + ICMKernelFactory, + _PureKernelFactory, +) +from baybe.surrogates.gaussian_process.components.mean import MeanFactoryProtocol +from baybe.surrogates.gaussian_process.presets.baybe import BayBEFitCriterionFactory + +if TYPE_CHECKING: + from gpytorch.kernels import Kernel as GPyTorchKernel + from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood + from gpytorch.means import Mean as GPyTorchMean + from torch import Tensor + + +@define +class BotorchKernelFactory(_PureKernelFactory): + """A factory providing BoTorch kernels.""" + + _uses_parameter_names: ClassVar[bool] = True + # See base class. + + _supported_parameter_kinds: ClassVar[_ParameterKind] = ( + _ParameterKind.REGULAR | _ParameterKind.TASK + ) + # See base class. + + @override + def _make( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel | GPyTorchKernel: + from botorch.models.kernels.positive_index import PositiveIndexKernel + from botorch.models.utils.gpytorch_modules import ( + get_covar_module_with_dim_scaled_prior, + ) + + parameter_names = self.get_parameter_names(searchspace) + + # Resolve parameter names to active dimension indices + active_dims: list[int] | None + if parameter_names is not None: + active_dims = list( + chain.from_iterable( + searchspace.get_comp_rep_parameter_indices(name) + for name in parameter_names + ) + ) + ard_num_dims = len(active_dims) + else: + active_dims = None + ard_num_dims = len(searchspace.comp_rep_columns) + + # Determine if the selected parameters include a task parameter + task_idx = searchspace.task_idx + is_multitask = task_idx is not None and ( + active_dims is None or task_idx in active_dims + ) + + if not is_multitask: + return get_covar_module_with_dim_scaled_prior( + ard_num_dims=ard_num_dims, active_dims=active_dims + ) + + assert task_idx is not None + base_idcs = [ + idx + for idx in (active_dims or range(len(searchspace.comp_rep_columns))) + if idx != task_idx + ] + base = get_covar_module_with_dim_scaled_prior( + ard_num_dims=len(base_idcs), active_dims=base_idcs + ) + index_kernel = PositiveIndexKernel( + num_tasks=searchspace.n_tasks, + rank=searchspace.n_tasks, + active_dims=[task_idx], + ) + return ICMKernelFactory(base, index_kernel)(searchspace, train_x, train_y) + + +class BotorchMeanFactory(MeanFactoryProtocol): + """A factory providing BoTorch mean functions.""" + + @override + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> GPyTorchMean: + from gpytorch.means import ConstantMean + + from baybe.surrogates.gaussian_process.components._gpytorch import ( + HadamardConstantMean, + ) + + if searchspace.n_tasks == 1: + return ConstantMean() + + assert searchspace.task_idx is not None + return HadamardConstantMean( + ConstantMean(), searchspace.n_tasks, searchspace.task_idx + ) + + +class BotorchLikelihoodFactory(LikelihoodFactoryProtocol): + """A factory providing BoTorch likelihoods.""" + + @override + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> GPyTorchLikelihood: + from botorch.models.utils.gpytorch_modules import ( + get_gaussian_likelihood_with_lognormal_prior, + ) + + if searchspace.n_tasks == 1: + return get_gaussian_likelihood_with_lognormal_prior() + + assert searchspace.task_idx is not None + return make_botorch_multitask_likelihood( + num_tasks=searchspace.n_tasks, task_feature=searchspace.task_idx + ) + + +# Collect leftover original slotted classes processed by `attrs.define` +gc.collect() + +# Aliases for generic preset imports +PRESET_KERNEL_FACTORY = BotorchKernelFactory() +PRESET_MEAN_FACTORY = BotorchMeanFactory() +PRESET_LIKELIHOOD_FACTORY = BotorchLikelihoodFactory() + +# Botorch dictates no speecific criterion, so we fill the preset with our default +PRESET_FIT_CRITERION_FACTORY = BayBEFitCriterionFactory() diff --git a/baybe/surrogates/gaussian_process/presets/chen.py b/baybe/surrogates/gaussian_process/presets/chen.py index 5e90c4aa5c..3bb6054606 100644 --- a/baybe/surrogates/gaussian_process/presets/chen.py +++ b/baybe/surrogates/gaussian_process/presets/chen.py @@ -18,6 +18,10 @@ to_parameter_selector, ) from baybe.priors.basic import GammaPrior +from baybe.surrogates.gaussian_process.components.fit_criterion import ( + FitCriterion, + PlainFitCriterionFactory, +) from baybe.surrogates.gaussian_process.components.kernel import ( _PureKernelFactory, ) @@ -68,10 +72,14 @@ def _make( ) +CHENFitCriterionFactory = PlainFitCriterionFactory(FitCriterion.MARGINAL_LOG_LIKELIHOOD) +"""A factory providing fitting criteria for the CHEN preset.""" + # Collect leftover original slotted classes processed by `attrs.define` gc.collect() -# Aliases for generic preset imports -PresetKernelFactory = CHENKernelFactory -PresetMeanFactory = LazyConstantMeanFactory -PresetLikelihoodFactory = LazyGaussianLikelihoodFactory +# Preset defaults +PRESET_KERNEL_FACTORY = CHENKernelFactory() +PRESET_MEAN_FACTORY = LazyConstantMeanFactory() +PRESET_LIKELIHOOD_FACTORY = LazyGaussianLikelihoodFactory() +PRESET_FIT_CRITERION_FACTORY = CHENFitCriterionFactory diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py index 5347cf85e5..70d5ac9a7c 100644 --- a/baybe/surrogates/gaussian_process/presets/core.py +++ b/baybe/surrogates/gaussian_process/presets/core.py @@ -11,6 +11,9 @@ class GaussianProcessPreset(Enum): BAYBE = "BAYBE" """The default BayBE settings of the Gaussian process surrogate class.""" + BOTORCH = "BOTORCH" + """The BoTorch settings.""" + CHEN = "CHEN" """The adaptive kernel hyperprior settings proposed by :cite:p:`Chen2026`.""" diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py index 539e2ef72c..d82d6f9d26 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo.py +++ b/baybe/surrogates/gaussian_process/presets/edbo.py @@ -21,6 +21,10 @@ from baybe.parameters.substance import SubstanceParameter from baybe.priors.basic import GammaPrior from baybe.searchspace.discrete import SubspaceDiscrete +from baybe.surrogates.gaussian_process.components.fit_criterion import ( + FitCriterion, + PlainFitCriterionFactory, +) from baybe.surrogates.gaussian_process.components.kernel import ( _PureKernelFactory, ) @@ -176,9 +180,13 @@ def __call__( # Collect leftover original slotted classes processed by `attrs.define` +EDBOFitCriterionFactory = PlainFitCriterionFactory(FitCriterion.MARGINAL_LOG_LIKELIHOOD) +"""A factory providing fitting criteria for the EDBO preset.""" + gc.collect() -# Aliases for generic preset imports -PresetKernelFactory = EDBOKernelFactory -PresetMeanFactory = EDBOMeanFactory -PresetLikelihoodFactory = EDBOLikelihoodFactory +# Preset defaults +PRESET_KERNEL_FACTORY = EDBOKernelFactory() +PRESET_MEAN_FACTORY = EDBOMeanFactory() +PRESET_LIKELIHOOD_FACTORY = EDBOLikelihoodFactory() +PRESET_FIT_CRITERION_FACTORY = EDBOFitCriterionFactory diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py index 904f711f31..a2eda692d1 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py +++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py @@ -18,6 +18,10 @@ to_parameter_selector, ) from baybe.priors.basic import GammaPrior +from baybe.surrogates.gaussian_process.components.fit_criterion import ( + FitCriterion, + PlainFitCriterionFactory, +) from baybe.surrogates.gaussian_process.components.kernel import ( _PureKernelFactory, ) @@ -126,10 +130,16 @@ def __call__( return likelihood +SmoothedEDBOFitCriterionFactory = PlainFitCriterionFactory( + FitCriterion.MARGINAL_LOG_LIKELIHOOD +) +"""A factory providing fitting criteria for the smoothed EDBO preset.""" + # Collect leftover original slotted classes processed by `attrs.define` gc.collect() -# Aliases for generic preset imports -PresetKernelFactory = SmoothedEDBOKernelFactory -PresetMeanFactory = SmoothedEDBOMeanFactory -PresetLikelihoodFactory = SmoothedEDBOLikelihoodFactory +# Preset defaults +PRESET_KERNEL_FACTORY = SmoothedEDBOKernelFactory() +PRESET_MEAN_FACTORY = SmoothedEDBOMeanFactory() +PRESET_LIKELIHOOD_FACTORY = SmoothedEDBOLikelihoodFactory() +PRESET_FIT_CRITERION_FACTORY = SmoothedEDBOFitCriterionFactory diff --git a/streamlit/surrogate_models.py b/streamlit/surrogate_models.py index 83415b91b9..c88d441164 100644 --- a/streamlit/surrogate_models.py +++ b/streamlit/surrogate_models.py @@ -19,11 +19,12 @@ from baybe.acquisition import qLogExpectedImprovement from baybe.acquisition.base import AcquisitionFunction from baybe.exceptions import IncompatibleSurrogateError -from baybe.parameters import NumericalDiscreteParameter +from baybe.parameters import NumericalDiscreteParameter, TaskParameter from baybe.recommenders import BotorchRecommender from baybe.searchspace import SearchSpace from baybe.surrogates import CustomONNXSurrogate, GaussianProcessSurrogate from baybe.surrogates.base import Surrogate +from baybe.surrogates.gaussian_process.presets import GaussianProcessPreset from baybe.targets import NumericalTarget from baybe.utils.basic import get_subclasses @@ -90,13 +91,36 @@ def main(): } acquisition_function_names = list(acquisition_function_classes.keys()) - # Streamlit simulation parameters + # >>>>> Sidebar options >>>>> + # Domain st.sidebar.markdown("# Domain") + st_enable_multitask = st.sidebar.toggle("Multi-task") + st_n_tasks = 2 if st_enable_multitask else 1 st_random_seed = int(st.sidebar.number_input("Random seed", value=1337)) st_function_name = st.sidebar.selectbox( "Test function", list(test_functions.keys()) ) st_minimize = st.sidebar.checkbox("Minimize") + + # Training data + st.sidebar.markdown("---") + st.sidebar.markdown("# Training Data") + st_n_training_points = st.sidebar.slider( + "Number of training points", + 0 if st_enable_multitask else 1, + 20, + 0 if st_enable_multitask else 5, + ) + if st_enable_multitask: + st_n_training_points_other = st.sidebar.slider( + "Number of source training points", 0, 20, 5 + ) + st_offtask_scale = st.sidebar.slider("Source scale factor", -20.0, 20.0, 1.0) + st_offtask_offset_factor = st.sidebar.slider( + "Source offset", -100.0, 100.0, 0.0 + ) + + # Model st.sidebar.markdown("---") st.sidebar.markdown("# Model") st_surrogate_name = st.sidebar.selectbox( @@ -104,13 +128,26 @@ def main(): surrogate_model_names, surrogate_model_names.index(GaussianProcessSurrogate.__name__), ) + + st_gp_preset = None + st_transfer_learning = False + if st_surrogate_name == GaussianProcessSurrogate.__name__: + preset_names = [preset.value for preset in GaussianProcessPreset] + st_gp_preset = st.sidebar.selectbox( + "GP Preset", + preset_names, + index=preset_names.index(GaussianProcessPreset.BAYBE.value), + ) + if st_enable_multitask: + st_transfer_learning = st.sidebar.checkbox("Transfer learning", value=True) st_acqf_name = st.sidebar.selectbox( "Acquisition function", acquisition_function_names, acquisition_function_names.index(qLogExpectedImprovement.__name__), ) - st_n_training_points = st.sidebar.slider("Number of training points", 1, 20, 5) st_n_recommendations = st.sidebar.slider("Number of recommendations", 1, 20, 5) + + # Surrogate consistency validation st.sidebar.markdown("---") st.sidebar.markdown("# Validation") st.sidebar.markdown( @@ -127,6 +164,10 @@ def main(): ) st_function_amplitude = st.sidebar.slider("Function amplitude", 1.0, 100.0, 1.0) st_function_bias = st.sidebar.slider("Function bias", -100.0, 100.0, 0.0) + # <<<<< Sidebar options <<<<< + + # Derived settings + st_use_separate_gps = st_n_tasks > 1 and not st_transfer_learning # Set the chosen random seed active_settings.random_seed = st_random_seed @@ -140,68 +181,184 @@ def main(): bias=st_function_bias, ) - # Create the training data - train_x = np.random.uniform( - st_lower_parameter_limit, st_upper_parameter_limit, st_n_training_points - ) - train_y = fun(train_x) - measurements = pd.DataFrame({"x": train_x, "y": train_y}) + # Generate task-specific transforms (scale and offset for each task) + task_names = ["target", "source"] if st_n_tasks > 1 else ["target"] + task_transforms = {} + for task_idx in range(st_n_tasks): + task_name = task_names[task_idx] + if task_idx == 0: + # Target: use original function values + task_transforms[task_name] = {"scale": 1.0, "offset": 0.0} + else: + # Source: use user-specified scale and offset + scale = st_offtask_scale + offset = st_offtask_offset_factor * st_function_amplitude + task_transforms[task_name] = {"scale": scale, "offset": offset} + + # Create training data + measurements_list = [] + for task_idx in range(st_n_tasks): + task_name = task_names[task_idx] + transform = task_transforms[task_name] + n_points = st_n_training_points if task_idx == 0 else st_n_training_points_other + train_x = np.random.uniform( + st_lower_parameter_limit, st_upper_parameter_limit, n_points + ) + task_measurements = pd.DataFrame( + { + "x": train_x, + "task": task_name, + "y": fun(train_x) * transform["scale"] + transform["offset"], + } + ) + measurements_list.append(task_measurements) + measurements = pd.concat(measurements_list, ignore_index=True) # Create the plotting grid and corresponding target values test_x = np.linspace( st_lower_parameter_limit, st_upper_parameter_limit, N_PARAMETER_VALUES ) - test_y = fun(test_x) - candidates = pd.DataFrame({"x": test_x, "y": test_y}) + candidates_list = [] + test_ys = {} + for task_idx in range(st_n_tasks): + task_name = task_names[task_idx] + transform = task_transforms[task_name] + test_ys[task_name] = fun(test_x) * transform["scale"] + transform["offset"] + task_candidates = pd.DataFrame( + {"x": test_x, "task": task_name, "y": test_ys[task_name]} + ) + candidates_list.append(task_candidates) + candidates = pd.concat(candidates_list, ignore_index=True) # Create the searchspace and objective - parameter = NumericalDiscreteParameter( - name="x", - values=np.linspace( - st_lower_parameter_limit, st_upper_parameter_limit, N_PARAMETER_VALUES - ), - ) - searchspace = SearchSpace.from_product(parameters=[parameter]) + parameters = [ + NumericalDiscreteParameter( + name="x", + values=np.linspace( + st_lower_parameter_limit, st_upper_parameter_limit, N_PARAMETER_VALUES + ), + ) + ] + if st_transfer_learning: + parameters.append( + TaskParameter( + name="task", + values=task_names, + active_values=["target"], + ) + ) + searchspace = SearchSpace.from_product(parameters=parameters) objective = NumericalTarget(name="y", minimize=st_minimize).to_objective() - # Create the acquisition function and the recommender + # Create the acquisition function acqf_cls = acquisition_function_classes[st_acqf_name] try: acqf = acqf_cls(maximize=not st_minimize) except TypeError: acqf = acqf_cls() - recommender = BotorchRecommender( - surrogate_model=surrogate_model_classes[st_surrogate_name](), - acquisition_function=acqf, - ) - # Get the recommendations and extract the posterior mean / standard deviation - try: - recommendations = recommender.recommend( - st_n_recommendations, searchspace, objective, measurements - ) - except IncompatibleSurrogateError: - st.error( - f"You requested {st_n_recommendations} recommendations but the selected " - f"surrogate class does not support recommending more than one candidate " - f"at a time." + def make_surrogate(): + if st_surrogate_name == GaussianProcessSurrogate.__name__: + assert st_gp_preset is not None + return GaussianProcessSurrogate.from_preset( + preset=GaussianProcessPreset[st_gp_preset] + ) + return surrogate_model_classes[st_surrogate_name]() + + if st_use_separate_gps: + # One independent GP per task, each trained without the task column + stats_by_task = {} + for task_name in task_names: + task_meas = measurements[measurements["task"] == task_name][ + ["x", "y"] + ].reset_index(drop=True) + task_recommender = BotorchRecommender( + surrogate_model=make_surrogate(), + acquisition_function=acqf, + ) + if task_name == "target": + try: + recommendations = task_recommender.recommend( + st_n_recommendations, searchspace, objective, task_meas + ) + except IncompatibleSurrogateError: + st.error( + f"You requested {st_n_recommendations} recommendations but " + f"the selected surrogate class does not support recommending " + f"more than one candidate at a time." + ) + st.stop() + task_surrogate = task_recommender.get_surrogate( + searchspace, objective, task_meas + ) + stats_by_task[task_name] = task_surrogate.posterior_stats( + pd.DataFrame({"x": test_x}) + ) + else: + # Single recommender (single task, or multi-task with transfer learning) + recommender = BotorchRecommender( + surrogate_model=make_surrogate(), + acquisition_function=acqf, ) - st.stop() - surrogate = recommender.get_surrogate(searchspace, objective, measurements) - stats = surrogate.posterior_stats(candidates) - mean, std = stats["y_mean"], stats["y_std"] + try: + recommendations = recommender.recommend( + st_n_recommendations, searchspace, objective, measurements + ) + except IncompatibleSurrogateError: + st.error( + f"You requested {st_n_recommendations} recommendations but the " + f"selected surrogate class does not support recommending more than " + f"one candidate at a time." + ) + st.stop() + surrogate = recommender.get_surrogate(searchspace, objective, measurements) + stats = surrogate.posterior_stats(candidates) # Visualize the test function, training points, model predictions, recommendations - fig = plt.figure() - plt.plot(test_x, test_y, color="tab:blue", label="Test function") - plt.plot(train_x, train_y, "o", color="tab:blue") - plt.plot(test_x, mean, color="tab:red", label="Surrogate model") - plt.fill_between(test_x, mean - std, mean + std, alpha=0.2, color="tab:red") - plt.vlines( - recommendations, *plt.gca().get_ylim(), color="k", label="Recommendations" - ) - plt.legend() - st.pyplot(fig) + if st_n_tasks > 1: + cols = st.columns(st_n_tasks) + + for task_idx in range(st_n_tasks): + task_name = task_names[task_idx] + task_mask = candidates["task"] == task_name if st_n_tasks > 1 else slice(None) + + if st_use_separate_gps: + task_stats = stats_by_task[task_name] + mean = task_stats["y_mean"].values + std = task_stats["y_std"].values + elif st_n_tasks > 1: + mean = stats["y_mean"][task_mask].values + std = stats["y_std"][task_mask].values + else: + mean = stats["y_mean"].values + std = stats["y_std"].values + + test_y = test_ys[task_name] + train_mask = ( + measurements["task"] == task_name if st_n_tasks > 1 else slice(None) + ) + train_y = measurements[train_mask]["y"].values + task_train_x = measurements[train_mask]["x"].values + + fig = plt.figure() + plt.plot(test_x, test_y, color="tab:blue", label="Test function") + plt.plot(task_train_x, train_y, "o", color="tab:blue") + plt.plot(test_x, mean, color="tab:red", label="Surrogate model") + plt.fill_between(test_x, mean - std, mean + std, alpha=0.2, color="tab:red") + if task_name == "target": + plt.vlines( + recommendations["x"] if st_n_tasks > 1 else recommendations, + *plt.gca().get_ylim(), + color="k", + label="Recommendations", + ) + plt.legend() + if st_n_tasks > 1: + plt.title(task_name.capitalize()) + with cols[task_idx]: + st.pyplot(fig) + else: + st.pyplot(fig) if __name__ == "__main__": diff --git a/tests/test_gp.py b/tests/test_gp.py index 30ca69bc0e..423790f060 100644 --- a/tests/test_gp.py +++ b/tests/test_gp.py @@ -1,6 +1,11 @@ """Tests for the Gaussian Process surrogate.""" +import pandas as pd import pytest +import torch +from botorch.fit import fit_gpytorch_mll +from botorch.models import MultiTaskGP, SingleTaskGP +from botorch.models.transforms import Normalize, Standardize from gpytorch.kernels import MaternKernel as GPyTorchMaternKernel from gpytorch.kernels import RBFKernel as GPyTorchRBFKernel from gpytorch.kernels import ScaleKernel as GPyTorchScaleKernel @@ -8,22 +13,35 @@ from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood from gpytorch.means import ConstantMean from gpytorch.means import Mean as GPyTorchMean +from gpytorch.mlls import ExactMarginalLogLikelihood from pandas.testing import assert_frame_equal from pytest import param +from baybe import active_settings from baybe.kernels.basic import MaternKernel, RBFKernel from baybe.kernels.composite import ScaleKernel +from baybe.parameters.categorical import TaskParameter from baybe.parameters.numerical import NumericalContinuousParameter +from baybe.searchspace.core import SearchSpace +from baybe.surrogates.gaussian_process.components.fit_criterion import FitCriterion from baybe.surrogates.gaussian_process.components.generic import PlainGPComponentFactory from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate from baybe.surrogates.gaussian_process.presets import GaussianProcessPreset from baybe.targets.numerical import NumericalTarget -from baybe.utils.dataframe import create_fake_input +from baybe.utils.dataframe import create_fake_input, to_tensor searchspace = NumericalContinuousParameter("p", (0, 1)).to_searchspace() +searchspace_mt = SearchSpace.from_product( + [ + NumericalContinuousParameter("p", (0, 1)), + TaskParameter("task", ["a", "b", "c"]), + ] +) objective = NumericalTarget("t").to_objective() measurements = create_fake_input(searchspace.parameters, objective.targets, n_rows=100) - +measurements_mt = create_fake_input( + searchspace_mt.parameters, objective.targets, n_rows=100 +) baybe_kernel = ScaleKernel(MaternKernel() + RBFKernel()) gpytorch_kernel = GPyTorchScaleKernel(GPyTorchMaternKernel() + GPyTorchRBFKernel()) @@ -36,6 +54,52 @@ def _dummy_likelihood_factory(*args, **kwargs) -> GPyTorchLikelihood: return GaussianLikelihood() +def _posterior_stats_botorch( + searchspace: SearchSpace, measurements: pd.DataFrame +) -> pd.DataFrame: + """The essential BoTorch steps to produce posterior estimates.""" + train_X = to_tensor(searchspace.transform(measurements, allow_extra=True)) + train_Y = to_tensor(objective.transform(measurements, allow_extra=True)) + + # >>>>> Code adapted from BoTorch landing page: https://botorch.org/ >>>>> + # NOTE: We normalize according to the searchspace bounds to ensure consistency with + # the BayBE GP implementation. + if searchspace.n_tasks == 1: + gp = SingleTaskGP( + train_X=train_X, + train_Y=train_Y, + input_transform=Normalize( + d=len(searchspace.comp_rep_columns), + bounds=to_tensor(searchspace.scaling_bounds), + ), + outcome_transform=Standardize(m=1), + ) + else: + assert searchspace.task_idx is not None + non_task_idcs = [ + i for i in range(train_X.shape[-1]) if i != searchspace.task_idx + ] + gp = MultiTaskGP( + train_X=train_X, + train_Y=train_Y, + task_feature=searchspace.task_idx, + input_transform=Normalize( + d=len(searchspace.comp_rep_columns), + indices=non_task_idcs, + bounds=to_tensor(searchspace.scaling_bounds), + ), + ) + mll = ExactMarginalLogLikelihood(gp.likelihood, gp) + fit_gpytorch_mll(mll) + # <<<<<<<<<< + + with torch.no_grad(): + posterior = gp.posterior(train_X) + mean = posterior.mean + std = posterior.variance.sqrt() + return pd.DataFrame({"t_mean": mean.numpy().ravel(), "t_std": std.numpy().ravel()}) + + @pytest.mark.parametrize( ("component_1", "component_2"), [ @@ -89,24 +153,32 @@ def test_presets(preset: GaussianProcessPreset): kernel = GPyTorchMaternKernel() mean = ConstantMean() likelihood = GaussianLikelihood() + criterion = FitCriterion.LEAVE_ONE_OUT_PSEUDOLIKELIHOOD # Works without overrides ... - GaussianProcessSurrogate.from_preset(preset) + gp1 = GaussianProcessSurrogate.from_preset(preset) # ... and with overrides - gp = GaussianProcessSurrogate.from_preset( + gp2 = GaussianProcessSurrogate.from_preset( preset, kernel_or_factory=kernel, mean_or_factory=mean, likelihood_or_factory=likelihood, + criterion_or_factory=criterion, ) - assert isinstance(gp.kernel_factory, PlainGPComponentFactory) - assert gp.kernel_factory.component is kernel - assert isinstance(gp.mean_factory, PlainGPComponentFactory) - assert gp.mean_factory.component is mean - assert isinstance(gp.likelihood_factory, PlainGPComponentFactory) - assert gp.likelihood_factory.component is likelihood - gp.fit(searchspace, objective, measurements) + + # Check that the overrides were applied correctly + assert isinstance(gp2.kernel_factory, PlainGPComponentFactory) + assert gp2.kernel_factory.component is kernel + assert isinstance(gp2.mean_factory, PlainGPComponentFactory) + assert gp2.mean_factory.component is mean + assert isinstance(gp2.likelihood_factory, PlainGPComponentFactory) + assert gp2.likelihood_factory.component is likelihood + assert isinstance(gp2.criterion_factory, PlainGPComponentFactory) + assert gp2.criterion_factory.component == criterion + assert gp2.criterion_factory != gp1.criterion_factory + + gp2.fit(searchspace, objective, measurements) def test_invalid_components(): @@ -116,4 +188,29 @@ def test_invalid_components(): with pytest.raises(TypeError, match="Component must be one of"): GaussianProcessSurrogate(mean_or_factory=GaussianLikelihood()) with pytest.raises(TypeError, match="Component must be one of"): - GaussianProcessSurrogate(likelihood_or_factory=MaternKernel()) + GaussianProcessSurrogate( + likelihood_or_factory=FitCriterion.LEAVE_ONE_OUT_PSEUDOLIKELIHOOD + ) + with pytest.raises(TypeError, match="Component must be one of"): + GaussianProcessSurrogate(criterion_or_factory=MaternKernel()) + + +@pytest.mark.parametrize("multitask", [False, True], ids=["single-task", "multi-task"]) +def test_botorch_preset(multitask: bool): + """The BoTorch preset exactly mimics BoTorch's behavior.""" + if multitask: + sp = searchspace_mt + data = measurements_mt + else: + sp = searchspace + data = measurements + + active_settings.random_seed = 1337 + gp = GaussianProcessSurrogate.from_preset("BOTORCH") + gp.fit(sp, objective, data) + posterior1 = gp.posterior_stats(data) + + active_settings.random_seed = 1337 + posterior2 = _posterior_stats_botorch(sp, data) + + assert_frame_equal(posterior1, posterior2)