From a7b96ac0c8009cc5b707ad33ab44b745a3620f6c Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Wed, 7 Jan 2026 21:33:19 +0000 Subject: [PATCH 01/14] Multi fidelity searchspaces and surrogate modelling --- baybe/kernels/__init__.py | 2 + baybe/recommenders/pure/bayesian/base.py | 7 + baybe/searchspace/core.py | 144 +++++++++ baybe/surrogates/bandit.py | 3 + baybe/surrogates/base.py | 17 ++ baybe/surrogates/custom.py | 3 + baybe/surrogates/gaussian_process/core.py | 23 ++ .../gaussian_process/multi_fidelity.py | 288 ++++++++++++++++++ .../gaussian_process/presets/core.py | 27 ++ .../gaussian_process/presets/fidelity.py | 54 ++++ baybe/surrogates/linear.py | 3 + baybe/surrogates/naive.py | 3 + baybe/surrogates/ngboost.py | 3 + baybe/surrogates/random_forest.py | 3 + 14 files changed, 580 insertions(+) create mode 100644 baybe/surrogates/gaussian_process/multi_fidelity.py create mode 100644 baybe/surrogates/gaussian_process/presets/fidelity.py diff --git a/baybe/kernels/__init__.py b/baybe/kernels/__init__.py index 9323a2b631..4e50c5728c 100644 --- a/baybe/kernels/__init__.py +++ b/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", diff --git a/baybe/recommenders/pure/bayesian/base.py b/baybe/recommenders/pure/bayesian/base.py index 4ac5c1eed2..8da8607f6f 100644 --- a/baybe/recommenders/pure/bayesian/base.py +++ b/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, diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 5510af704f..c2a47cbdd0 100644 --- a/baybe/searchspace/core.py +++ b/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. diff --git a/baybe/surrogates/bandit.py b/baybe/surrogates/bandit.py index ad6563cc43..3437e2b945 100644 --- a/baybe/surrogates/bandit.py +++ b/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.""" diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py index 205e32f703..244b0320e5 100644 --- a/baybe/surrogates/base.py +++ b/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) diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py index 79c4c6ea1c..2b65b08a5f 100644 --- a/baybe/surrogates/custom.py +++ b/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.""" diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 2b1a3f361e..fd11cb9939 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -53,6 +53,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`.""" @@ -81,6 +83,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?""" + self.n_fidelity_dimensions > 0 + + @property + def fidelity_idx(self) -> int: + """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.""" diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py new file mode 100644 index 0000000000..efff859941 --- /dev/null +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -0,0 +1,288 @@ +"""Multi-fidelity Gaussian process surrogates.""" + +from __future__ import annotations + +import gc +from typing import TYPE_CHECKING, ClassVar + +from attrs import define, field +from typing_extensions import override + +from baybe.parameters.base import Parameter +from baybe.surrogates.base import Surrogate +from baybe.surrogates.gaussian_process.core import ( + GaussianProcessSurrogate, + _ModelContext, +) +from baybe.surrogates.gaussian_process.kernel_factory import ( + KernelFactory, + to_kernel_factory, +) +from baybe.surrogates.gaussian_process.presets import ( + GaussianProcessPreset, + make_gp_from_preset, +) +from baybe.surrogates.gaussian_process.presets.default import ( + DefaultKernelFactory, + _default_noise_factory, +) +from baybe.surrogates.gaussian_process.presets.fidelity import ( + DefaultFidelityKernelFactory, +) +from baybe.utils.conversion import to_string + +if TYPE_CHECKING: + from botorch.models.gpytorch import GPyTorchModel + from botorch.models.transforms.input import InputTransform + from botorch.models.transforms.outcome import OutcomeTransform + from botorch.posteriors import Posterior + from torch import Tensor + + +@define +class MultiFidelityGaussianProcessSurrogate(Surrogate): + """Multi fidelity Gaussian process with customisable kernel.""" + + supports_transfer_learning: ClassVar[bool] = False + # See base class. + + supports_multi_fidelity: ClassVar[bool] = True + # See base class. + + kernel_factory: KernelFactory = field( + alias="kernel_or_factory", + factory=DefaultKernelFactory, + converter=to_kernel_factory, + ) + """The factory used to create the kernel of the Gaussian process. + Accepts either a :class:`baybe.kernels.base.Kernel` or a + :class:`.kernel_factory.KernelFactory`. + When passing a :class:`baybe.kernels.base.Kernel`, it gets automatically wrapped + into a :class:`.kernel_factory.PlainKernelFactory`.""" + + fidelity_kernel_factory: KernelFactory = field( + alias="fidelity_kernel_or_factory", + factory=DefaultFidelityKernelFactory, + converter=to_kernel_factory, + ) + """The factory used to create the fidelity kernel of the Gaussian process. + Accepts either a :class:`baybe.kernels.base.Kernel` or a + :class:`.kernel_factory.KernelFactory`. + When passing a :class:`baybe.kernels.base.Kernel`, it gets automatically wrapped + into a :class:`.kernel_factory.PlainKernelFactory`.""" + + _model = field(init=False, default=None, eq=False) + """The actual model.""" + + @staticmethod + def from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate: + """Create a Gaussian process surrogate from one of the defined presets.""" + return make_gp_from_preset(preset) + + @override + def to_botorch(self) -> GPyTorchModel: + return self._model + + @override + @staticmethod + def _make_parameter_scaler_factory( + parameter: Parameter, + ) -> type[InputTransform] | None: + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + return None + + @override + @staticmethod + def _make_target_scaler_factory() -> type[OutcomeTransform] | None: + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + return None + + @override + def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior: + return self._model.posterior(candidates_comp_scaled) + + @override + def _fit(self, train_x: Tensor, train_y: Tensor) -> None: + import botorch + import gpytorch + import torch + from botorch.models.transforms import Normalize, Standardize + + # FIXME[typing]: It seems there is currently no better way to inform the type + # checker that the attribute is available at the time of the function call + assert self._searchspace is not None + + context = _ModelContext(self._searchspace) + + numerical_idxs = context.get_numerical_indices(train_x.shape[-1]) + + numerical_design_idxs = tuple( + idx for idx in numerical_idxs if idx != context.fidelity_idx + ) + + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + input_transform = Normalize( + train_x.shape[-1], + bounds=context.parameter_bounds, + indices=list(numerical_idxs), + ) + outcome_transform = Standardize(train_y.shape[-1]) + + # extract the batch shape of the training data + batch_shape = train_x.shape[:-2] + + # create GP mean + mean_module = gpytorch.means.ConstantMean(batch_shape=batch_shape) + + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + input_transform = botorch.models.transforms.Normalize( + train_x.shape[-1], + bounds=context.parameter_bounds, + indices=list(numerical_design_idxs), + ) + outcome_transform = botorch.models.transforms.Standardize(train_y.shape[-1]) + + base_covar_module = self.kernel_factory( + context.searchspace, train_x, train_y + ).to_gpytorch( + ard_num_dims=train_x.shape[-1] - context.n_fidelity_dimensions, + active_dims=numerical_design_idxs, + batch_shape=batch_shape, + ) + + fidelity_covar_module = self.fidelity_kernel_factory( + num_tasks=context.n_fidelities, + active_dims=context.fidelity_idx, + rank=context.n_fidelities, # TODO: make controllable + ).to_gpytorch( + ard_num_dims=1, + active_dims=(context.fidelity_idx,), + batch_shape=batch_shape, + ) + + covar_module = base_covar_module * fidelity_covar_module + + # create GP likelihood + noise_prior = _default_noise_factory(context.searchspace, train_x, train_y) + likelihood = gpytorch.likelihoods.GaussianLikelihood( + noise_prior=noise_prior[0].to_gpytorch(), batch_shape=batch_shape + ) + likelihood.noise = torch.tensor([noise_prior[1]]) + + # construct and fit the Gaussian process + self._model = botorch.models.SingleTaskGP( + train_x, + train_y, + input_transform=input_transform, + outcome_transform=outcome_transform, + mean_module=mean_module, + covar_module=covar_module, + likelihood=likelihood, + ) + + mll = gpytorch.ExactMarginalLogLikelihood(self._model.likelihood, self._model) + + botorch.fit.fit_gpytorch_mll(mll) + + @override + def __str__(self) -> str: + fields = [ + to_string( + to_string("Kernel factory", self.kernel_factory, single_line=True), + "Fidelity kernel factory", + self.fidelity_kernel_factory, + single_line=True, + ), + ] + return to_string(super().__str__(), *fields) + + +@define +class GaussianProcessSurrogateSTMF(GaussianProcessSurrogate): + """Botorch's single task multi fidelity Gaussian process.""" + + supports_transfer_learning: ClassVar[bool] = False + # See base class. + + supports_multi_fidelity: ClassVar[bool] = True + # See base class. + + kernel_factory: KernelFactory = field(init=False, default=None) + """Design kernel is set to Matern within SingleTaskMultiFidelityGP.""" + + # 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) + """The actual model.""" + + @staticmethod + def from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate: + """Create a Gaussian process surrogate from one of the defined presets.""" + return make_gp_from_preset(preset) + + @override + def to_botorch(self) -> GPyTorchModel: + return self._model + + @override + @staticmethod + def _make_parameter_scaler_factory( + parameter: Parameter, + ) -> type[InputTransform] | None: + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + return None + + @override + @staticmethod + def _make_target_scaler_factory() -> type[OutcomeTransform] | None: + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + return None + + @override + def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior: + return self._model.posterior(candidates_comp_scaled) + + @override + def _fit(self, train_x: Tensor, train_y: Tensor) -> None: + import botorch + import gpytorch + + assert self._searchspace is not None + + context = _ModelContext(self._searchspace) + + numerical_design_idxs = context.get_numerical_indices(train_x.shape[-1]) + + if context.is_multi_fidelity: + numerical_design_idxs = tuple( + idx for idx in numerical_design_idxs if idx != context.fidelity_idx + ) + + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + input_transform = botorch.models.transforms.Normalize( + train_x.shape[-1], + bounds=context.parameter_bounds, + indices=numerical_design_idxs, + ) + outcome_transform = botorch.models.transforms.Standardize(train_y.shape[-1]) + + # construct and fit the Gaussian process + self._model = botorch.models.SingleTaskMultiFidelityGP( + train_x, + train_y, + input_transform=input_transform, + outcome_transform=outcome_transform, + data_fidelities=[context.fidelity_idx], + ) + + mll = gpytorch.ExactMarginalLogLikelihood(self._model.likelihood, self._model) + + botorch.fit.fit_gpytorch_mll(mll) + + @override + def __str__(self) -> str: + return "SingleTaskMultiFidelityGP with Botorch defaults." + + +# Collect leftover original slotted classes processed by `attrs.define` +gc.collect() diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py index ad77df0b4d..aa9fd62900 100644 --- a/baybe/surrogates/gaussian_process/presets/core.py +++ b/baybe/surrogates/gaussian_process/presets/core.py @@ -4,6 +4,8 @@ from enum import Enum +from surrogates.gaussian_process.core import GaussianProcessSurrogate + class GaussianProcessPreset(Enum): """Available Gaussian process surrogate presets.""" @@ -16,3 +18,28 @@ class GaussianProcessPreset(Enum): EDBO_SMOOTHED = "EDBO_SMOOTHED" """A smoothed version of the EDBO settings.""" + + BOTORCH_STMF = "BOTORCH_STMF" + """Recreates the default settings of the BOTORCH SingleTaskMultiFidelityGP.""" + + +def make_gp_from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate: + """Create a :class:`GaussianProcessSurrogate` from a :class:`GaussianProcessPreset.""" # noqa: E501 + from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate + from baybe.surrogates.gaussian_process.multi_fidelity import ( + GaussianProcessSurrogateSTMF, + MultiFidelityGaussianProcessSurrogate, + ) + + if preset is GaussianProcessPreset.BAYBE: + return GaussianProcessSurrogate() + + if preset is GaussianProcessPreset.MFGP: + return MultiFidelityGaussianProcessSurrogate() + + if preset is GaussianProcessPreset.BOTORCH_STMF: + return GaussianProcessSurrogateSTMF() + + raise ValueError( + f"Unknown '{GaussianProcessPreset.__name__}' with name '{preset.name}'." + ) diff --git a/baybe/surrogates/gaussian_process/presets/fidelity.py b/baybe/surrogates/gaussian_process/presets/fidelity.py new file mode 100644 index 0000000000..16ba64f316 --- /dev/null +++ b/baybe/surrogates/gaussian_process/presets/fidelity.py @@ -0,0 +1,54 @@ +"""Kernels for Gaussian process fidelity surrogates.""" + +from __future__ import annotations + +import gc +from typing import TYPE_CHECKING + +from attrs import define +from typing_extensions import override + +from baybe.kernels.basic import IndexKernel +from baybe.surrogates.gaussian_process.kernel_factory import KernelFactory + +if TYPE_CHECKING: + from torch import Tensor + + from baybe.kernels.base import Kernel + from baybe.searchspace.core import SearchSpace + + +@define +class IndependentFidelityKernelFactory(KernelFactory): + """Rank 0 index kernel treating fidelities as independent.""" + + @override + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel: + return IndexKernel( + num_tasks=searchspace.n_fidelities, + active_dims=searchspace.fidelity_idx, + rank=0, + ) + + +@define +class IndexFidelityKernelFactory(KernelFactory): + """Full rank index kernel modelling dependent fidelities.""" + + @override + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel: + return IndexKernel( + num_tasks=searchspace.n_fidelities, + active_dims=searchspace.fidelity_idx, + rank=searchspace.n_fidelities, + ) + + +DefaultFidelityKernelFactory = IndexFidelityKernelFactory + +# Collect leftover original slotted classes processed by `attrs.define` +gc.collect() diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py index ba1fed5b2c..94746a16b9 100644 --- a/baybe/surrogates/linear.py +++ b/baybe/surrogates/linear.py @@ -44,6 +44,9 @@ class BayesianLinearSurrogate(IndependentGaussianSurrogate): supports_transfer_learning: ClassVar[bool] = False # See base class. + supports_multi_fidelity: ClassVar[bool] = False + # See base class. + model_params: _ARDRegressionParams = field( factory=dict, converter=dict, diff --git a/baybe/surrogates/naive.py b/baybe/surrogates/naive.py index 3912c6b128..b407b48f08 100644 --- a/baybe/surrogates/naive.py +++ b/baybe/surrogates/naive.py @@ -26,6 +26,9 @@ class MeanPredictionSurrogate(IndependentGaussianSurrogate): supports_transfer_learning: ClassVar[bool] = False # See base class. + supports_multi_fidelity: ClassVar[bool] = False + # See base class. + _model: float | None = field(init=False, default=None, eq=False) """The estimated posterior mean value of the training targets.""" diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py index 6cb1cee135..f05a9ebc0d 100644 --- a/baybe/surrogates/ngboost.py +++ b/baybe/surrogates/ngboost.py @@ -49,6 +49,9 @@ class NGBoostSurrogate(IndependentGaussianSurrogate): supports_transfer_learning: ClassVar[bool] = False # See base class. + supports_multi_fidelity: ClassVar[bool] = False + # See base class. + _default_model_params: ClassVar[dict] = {"n_estimators": 25, "verbose": False} """Class variable encoding the default model parameters.""" diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py index 91ad599bb1..6ee1f7ed70 100644 --- a/baybe/surrogates/random_forest.py +++ b/baybe/surrogates/random_forest.py @@ -64,6 +64,9 @@ class RandomForestSurrogate(Surrogate): supports_transfer_learning: ClassVar[bool] = False # See base class. + supports_multi_fidelity: ClassVar[bool] = False + # See base class. + model_params: _RandomForestRegressorParams = field( factory=dict, converter=dict, From 748ffa75e1e2b2b507919e7426c4702c232445d1 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Tue, 13 Jan 2026 18:06:16 +0000 Subject: [PATCH 02/14] Typing fixes --- baybe/surrogates/gaussian_process/core.py | 4 ++-- baybe/surrogates/gaussian_process/multi_fidelity.py | 8 ++++---- baybe/surrogates/gaussian_process/presets/core.py | 6 +++++- baybe/surrogates/gaussian_process/presets/fidelity.py | 2 -- 4 files changed, 11 insertions(+), 9 deletions(-) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index fd11cb9939..910f9f5910 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -92,10 +92,10 @@ def n_fidelity_dimensions(self) -> int: @property def is_multi_fidelity(self) -> bool: """Are there any fidelity dimensions?""" - self.n_fidelity_dimensions > 0 + return self.n_fidelity_dimensions > 0 @property - def fidelity_idx(self) -> int: + def fidelity_idx(self) -> int | None: """The computational column index of the task parameter, if available.""" return self.searchspace.fidelity_idx diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index efff859941..9c972d0ae7 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -15,6 +15,7 @@ _ModelContext, ) from baybe.surrogates.gaussian_process.kernel_factory import ( + DiscreteFidelityKernelFactory, KernelFactory, to_kernel_factory, ) @@ -60,7 +61,7 @@ class MultiFidelityGaussianProcessSurrogate(Surrogate): When passing a :class:`baybe.kernels.base.Kernel`, it gets automatically wrapped into a :class:`.kernel_factory.PlainKernelFactory`.""" - fidelity_kernel_factory: KernelFactory = field( + fidelity_kernel_factory: DiscreteFidelityKernelFactory = field( alias="fidelity_kernel_or_factory", factory=DefaultFidelityKernelFactory, converter=to_kernel_factory, @@ -152,7 +153,6 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: fidelity_covar_module = self.fidelity_kernel_factory( num_tasks=context.n_fidelities, - active_dims=context.fidelity_idx, rank=context.n_fidelities, # TODO: make controllable ).to_gpytorch( ard_num_dims=1, @@ -198,7 +198,7 @@ def __str__(self) -> str: @define -class GaussianProcessSurrogateSTMF(GaussianProcessSurrogate): +class GaussianProcessSurrogateSTMF(Surrogate): """Botorch's single task multi fidelity Gaussian process.""" supports_transfer_learning: ClassVar[bool] = False @@ -262,7 +262,7 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: input_transform = botorch.models.transforms.Normalize( train_x.shape[-1], bounds=context.parameter_bounds, - indices=numerical_design_idxs, + indices=list(numerical_design_idxs), ) outcome_transform = botorch.models.transforms.Standardize(train_y.shape[-1]) diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py index aa9fd62900..2fbf4b000b 100644 --- a/baybe/surrogates/gaussian_process/presets/core.py +++ b/baybe/surrogates/gaussian_process/presets/core.py @@ -3,6 +3,10 @@ from __future__ import annotations from enum import Enum +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from baybe.surrogates.base import Surrogate from surrogates.gaussian_process.core import GaussianProcessSurrogate @@ -23,7 +27,7 @@ class GaussianProcessPreset(Enum): """Recreates the default settings of the BOTORCH SingleTaskMultiFidelityGP.""" -def make_gp_from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate: +def make_gp_from_preset(preset: GaussianProcessPreset) -> Surrogate: """Create a :class:`GaussianProcessSurrogate` from a :class:`GaussianProcessPreset.""" # noqa: E501 from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate from baybe.surrogates.gaussian_process.multi_fidelity import ( diff --git a/baybe/surrogates/gaussian_process/presets/fidelity.py b/baybe/surrogates/gaussian_process/presets/fidelity.py index 16ba64f316..328705411d 100644 --- a/baybe/surrogates/gaussian_process/presets/fidelity.py +++ b/baybe/surrogates/gaussian_process/presets/fidelity.py @@ -28,7 +28,6 @@ def __call__( ) -> Kernel: return IndexKernel( num_tasks=searchspace.n_fidelities, - active_dims=searchspace.fidelity_idx, rank=0, ) @@ -43,7 +42,6 @@ def __call__( ) -> Kernel: return IndexKernel( num_tasks=searchspace.n_fidelities, - active_dims=searchspace.fidelity_idx, rank=searchspace.n_fidelities, ) From 1a7e6c9059489b8fd03a3c0ec016aa7998b739ca Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Wed, 14 Jan 2026 11:44:41 +0000 Subject: [PATCH 03/14] More typing fixes --- .../surrogates/gaussian_process/multi_fidelity.py | 14 ++++++++++---- .../gaussian_process/presets/fidelity.py | 8 +++++--- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index 9c972d0ae7..bb9009887e 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -11,7 +11,6 @@ from baybe.parameters.base import Parameter from baybe.surrogates.base import Surrogate from baybe.surrogates.gaussian_process.core import ( - GaussianProcessSurrogate, _ModelContext, ) from baybe.surrogates.gaussian_process.kernel_factory import ( @@ -76,7 +75,7 @@ class MultiFidelityGaussianProcessSurrogate(Surrogate): """The actual model.""" @staticmethod - def from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate: + def from_preset(preset: GaussianProcessPreset) -> Surrogate: """Create a Gaussian process surrogate from one of the defined presets.""" return make_gp_from_preset(preset) @@ -156,7 +155,9 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: rank=context.n_fidelities, # TODO: make controllable ).to_gpytorch( ard_num_dims=1, - active_dims=(context.fidelity_idx,), + active_dims=None + if context.fidelity_idx is None + else (context.fidelity_idx,), batch_shape=batch_shape, ) @@ -216,7 +217,7 @@ class GaussianProcessSurrogateSTMF(Surrogate): """The actual model.""" @staticmethod - def from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate: + def from_preset(preset: GaussianProcessPreset) -> Surrogate: """Create a Gaussian process surrogate from one of the defined presets.""" return make_gp_from_preset(preset) @@ -253,6 +254,11 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: numerical_design_idxs = context.get_numerical_indices(train_x.shape[-1]) + assert context.is_multi_fidelity, ( + "GaussianProcessSurrogateSTMF can only " + "be fit on multi fidelity searchspaces." + ) + if context.is_multi_fidelity: numerical_design_idxs = tuple( idx for idx in numerical_design_idxs if idx != context.fidelity_idx diff --git a/baybe/surrogates/gaussian_process/presets/fidelity.py b/baybe/surrogates/gaussian_process/presets/fidelity.py index 328705411d..366ca2fae8 100644 --- a/baybe/surrogates/gaussian_process/presets/fidelity.py +++ b/baybe/surrogates/gaussian_process/presets/fidelity.py @@ -9,7 +9,9 @@ from typing_extensions import override from baybe.kernels.basic import IndexKernel -from baybe.surrogates.gaussian_process.kernel_factory import KernelFactory +from baybe.surrogates.gaussian_process.kernel_factory import ( + DefaultFidelityKernelFactory, +) if TYPE_CHECKING: from torch import Tensor @@ -19,7 +21,7 @@ @define -class IndependentFidelityKernelFactory(KernelFactory): +class IndependentFidelityKernelFactory(DefaultFidelityKernelFactory): """Rank 0 index kernel treating fidelities as independent.""" @override @@ -33,7 +35,7 @@ def __call__( @define -class IndexFidelityKernelFactory(KernelFactory): +class IndexFidelityKernelFactory(DefaultFidelityKernelFactory): """Full rank index kernel modelling dependent fidelities.""" @override From 078f16002801624192303fed848c41afe670e7f9 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 16 Jan 2026 12:23:33 +0000 Subject: [PATCH 04/14] More typing fixes with some unresolved --- baybe/surrogates/gaussian_process/multi_fidelity.py | 4 +++- baybe/surrogates/gaussian_process/presets/fidelity.py | 6 +++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index bb9009887e..0ff8f63f02 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -278,7 +278,9 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: train_y, input_transform=input_transform, outcome_transform=outcome_transform, - data_fidelities=[context.fidelity_idx], + data_fidelities=None + if context.fidelity_idx is None + else (context.fidelity_idx,), ) mll = gpytorch.ExactMarginalLogLikelihood(self._model.likelihood, self._model) diff --git a/baybe/surrogates/gaussian_process/presets/fidelity.py b/baybe/surrogates/gaussian_process/presets/fidelity.py index 366ca2fae8..9ecf9b5bd5 100644 --- a/baybe/surrogates/gaussian_process/presets/fidelity.py +++ b/baybe/surrogates/gaussian_process/presets/fidelity.py @@ -10,7 +10,7 @@ from baybe.kernels.basic import IndexKernel from baybe.surrogates.gaussian_process.kernel_factory import ( - DefaultFidelityKernelFactory, + DiscreteFidelityKernelFactory, ) if TYPE_CHECKING: @@ -21,7 +21,7 @@ @define -class IndependentFidelityKernelFactory(DefaultFidelityKernelFactory): +class IndependentFidelityKernelFactory(DiscreteFidelityKernelFactory): """Rank 0 index kernel treating fidelities as independent.""" @override @@ -35,7 +35,7 @@ def __call__( @define -class IndexFidelityKernelFactory(DefaultFidelityKernelFactory): +class IndexFidelityKernelFactory(DiscreteFidelityKernelFactory): """Full rank index kernel modelling dependent fidelities.""" @override From 120e0e5b833b8ec7c85d61413c27abdde47aee08 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Thu, 5 Feb 2026 08:11:17 +0000 Subject: [PATCH 05/14] Typo fix --- baybe/surrogates/gaussian_process/multi_fidelity.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index 0ff8f63f02..ac16057d7c 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -151,8 +151,7 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: ) fidelity_covar_module = self.fidelity_kernel_factory( - num_tasks=context.n_fidelities, - rank=context.n_fidelities, # TODO: make controllable + searchspace=self._searchspace ).to_gpytorch( ard_num_dims=1, active_dims=None From 385327339de58059188b396ba4f518d7f1241f3c Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 27 Feb 2026 07:24:44 +0000 Subject: [PATCH 06/14] Integrating multi fidelity surrogate models with multitask refactor --- .../gaussian_process/components/kernel.py | 21 ++- .../gaussian_process/multi_fidelity.py | 174 ------------------ .../gaussian_process/presets/core.py | 10 - .../gaussian_process/presets/fidelity.py | 54 ------ 4 files changed, 19 insertions(+), 240 deletions(-) delete mode 100644 baybe/surrogates/gaussian_process/presets/fidelity.py diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 8b6ddff280..3837173c7b 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -16,6 +16,7 @@ from baybe.kernels.base import Kernel from baybe.parameters.categorical import TaskParameter from baybe.parameters.enum import _ParameterKind +from baybe.parameters.fidelity import CategoricalFidelityParameter from baybe.parameters.selectors import ( ParameterSelectorProtocol, TypeSelector, @@ -28,6 +29,7 @@ PlainGPComponentFactory, to_component_factory, ) +from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol if TYPE_CHECKING: from gpytorch.kernels import Kernel as GPyTorchKernel @@ -157,7 +159,15 @@ 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: @@ -165,7 +175,14 @@ def _default_task_kernel_factory(self) -> KernelFactoryProtocol: BayBETaskKernelFactory, ) - return BayBETaskKernelFactory(TypeSelector((TaskParameter,))) + return BayBETaskKernelFactory( + TypeSelector( + ( + TaskParameter, + CategoricalFidelityParameter, + ) + ) + ) @override def __call__( diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index ac16057d7c..2d6783d2eb 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -13,23 +13,10 @@ from baybe.surrogates.gaussian_process.core import ( _ModelContext, ) -from baybe.surrogates.gaussian_process.kernel_factory import ( - DiscreteFidelityKernelFactory, - KernelFactory, - to_kernel_factory, -) from baybe.surrogates.gaussian_process.presets import ( GaussianProcessPreset, make_gp_from_preset, ) -from baybe.surrogates.gaussian_process.presets.default import ( - DefaultKernelFactory, - _default_noise_factory, -) -from baybe.surrogates.gaussian_process.presets.fidelity import ( - DefaultFidelityKernelFactory, -) -from baybe.utils.conversion import to_string if TYPE_CHECKING: from botorch.models.gpytorch import GPyTorchModel @@ -39,164 +26,6 @@ from torch import Tensor -@define -class MultiFidelityGaussianProcessSurrogate(Surrogate): - """Multi fidelity Gaussian process with customisable kernel.""" - - supports_transfer_learning: ClassVar[bool] = False - # See base class. - - supports_multi_fidelity: ClassVar[bool] = True - # See base class. - - kernel_factory: KernelFactory = field( - alias="kernel_or_factory", - factory=DefaultKernelFactory, - converter=to_kernel_factory, - ) - """The factory used to create the kernel of the Gaussian process. - Accepts either a :class:`baybe.kernels.base.Kernel` or a - :class:`.kernel_factory.KernelFactory`. - When passing a :class:`baybe.kernels.base.Kernel`, it gets automatically wrapped - into a :class:`.kernel_factory.PlainKernelFactory`.""" - - fidelity_kernel_factory: DiscreteFidelityKernelFactory = field( - alias="fidelity_kernel_or_factory", - factory=DefaultFidelityKernelFactory, - converter=to_kernel_factory, - ) - """The factory used to create the fidelity kernel of the Gaussian process. - Accepts either a :class:`baybe.kernels.base.Kernel` or a - :class:`.kernel_factory.KernelFactory`. - When passing a :class:`baybe.kernels.base.Kernel`, it gets automatically wrapped - into a :class:`.kernel_factory.PlainKernelFactory`.""" - - _model = field(init=False, default=None, eq=False) - """The actual model.""" - - @staticmethod - def from_preset(preset: GaussianProcessPreset) -> Surrogate: - """Create a Gaussian process surrogate from one of the defined presets.""" - return make_gp_from_preset(preset) - - @override - def to_botorch(self) -> GPyTorchModel: - return self._model - - @override - @staticmethod - def _make_parameter_scaler_factory( - parameter: Parameter, - ) -> type[InputTransform] | None: - # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. - return None - - @override - @staticmethod - def _make_target_scaler_factory() -> type[OutcomeTransform] | None: - # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. - return None - - @override - def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior: - return self._model.posterior(candidates_comp_scaled) - - @override - def _fit(self, train_x: Tensor, train_y: Tensor) -> None: - import botorch - import gpytorch - import torch - from botorch.models.transforms import Normalize, Standardize - - # FIXME[typing]: It seems there is currently no better way to inform the type - # checker that the attribute is available at the time of the function call - assert self._searchspace is not None - - context = _ModelContext(self._searchspace) - - numerical_idxs = context.get_numerical_indices(train_x.shape[-1]) - - numerical_design_idxs = tuple( - idx for idx in numerical_idxs if idx != context.fidelity_idx - ) - - # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. - input_transform = Normalize( - train_x.shape[-1], - bounds=context.parameter_bounds, - indices=list(numerical_idxs), - ) - outcome_transform = Standardize(train_y.shape[-1]) - - # extract the batch shape of the training data - batch_shape = train_x.shape[:-2] - - # create GP mean - mean_module = gpytorch.means.ConstantMean(batch_shape=batch_shape) - - # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. - input_transform = botorch.models.transforms.Normalize( - train_x.shape[-1], - bounds=context.parameter_bounds, - indices=list(numerical_design_idxs), - ) - outcome_transform = botorch.models.transforms.Standardize(train_y.shape[-1]) - - base_covar_module = self.kernel_factory( - context.searchspace, train_x, train_y - ).to_gpytorch( - ard_num_dims=train_x.shape[-1] - context.n_fidelity_dimensions, - active_dims=numerical_design_idxs, - batch_shape=batch_shape, - ) - - fidelity_covar_module = self.fidelity_kernel_factory( - searchspace=self._searchspace - ).to_gpytorch( - ard_num_dims=1, - active_dims=None - if context.fidelity_idx is None - else (context.fidelity_idx,), - batch_shape=batch_shape, - ) - - covar_module = base_covar_module * fidelity_covar_module - - # create GP likelihood - noise_prior = _default_noise_factory(context.searchspace, train_x, train_y) - likelihood = gpytorch.likelihoods.GaussianLikelihood( - noise_prior=noise_prior[0].to_gpytorch(), batch_shape=batch_shape - ) - likelihood.noise = torch.tensor([noise_prior[1]]) - - # construct and fit the Gaussian process - self._model = botorch.models.SingleTaskGP( - train_x, - train_y, - input_transform=input_transform, - outcome_transform=outcome_transform, - mean_module=mean_module, - covar_module=covar_module, - likelihood=likelihood, - ) - - mll = gpytorch.ExactMarginalLogLikelihood(self._model.likelihood, self._model) - - botorch.fit.fit_gpytorch_mll(mll) - - @override - def __str__(self) -> str: - fields = [ - to_string( - to_string("Kernel factory", self.kernel_factory, single_line=True), - "Fidelity kernel factory", - self.fidelity_kernel_factory, - single_line=True, - ), - ] - return to_string(super().__str__(), *fields) - - @define class GaussianProcessSurrogateSTMF(Surrogate): """Botorch's single task multi fidelity Gaussian process.""" @@ -207,9 +36,6 @@ class GaussianProcessSurrogateSTMF(Surrogate): supports_multi_fidelity: ClassVar[bool] = True # See base class. - kernel_factory: KernelFactory = field(init=False, default=None) - """Design kernel is set to Matern within SingleTaskMultiFidelityGP.""" - # 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) diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py index 2fbf4b000b..5f659c0e01 100644 --- a/baybe/surrogates/gaussian_process/presets/core.py +++ b/baybe/surrogates/gaussian_process/presets/core.py @@ -8,8 +8,6 @@ if TYPE_CHECKING: from baybe.surrogates.base import Surrogate -from surrogates.gaussian_process.core import GaussianProcessSurrogate - class GaussianProcessPreset(Enum): """Available Gaussian process surrogate presets.""" @@ -29,18 +27,10 @@ class GaussianProcessPreset(Enum): def make_gp_from_preset(preset: GaussianProcessPreset) -> Surrogate: """Create a :class:`GaussianProcessSurrogate` from a :class:`GaussianProcessPreset.""" # noqa: E501 - from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate from baybe.surrogates.gaussian_process.multi_fidelity import ( GaussianProcessSurrogateSTMF, - MultiFidelityGaussianProcessSurrogate, ) - if preset is GaussianProcessPreset.BAYBE: - return GaussianProcessSurrogate() - - if preset is GaussianProcessPreset.MFGP: - return MultiFidelityGaussianProcessSurrogate() - if preset is GaussianProcessPreset.BOTORCH_STMF: return GaussianProcessSurrogateSTMF() diff --git a/baybe/surrogates/gaussian_process/presets/fidelity.py b/baybe/surrogates/gaussian_process/presets/fidelity.py deleted file mode 100644 index 9ecf9b5bd5..0000000000 --- a/baybe/surrogates/gaussian_process/presets/fidelity.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Kernels for Gaussian process fidelity surrogates.""" - -from __future__ import annotations - -import gc -from typing import TYPE_CHECKING - -from attrs import define -from typing_extensions import override - -from baybe.kernels.basic import IndexKernel -from baybe.surrogates.gaussian_process.kernel_factory import ( - DiscreteFidelityKernelFactory, -) - -if TYPE_CHECKING: - from torch import Tensor - - from baybe.kernels.base import Kernel - from baybe.searchspace.core import SearchSpace - - -@define -class IndependentFidelityKernelFactory(DiscreteFidelityKernelFactory): - """Rank 0 index kernel treating fidelities as independent.""" - - @override - def __call__( - self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor - ) -> Kernel: - return IndexKernel( - num_tasks=searchspace.n_fidelities, - rank=0, - ) - - -@define -class IndexFidelityKernelFactory(DiscreteFidelityKernelFactory): - """Full rank index kernel modelling dependent fidelities.""" - - @override - def __call__( - self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor - ) -> Kernel: - return IndexKernel( - num_tasks=searchspace.n_fidelities, - rank=searchspace.n_fidelities, - ) - - -DefaultFidelityKernelFactory = IndexFidelityKernelFactory - -# Collect leftover original slotted classes processed by `attrs.define` -gc.collect() From 6c13232b36535b3689bab4a58fab07270de39b4c Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 27 Feb 2026 08:08:33 +0000 Subject: [PATCH 07/14] Integrate typing --- baybe/parameters/fidelity.py | 10 +++++++--- baybe/surrogates/gaussian_process/core.py | 2 +- .../surrogates/gaussian_process/multi_fidelity.py | 15 ++++----------- 3 files changed, 12 insertions(+), 15 deletions(-) diff --git a/baybe/parameters/fidelity.py b/baybe/parameters/fidelity.py index 15235c9c78..14ecf1a51a 100644 --- a/baybe/parameters/fidelity.py +++ b/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, ) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 910f9f5910..376d4832c8 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -117,7 +117,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) ] diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index 2d6783d2eb..2cdca81801 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -13,7 +13,7 @@ from baybe.surrogates.gaussian_process.core import ( _ModelContext, ) -from baybe.surrogates.gaussian_process.presets import ( +from baybe.surrogates.gaussian_process.presets.core import ( GaussianProcessPreset, make_gp_from_preset, ) @@ -77,25 +77,18 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: context = _ModelContext(self._searchspace) - numerical_design_idxs = context.get_numerical_indices(train_x.shape[-1]) - assert context.is_multi_fidelity, ( "GaussianProcessSurrogateSTMF can only " "be fit on multi fidelity searchspaces." ) - if context.is_multi_fidelity: - numerical_design_idxs = tuple( - idx for idx in numerical_design_idxs if idx != context.fidelity_idx - ) - # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. - input_transform = botorch.models.transforms.Normalize( + input_transform = botorch.models.transforms.Normalize( # type: ignore[attr-defined] train_x.shape[-1], bounds=context.parameter_bounds, - indices=list(numerical_design_idxs), + indices=context.numerical_indices, ) - outcome_transform = botorch.models.transforms.Standardize(train_y.shape[-1]) + outcome_transform = botorch.models.transforms.Standardize(train_y.shape[-1]) # type: ignore[attr-defined] # construct and fit the Gaussian process self._model = botorch.models.SingleTaskMultiFidelityGP( From 782754e2a4c8c5dd9b9f75c9d64adbfaad82fce2 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 6 Mar 2026 12:22:56 +0000 Subject: [PATCH 08/14] Integrating kernel factories with multi fidelity --- baybe/surrogates/gaussian_process/components/kernel.py | 1 - 1 file changed, 1 deletion(-) diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 3837173c7b..f69130dab8 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -29,7 +29,6 @@ PlainGPComponentFactory, to_component_factory, ) -from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol if TYPE_CHECKING: from gpytorch.kernels import Kernel as GPyTorchKernel From a85113db7e92e6bc52d53d2f16a5f04ac4522a36 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 20 Mar 2026 04:32:00 +0000 Subject: [PATCH 09/14] Move _ModelContext --- baybe/surrogates/gaussian_process/core.py | 69 +---------------- .../gaussian_process/multi_fidelity.py | 4 +- baybe/surrogates/gaussian_process/utils.py | 77 +++++++++++++++++++ 3 files changed, 79 insertions(+), 71 deletions(-) create mode 100644 baybe/surrogates/gaussian_process/utils.py diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 376d4832c8..de9b517425 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -39,6 +39,7 @@ BayBELikelihoodFactory, BayBEMeanFactory, ) +from baybe.surrogates.gaussian_process.utils import _ModelContext from baybe.utils.boolean import strtobool from baybe.utils.conversion import to_string @@ -53,74 +54,6 @@ 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`.""" - - searchspace: SearchSpace = field(validator=instance_of(SearchSpace)) - """The search space the model is trained on.""" - - @property - def task_idx(self) -> int | None: - """The computational column index of the task parameter, if available.""" - return self.searchspace.task_idx - - @property - def is_multitask(self) -> bool: - """Indicates if model is to be operated in a multi-task context.""" - return self.n_task_dimensions > 0 - - @property - def n_task_dimensions(self) -> int: - """The number of task dimensions.""" - # TODO: Generalize to multiple task parameters - return 1 if self.task_idx is not None else 0 - - @property - 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.""" - import torch - - return torch.from_numpy(self.searchspace.scaling_bounds.values) - - @property - def numerical_indices(self) -> list[int]: - """The indices of the regular numerical model inputs.""" - return [ - i - for i in range(len(self.searchspace.comp_rep_columns)) - if i not in (self.task_idx, self.fidelity_idx) - ] - - def _mark_custom_kernel( value: Kernel | KernelFactoryProtocol, self: GaussianProcessSurrogate ) -> Kernel | KernelFactoryProtocol: diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index 2cdca81801..2af3aa4384 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -10,13 +10,11 @@ from baybe.parameters.base import Parameter from baybe.surrogates.base import Surrogate -from baybe.surrogates.gaussian_process.core import ( - _ModelContext, -) from baybe.surrogates.gaussian_process.presets.core import ( GaussianProcessPreset, make_gp_from_preset, ) +from baybe.surrogates.gaussian_process.utils import _ModelContext if TYPE_CHECKING: from botorch.models.gpytorch import GPyTorchModel diff --git a/baybe/surrogates/gaussian_process/utils.py b/baybe/surrogates/gaussian_process/utils.py new file mode 100644 index 0000000000..5615a0fb84 --- /dev/null +++ b/baybe/surrogates/gaussian_process/utils.py @@ -0,0 +1,77 @@ +"""Gaussian process utilities.""" + +from typing import TYPE_CHECKING + +from attrs import define, field +from attrs.validators import instance_of + +from baybe.searchspace.core import SearchSpace + +if TYPE_CHECKING: + from torch import Tensor + + +@define +class _ModelContext: + """Model context for Gaussian process surrogates.""" + + searchspace: SearchSpace = field(validator=instance_of(SearchSpace)) + """The search space the model is trained on.""" + + @property + def task_idx(self) -> int | None: + """The computational column index of the task parameter, if available.""" + return self.searchspace.task_idx + + @property + def is_multitask(self) -> bool: + """Indicates if model is to be operated in a multi-task context.""" + return self.n_task_dimensions > 0 + + @property + def n_task_dimensions(self) -> int: + """The number of task dimensions.""" + # TODO: Generalize to multiple task parameters + return 1 if self.task_idx is not None else 0 + + @property + 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.""" + # TODO: Generalize to multiple fidelity parameters + 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.""" + import torch + + return torch.from_numpy(self.searchspace.scaling_bounds.values) + + @property + def numerical_indices(self) -> list[int]: + """The indices of the regular numerical model inputs.""" + return [ + i + for i in range(len(self.searchspace.comp_rep_columns)) + if i not in (self.task_idx, self.fidelity_idx) + ] From 3896724c3637c7acf26d89dbf458ef68a469f908 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 20 Mar 2026 04:51:30 +0000 Subject: [PATCH 10/14] Remove deprecated/unused from_preset from STMF GP. --- .../gaussian_process/multi_fidelity.py | 9 ----- .../gaussian_process/presets/__init__.py | 3 -- .../gaussian_process/presets/core.py | 39 ------------------- 3 files changed, 51 deletions(-) delete mode 100644 baybe/surrogates/gaussian_process/presets/core.py diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index 2af3aa4384..7d0aaab50f 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -10,10 +10,6 @@ from baybe.parameters.base import Parameter from baybe.surrogates.base import Surrogate -from baybe.surrogates.gaussian_process.presets.core import ( - GaussianProcessPreset, - make_gp_from_preset, -) from baybe.surrogates.gaussian_process.utils import _ModelContext if TYPE_CHECKING: @@ -39,11 +35,6 @@ class GaussianProcessSurrogateSTMF(Surrogate): _model = field(init=False, default=None, eq=False) """The actual model.""" - @staticmethod - def from_preset(preset: GaussianProcessPreset) -> Surrogate: - """Create a Gaussian process surrogate from one of the defined presets.""" - return make_gp_from_preset(preset) - @override def to_botorch(self) -> GPyTorchModel: return self._model diff --git a/baybe/surrogates/gaussian_process/presets/__init__.py b/baybe/surrogates/gaussian_process/presets/__init__.py index deb7de9e64..017e673550 100644 --- a/baybe/surrogates/gaussian_process/presets/__init__.py +++ b/baybe/surrogates/gaussian_process/presets/__init__.py @@ -7,9 +7,6 @@ BayBEMeanFactory, ) -# Core -from baybe.surrogates.gaussian_process.presets.core import GaussianProcessPreset - # EDBO preset from baybe.surrogates.gaussian_process.presets.edbo import ( EDBOKernelFactory, diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py deleted file mode 100644 index 5f659c0e01..0000000000 --- a/baybe/surrogates/gaussian_process/presets/core.py +++ /dev/null @@ -1,39 +0,0 @@ -"""Preset configurations for Gaussian process surrogates.""" - -from __future__ import annotations - -from enum import Enum -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from baybe.surrogates.base import Surrogate - - -class GaussianProcessPreset(Enum): - """Available Gaussian process surrogate presets.""" - - BAYBE = "BAYBE" - """The default BayBE settings of the Gaussian process surrogate class.""" - - EDBO = "EDBO" - """The EDBO settings.""" - - EDBO_SMOOTHED = "EDBO_SMOOTHED" - """A smoothed version of the EDBO settings.""" - - BOTORCH_STMF = "BOTORCH_STMF" - """Recreates the default settings of the BOTORCH SingleTaskMultiFidelityGP.""" - - -def make_gp_from_preset(preset: GaussianProcessPreset) -> Surrogate: - """Create a :class:`GaussianProcessSurrogate` from a :class:`GaussianProcessPreset.""" # noqa: E501 - from baybe.surrogates.gaussian_process.multi_fidelity import ( - GaussianProcessSurrogateSTMF, - ) - - if preset is GaussianProcessPreset.BOTORCH_STMF: - return GaussianProcessSurrogateSTMF() - - raise ValueError( - f"Unknown '{GaussianProcessPreset.__name__}' with name '{preset.name}'." - ) From 76ff8f2bd71e7bf7dd2ec7b43b9fbdab1134549f Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 20 Mar 2026 07:44:03 +0000 Subject: [PATCH 11/14] Mainly docstring fixes --- baybe/recommenders/pure/bayesian/base.py | 7 ---- baybe/searchspace/core.py | 35 ++++++++++--------- baybe/surrogates/base.py | 4 +-- baybe/surrogates/gaussian_process/core.py | 3 ++ .../gaussian_process/multi_fidelity.py | 15 ++++---- baybe/surrogates/gaussian_process/utils.py | 7 +--- 6 files changed, 32 insertions(+), 39 deletions(-) diff --git a/baybe/recommenders/pure/bayesian/base.py b/baybe/recommenders/pure/bayesian/base.py index 8da8607f6f..4ac5c1eed2 100644 --- a/baybe/recommenders/pure/bayesian/base.py +++ b/baybe/recommenders/pure/bayesian/base.py @@ -44,13 +44,6 @@ 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, diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index c2a47cbdd0..760e2178dc 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -55,8 +55,8 @@ class SearchSpaceType(Enum): class SearchSpaceTaskType(Enum): """Enum class for different types of task and/or fidelity subspaces.""" - SINGLETASK = "SINGLETASK" - """Flag for search spaces with no task parameters.""" + NOTASK = "NOTASK" + """Flag for search spaces with a single task, meaning no task parameter.""" CATEGORICALTASK = "CATEGORICALTASK" """Flag for search spaces with a categorical task parameter.""" @@ -304,9 +304,8 @@ def task_idx(self) -> int | None: @property def fidelity_idx(self) -> int | None: - """The column index of the task parameter in computational representation.""" + """Column index of the fidelity parameter in computational representation.""" try: - # See TODO [16932] and TODO [11611] fidelity_param = next( p for p in self.parameters @@ -334,8 +333,7 @@ def n_tasks(self) -> int: @property def n_fidelities(self) -> int: - """The number of tasks encoded in the search space.""" - # See TODO [16932] + """The number of fidelities encoded in the search space.""" try: fidelity_param = next( p @@ -355,20 +353,23 @@ def n_fidelities(self) -> int: 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,)) + task_param = next( + p + for p in self.parameters + if isinstance( + p, + (TaskParameter), + ) ) except StopIteration: - fidelity_param = None + task_param = None - return 1 if fidelity_param is not None else 0 + return 1 if task_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 @@ -387,8 +388,8 @@ 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. + ValueError: If search space contains more than one task/fidelity parameter. + ValueError: An unrecognised fidelity parameter type is in search space. """ task_like_parameters = ( TaskParameter, @@ -401,14 +402,14 @@ def task_type(self) -> SearchSpaceTaskType: ) if n_task_like_parameters == 0: - return SearchSpaceTaskType.SINGLETASK + return SearchSpaceTaskType.NOTASK 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." + "Search space must not contain more than one task/fidelity parameter." ) return SearchSpaceTaskType.MULTIPLETASKPARAMETER @@ -530,7 +531,7 @@ def transform( @property def constraints_augmentable(self) -> tuple[Constraint, ...]: - """The searchspace constraints that can be considered during augmentation.""" + """The search space constraints that can be considered during augmentation.""" return tuple(c for c in self.constraints if c.eval_during_augmentation) def get_parameters_by_name(self, names: Sequence[str]) -> tuple[Parameter, ...]: diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py index 244b0320e5..fc9d4755a6 100644 --- a/baybe/surrogates/base.py +++ b/baybe/surrogates/base.py @@ -82,11 +82,11 @@ def to_botorch(self) -> Model: class Surrogate(ABC, SurrogateProtocol, SerialMixin): """Abstract base class for all surrogate models.""" - supports_transfer_learning: ClassVar[bool] + supports_transfer_learning: ClassVar[bool] = False """Class variable encoding whether or not the surrogate supports transfer learning.""" - supports_multi_fidelity: ClassVar[bool] + supports_multi_fidelity: ClassVar[bool] = False """Class variable encoding whether or not the surrogate supports multi fidelity Bayesian optimization.""" diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index de9b517425..7171a57f79 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -86,6 +86,9 @@ class GaussianProcessSurrogate(Surrogate): supports_transfer_learning: ClassVar[bool] = True # See base class. + supports_multi_fidelity = True + # See base class. + _custom_kernel: bool = field(init=False, default=False, repr=False, eq=False) # For deprecation only! diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index 7d0aaab50f..3d558d5b1a 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -30,8 +30,6 @@ class GaussianProcessSurrogateSTMF(Surrogate): supports_multi_fidelity: ClassVar[bool] = True # See base class. - # 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) """The actual model.""" @@ -44,7 +42,6 @@ def to_botorch(self) -> GPyTorchModel: def _make_parameter_scaler_factory( parameter: Parameter, ) -> type[InputTransform] | None: - # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. return None @override @@ -66,9 +63,8 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: context = _ModelContext(self._searchspace) - assert context.is_multi_fidelity, ( - "GaussianProcessSurrogateSTMF can only " - "be fit on multi fidelity searchspaces." + assert context.n_fidelity_dimensions > 0, ( + f"{self.__class__.__name__} can only be fit on multi fidelity searchspaces." ) # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. @@ -96,7 +92,12 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: @override def __str__(self) -> str: - return "SingleTaskMultiFidelityGP with Botorch defaults." + return ( + "Wrapper for a" + ":class:`~botorch.models.gp_regression_fidelity.SingleTaskMultiFidelityGP`," + "used as the default GP for discrete numerical fidelity parameters in," + "e.g., multi fidelity knowledge gradient." + ) # Collect leftover original slotted classes processed by `attrs.define` diff --git a/baybe/surrogates/gaussian_process/utils.py b/baybe/surrogates/gaussian_process/utils.py index 5615a0fb84..a8b35bf6fe 100644 --- a/baybe/surrogates/gaussian_process/utils.py +++ b/baybe/surrogates/gaussian_process/utils.py @@ -45,14 +45,9 @@ def n_fidelity_dimensions(self) -> int: # TODO: Generalize to multiple fidelity parameters 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.""" + """The computational column index of the fidelity parameter, if available.""" return self.searchspace.fidelity_idx @property From 4ee22e808f01c1fa8e9c68524ec18de9b6ddd135 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 27 Mar 2026 10:16:52 +0000 Subject: [PATCH 12/14] Separating task and fidelity types in search space. --- baybe/searchspace/core.py | 173 ++++++++++++++------------------------ 1 file changed, 64 insertions(+), 109 deletions(-) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 760e2178dc..b0a979810b 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -55,25 +55,25 @@ class SearchSpaceType(Enum): class SearchSpaceTaskType(Enum): """Enum class for different types of task and/or fidelity subspaces.""" - NOTASK = "NOTASK" + SINGLETASK = "SINGLETASK" """Flag for search spaces with a single task, meaning no task parameter.""" - CATEGORICALTASK = "CATEGORICALTASK" + CATEGORICALMULTITASK = "CATEGORICALMULTITASK" """Flag for search spaces with a categorical task parameter.""" - NUMERICALFIDELITY = "NUMERICALFIDELITY" + +class SearchSpaceFidelityType(Enum): + """Enum class for different types of task and/or fidelity subspaces.""" + + SINGLEFIDELITY = "SINGLEFIDELITY" + """Flag for search spaces with a single fidelity, meaning no fidelity parameter.""" + + NUMERICALDISCRETEMULTIFIDELITY = "NUMERICALDISCRETEMULTIFIDELITY" """Flag for search spaces with a discrete numerical (ordered) fidelity parameter.""" - CATEGORICALFIDELITY = "CATEGORICALFIDELITY" + CATEGORICALMULTIFIDELITY = "CATEGORICALMULTIFIDELITY" """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): @@ -285,15 +285,32 @@ def _task_parameter(self) -> TaskParameter | None: if not params: return None - assert len(params) == 1 # currently ensured by parameter validation step + return params[0] + + @property + def _fidelity_parameter( + self, + ) -> NumericalDiscreteFidelityParameter | CategoricalFidelityParameter | None: + """The (single) fidelity parameter of the space, if it exists.""" + # Currently private, see comment above + fidelity_parameters = ( + NumericalDiscreteFidelityParameter, + CategoricalFidelityParameter, + ) + + params = [p for p in self.parameters if isinstance(p, fidelity_parameters)] + + if not params: + return None + return params[0] @property def task_idx(self) -> int | None: - """The column index of the task parameter in computational representation.""" + """Column index of the task parameter in computational representation.""" if (task_param := self._task_parameter) is None: return None - # TODO[11611]: The current approach has three limitations: + # TODO [11611]: The current approach has three limitations: # 1. It matches by column name and thus assumes that the parameter name # is used as the column name. # 2. It relies on the current implementation detail that discrete parameters @@ -305,18 +322,9 @@ def task_idx(self) -> int | None: @property def fidelity_idx(self) -> int | None: """Column index of the fidelity parameter in computational representation.""" - try: - fidelity_param = next( - p - for p in self.parameters - if isinstance( - p, - (CategoricalFidelityParameter, NumericalDiscreteFidelityParameter), - ) - ) - except StopIteration: + if (fidelity_param := self._fidelity_parameter) is None: return None - + # See TODO [11611] above return cast(int, self.discrete.comp_rep.columns.get_loc(fidelity_param.name)) @property @@ -334,103 +342,50 @@ def n_tasks(self) -> int: @property def n_fidelities(self) -> int: """The number of fidelities encoded in the search space.""" - 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: + # See TODO [16932] above + if (fidelity_param := self._fidelity_parameter) is None: + # When there are no task parameters, we effectively have a single task return 1 + return len(fidelity_param.values) @property - def n_task_dimensions(self) -> int: - """The number of task dimensions.""" - try: - task_param = next( - p - for p in self.parameters - if isinstance( - p, - (TaskParameter), - ) - ) - except StopIteration: - task_param = None - - return 1 if task_param is not None else 0 + def task_type(self) -> SearchSpaceTaskType: + """Return the task type of the search space.""" + task_parameters = (p for p in self.parameters if isinstance(p, TaskParameter)) - @property - def n_fidelity_dimensions(self) -> int: - """The number of fidelity dimensions.""" - try: - fidelity_param = next( - p - for p in self.parameters - if isinstance( - p, - (CategoricalFidelityParameter, NumericalDiscreteFidelityParameter), - ) + if len(task_parameters) == 0: + return SearchSpaceTaskType.SINGLETASK + elif len(task_parameters) == 1: + return SearchSpaceTaskType.CATEGORICALMULTITASK + else: + raise NotImplementedError( + "BayBE does not currently support search" + "spaces with multiple task parameters." ) - 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 search space contains more than one task/fidelity parameter. - ValueError: An unrecognised fidelity parameter type is in search space. - """ - task_like_parameters = ( - TaskParameter, + def fidelity_type(self) -> SearchSpaceFidelityType: + """Return the fidelity type of the search space.""" + fidelity_parameters = ( CategoricalFidelityParameter, NumericalDiscreteFidelityParameter, ) - n_task_like_parameters = sum( - isinstance(p, (task_like_parameters)) for p in self.parameters + fidelity_parameters = ( + p for p in self.parameters if isinstance(p, fidelity_parameters) ) - if n_task_like_parameters == 0: - return SearchSpaceTaskType.NOTASK - 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( - "Search space 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 len(fidelity_parameters) == 0: + return SearchSpaceFidelityType.SINGLEFIDELITY + elif len(fidelity_parameters) == 1: + if isinstance(fidelity_parameters[0], CategoricalFidelityParameter): + return SearchSpaceFidelityType.CATEGORICALMULTIFIDELITY + if isinstance(fidelity_parameters[0], NumericalDiscreteFidelityParameter): + return SearchSpaceFidelityType.NUMERICALDISCRETEMULTIFIDELITY + else: + raise NotImplementedError( + "BayBE does not currently support search" + "spaces with multiple fidelity 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. From dc3b0f4aea9611aeefe4a519f7476959c0c266a3 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Tue, 31 Mar 2026 14:41:55 +0100 Subject: [PATCH 13/14] Tidying supports_ inheritance --- baybe/surrogates/bandit.py | 8 +------- baybe/surrogates/custom.py | 8 +------- baybe/surrogates/gaussian_process/multi_fidelity.py | 3 --- baybe/surrogates/linear.py | 8 +------- baybe/surrogates/naive.py | 8 +------- baybe/surrogates/ngboost.py | 6 ------ baybe/surrogates/random_forest.py | 8 +------- 7 files changed, 5 insertions(+), 44 deletions(-) diff --git a/baybe/surrogates/bandit.py b/baybe/surrogates/bandit.py index 3437e2b945..f003f44c64 100644 --- a/baybe/surrogates/bandit.py +++ b/baybe/surrogates/bandit.py @@ -3,7 +3,7 @@ from __future__ import annotations import gc -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any from attrs import define, field from typing_extensions import override @@ -29,12 +29,6 @@ class BetaBernoulliMultiArmedBanditSurrogate(Surrogate): """A multi-armed bandit model with Bernoulli likelihood and beta prior.""" - 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.""" diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py index 2b65b08a5f..00e346319d 100644 --- a/baybe/surrogates/custom.py +++ b/baybe/surrogates/custom.py @@ -11,7 +11,7 @@ from __future__ import annotations import gc -from typing import TYPE_CHECKING, Any, ClassVar, NoReturn +from typing import TYPE_CHECKING, Any, NoReturn import cattrs from attrs import define, field, validators @@ -67,12 +67,6 @@ class CustomONNXSurrogate(IndependentGaussianSurrogate): Note that these surrogates cannot be retrained. """ - 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.""" diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py index 3d558d5b1a..6e5fa3cdef 100644 --- a/baybe/surrogates/gaussian_process/multi_fidelity.py +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -24,9 +24,6 @@ class GaussianProcessSurrogateSTMF(Surrogate): """Botorch's single task multi fidelity Gaussian process.""" - supports_transfer_learning: ClassVar[bool] = False - # See base class. - supports_multi_fidelity: ClassVar[bool] = True # See base class. diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py index 94746a16b9..3c4ded8f45 100644 --- a/baybe/surrogates/linear.py +++ b/baybe/surrogates/linear.py @@ -3,7 +3,7 @@ from __future__ import annotations import gc -from typing import TYPE_CHECKING, ClassVar, TypedDict +from typing import TYPE_CHECKING, TypedDict from attrs import define, field from typing_extensions import override @@ -41,12 +41,6 @@ class _ARDRegressionParams(TypedDict, total=False): class BayesianLinearSurrogate(IndependentGaussianSurrogate): """A Bayesian linear regression surrogate model.""" - supports_transfer_learning: ClassVar[bool] = False - # See base class. - - supports_multi_fidelity: ClassVar[bool] = False - # See base class. - model_params: _ARDRegressionParams = field( factory=dict, converter=dict, diff --git a/baybe/surrogates/naive.py b/baybe/surrogates/naive.py index b407b48f08..ef7d624c5b 100644 --- a/baybe/surrogates/naive.py +++ b/baybe/surrogates/naive.py @@ -3,7 +3,7 @@ from __future__ import annotations import gc -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING from attrs import define, field from typing_extensions import override @@ -23,12 +23,6 @@ class MeanPredictionSurrogate(IndependentGaussianSurrogate): as posterior mean and a (data-independent) constant posterior variance. """ - supports_transfer_learning: ClassVar[bool] = False - # See base class. - - supports_multi_fidelity: ClassVar[bool] = False - # See base class. - _model: float | None = field(init=False, default=None, eq=False) """The estimated posterior mean value of the training targets.""" diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py index f05a9ebc0d..d4a4363e3f 100644 --- a/baybe/surrogates/ngboost.py +++ b/baybe/surrogates/ngboost.py @@ -46,12 +46,6 @@ class _NGBRegressorParams(TypedDict, total=False): class NGBoostSurrogate(IndependentGaussianSurrogate): """A natural-gradient-boosting surrogate model.""" - supports_transfer_learning: ClassVar[bool] = False - # See base class. - - supports_multi_fidelity: ClassVar[bool] = False - # See base class. - _default_model_params: ClassVar[dict] = {"n_estimators": 25, "verbose": False} """Class variable encoding the default model parameters.""" diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py index 6ee1f7ed70..f7726828be 100644 --- a/baybe/surrogates/random_forest.py +++ b/baybe/surrogates/random_forest.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Collection -from typing import TYPE_CHECKING, ClassVar, Literal, Protocol, TypedDict +from typing import TYPE_CHECKING, Literal, Protocol, TypedDict import numpy as np import numpy.typing as npt @@ -61,12 +61,6 @@ def predict(self, x: np.ndarray, /) -> np.ndarray: ... class RandomForestSurrogate(Surrogate): """A random forest surrogate model.""" - supports_transfer_learning: ClassVar[bool] = False - # See base class. - - supports_multi_fidelity: ClassVar[bool] = False - # See base class. - model_params: _RandomForestRegressorParams = field( factory=dict, converter=dict, From 6451eb7d2fc1e47be8c1333a0a16cc3de23c2090 Mon Sep 17 00:00:00 2001 From: "Alexander V. Hopp" Date: Fri, 24 Apr 2026 16:38:14 +0200 Subject: [PATCH 14/14] Remove unused imports after ModelContext move --- baybe/surrogates/gaussian_process/core.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 7171a57f79..a5e43b145c 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -10,14 +10,13 @@ from attrs import Converter, define, field from attrs.converters import pipe -from attrs.validators import instance_of, is_callable +from attrs.validators import is_callable from typing_extensions import Self, override from baybe.exceptions import DeprecationError from baybe.kernels.base import Kernel from baybe.parameters.base import Parameter from baybe.parameters.categorical import TaskParameter -from baybe.searchspace.core import SearchSpace from baybe.surrogates.base import Surrogate from baybe.surrogates.gaussian_process.components.generic import ( GPComponentType,