diff --git a/CHANGELOG.md b/CHANGELOG.md index b2790450b..03dba78b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,8 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `mypy` for search space and objectives - Class hierarchy for objectives - Deserialization is now also possible from optional class name abbreviations -- `Kernel`, `MaternKernel`, `AdditiveKernel`, `ProductKernel` and `ScaleKernel` - classes for specifying kernels +- `AdditiveKernel`, `LinearKernel`, `MaternKernel`, `PeriodicKernel`, + `PiecewisePolynomialKernel`, `PolynomialKernel`, `ProductKernel`, `RBFKernel`, + `RFFKernel`, `RQKernel`, `ScaleKernel` classes for specifying kernels +- `GammaPrior`, `HalfCauchyPrior`, `NormalPrior`, `HalfNormalPrior`, `LogNormalPrior` + and `SmoothedBoxPrior` classes for specifying priors - `KernelFactory` protocol enabling context-dependent construction of kernels - Preset mechanism for `GaussianProcessSurrogate` - `hypothesis` strategies and roundtrip test for kernels, constraints, objectives, @@ -18,8 +21,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - New acquisition functions: `qSR`, `qNEI`, `LogEI`, `qLogEI`, `qLogNEI` - Serialization user guide - Basic deserialization tests using different class type specifiers -- `GammaPrior`, `HalfCauchyPrior`, `NormalPrior`, `HalfNormalPrior`, `LogNormalPrior` - and `SmoothedBoxPrior` can now be chosen as lengthscale prior - Environment variables user guide - Utility for estimating memory requirements of discrete product search space diff --git a/baybe/kernels/__init__.py b/baybe/kernels/__init__.py index 7eb009e2e..9323a2b63 100644 --- a/baybe/kernels/__init__.py +++ b/baybe/kernels/__init__.py @@ -1,11 +1,31 @@ -"""Kernels for Gaussian process surrogate models.""" +"""Kernels for Gaussian process surrogate models. -from baybe.kernels.basic import MaternKernel +The kernel classes mimic classes from GPyTorch. For details on specification and +arguments see https://docs.gpytorch.ai/en/stable/kernels.html. +""" + +from baybe.kernels.basic import ( + LinearKernel, + MaternKernel, + PeriodicKernel, + PiecewisePolynomialKernel, + PolynomialKernel, + RBFKernel, + RFFKernel, + RQKernel, +) from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel __all__ = [ "AdditiveKernel", + "LinearKernel", "MaternKernel", + "PeriodicKernel", + "PiecewisePolynomialKernel", + "PolynomialKernel", "ProductKernel", + "RBFKernel", + "RFFKernel", + "RQKernel", "ScaleKernel", ] diff --git a/baybe/kernels/basic.py b/baybe/kernels/basic.py index f4805e69a..4bfabaf04 100644 --- a/baybe/kernels/basic.py +++ b/baybe/kernels/basic.py @@ -4,7 +4,7 @@ from attrs import define, field from attrs.converters import optional as optional_c -from attrs.validators import in_, instance_of +from attrs.validators import ge, in_, instance_of from attrs.validators import optional as optional_v from baybe.kernels.base import Kernel @@ -13,6 +13,34 @@ from baybe.utils.validation import finite_float +@define(frozen=True) +class LinearKernel(Kernel): + """A linear kernel.""" + + variance_prior: Optional[Prior] = field( + default=None, validator=optional_v(instance_of(Prior)) + ) + """An optional prior on the kernel variance parameter.""" + + variance_initial_value: Optional[float] = field( + default=None, converter=optional_c(float), validator=optional_v(finite_float) + ) + """An optional initial value for the kernel variance parameter.""" + + def to_gpytorch(self, *args, **kwargs): # noqa: D102 + # See base class. + import torch + + from baybe.utils.torch import DTypeFloatTorch + + gpytorch_kernel = super().to_gpytorch(*args, **kwargs) + if (initial_value := self.variance_initial_value) is not None: + gpytorch_kernel.variance = torch.tensor( + initial_value, dtype=DTypeFloatTorch + ) + return gpytorch_kernel + + @define(frozen=True) class MaternKernel(Kernel): """A Matern kernel using a smoothness parameter.""" @@ -34,3 +62,138 @@ class MaternKernel(Kernel): default=None, converter=optional_c(float), validator=optional_v(finite_float) ) """An optional initial value for the kernel lengthscale.""" + + +@define(frozen=True) +class PeriodicKernel(Kernel): + """A periodic kernel.""" + + lengthscale_prior: Optional[Prior] = field( + default=None, validator=optional_v(instance_of(Prior)) + ) + """An optional prior on the kernel lengthscale.""" + + lengthscale_initial_value: Optional[float] = field( + default=None, converter=optional_c(float), validator=optional_v(finite_float) + ) + """An optional initial value for the kernel lengthscale.""" + + period_length_prior: Optional[Prior] = field( + default=None, validator=optional_v(instance_of(Prior)) + ) + """An optional prior on the kernel period length.""" + + period_length_initial_value: Optional[float] = field( + default=None, converter=optional_c(float), validator=optional_v(finite_float) + ) + """An optional initial value for the kernel period length.""" + + def to_gpytorch(self, *args, **kwargs): # noqa: D102 + # See base class. + import torch + + from baybe.utils.torch import DTypeFloatTorch + + gpytorch_kernel = super().to_gpytorch(*args, **kwargs) + # lengthscale is handled by the base class + + if (initial_value := self.period_length_initial_value) is not None: + gpytorch_kernel.period_length = torch.tensor( + initial_value, dtype=DTypeFloatTorch + ) + return gpytorch_kernel + + +@define(frozen=True) +class PiecewisePolynomialKernel(Kernel): + """A piecewise polynomial kernel.""" + + q: int = field(validator=in_([0, 1, 2, 3]), default=2) + """A smoothness parameter.""" + + lengthscale_prior: Optional[Prior] = field( + default=None, validator=optional_v(instance_of(Prior)) + ) + """An optional prior on the kernel lengthscale.""" + + lengthscale_initial_value: Optional[float] = field( + default=None, converter=optional_c(float), validator=optional_v(finite_float) + ) + """An optional initial value for the kernel lengthscale.""" + + +@define(frozen=True) +class PolynomialKernel(Kernel): + """A polynomial kernel.""" + + power: int = field(validator=[instance_of(int), ge(0)]) + """The power of the polynomial term.""" + + offset_prior: Optional[Prior] = field( + default=None, validator=optional_v(instance_of(Prior)) + ) + """An optional prior on the kernel offset.""" + + offset_initial_value: Optional[float] = field( + default=None, converter=optional_c(float), validator=optional_v(finite_float) + ) + """An optional initial value for the kernel offset.""" + + def to_gpytorch(self, *args, **kwargs): # noqa: D102 + # See base class. + import torch + + from baybe.utils.torch import DTypeFloatTorch + + gpytorch_kernel = super().to_gpytorch(*args, **kwargs) + if (initial_value := self.offset_initial_value) is not None: + gpytorch_kernel.offset = torch.tensor(initial_value, dtype=DTypeFloatTorch) + return gpytorch_kernel + + +@define(frozen=True) +class RBFKernel(Kernel): + """A radial basis function (RBF) kernel.""" + + lengthscale_prior: Optional[Prior] = field( + default=None, validator=optional_v(instance_of(Prior)) + ) + """An optional prior on the kernel lengthscale.""" + + lengthscale_initial_value: Optional[float] = field( + default=None, converter=optional_c(float), validator=optional_v(finite_float) + ) + """An optional initial value for the kernel lengthscale.""" + + +@define(frozen=True) +class RFFKernel(Kernel): + """A random Fourier features (RFF) kernel.""" + + num_samples: int = field(validator=[instance_of(int), ge(1)]) + """The number of frequencies to draw.""" + + lengthscale_prior: Optional[Prior] = field( + default=None, validator=optional_v(instance_of(Prior)) + ) + """An optional prior on the kernel lengthscale.""" + + lengthscale_initial_value: Optional[float] = field( + default=None, converter=optional_c(float), validator=optional_v(finite_float) + ) + """An optional initial value for the kernel lengthscale.""" + + +@define(frozen=True) +class RQKernel(Kernel): + """A rational quadratic (RQ) kernel.""" + + lengthscale_prior: Optional[Prior] = field( + default=None, validator=optional_v(instance_of(Prior)) + ) + """An optional prior on the kernel lengthscale.""" + + lengthscale_initial_value: Optional[float] = field( + default=None, converter=optional_c(float), validator=optional_v(finite_float) + ) + """An optional initial value for the kernel lengthscale.""" diff --git a/baybe/kernels/composite.py b/baybe/kernels/composite.py index e9d62bb82..07a04c988 100644 --- a/baybe/kernels/composite.py +++ b/baybe/kernels/composite.py @@ -1,5 +1,5 @@ """Composite kernels (that is, kernels composed of other kernels).""" - +from functools import reduce from operator import add, mul from typing import Optional @@ -56,7 +56,7 @@ class AdditiveKernel(Kernel): def to_gpytorch(self, *args, **kwargs): # noqa: D102 # See base class. - return add(*(k.to_gpytorch(*args, **kwargs) for k in self.base_kernels)) + return reduce(add, (k.to_gpytorch(*args, **kwargs) for k in self.base_kernels)) @define(frozen=True) @@ -71,4 +71,4 @@ class ProductKernel(Kernel): def to_gpytorch(self, *args, **kwargs): # noqa: D102 # See base class. - return mul(*(k.to_gpytorch(*args, **kwargs) for k in self.base_kernels)) + return reduce(mul, (k.to_gpytorch(*args, **kwargs) for k in self.base_kernels)) diff --git a/baybe/priors/__init__.py b/baybe/priors/__init__.py index b133c61f9..77f71ae67 100644 --- a/baybe/priors/__init__.py +++ b/baybe/priors/__init__.py @@ -1,4 +1,8 @@ -"""Prior distributions.""" +"""Prior distributions. + +The prior classes mimic classes from GPyTorch. For details on specification and +arguments see https://docs.gpytorch.ai/en/stable/priors.html. +""" from baybe.priors.basic import ( GammaPrior, diff --git a/tests/hypothesis_strategies/kernels.py b/tests/hypothesis_strategies/kernels.py index d47fb179c..479f6233f 100644 --- a/tests/hypothesis_strategies/kernels.py +++ b/tests/hypothesis_strategies/kernels.py @@ -4,7 +4,16 @@ import hypothesis.strategies as st -from baybe.kernels.basic import MaternKernel +from baybe.kernels.basic import ( + LinearKernel, + MaternKernel, + PeriodicKernel, + PiecewisePolynomialKernel, + PolynomialKernel, + RBFKernel, + RFFKernel, + RQKernel, +) from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel from ..hypothesis_strategies.basic import finite_floats @@ -19,6 +28,13 @@ class KernelType(Enum): PRODUCT = "PRODUCT" +linear_kernels = st.builds( + LinearKernel, + variance_prior=st.one_of(st.none(), priors), + variance_initial_value=st.one_of(st.none(), finite_floats()), +) +"""A strategy that generates linear kernels.""" + matern_kernels = st.builds( MaternKernel, nu=st.sampled_from((0.5, 1.5, 2.5)), @@ -27,8 +43,65 @@ class KernelType(Enum): ) """A strategy that generates Matern kernels.""" +periodic_kernels = st.builds( + PeriodicKernel, + lengthscale_prior=st.one_of(st.none(), priors), + lengthscale_initial_value=st.one_of(st.none(), finite_floats()), + period_length_prior=st.one_of(st.none(), priors), + period_length_initial_value=st.one_of(st.none(), finite_floats()), +) +"""A strategy that generates periodic kernels.""" + +piecewise_polynomial_kernels = st.builds( + PiecewisePolynomialKernel, + q=st.integers(min_value=0, max_value=3), + lengthscale_prior=st.one_of(st.none(), priors), + lengthscale_initial_value=st.one_of(st.none(), finite_floats()), +) +"""A strategy that generates piecewise polynomial kernels.""" + +polynomial_kernels = st.builds( + PolynomialKernel, + power=st.integers(min_value=0), + offset_prior=st.one_of(st.none(), priors), + offset_initial_value=st.one_of(st.none(), finite_floats()), +) +"""A strategy that generates polynomial kernels.""" -base_kernels = st.one_of([matern_kernels]) +rbf_kernels = st.builds( + RBFKernel, + lengthscale_prior=st.one_of(st.none(), priors), + lengthscale_initial_value=st.one_of(st.none(), finite_floats()), +) +"""A strategy that generates radial basis function (RBF) kernels.""" + +rff_kernels = st.builds( + RFFKernel, + num_samples=st.integers(min_value=1), + lengthscale_prior=st.one_of(st.none(), priors), + lengthscale_initial_value=st.one_of(st.none(), finite_floats()), +) +"""A strategy that generates random Fourier features (RFF) kernels.""" + +rq_kernels = st.builds( + RQKernel, + lengthscale_prior=st.one_of(st.none(), priors), + lengthscale_initial_value=st.one_of(st.none(), finite_floats()), +) +"""A strategy that generates rational quadratic (RQ) kernels.""" + +base_kernels = st.one_of( + [ + matern_kernels, # on top because it is the default for many use cases + linear_kernels, + rbf_kernels, + rq_kernels, + rff_kernels, + piecewise_polynomial_kernels, + polynomial_kernels, + periodic_kernels, + ] +) """A strategy that generates base kernels to be used within more complex kernels.""" diff --git a/tests/test_iterations.py b/tests/test_iterations.py index 094cb4833..2105f76b1 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -4,7 +4,17 @@ import pytest from baybe.acquisition.base import AcquisitionFunction -from baybe.kernels.basic import MaternKernel +from baybe.kernels.base import Kernel +from baybe.kernels.basic import ( + LinearKernel, + MaternKernel, + PeriodicKernel, + PiecewisePolynomialKernel, + PolynomialKernel, + RBFKernel, + RFFKernel, + RQKernel, +) from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel from baybe.priors import ( GammaPrior, @@ -131,17 +141,40 @@ SmoothedBoxPrior(0, 3, 0.1), ] -valid_base_kernels = [MaternKernel(lengthscale_prior=prior) for prior in valid_priors] +valid_base_kernels: list[Kernel] = [ + cls(**arg_dict) + for prior in valid_priors + for cls, arg_dict in [ + (MaternKernel, {"lengthscale_prior": prior}), + (LinearKernel, {"variance_prior": prior}), + (PeriodicKernel, {"period_length_prior": prior}), + (PeriodicKernel, {"lengthscale_prior": prior}), + (PiecewisePolynomialKernel, {"lengthscale_prior": prior}), + (PolynomialKernel, {"offset_prior": prior, "power": 2}), + (RBFKernel, {"lengthscale_prior": prior}), + (RQKernel, {"lengthscale_prior": prior}), + (RFFKernel, {"lengthscale_prior": prior, "num_samples": 5}), + ] +] valid_scale_kernels = [ - ScaleKernel(base_kernel=base_kernel, outputscale_prior=prior) + ScaleKernel(base_kernel=base_kernel, outputscale_prior=HalfCauchyPrior(scale=1)) for base_kernel in valid_base_kernels - for prior in valid_priors ] valid_composite_kernels = [ AdditiveKernel([MaternKernel(1.5), MaternKernel(2.5)]), + AdditiveKernel([PolynomialKernel(1), PolynomialKernel(2), PolynomialKernel(3)]), + AdditiveKernel([RBFKernel(), RQKernel(), PolynomialKernel(1)]), ProductKernel([MaternKernel(1.5), MaternKernel(2.5)]), + ProductKernel([RBFKernel(), RQKernel(), PolynomialKernel(1)]), + ProductKernel([PolynomialKernel(1), PolynomialKernel(2), PolynomialKernel(3)]), + AdditiveKernel( + [ + ProductKernel([MaternKernel(1.5), MaternKernel(2.5)]), + AdditiveKernel([MaternKernel(1.5), MaternKernel(2.5)]), + ] + ), ] valid_kernels = valid_base_kernels + valid_scale_kernels + valid_composite_kernels @@ -178,15 +211,6 @@ def test_iter_nonmc_acquisition_function(campaign, n_iterations, batch_size): run_iterations(campaign, n_iterations, batch_size) -@pytest.mark.slow -@pytest.mark.parametrize( - "lengthscale_prior", valid_priors, ids=[c.__class__ for c in valid_priors] -) -@pytest.mark.parametrize("n_iterations", [3], ids=["i3"]) -def test_iter_prior(campaign, n_iterations, batch_size): - run_iterations(campaign, n_iterations, batch_size) - - @pytest.mark.slow @pytest.mark.parametrize( "kernel", valid_kernels, ids=[c.__class__ for c in valid_kernels]