Implement Priors from "Vanilla Bayesian Optimization Performs Great i…

…n High Dimensions" (#402)

bofire/data_models/kernels/aggregation.py

@@ -4,7 +4,7 @@
 from bofire.data_models.kernels.continuous import LinearKernel, MaternKernel, RBFKernel
 from bofire.data_models.kernels.kernel import Kernel
 from bofire.data_models.kernels.molecular import TanimotoKernel
-from bofire.data_models.priors.api import AnyPrior
+from bofire.data_models.priors.api import AnyGeneralPrior
 
 
 class AdditiveKernel(Kernel):
@@ -52,7 +52,7 @@ class ScaleKernel(Kernel):
         TanimotoKernel,
         "ScaleKernel",
     ]
-    outputscale_prior: Optional[AnyPrior] = None
+    outputscale_prior: Optional[AnyGeneralPrior] = None
 
 
 AdditiveKernel.model_rebuild()

bofire/data_models/kernels/continuous.py

@@ -3,7 +3,7 @@
 from pydantic import field_validator
 
 from bofire.data_models.kernels.kernel import Kernel
-from bofire.data_models.priors.api import AnyPrior
+from bofire.data_models.priors.api import AnyGeneralPrior, AnyPrior
 
 
 class ContinuousKernel(Kernel):
@@ -31,10 +31,10 @@ def validate_nu(cls, nu):
 
 class LinearKernel(ContinuousKernel):
     type: Literal["LinearKernel"] = "LinearKernel"
-    variance_prior: Optional[AnyPrior] = None
+    variance_prior: Optional[AnyGeneralPrior] = None
 
 
 class PolynomialKernel(ContinuousKernel):
     type: Literal["PolynomialKernel"] = "PolynomialKernel"
-    offset_prior: Optional[AnyPrior] = None
+    offset_prior: Optional[AnyGeneralPrior] = None
     power: int = 2

bofire/data_models/priors/api.py

@@ -3,7 +3,11 @@
 
 from bofire.data_models.priors.gamma import GammaPrior
 from bofire.data_models.priors.lkj import LKJPrior
-from bofire.data_models.priors.normal import NormalPrior
+from bofire.data_models.priors.normal import (
+    DimensionalityScaledLogNormalPrior,
+    LogNormalPrior,
+    NormalPrior,
+)
 from bofire.data_models.priors.prior import Prior
 
 AbstractPrior = Prior
@@ -12,8 +16,14 @@
     GammaPrior,
     NormalPrior,
     LKJPrior,
+    LogNormalPrior,
+    DimensionalityScaledLogNormalPrior,
 ]
 
+# these are priors that are generally applicable
+# and do not depend on problem specific extra parameters
+AnyGeneralPrior = Union[GammaPrior, NormalPrior, LKJPrior, LogNormalPrior]
+
 # default priors of interest
 # botorch defaults
 BOTORCH_LENGTHCALE_PRIOR = partial(GammaPrior, concentration=3.0, rate=6.0)
@@ -32,3 +42,7 @@
 LKJ_PRIOR = partial(
     LKJPrior, shape=2.0, sd_prior=GammaPrior(concentration=2.0, rate=0.15)
 )
+
+# Hvarfner priors
+HVARFNER_NOISE_PRIOR = partial(LogNormalPrior, loc=-4, scale=1)
+HVARFNER_LENGTHSCALE_PRIOR = DimensionalityScaledLogNormalPrior

bofire/data_models/priors/normal.py

@@ -1,5 +1,6 @@
 from typing import Literal
 
+import numpy as np
 from pydantic import PositiveFloat
 
 from bofire.data_models.priors.prior import Prior
@@ -16,3 +17,31 @@ class NormalPrior(Prior):
     type: Literal["NormalPrior"] = "NormalPrior"
     loc: float
     scale: PositiveFloat
