botorch_gaussian_process_fitting_with_ucb_acquisition_optimization.py

1import torch
2from botorch.models import SingleTaskGP
3from botorch.fit import fit_gpytorch_mll
4from botorch.acquisition import UpperConfidenceBound
5from botorch.optim import optimize_acqf
6from gpytorch.mlls import ExactMarginalLogLikelihood
7
8# 1. Generate some training data
9train_X = torch.rand(10, 2)
10# Objective function: y = sin(2*pi*x[0]) * cos(2*pi*x[1])
11train_Y = torch.sin(2 * torch.pi * train_X[:, [0]]) * torch.cos(2 * torch.pi * train_X[:, [1]])
12
13# 2. Define and fit a Gaussian Process model
14gp = SingleTaskGP(train_X, train_Y)
15mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
16fit_gpytorch_mll(mll)
17
18# 3. Define the acquisition function (Upper Confidence Bound)
19UCB = UpperConfidenceBound(gp, beta=0.1)
20
21# 4. Optimize the acquisition function to find the next candidate point
22bounds = torch.stack([torch.zeros(2), torch.ones(2)])
23candidate, acq_value = optimize_acqf(
24    UCB, bounds=bounds, q=1, num_restarts=5, raw_samples=20,
25)
26
27print(f"Next candidate point to sample: {candidate}")