Conversation
* Define mean functions for the available presets * Pass both mean and likelihood when instantiating preset
There was a problem hiding this comment.
Pull request overview
Generalizes Gaussian Process (GP) surrogate configuration into reusable, generic “component” factories and adds new GP presets, including EDBO variants, while improving serialization behavior and adding tests for GPyTorch kernel support.
Changes:
- Introduces generic GP component factory infrastructure (kernel/mean/likelihood) and wires it into
GaussianProcessSurrogate. - Adds
EDBOand smoothed EDBO preset implementations including mean/likelihood factories. - Updates serialization hooks to better support generic base-class dispatch and clearer error messages; adds tests for GPyTorch kernel usage + serialization blocking.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_gp.py | Adds tests for using GPyTorch kernels and blocking serialization. |
| baybe/utils/boolean.py | Extends is_abstract to handle parameterized generics. |
| baybe/utils/basic.py | Extends get_subclasses to handle parameterized generics. |
| baybe/surrogates/gaussian_process/presets/utils.py | Adds lazy-loaded constant mean factory for presets. |
| baybe/surrogates/gaussian_process/presets/edbo_smoothed.py | Implements smoothed EDBO kernel + likelihood factories. |
| baybe/surrogates/gaussian_process/presets/edbo.py | Refactors EDBO logic into kernel/mean/likelihood factories. |
| baybe/surrogates/gaussian_process/presets/default.py | Re-exports default factories via smoothed EDBO + constant mean. |
| baybe/surrogates/gaussian_process/presets/core.py | Adds EDBO presets and constructs GP from preset factories. |
| baybe/surrogates/gaussian_process/presets/init.py | Exposes new preset factories in public API. |
| baybe/surrogates/gaussian_process/kernel_factory.py | Removes legacy kernel-factory module in favor of generic components. |
| baybe/surrogates/gaussian_process/core.py | Adds mean/likelihood factories and generic component conversion. |
| baybe/surrogates/gaussian_process/components/mean.py | Adds mean component type aliases around generic factory base. |
| baybe/surrogates/gaussian_process/components/likelihood.py | Adds likelihood component type aliases around generic factory base. |
| baybe/surrogates/gaussian_process/components/kernel.py | Adds kernel component type aliases around generic factory base. |
| baybe/surrogates/gaussian_process/components/generic.py | Introduces generic component factory + serialization blocking for GPyTorch kernels. |
| baybe/surrogates/gaussian_process/components/init.py | Adds package exports for component factory types. |
| baybe/serialization/core.py | Updates abstract-base hooks (generic-aware) + improves error messages. |
| baybe/kernels/base.py | Updates import path for PlainKernelFactory after refactor. |
| CHANGELOG.md | Documents new GP component support and presets. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| from baybe.surrogates.gaussian_process.components.kernel import ( | ||
| KernelFactory, | ||
| PlainComponentFactory, |
There was a problem hiding this comment.
PlainComponentFactory is not defined/exported by components.kernel (that module defines PlainKernelFactory). This import will fail at runtime. Import PlainComponentFactory from baybe.surrogates.gaussian_process.components.generic, or change the export to PlainKernelFactory (and update __all__ accordingly).
| from baybe.surrogates.gaussian_process.components.kernel import ( | |
| KernelFactory, | |
| PlainComponentFactory, | |
| from baybe.surrogates.gaussian_process.components.generic import ( | |
| PlainComponentFactory, | |
| ) | |
| from baybe.surrogates.gaussian_process.components.kernel import ( | |
| KernelFactory, |
| def _validate_component(instance, attribute: Attribute, value: Any): | ||
| """Validate that an object is a BayBE or a GPyTorch GP component.""" | ||
| if isinstance(value, Kernel) or _is_gpytorch_kernel_class(type(value)): | ||
| return | ||
|
|
||
| raise TypeError( | ||
| f"The object provided for '{attribute.alias}' of " | ||
| f"'{instance.__class__.__name__}' must be a BayBE or a GPyTorch GP component. " | ||
| f"Got: {type(value)}" | ||
| ) |
There was a problem hiding this comment.
This validator currently only accepts BayBE kernels or GPyTorch kernels, but the PR also introduces MeanFactory and LikelihoodFactory. Passing a GPyTorch mean/likelihood instance (which GaussianProcessSurrogate claims to accept) will raise TypeError. Expand validation to also accept GPyTorch mean and likelihood instances (ideally via lazy issubclass(type(value), gpytorch.means.Mean) / gpytorch.likelihoods.Likelihood checks similar to _is_gpytorch_kernel_class).
| def to_component_factory(x: Component | ComponentFactory, /) -> ComponentFactory: | ||
| """Wrap a component into a plain component factory (with factory passthrough).""" | ||
| if isinstance(x, Component) or _is_gpytorch_kernel_class(type(x)): | ||
| return PlainComponentFactory(x) | ||
| return x |
There was a problem hiding this comment.
At runtime, Component is set to Kernel only (to preserve serialization compatibility), so isinstance(x, Component) will not recognize GPyTorch mean/likelihood objects. As a result, providing mean_or_factory=<gpytorch Mean instance> / likelihood_or_factory=<gpytorch Likelihood instance> will fall through and return the raw object (not a factory), which will fail when called in _fit. Consider using explicit runtime checks for GPyTorch Mean/Likelihood instances (lazy-loaded) in addition to kernels, and wrap them in PlainComponentFactory.
| EDBO_SMOOTHED = "EDBO_SMOOTHED" | ||
| """A smoothed version of the EDBO settings.""" | ||
|
|
||
|
|
There was a problem hiding this comment.
The PR description and CHANGELOG.md mention a SMOOTHED_EDBO preset, but the enum value is EDBO_SMOOTHED. Since this is a user-facing API (and affects serialization/config), please align naming across the enum/make_gp_from_preset/changelog/docs. If the external name must be SMOOTHED_EDBO, consider renaming the enum member (or providing a backward-compatible alias if EDBO_SMOOTHED was already released).
| EDBO_SMOOTHED = "EDBO_SMOOTHED" | |
| """A smoothed version of the EDBO settings.""" | |
| SMOOTHED_EDBO = "SMOOTHED_EDBO" | |
| """A smoothed version of the EDBO settings.""" | |
| # Backwards-compatible alias; preferred name is ``SMOOTHED_EDBO``. | |
| EDBO_SMOOTHED = SMOOTHED_EDBO | |
| """Deprecated alias for :attr:`SMOOTHED_EDBO`.""" |
| from baybe.surrogates.gaussian_process.presets.utils import LazyConstantMeanFactory | ||
|
|
||
| if TYPE_CHECKING: | ||
| from gpytorch.likelihoods import Likelihood as GPyTorchLikelihhood |
There was a problem hiding this comment.
Correct the spelling of GPyTorchLikelihhood to GPyTorchLikelihood.
| from baybe.surrogates.gaussian_process.presets.utils import LazyConstantMeanFactory | ||
|
|
||
| if TYPE_CHECKING: | ||
| from gpytorch.likelihoods import Likelihood as GPyTorchLikelihhood |
There was a problem hiding this comment.
Correct the spelling of GPyTorchLikelihhood to GPyTorchLikelihood.
| from gpytorch.likelihoods import Likelihood as GPyTorchLikelihhood | |
| from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood |
| @@ -26,7 +26,9 @@ class Kernel(ABC, SerialMixin): | |||
|
|
|||
| def to_factory(self) -> PlainKernelFactory: | |||
| """Wrap the kernel in a :class:`baybe.surrogates.gaussian_process.kernel_factory.PlainKernelFactory`.""" # noqa: E501 | |||
There was a problem hiding this comment.
The docstring still references the removed module path baybe.surrogates.gaussian_process.kernel_factory.PlainKernelFactory. Update it to baybe.surrogates.gaussian_process.components.kernel.PlainKernelFactory (or whatever the intended public path is) to avoid stale documentation.
| """Wrap the kernel in a :class:`baybe.surrogates.gaussian_process.kernel_factory.PlainKernelFactory`.""" # noqa: E501 | |
| """Wrap the kernel in a :class:`baybe.surrogates.gaussian_process.components.kernel.PlainKernelFactory`.""" # noqa: E501 |
| Accepts: | ||
| * :class:`baybe.kernels.base.Kernel` | ||
| * :class:`.kernel_factory.KernelFactory` | ||
| * :class:`gpytorch.kernels.Kernel` |
There was a problem hiding this comment.
These docstring references (.kernel_factory.KernelFactory, similarly .mean_factory.* / .likelihood_factory.* below) no longer match the refactored module layout (components.kernel, components.mean, components.likelihood). Update the paths so docs/rendered references point to the correct types.
| def test_gpytorch_kernel(): | ||
| """The GP accepts GPyTorch kernels and produces the same result as with BayBE kernels.""" # noqa: E501 | ||
| measurements = create_fake_input(searchspace.parameters, objective.targets) | ||
| k1 = GPyTorchScaleKernel(GPyTorchMaternKernel() + GPyTorchRBFKernel()) | ||
| k2 = ScaleKernel(AdditiveKernel([MaternKernel(), RBFKernel()])) | ||
| gp1 = GaussianProcessSurrogate(k1) | ||
| gp2 = GaussianProcessSurrogate(k2) |
There was a problem hiding this comment.
This PR also adds support for configuring GPyTorch means and likelihoods, but tests only cover kernel acceptance + kernel serialization blocking. Add tests that (1) passing a GPyTorch Mean and Likelihood instance through GaussianProcessSurrogate(mean_or_factory=..., likelihood_or_factory=...) trains successfully, and (2) serialization is either explicitly blocked (raising NotImplementedError) or supported in a defined way for these components.
EDBOandSMOOTHED_EDBOGP preset