Add generic GP components by AdrianSosic · Pull Request #746 · emdgroup/baybe

AdrianSosic · 2026-02-11T14:29:37Z

Generalizes the existing GP kernel into multi-purpose component logic using generics
Adds support for GPyTorch GP components
Enables configuration of GP mean and likelihood
Adds EDBO and SMOOTHED_EDBO GP preset
Completes preset logic by including mean and likelihood components

* Define mean functions for the available presets * Pass both mean and likelihood when instantiating preset

Copilot

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 EDBO and 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.

Copilot · 2026-02-11T14:32:22Z

baybe/surrogates/gaussian_process/components/__init__.py

+from baybe.surrogates.gaussian_process.components.kernel import (
+    KernelFactory,
+    PlainComponentFactory,


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).

Suggested change

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,

Copilot · 2026-02-11T14:32:22Z

baybe/surrogates/gaussian_process/components/generic.py

+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)}"
+    )


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).

Copilot · 2026-02-11T14:32:23Z

baybe/surrogates/gaussian_process/components/generic.py

+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


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.

Copilot · 2026-02-11T14:32:23Z

baybe/surrogates/gaussian_process/presets/core.py

+    EDBO_SMOOTHED = "EDBO_SMOOTHED"
+    """A smoothed version of the EDBO settings."""




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).

Suggested change

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`."""

Copilot · 2026-02-11T14:32:23Z

baybe/surrogates/gaussian_process/presets/edbo_smoothed.py

+from baybe.surrogates.gaussian_process.presets.utils import LazyConstantMeanFactory
+
+if TYPE_CHECKING:
+    from gpytorch.likelihoods import Likelihood as GPyTorchLikelihhood


Correct the spelling of GPyTorchLikelihhood to GPyTorchLikelihood.

Copilot · 2026-02-11T14:32:24Z

baybe/surrogates/gaussian_process/presets/edbo.py

+from baybe.surrogates.gaussian_process.presets.utils import LazyConstantMeanFactory

 if TYPE_CHECKING:
+    from gpytorch.likelihoods import Likelihood as GPyTorchLikelihhood


Correct the spelling of GPyTorchLikelihhood to GPyTorchLikelihood.

Suggested change

from gpytorch.likelihoods import Likelihood as GPyTorchLikelihhood

from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood

Copilot · 2026-02-11T14:32:24Z

baybe/kernels/base.py

@@ -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


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.

Suggested change

"""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

Copilot · 2026-02-11T14:32:24Z

baybe/surrogates/gaussian_process/core.py

+    Accepts:
+        * :class:`baybe.kernels.base.Kernel`
+        * :class:`.kernel_factory.KernelFactory`
+        * :class:`gpytorch.kernels.Kernel`


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.

Copilot · 2026-02-11T14:32:25Z

tests/test_gp.py

+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)


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.

AdrianSosic added 18 commits February 10, 2026 17:54

Extract current GP defaults in into separate EDBO module

5e8c7de

Add EDBO presets

6b48716

Introduce generic component factories

577abd2

Make serialization utilties handle generics

19220bd

Correctly (un)structure generic classes

b0cd172

Add execution path for non-generic classes

50e531e

Add support for GPyTorch kernels

23a71a7

Block serialization of GPyTorch kernels

f669a21

Enable configuration of GP mean

8234082

Enable configuration of GP likelihood

e15e243

Complete the current preset framework

108e082

* Define mean functions for the available presets * Pass both mean and likelihood when instantiating preset

Update CHANGELOG.md

50d0a2a

Update attribute docstrings

aeb5d8a

Reorganize modules into subpackage

ffb8702

Add missing components to TypeVar

d908fcb

Add missing entries to __init__.py

9df998f

Deduplicate code in default preset

7a509d4

Add execution path for non-generic classes

b1f6215

AdrianSosic self-assigned this Feb 11, 2026

AdrianSosic added enhancement Expand / change existing functionality new feature New functionality refactor labels Feb 11, 2026

AdrianSosic marked this pull request as ready for review February 11, 2026 14:29

AdrianSosic requested review from AVHopp and Scienfitz as code owners February 11, 2026 14:29

Copilot AI review requested due to automatic review settings February 11, 2026 14:29

AdrianSosic added the dev label Feb 11, 2026

AdrianSosic mentioned this pull request Feb 11, 2026

Beef up GaussianProcessSurrogate #745

Draft

15 tasks

Copilot AI reviewed Feb 11, 2026

View reviewed changes

Fix typing issues

0056011

Fix assert_never import

45d0486

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add generic GP components#746

Add generic GP components#746
AdrianSosic wants to merge 20 commits intodev/gpfrom
refactor/gp_components

AdrianSosic commented Feb 11, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Copilot AI Feb 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		EDBO_SMOOTHED = "EDBO_SMOOTHED"
		"""A smoothed version of the EDBO settings."""

-    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 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

Conversation

AdrianSosic commented Feb 11, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant