diff --git a/CHANGELOG.md b/CHANGELOG.md index 78b2a663f4..450bffa1e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] +### Added +- `ProjectionKernel` class and corresponding `ProjectionKernelFactory` class + ## [0.14.2] - 2026-01-14 ### Added - `NumericalTarget.match_*` constructors now accept a `mismatch_instead` argument. If diff --git a/baybe/kernels/__init__.py b/baybe/kernels/__init__.py index 9323a2b631..1e5393f1b9 100644 --- a/baybe/kernels/__init__.py +++ b/baybe/kernels/__init__.py @@ -14,7 +14,12 @@ RFFKernel, RQKernel, ) -from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel +from baybe.kernels.composite import ( + AdditiveKernel, + ProductKernel, + ProjectionKernel, + ScaleKernel, +) __all__ = [ "AdditiveKernel", @@ -24,6 +29,7 @@ "PiecewisePolynomialKernel", "PolynomialKernel", "ProductKernel", + "ProjectionKernel", "RBFKernel", "RFFKernel", "RQKernel", diff --git a/baybe/kernels/_gpytorch.py b/baybe/kernels/_gpytorch.py new file mode 100644 index 0000000000..86c7f4d8b5 --- /dev/null +++ b/baybe/kernels/_gpytorch.py @@ -0,0 +1,40 @@ +"""GPyTorch kernel implementations.""" + +from typing import Any + +import torch +from gpytorch.kernels import Kernel +from torch import Tensor + +from baybe.utils.torch import DTypeFloatTorch + +_ConvertibleToTensor = Any +"""A type alias for objects convertible to tensors.""" + + +class GPyTorchProjectionKernel(Kernel): + """GPyTorch implementation of :class:`baybe.kernels.composite.ProjectionKernel`.""" + + def __init__( + self, + base_kernel: Kernel, + projection_matrix: _ConvertibleToTensor, + *, + learn_projection: bool = False, + ): + super().__init__() + + self.base_kernel = base_kernel + self.learn_projection = learn_projection + + matrix = torch.tensor(projection_matrix, dtype=DTypeFloatTorch) + if self.learn_projection: + self.register_parameter("projection_matrix", torch.nn.Parameter(matrix)) + else: + self.register_buffer("projection_matrix", matrix) + + def forward(self, x1: Tensor, x2: Tensor, **kwargs): + """Apply the base kernel to the projected input tensors.""" + x1_proj = x1 @ self.projection_matrix + x2_proj = x2 @ self.projection_matrix + return self.base_kernel(x1_proj, x2_proj, **kwargs) diff --git a/baybe/kernels/composite.py b/baybe/kernels/composite.py index dab7abbf58..e769fd4470 100644 --- a/baybe/kernels/composite.py +++ b/baybe/kernels/composite.py @@ -4,7 +4,8 @@ from functools import reduce from operator import add, mul -from attrs import define, field +import numpy as np +from attrs import Attribute, cmp_using, define, field from attrs.converters import optional as optional_c from attrs.validators import deep_iterable, gt, instance_of, min_len from attrs.validators import optional as optional_v @@ -83,5 +84,45 @@ def to_gpytorch(self, *args, **kwargs): return reduce(mul, (k.to_gpytorch(*args, **kwargs) for k in self.base_kernels)) +@define(frozen=True) +class ProjectionKernel(CompositeKernel): + """A projection kernel for dimensionality reduction.""" + + base_kernel: Kernel = field(validator=instance_of(Kernel)) + """The kernel applied to the projected inputs.""" + + projection_matrix: np.ndarray = field( + converter=np.asarray, eq=cmp_using(eq=np.array_equal) + ) + """The projection matrix.""" + + learn_projection: bool = field( + default=True, validator=instance_of(bool), kw_only=True + ) + """Boolean specifying if the projection matrix should be learned. + + If ``True``, the provided projection matrix is used as initial value.""" + + @projection_matrix.validator + def _validate_projection_matrix(self, _: Attribute, value: np.ndarray): + if value.ndim != 2: + raise ValueError( + f"The projection matrix must be 2-dimensional, " + f"but has shape {value.shape}." + ) + + @override + def to_gpytorch(self, **kwargs): + from baybe.kernels._gpytorch import GPyTorchProjectionKernel + + n_projections = self.projection_matrix.shape[-1] + gpytorch_kernel = self.base_kernel.to_gpytorch(ard_num_dims=n_projections) + return GPyTorchProjectionKernel( + gpytorch_kernel, + projection_matrix=self.projection_matrix, + learn_projection=self.learn_projection, + ) + + # Collect leftover original slotted classes processed by `attrs.define` gc.collect() diff --git a/baybe/serialization/core.py b/baybe/serialization/core.py index 6381bac26c..b8ff3d44e8 100644 --- a/baybe/serialization/core.py +++ b/baybe/serialization/core.py @@ -9,6 +9,7 @@ import attrs import cattrs +import numpy as np import pandas as pd from cattrs.strategies import configure_union_passthrough @@ -104,6 +105,15 @@ def _unstructure_dataframe_hook(df: pd.DataFrame) -> str: return base64.b64encode(pickled_df).decode("utf-8") +_unstructure_ndarray_hook = _unstructure_dataframe_hook + + +def _structure_ndarray_hook(obj: str, _) -> np.ndarray: + """Deserialize a numpy ndarray.""" + pickled_array = base64.b64decode(obj.encode("utf-8")) + return pickle.loads(pickled_array) + + def block_serialization_hook(obj: Any) -> NoReturn: # noqa: DOC101, DOC103 """Prevent serialization of the passed object. @@ -163,6 +173,8 @@ def select_constructor_hook(specs: dict, cls: type[_T]) -> _T: ) converter.register_unstructure_hook(pd.DataFrame, _unstructure_dataframe_hook) converter.register_structure_hook(pd.DataFrame, _structure_dataframe_hook) +converter.register_unstructure_hook(np.ndarray, _unstructure_ndarray_hook) +converter.register_structure_hook(np.ndarray, _structure_ndarray_hook) converter.register_unstructure_hook(datetime, lambda x: x.isoformat()) converter.register_structure_hook(datetime, lambda x, _: datetime.fromisoformat(x)) converter.register_unstructure_hook(timedelta, lambda x: f"{x.total_seconds()}s") diff --git a/baybe/serialization/utils.py b/baybe/serialization/utils.py index d5a5fd6a7f..37435464b8 100644 --- a/baybe/serialization/utils.py +++ b/baybe/serialization/utils.py @@ -2,18 +2,33 @@ from typing import Any +import numpy as np import pandas as pd +def serialize_ndarray(array: np.ndarray, /) -> Any: + """Serialize a numpy ndarray.""" + from baybe.serialization import converter + + return converter.unstructure(array) + + +def deserialize_ndarray(serialized_array: Any, /) -> np.ndarray: + """Deserialize a numpy ndarray.""" + from baybe.serialization import converter + + return converter.structure(serialized_array, np.ndarray) + + def serialize_dataframe(df: pd.DataFrame, /) -> Any: """Serialize a pandas dataframe.""" - from baybe.searchspace.core import converter + from baybe.serialization import converter return converter.unstructure(df) def deserialize_dataframe(serialized_df: Any, /) -> pd.DataFrame: """Deserialize a pandas dataframe.""" - from baybe.searchspace.core import converter + from baybe.serialization import converter return converter.structure(serialized_df, pd.DataFrame) diff --git a/baybe/surrogates/gaussian_process/kernel_factory.py b/baybe/surrogates/gaussian_process/kernel_factory.py index b3ae4983eb..01aadd9123 100644 --- a/baybe/surrogates/gaussian_process/kernel_factory.py +++ b/baybe/surrogates/gaussian_process/kernel_factory.py @@ -3,13 +3,16 @@ from __future__ import annotations import gc +from enum import Enum from typing import TYPE_CHECKING, Protocol +import numpy as np from attrs import define, field from attrs.validators import instance_of -from typing_extensions import override +from typing_extensions import assert_never, override from baybe.kernels.base import Kernel +from baybe.kernels.composite import AdditiveKernel, ProjectionKernel from baybe.searchspace import SearchSpace from baybe.serialization.mixin import SerialMixin @@ -17,6 +20,101 @@ from torch import Tensor +def to_kernel_factory(x: Kernel | KernelFactory, /) -> KernelFactory: + """Wrap a kernel into a plain kernel factory (with factory passthrough).""" + return x.to_factory() if isinstance(x, Kernel) else x + + +class ProjectionMatrixInitialization(Enum): + """Initialization strategies for kernel projection matrices.""" + + MASKING = "MASKING" + """Axis-aligned masking (random selection of input dimensions).""" + + ORTHONORMAL = "ORTHONORMAL" + """Random orthonormal basis.""" + + PLS = "PLS" + """Partial Least Squares (PLS) directions.""" + + SPHERICAL = "SPHERICAL" + """Uniform random sampling on the unit sphere.""" + + +def _make_projection_matrices( + n_projections: int, + n_matrices: int, + initialization: ProjectionMatrixInitialization, + train_x: np.ndarray, + train_y: np.ndarray, +) -> np.ndarray: + """Create a collection of projection matrices. + + Args: + n_projections: The number of projections in each matrix. + n_matrices: The number of projection matrices to create. + initialization: The initialization strategy to use. + train_x: The training inputs. + train_y: The training outputs. + + Returns: + An array of shape ``(n_matrices, n_input_dims, n_projections)`` containing the + created matrices. + """ + n_input_dims = train_x.shape[-1] + + if n_matrices == 0: + return np.empty((0, n_input_dims, n_projections)) + + if initialization is ProjectionMatrixInitialization.MASKING: + matrices = [] + for _ in range(n_matrices): + matrix = np.eye(n_input_dims) + matrix = matrix[ + :, np.random.choice(n_input_dims, n_projections, replace=False) + ] + matrices.append(matrix) + + elif initialization is ProjectionMatrixInitialization.ORTHONORMAL: + matrices = [] + for _ in range(n_matrices): + random_matrix = np.random.randn(n_input_dims, n_projections) + q, _ = np.linalg.qr(random_matrix) + matrices.append(q[:, :n_projections]) + + elif initialization is ProjectionMatrixInitialization.PLS: + from sklearn.cross_decomposition import PLSRegression + + pls = PLSRegression(n_components=n_projections) + pls.fit(train_x, train_y) + M = pls.x_rotations_ + + # IMPROVE: One could use the remaining PLS directions for the next matrices + # until they are exhausted, then switch to orthonormal. + matrices = [ + M, + *_make_projection_matrices( + n_projections=n_projections, + n_matrices=n_matrices - 1, + initialization=ProjectionMatrixInitialization.ORTHONORMAL, + train_x=train_x, + train_y=train_y, + ), + ] + + elif initialization is ProjectionMatrixInitialization.SPHERICAL: + matrices = [] + for _ in range(n_matrices): + matrix = np.random.randn(n_input_dims, n_projections) + matrix = matrix / np.linalg.norm(matrix, axis=0, keepdims=True) + matrices.append(matrix) + + else: + assert_never(initialization) + + return np.stack(matrices) if n_matrices > 1 else matrices[0][None, ...] + + class KernelFactory(Protocol): """A protocol defining the interface expected for kernel factories.""" @@ -41,9 +139,61 @@ def __call__( return self.kernel -def to_kernel_factory(x: Kernel | KernelFactory, /) -> KernelFactory: - """Wrap a kernel into a plain kernel factory (with factory passthrough).""" - return x.to_factory() if isinstance(x, Kernel) else x +@define(frozen=True) +class ProjectionKernelFactory(KernelFactory, SerialMixin): + """A factory producing projected kernels.""" + + n_projections: int = field(validator=instance_of(int)) + """The number of projections to be used in each projection matrix.""" + + n_matrices: int = field(validator=instance_of(int)) + """The number of projection matrices to be used.""" + + initialization: ProjectionMatrixInitialization = field( + converter=ProjectionMatrixInitialization + ) + """The initialization strategy for the projection matrices.""" + + base_kernel_factory: KernelFactory = field( + alias="kernel_or_factory", + converter=to_kernel_factory, + ) + """The factory creating the base kernel to be applied to the projected inputs.""" + + learn_projection: bool = field( + default=True, validator=instance_of(bool), kw_only=True + ) + """See :attr:`baybe.kernels.composite.ProjectionKernel.learn_projection`.""" + + @base_kernel_factory.default + def _default_base_kernel_factory(self) -> KernelFactory: + from baybe.surrogates.gaussian_process.presets.default import ( + DefaultKernelFactory, + ) + + return DefaultKernelFactory() + + @override + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel: + base_kernel = self.base_kernel_factory(searchspace, train_x, train_y) + projection_matrices = _make_projection_matrices( + n_projections=self.n_projections, + n_matrices=self.n_matrices, + initialization=self.initialization, + train_x=train_x.numpy(), + train_y=train_y.numpy(), + ) + kernels = [ + ProjectionKernel( + base_kernel=base_kernel, + projection_matrix=m, + learn_projection=self.learn_projection, + ) + for m in projection_matrices + ] + return AdditiveKernel(kernels) if self.n_matrices > 1 else kernels[0] # Collect leftover original slotted classes processed by `attrs.define` diff --git a/docs/references.bib b/docs/references.bib index 307529d299..a28f0a728c 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -8,3 +8,19 @@ @inproceedings{NIPS2007_66368270 volume = {20}, year = {2007} } +@InProceedings{pmlr-v235-hvarfner24a, + title = {Vanilla {B}ayesian Optimization Performs Great in High Dimensions}, + author = {Hvarfner, Carl and Hellsten, Erik Orm and Nardi, Luigi}, + booktitle = {Proceedings of the 41st International Conference on Machine Learning}, + pages = {20793--20817}, + year = {2024}, + editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, + volume = {235}, + series = {Proceedings of Machine Learning Research}, + month = {21--27 Jul}, + publisher = {PMLR}, + pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/hvarfner24a/hvarfner24a.pdf}, + url = {https://proceedings.mlr.press/v235/hvarfner24a.html}, + abstract = {High-dimensional optimization problems have long been considered the Achilles’ heel of Bayesian optimization algorithms. Spurred by the curse of dimensionality, a large collection of algorithms aim to make BO more performant in this setting, commonly by imposing various simplifying assumptions on the objective, thereby decreasing its presumed complexity. In this paper, we identify the degeneracies that make vanilla BO poorly suited to high-dimensional tasks, and further show how existing algorithms address these degeneracies through the lens of model complexity. Motivated by the model complexity measure, we derive an enhancement to the prior assumptions that are typical of the vanilla BO algorithm, which reduces the complexity to manageable levels without imposing structural restrictions on the objective. Our modification - a simple scaling of the Gaussian process lengthscale prior in the dimensionality - reveals that standard BO works drastically better than previously thought in high dimensions. Our insights are supplemented by substantial out-performance of existing state-of-the-art on multiple commonly considered real-world high-dimensional tasks.} +} + diff --git a/docs/userguide/transfer_learning.md b/docs/userguide/transfer_learning.md index 6f509df31b..83b20a4f27 100644 --- a/docs/userguide/transfer_learning.md +++ b/docs/userguide/transfer_learning.md @@ -186,6 +186,7 @@ on the optimization: ``` ```{bibliography} +# :filter: docname in docnames ``` [`TaskParameter`]: baybe.parameters.categorical.TaskParameter \ No newline at end of file diff --git a/examples/Custom_Surrogates/projection_kernel.py b/examples/Custom_Surrogates/projection_kernel.py new file mode 100644 index 0000000000..bd00c9db71 --- /dev/null +++ b/examples/Custom_Surrogates/projection_kernel.py @@ -0,0 +1,201 @@ +# # Gaussian Processes in High-Dimensional Spaces + +# A common challenge when applying Gaussian process models to high-dimensional spaces is +# to avoid model overfitting, which can easily happen when the hyperparameters of the +# model are not chosen with sufficient care {cite:p}`pmlr-v235-hvarfner24a`. While there +# exist various strategies to mitigate this issue, a simple yet effective approach is +# to model the data on a lower-dimensional subspace of the original parameter space. +# By reducing the effective dimensionality, this technique improves sample efficiency +# and reduces the risk of overfitting. Here, we demonstrate how to use BayBE's +# {class}`~baybe.kernels.composite.ProjectionKernel` to implement this idea. + + +# ## Imports + +import os +from typing import Any + +import numpy as np +import pandas as pd +import seaborn as sns +from matplotlib import pyplot as plt +from scipy.special import expit +from scipy.stats import spearmanr +from sklearn.metrics import r2_score + +from baybe.kernels import MaternKernel, ProjectionKernel +from baybe.parameters import NumericalContinuousParameter +from baybe.searchspace import SearchSpace +from baybe.surrogates import GaussianProcessSurrogate +from baybe.surrogates.base import Surrogate +from baybe.surrogates.gaussian_process.kernel_factory import ProjectionKernelFactory +from baybe.surrogates.gaussian_process.presets.default import DefaultKernelFactory +from baybe.targets import NumericalTarget +from baybe.utils.dataframe import arrays_to_dataframes +from baybe.utils.random import set_random_seed + +# ## Settings + +# Before we start, let us define some general settings for this example: + +# Simulation Settings +SMOKE_TEST = "SMOKE_TEST" in os.environ +N_MC = 1 if SMOKE_TEST else 10 + +# Data Settings +N_TRAIN_DATA = [2] if SMOKE_TEST else [20, 50, 100, 200] +N_TEST_DATA = 100 + +# Problem Settings +N_DIMENSIONS_TOTAL = 3 if SMOKE_TEST else 20 +N_DIMENSIONS_SUBSPACE = 2 + +# Projection Kernel Settings +N_PROJECTIONS = 2 # Dimensionality of the assumed subspace +N_MATRICES = 1 if SMOKE_TEST else 3 # Number of projection matrices to use + +set_random_seed(1337) + +# ## The Scenario + +# To keep things simple, we consider a single target that needs to be modeled: + +objective = NumericalTarget("t").to_objective() + +# We consider a situation where the available regression parameters span a rather +# high-dimensional space: + +parameters = [ + NumericalContinuousParameter(f"p{i}", (-1, 1)) for i in range(N_DIMENSIONS_TOTAL) +] +searchspace = SearchSpace.from_product(parameters) + +# However, we assume that the actual relationship between the regression parameters and +# the target can be expressed via a low-dimensional function that is embedded into the +# higher-dimensional parameter space. To benchmark the proposed approach, let us +# generate this ground truth subspace by randomly drawing some unit-length vectors +# spanning it: + +subspace = np.random.randn(N_DIMENSIONS_TOTAL, N_DIMENSIONS_SUBSPACE) +subspace /= np.linalg.norm(subspace, axis=0, keepdims=True) + + +# As a last component to describe the modeling scenario, we define the corresponding +# ground truth function that operates on this subspace: + + +@arrays_to_dataframes(searchspace.parameter_names, ["t"]) +def low_dimensional_function(array: np.ndarray, /) -> np.ndarray: + """The low-dimensional ground truth function. + + Computes the sigmoid transformation of the norm of its input. + """ + assert array.shape[-1] == N_DIMENSIONS_TOTAL + return expit(np.linalg.norm(array @ subspace, axis=1, keepdims=True)) + + +# ## Subspace Modeling + +# Now, let us turn to the modeling side. First, we define a little helper that lets +# us easily train and evaluate different surrogate models: + + +def predict( + model: Surrogate, train_data: pd.DataFrame, test_data: pd.DataFrame +) -> pd.DataFrame: + """Fit and evaluate a surrogate model.""" + model.fit(searchspace, objective, train_data) + return model.posterior_stats(test_data[list(searchspace.parameter_names)]) + + +# We compare the following models: +# - **Vanilla:** +# The vanilla Gaussian process model operating on the **full** parameter space. +# - **Learned Projection:** +# A Gaussian process model operating on **learned** low-dimensional subspaces. +# - **Ideal Projection:** +# A Gaussian process model operating on the **ground truth** low-dimensional subspace. + +models: dict[str, Surrogate] = { + "Vanilla": GaussianProcessSurrogate(), + "Learned Projection": GaussianProcessSurrogate( + ProjectionKernelFactory( + n_projections=N_PROJECTIONS, + n_matrices=N_MATRICES, + initialization="PLS", + kernel_or_factory=DefaultKernelFactory(), + learn_projection=True, + ) + ), + "Ideal Projection": GaussianProcessSurrogate( + ProjectionKernel( + MaternKernel(), + projection_matrix=subspace, + learn_projection=False, + ) + ), +} + +# We can now evaluate the regression performance of these models for different training +# data set sizes. The entire process is repeated for several Monte Carlo iterations: + +metrics: list[dict[str, Any]] = [] +for mc in range(N_MC): + test_data = searchspace.continuous.sample_uniform(N_TEST_DATA) + test_data["t"] = low_dimensional_function(test_data) + for n_train_data in N_TRAIN_DATA: + train_data = searchspace.continuous.sample_uniform(n_train_data) + train_data["t"] = low_dimensional_function(train_data) + for model_name, model in models.items(): + predictions = predict(model, train_data, test_data) + r2 = r2_score(test_data["t"], predictions["t_mean"]) + correlation = spearmanr(test_data["t"], predictions["t_mean"]).correlation + metrics.append( + { + "mc_iteration": mc, + "model": model_name, + "n_train_data": n_train_data, + "rank_correlation": correlation, + "r2_score": r2, + } + ) + +df = pd.DataFrame.from_records(metrics) + +# ## Results + +# Finally, we visualize the resulting regression metrics. As expected, when operating on +# the ground truth subspace, the Gaussian process can easily capture the structure of +# the data, even for relatively small training set sizes. More importantly, even when +# learning the subspace from the training data itself (the realistic scenario where the +# ground truth is unknown), the projection approach identifies patterns in the data +# significantly earlier compared to the vanilla approach. This effect is more prominent +# for the underlying ordering of the target data (as reflected by the rank correlation) +# than for the actual prediction values. Nevertheless, when applied in a Bayesian +# optimization context, rank correlation is oftentimes the more important metric, since +# data-driven optimization heavily relies on correctly ranking candidates. + +fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) + +# fmt: off +sns.lineplot(df, x="n_train_data", y="rank_correlation", hue="model", ax=ax1) +ax1.set_title("Spearman Correlation") +ax1.set_xlabel("Number of Training Points") +ax1.set_ylabel("Rank Correlation") +# fmt: on + +# fmt: off +sns.lineplot(df, x="n_train_data", y="r2_score", hue="model", ax=ax2) +ax2.set_title("R² Score") +ax2.set_xlabel("Number of Training Points") +ax2.set_ylabel("R² Score") +# fmt: on + +plt.tight_layout() +if not SMOKE_TEST: + plt.savefig("projection_kernel.svg") +plt.show() + +# ```{bibliography} +# :filter: docname in docnames +# ``` diff --git a/examples/Custom_Surrogates/projection_kernel.svg b/examples/Custom_Surrogates/projection_kernel.svg new file mode 100644 index 0000000000..0ef571aabb --- /dev/null +++ b/examples/Custom_Surrogates/projection_kernel.svg @@ -0,0 +1,1834 @@ + + + + + + + + 2025-11-27T09:53:17.925349 + image/svg+xml + + + Matplotlib v3.10.7, https://matplotlib.org/ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/tests/hypothesis_strategies/kernels.py b/tests/hypothesis_strategies/kernels.py index 2f2ac87432..1cf5a6c847 100644 --- a/tests/hypothesis_strategies/kernels.py +++ b/tests/hypothesis_strategies/kernels.py @@ -3,6 +3,8 @@ from enum import Enum import hypothesis.strategies as st +from hypothesis.extra.numpy import arrays +from typing_extensions import assert_never from baybe.kernels.basic import ( LinearKernel, @@ -14,7 +16,12 @@ RFFKernel, RQKernel, ) -from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel +from baybe.kernels.composite import ( + AdditiveKernel, + ProductKernel, + ProjectionKernel, + ScaleKernel, +) from tests.hypothesis_strategies.basic import positive_finite_floats from tests.hypothesis_strategies.priors import priors @@ -25,6 +32,7 @@ class KernelType(Enum): SINGLE = "SINGLE" ADDITIVE = "ADDITIVE" PRODUCT = "PRODUCT" + PROJECTION = "PROJECTION" linear_kernels = st.builds( @@ -121,6 +129,26 @@ def single_kernels(draw: st.DrawFn): return base_kernel +@st.composite +def projection_kernels(draw: st.DrawFn): + """Generate :class:`baybe.kernels.composite.ProjectionKernel`.""" + MAX_DIM = 5 + shape = draw( + st.tuples( + st.integers(min_value=1, max_value=MAX_DIM), + st.integers(min_value=1, max_value=MAX_DIM), + ) + ) + projection_matrix = draw( + arrays(float, shape, elements=st.floats(allow_nan=False, allow_infinity=False)) + ) + base_kernel = draw(single_kernels()) + learn_projection = draw(st.booleans()) + return ProjectionKernel( + base_kernel, projection_matrix, learn_projection=learn_projection + ) + + @st.composite def kernels(draw: st.DrawFn): """Generate :class:`baybe.kernels.base.Kernel`.""" @@ -128,9 +156,12 @@ def kernels(draw: st.DrawFn): if kernel_type is KernelType.SINGLE: return draw(single_kernels()) - - base_kernels = draw(st.lists(single_kernels(), min_size=2)) if kernel_type is KernelType.ADDITIVE: + base_kernels = draw(st.lists(single_kernels(), min_size=2)) return AdditiveKernel(base_kernels) if kernel_type is KernelType.PRODUCT: + base_kernels = draw(st.lists(single_kernels(), min_size=2)) return ProductKernel(base_kernels) + if kernel_type is KernelType.PROJECTION: + return draw(projection_kernels()) + assert_never(kernel_type)