From 995cbf889ffdd267c0982a35d56e0ed3369e9abd Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Wed, 7 Jan 2026 21:33:19 +0000 Subject: [PATCH 1/8] 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 617a4a247c..eaf07b5be5 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -52,6 +52,8 @@ from torch import Tensor +# TODO Jordan MHS: _ModelContext is used by fidelity surrogate models now so may deserve +# its own file. @define class _ModelContext: """Model context for :class:`GaussianProcessSurrogate`.""" @@ -80,6 +82,27 @@ def n_tasks(self) -> int: """The number of tasks.""" return self.searchspace.n_tasks + @property + def n_fidelity_dimensions(self) -> int: + """The number of fidelity dimensions.""" + # Possible TODO: Generalize to multiple fidelity dimensions + return 1 if self.searchspace.fidelity_idx is not None else 0 + + @property + def is_multi_fidelity(self) -> bool: + """Are there any fidelity dimensions?""" + 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 686fc86da1d53eaf8c574ecc6e988f9258c433cf Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Tue, 13 Jan 2026 18:06:16 +0000 Subject: [PATCH 2/8] 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 eaf07b5be5..6054e52feb 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -91,10 +91,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 18ec63e6c36d493b759945e0bea87504ac21ad22 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Wed, 14 Jan 2026 11:44:41 +0000 Subject: [PATCH 3/8] 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 d3d0c2fc704e0a40c139a6f6ca36d37266be6fe2 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 16 Jan 2026 12:23:33 +0000 Subject: [PATCH 4/8] 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 585ad169821e4511b2f1170e33202fb0c21738f0 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Thu, 5 Feb 2026 08:11:17 +0000 Subject: [PATCH 5/8] 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 e4b28a7903057a0be2bcbab0d320e05ce591f65e Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 27 Feb 2026 07:24:44 +0000 Subject: [PATCH 6/8] 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 1f038345ff..e86feb41ce 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -10,6 +10,7 @@ from baybe.kernels.base import Kernel from baybe.kernels.composite import ProductKernel from baybe.parameters.categorical import TaskParameter +from baybe.parameters.fidelity import CategoricalFidelityParameter from baybe.parameters.selector import ( ParameterSelectorProtocol, TypeSelector, @@ -19,6 +20,7 @@ GPComponentFactoryProtocol, PlainGPComponentFactory, ) +from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol if TYPE_CHECKING: from gpytorch.kernels import Kernel as GPyTorchKernel @@ -79,7 +81,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: @@ -87,7 +97,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 eca58ef9b83f063e16c4bf78a549fe2abf088697 Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 27 Feb 2026 08:08:33 +0000 Subject: [PATCH 7/8] 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 6054e52feb..2c70a09151 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -116,7 +116,7 @@ def numerical_indices(self) -> list[int]: return [ i for i in range(len(self.searchspace.comp_rep_columns)) - if i != self.task_idx + if i not in (self.task_idx, self.fidelity_idx) ] 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 88db97f5bc30e22cbd40f1ae1793542c404d371e Mon Sep 17 00:00:00 2001 From: Jordan Penn Date: Fri, 6 Mar 2026 12:22:56 +0000 Subject: [PATCH 8/8] 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 e86feb41ce..eb70ded893 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -20,7 +20,6 @@ GPComponentFactoryProtocol, PlainGPComponentFactory, ) -from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol if TYPE_CHECKING: from gpytorch.kernels import Kernel as GPyTorchKernel