Source code for baybe.surrogates.gaussian_process.presets.chen

"""Preset for adaptive kernel hyperpriors proposed by :cite:p:`Chen2026`."""

from __future__ import annotations

import gc
import math
from typing import TYPE_CHECKING, ClassVar

import pandas as pd
from attrs import define
from typing_extensions import override

from baybe.kernels.basic import MaternKernel
from baybe.kernels.composite import ScaleKernel
from baybe.objectives.base import Objective
from baybe.priors.basic import GammaPrior
from baybe.surrogates.gaussian_process.components.fit_criterion import (
    _MLLForNonTLFitCriterionFactory,
)
from baybe.surrogates.gaussian_process.components.kernel import (
    _enable_transfer_learning,
    _PureKernelFactory,
)
from baybe.surrogates.gaussian_process.components.likelihood import (
    LazyGaussianLikelihoodFactory,
)
from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory

if TYPE_CHECKING:
    from baybe.kernels.base import Kernel
    from baybe.searchspace.core import SearchSpace


@define
class _ChenNumericalKernelFactory(_PureKernelFactory):
    """A factory providing the core numerical kernel for the Chen preset."""

    _uses_parameter_names: ClassVar[bool] = True
    # See base class.

    @override
    def _make(
        self, searchspace: SearchSpace, objective: Objective, measurements: pd.DataFrame
    ) -> Kernel:
        n_dimensions = self._get_effective_dimensionality(searchspace)
        lengthscale = 0.4 * math.sqrt(n_dimensions) + 4.0
        lengthscale_prior = GammaPrior(2.0 * lengthscale, 2.0)
        lengthscale_initial_value = lengthscale
        outputscale_prior = GammaPrior(1.0 * lengthscale, 1.0)
        outputscale_initial_value = lengthscale

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


ChenKernelFactory = _enable_transfer_learning(
    _ChenNumericalKernelFactory, "ChenKernelFactory"
)
"""A factory providing adaptive hyperprior kernels as proposed by :cite:p:`Chen2026`.

Takes a dimension-aware approach where kernel hyperpriors scale with the square root
of the effective dimensionality of the search space.
"""  # noqa: E501




class ChenMeanFactory(LazyConstantMeanFactory):
    """A factory providing mean functions for the Chen preset."""





class ChenLikelihoodFactory(LazyGaussianLikelihoodFactory):
    """A factory providing likelihoods for the Chen preset."""



# Collect leftover original slotted classes processed by `attrs.define`
gc.collect()

# Preset defaults
KERNEL_FACTORY = ChenKernelFactory()
MEAN_FACTORY = ChenMeanFactory()
LIKELIHOOD_FACTORY = ChenLikelihoodFactory()
FIT_CRITERION_FACTORY = _MLLForNonTLFitCriterionFactory()