"""Smoothed EDBO preset (adapted from :cite:p:`Shields2021`)."""
from __future__ import annotations
import gc
from typing import TYPE_CHECKING, ClassVar
import numpy as np
from attrs import define, field
from typing_extensions import override
from baybe.kernels.basic import MaternKernel
from baybe.kernels.composite import ScaleKernel
from baybe.parameters import TaskParameter
from baybe.parameters.selectors import (
ParameterSelectorProtocol,
TypeSelector,
to_parameter_selector,
)
from baybe.priors.basic import GammaPrior
from baybe.surrogates.gaussian_process.components.fit_criterion import (
_MLLForNonTLFitCriterionFactory,
)
from baybe.surrogates.gaussian_process.components.kernel import (
_PureKernelFactory,
)
from baybe.surrogates.gaussian_process.components.likelihood import (
LikelihoodFactoryProtocol,
)
from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory
if TYPE_CHECKING:
from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood
from torch import Tensor
from baybe.kernels.base import Kernel
from baybe.searchspace.core import SearchSpace
# Boundaries for low and high dimension limits
_DIM_LIMITS = (8, 75)
[docs]
@define
class SmoothedEDBOKernelFactory(_PureKernelFactory):
"""A factory providing smoothed versions of EDBO kernels (adapted from :cite:p:`Shields2021`).
Takes the low and high dimensional limits of
:class:`baybe.surrogates.gaussian_process.presets.edbo.EDBOKernelFactory`
and interpolates the prior moments linearly in between.
""" # noqa: E501
_uses_parameter_names: ClassVar[bool] = True
# See base class.
parameter_selector: ParameterSelectorProtocol | None = field(
factory=lambda: TypeSelector([TaskParameter], exclude=True),
converter=to_parameter_selector,
)
# TODO: Reuse base attribute (https://github.com/python-attrs/attrs/pull/1429)
@override
def _make(
self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
) -> Kernel:
effective_dims = train_x.shape[-1]
# Interpolate prior moments linearly between low D and high D regime.
# The high D regime itself is the average of the EDBO OHE and Mordred regime.
# Values outside the dimension limits will get the border value assigned.
lengthscale_prior = GammaPrior(
np.interp(effective_dims, _DIM_LIMITS, [1.2, 2.5]),
np.interp(effective_dims, _DIM_LIMITS, [1.1, 0.55]),
)
lengthscale_initial_value = np.interp(effective_dims, _DIM_LIMITS, [0.2, 6.0])
outputscale_prior = GammaPrior(
np.interp(effective_dims, _DIM_LIMITS, [5.0, 3.5]),
np.interp(effective_dims, _DIM_LIMITS, [0.5, 0.15]),
)
outputscale_initial_value = np.interp(effective_dims, _DIM_LIMITS, [8.0, 15.0])
return ScaleKernel(
MaternKernel(
nu=2.5,
lengthscale_prior=lengthscale_prior,
lengthscale_initial_value=lengthscale_initial_value,
parameter_names=self.get_parameter_names(searchspace),
),
outputscale_prior=outputscale_prior,
outputscale_initial_value=outputscale_initial_value,
)
SmoothedEDBOMeanFactory = LazyConstantMeanFactory
"""A factory providing mean functions for the smoothed EDBO preset."""
[docs]
@define
class SmoothedEDBOLikelihoodFactory(LikelihoodFactoryProtocol):
"""A factory providing smoothed versions of EDBO likelihoods (adapted from :cite:p:`Shields2021`).
Takes the low and high dimensional limits of
:class:`baybe.surrogates.gaussian_process.presets.edbo.EDBOLikelihoodFactory`
and interpolates the prior moments linearly in between.
""" # noqa: E501
@override
def __call__(
self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
) -> GPyTorchLikelihood:
import torch
from gpytorch.likelihoods import GaussianLikelihood
# Interpolate prior moments linearly between low D and high D regime.
# The high D regime itself is the average of the EDBO OHE and Mordred regime.
# Values outside the dimension limits will get the border value assigned.
effective_dims = train_x.shape[-1] - len(
[p for p in searchspace.parameters if isinstance(p, TaskParameter)]
)
prior = GammaPrior(
np.interp(effective_dims, _DIM_LIMITS, [1.05, 1.5]),
np.interp(effective_dims, _DIM_LIMITS, [0.5, 0.1]),
)
initial_value = np.interp(effective_dims, _DIM_LIMITS, [0.1, 5.0]).item()
likelihood = GaussianLikelihood(prior.to_gpytorch())
likelihood.noise = torch.tensor([initial_value])
return likelihood
SmoothedEDBOFitCriterionFactory = _MLLForNonTLFitCriterionFactory()
"""A factory providing fitting criteria for the smoothed EDBO preset."""
# Collect leftover original slotted classes processed by `attrs.define`
gc.collect()
# Preset defaults
KERNEL_FACTORY = SmoothedEDBOKernelFactory()
MEAN_FACTORY = SmoothedEDBOMeanFactory()
LIKELIHOOD_FACTORY = SmoothedEDBOLikelihoodFactory()
FIT_CRITERION_FACTORY = SmoothedEDBOFitCriterionFactory