+
+
+class LogNormalPrior(Prior):
+    """Log-normal prior based on the log-normal distribution
+
+    Attributes:
+        loc(float): mean/center of the log-normal distribution
+        scale(PositiveFloat): width of the log-normal distribution
+    """
+
+    type: Literal["LogNormalPrior"] = "LogNormalPrior"
+    loc: float
+    scale: float
+
+
+class DimensionalityScaledLogNormalPrior(Prior):
+    """This prior is a log-normal prior where loc and scale are scaled by the dimensionaly of the problem.
+    It was introduced by Hvarfner et al. in their paper https://arxiv.org/abs/2402.02229. More can be read in
+    this excellent blogpost: https://www.miguelgondu.com/blogposts/2024-03-16/when-does-vanilla-gpr-fail/
+    """
+
+    type: Literal[
+        "DimensionalityScaledLogNormalPrior"
+    ] = "DimensionalityScaledLogNormalPrior"
+    loc: PositiveFloat = np.sqrt(2)
+    loc_scaling: PositiveFloat = 0.5
+    scale: PositiveFloat = np.sqrt(3)
+    scale_scaling: float = 0.0

bofire/kernels/mapper.py

@@ -21,7 +21,7 @@ def map_RBFKernel(
         ard_num_dims=len(active_dims) if data_model.ard else None,
         active_dims=active_dims,  # type: ignore
         lengthscale_prior=(
-            priors.map(data_model.lengthscale_prior)
+            priors.map(data_model.lengthscale_prior, d=len(active_dims))
             if data_model.lengthscale_prior is not None
             else None
         ),
@@ -40,7 +40,7 @@ def map_MaternKernel(
         active_dims=active_dims,
         nu=data_model.nu,
         lengthscale_prior=(
-            priors.map(data_model.lengthscale_prior)
+            priors.map(data_model.lengthscale_prior, d=len(active_dims))
             if data_model.lengthscale_prior is not None
             else None
         ),

bofire/plot/prior.py

@@ -1,39 +1,37 @@
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 import numpy as np
-import plotly.express as px
+import plotly.graph_objects as go
 import torch
-
-import bofire.priors.api as priors
-from bofire.data_models.priors.api import AnyPrior
+from gpytorch.priors import Prior
 
 
 def plot_prior_pdf_plotly(
-    prior: AnyPrior,
+    priors: List[Prior],
     lower: float,
     upper: float,
     layout_options: Optional[Dict] = None,
+    labels: Optional[List[str]] = None,
 ):
-    """Plot the probability density function of the prior with plotly.
+    """Plot the probability density function of a gyptorch prior with plotly.
 
     Args:
-        prior (AnyPrior): The prior that should be plotted.
-        lower (float): lower bound for computing the prior pdf.
-        upper (float): upper bound for computing the prior pdf.
-        layout_options (Dict, optional): Layout options passed to plotly. Defaults to {}.
+        prior: The prior that should be plotted.
+        lower: lower bound for computing the prior pdf.
+        upper: upper bound for computing the prior pdf.
+        layout_options: Layout options passed to plotly. Defaults to None.
+        labels: Labels for the priors, that are shown in the plot. Defaults to None.
 
     Returns:
         fig, ax objects of the plot.
     """
-
+    use_labels = labels is not None and len(labels) == len(priors)
     x = np.linspace(lower, upper, 1000)
-
-    fig = px.line(
-        x=x,
-        y=np.exp(priors.map(prior).log_prob(torch.from_numpy(x)).numpy()),
-    )
-
+    fig = go.Figure()
+    for i, prior in enumerate(priors):
+        y = np.exp(prior.log_prob(torch.from_numpy(x)).numpy())
+        label = labels[i] if use_labels else prior.__class__.__name__  # type: ignore
+        fig.add_trace(go.Scatter(x=x, y=y, mode="lines", name=label))
     if layout_options is not None:
         fig.update_layout(layout_options)
-
     return fig

bofire/priors/api.py

@@ -1 +1 @@
-from bofire.priors.mapper import map  # noqa: F401
+from bofire.priors.mapper import map

bofire/priors/mapper.py

@@ -1,32 +1,56 @@
+import math
+
 import gpytorch
 
 import bofire.data_models.priors.api as data_models
 
 
-def map_NormalPrior(data_model: data_models.NormalPrior) -> gpytorch.priors.NormalPrior:
+def map_NormalPrior(
+    data_model: data_models.NormalPrior, **kwargs
+) -> gpytorch.priors.NormalPrior:
     return gpytorch.priors.NormalPrior(loc=data_model.loc, scale=data_model.scale)
 
 
-def map_GammaPrior(data_model: data_models.GammaPrior) -> gpytorch.priors.GammaPrior:
+def map_GammaPrior(
+    data_model: data_models.GammaPrior, **kwargs
+) -> gpytorch.priors.GammaPrior:
     return gpytorch.priors.GammaPrior(
         concentration=data_model.concentration, rate=data_model.rate
     )
 
 
-def map_LKJPrior(data_model: data_models.LKJPrior) -> gpytorch.priors.LKJPrior:
+def map_LKJPrior(
+    data_model: data_models.LKJPrior, **kwargs
+) -> gpytorch.priors.LKJPrior:
     return gpytorch.priors.LKJCovariancePrior(
         n=data_model.n_tasks, eta=data_model.shape, sd_prior=map(data_model.sd_prior)
     )
 
 
+def map_LogNormalPrior(
+    data_model: data_models.LogNormalPrior,
+    **kwargs,
+) -> gpytorch.priors.LogNormalPrior:
+    return gpytorch.priors.LogNormalPrior(loc=data_model.loc, scale=data_model.scale)
+
+
+def map_DimensionalityScaledLogNormalPrior(
+    data_model: data_models.DimensionalityScaledLogNormalPrior, d: int
+) -> gpytorch.priors.LogNormalPrior:
+    return gpytorch.priors.LogNormalPrior(
+        loc=data_model.loc + math.log(d) * data_model.loc_scaling,
+        scale=(data_model.scale**2 + math.log(d) * data_model.scale_scaling) ** 0.5,
+    )
+
+
 PRIOR_MAP = {
     data_models.NormalPrior: map_NormalPrior,
     data_models.GammaPrior: map_GammaPrior,
     data_models.LKJPrior: map_LKJPrior,
+    data_models.LogNormalPrior: map_LogNormalPrior,
+    data_models.DimensionalityScaledLogNormalPrior: map_DimensionalityScaledLogNormalPrior,
 }
 
 
-def map(
-    data_model: data_models.AnyPrior,
-) -> gpytorch.priors.Prior:
-    return PRIOR_MAP[data_model.__class__](data_model)
+def map(data_model: data_models.AnyPrior, **kwargs) -> gpytorch.priors.Prior:
+    return PRIOR_MAP[data_model.__class__](data_model, **kwargs)

bofire/strategies/random.py

@@ -286,18 +286,41 @@ def _sample_from_polytope(
             combined_eqs = unfixed_eqs + unfixed_interpoints  # type: ignore
 
             # now use the hit and run sampler
-            candidates = sample_q_batches_from_polytope(
-                n=1,
-                q=n,
-                bounds=bounds.to(**tkwargs),
-                inequality_constraints=(
-                    unfixed_ineqs if len(unfixed_ineqs) > 0 else None  # type: ignore
-                ),
-                equality_constraints=combined_eqs if len(combined_eqs) > 0 else None,
-                n_burnin=n_burnin,
-                thinning=n_thinning,
-                seed=seed,
-            ).squeeze(dim=0)
+            # this try except is needed as in the main branch of botorch the keyword argument
+            # `thinning` was changed to `n_thinning`, which is not yet in the latest release
+            # so we need to catch the TypeError and use the old keyword argument. As soon as the
+            # new release is out, we can remove this try except block.
+            # TODO: remove this try except block when the new release of botorch is out
+            try:
+                candidates = sample_q_batches_from_polytope(
+                    n=1,
+                    q=n,
+                    bounds=bounds.to(**tkwargs),
+                    inequality_constraints=(
+                        unfixed_ineqs if len(unfixed_ineqs) > 0 else None  # type: ignore
+                    ),
+                    equality_constraints=combined_eqs
+                    if len(combined_eqs) > 0
+                    else None,
+                    n_burnin=n_burnin,
+                    thinning=n_thinning,
+                    seed=seed,
+                ).squeeze(dim=0)
+            except TypeError:
+                candidates = sample_q_batches_from_polytope(
+                    n=1,
+                    q=n,
+                    bounds=bounds.to(**tkwargs),
+                    inequality_constraints=(
+                        unfixed_ineqs if len(unfixed_ineqs) > 0 else None  # type: ignore
+                    ),
+                    equality_constraints=combined_eqs
+                    if len(combined_eqs) > 0
+                    else None,
+                    n_burnin=n_burnin,
+                    n_thinning=n_thinning,  # type: ignore
+                    seed=seed,
+                ).squeeze(dim=0)
 
             # check that the random generated candidates are not always the same
             if (candidates.unique(dim=0).shape[0] != n) and (n > 1):

pyproject.toml

@@ -22,6 +22,8 @@ max-complexity = 18
 
 [tool.ruff.per-file-ignores]
 "bofire/surrogates/api.py" = ["F401"]
+"bofire/data_models/priors/api.py" = ["F401"]
+"bofire/priors/api.py" = ["F401"]
 "bofire/utils/annotated.py" = ["F401"]
 "bofire/data_models/outlier_detection/api.py" = ["F401"]
 "bofire/outlier_detection/api.py" = ["F401"]

tests/bofire/data_models/specs/kernels.py

@@ -1,4 +1,5 @@
 import bofire.data_models.kernels.api as kernels
+from bofire.data_models.priors.api import GammaPrior, LogNormalPrior
 from tests.bofire.data_models.specs.priors import specs as priors
 from tests.bofire.data_models.specs.specs import Specs
 
@@ -12,7 +13,7 @@
 )
 specs.add_valid(
     kernels.LinearKernel,
-    lambda: {"variance_prior": priors.valid().obj().model_dump()},
+    lambda: {"variance_prior": priors.valid(GammaPrior).obj().model_dump()},
 )
 specs.add_valid(
     kernels.MaternKernel,
@@ -44,7 +45,7 @@
     kernels.ScaleKernel,
     lambda: {
         "base_kernel": specs.valid(kernels.LinearKernel).obj().model_dump(),
-        "outputscale_prior": priors.valid().obj().model_dump(),
+        "outputscale_prior": priors.valid(LogNormalPrior).obj().model_dump(),
     },
 )
 specs.add_valid(

tests/bofire/data_models/specs/priors.py

@@ -87,3 +87,19 @@
             },
             error=ValidationError,
         )
+
+
+specs.add_valid(
+    priors.LogNormalPrior, lambda: {"loc": random.random(), "scale": random.random()}
+)
+
+
+specs.add_valid(
+    priors.DimensionalityScaledLogNormalPrior,
+    lambda: {
+        "loc": random.random(),
+        "loc_scaling": random.random(),
+        "scale": random.random(),
+        "scale_scaling": random.random(),
+    },
+)

tests/bofire/plot/test_plot_prior.py

@@ -1,6 +1,7 @@
+import bofire.priors.api as priors
 from bofire.data_models.priors.api import BOTORCH_LENGTHCALE_PRIOR
 from bofire.plot.api import plot_prior_pdf_plotly
 
 
 def test_plot_prior_pdf_plotly():
-    plot_prior_pdf_plotly(BOTORCH_LENGTHCALE_PRIOR(), lower=0, upper=10)
+    plot_prior_pdf_plotly([priors.map(BOTORCH_LENGTHCALE_PRIOR())], lower=0, upper=10)