Go back to wrap_kernel

Author: till-m

bayes_opt/bayesian_optimization.py

@@ -17,7 +17,7 @@
 from bayes_opt.domain_reduction import DomainTransformer
 from bayes_opt.event import DEFAULT_EVENTS, Events
 from bayes_opt.logger import _get_default_logger
-from bayes_opt.parameter import WrappedKernel
+from bayes_opt.parameter import wrap_kernel
 from bayes_opt.target_space import TargetSpace
 from bayes_opt.util import ensure_rng
 
@@ -153,7 +153,7 @@ def __init__(
 
         # Internal GP regressor
         self._gp = GaussianProcessRegressor(
-            kernel=WrappedKernel(Matern(nu=2.5), transform=self._space.kernel_transform),
+            kernel=wrap_kernel(Matern(nu=2.5), transform=self._space.kernel_transform),
             alpha=1e-6,
             normalize_y=True,
             n_restarts_optimizer=5,
@@ -330,7 +330,5 @@ def set_bounds(self, new_bounds: BoundsMapping) -> None:
     def set_gp_params(self, **params: Any) -> None:
         """Set parameters of the internal Gaussian Process Regressor."""
         if "kernel" in params:
-            params["kernel"] = WrappedKernel(
-                base_kernel=params["kernel"], transform=self._space.kernel_transform
-            )
+            params["kernel"] = wrap_kernel(kernel=params["kernel"], transform=self._space.kernel_transform)
         self._gp.set_params(**params)

bayes_opt/constraint.py

@@ -9,7 +9,7 @@
 from sklearn.gaussian_process import GaussianProcessRegressor
 from sklearn.gaussian_process.kernels import Matern
 
-from bayes_opt.parameter import WrappedKernel
+from bayes_opt.parameter import wrap_kernel
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -71,7 +71,7 @@ def __init__(
 
         self._model = [
             GaussianProcessRegressor(
-                kernel=WrappedKernel(Matern(nu=2.5), transform) if transform is not None else Matern(nu=2.5),
+                kernel=wrap_kernel(Matern(nu=2.5), transform) if transform is not None else Matern(nu=2.5),
                 alpha=1e-6,
                 normalize_y=True,
                 n_restarts_optimizer=5,

bayes_opt/parameter.py

@@ -4,6 +4,7 @@
 
 import abc
 from collections.abc import Sequence
+from inspect import signature
 from numbers import Number
 from typing import TYPE_CHECKING, Any, Callable, Union
 
@@ -458,69 +459,53 @@ def dim(self) -> int:
         return len(self.categories)
 
 
-class WrappedKernel(kernels.Kernel):
-    """Wrap a kernel with a parameter transformation.
-
-    The transform function is applied to the input before passing it to the base kernel.
+def wrap_kernel(kernel: kernels.Kernel, transform: Callable[[Any], Any]) -> kernels.Kernel:
+    """Wrap a kernel to transform input data before passing it to the kernel.
 
     Parameters
     ----------
-    base_kernel : kernels.Kernel
+    kernel : kernels.Kernel
+        The kernel to wrap.
 
-    transform : Callable[[Any], Any]
-    """
+    transform : Callable
+        The transformation function to apply to the input data.
 
-    def __init__(self, base_kernel: kernels.Kernel, transform: Callable[[Any], Any]) -> None:
-        super().__init__()
-        self.base_kernel = base_kernel
-        self.transform = transform
+    Returns
+    -------
+    kernels.Kernel
+        The wrapped kernel.
 
-    def __call__(self, X: NDArray[Float], Y: NDArray[Float] = None, eval_gradient: bool = False) -> Any:
-        """Return the kernel k(X, Y) and optionally its gradient after applying the transform.
+    Notes
+    -----
+    See https://arxiv.org/abs/1805.03463 for more information.
+    """
+    kernel_type = type(kernel)
 
-        For details, see the documentation of the base kernel.
+    class WrappedKernel(kernel_type):
+        @copy_signature(getattr(kernel_type.__init__, "deprecated_original", kernel_type.__init__))
+        def __init__(self, **kwargs: Any) -> None:
+            super().__init__(**kwargs)
 
-        Parameters
-        ----------
-        X :  ndarray of shape (n_samples_X, n_features)
-            Left argument of the returned kernel k(X, Y).
+        def __call__(self, X: Any, Y: Any = None, eval_gradient: bool = False) -> Any:
+            X = transform(X)
+            Y = transform(Y) if Y is not None else None
+            return super().__call__(X, Y, eval_gradient)
 
-        Y : ndarray of shape (n_samples_Y, n_features), default=None
-            Right argument of the returned kernel k(X, Y). If None, k(X, X) is evaluated.
+        def __reduce__(self) -> str | tuple[Any, ...]:
+            return (wrap_kernel, (kernel, transform))
 
-        eval_gradient : bool, default=False
-            Determines whether the gradient with respect to the kernel hyperparameter is calculated.
+    return WrappedKernel(**kernel.get_params())
 
-        Returns
-        -------
-        K : ndarray of shape (n_samples_X, n_samples_Y)
-
-        K_gradient : ndarray of shape (n_samples_X, n_samples_X, n_dims)
-        """
-        X = self.transform(X)
-        Y = self.transform(Y) if Y is not None else None
-        return self.base_kernel(X, Y, eval_gradient)
 
-    def is_stationary(self):
-        """Return whether the kernel is stationary."""
-        return self.base_kernel.is_stationary()
-
-    def diag(self, X: NDArray[Float]) -> NDArray[Float]:
-        """Return the diagonal of k(X, X).
-
-        This method allows for more efficient calculations than calling
-        np.diag(self(X)).
+def copy_signature(source_fct: Callable[..., Any]) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Clones a signature from a source function to a target function.
 
+    via
+    https://stackoverflow.com/a/58989918/
+    """
 
-        Parameters
-        ----------
-        X : array-like of shape (n_samples,)
-            Left argument of the returned kernel k(X, Y)
+    def copy(target_fct: Callable[..., Any]) -> Callable[..., Any]:
+        target_fct.__signature__ = signature(source_fct)
+        return target_fct
 
-        Returns
-        -------
-        K_diag : ndarray of shape (n_samples_X,)
-            Diagonal of kernel k(X, X)
-        """
-        X = self.transform(X)
-        return self.base_kernel.diag(X)
+    return copy