-
Notifications
You must be signed in to change notification settings - Fork 3
/
botorch_branin_basic_bayesian_optimization_oneloop.py
140 lines (111 loc) · 4.52 KB
/
botorch_branin_basic_bayesian_optimization_oneloop.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import torch
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_model
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.acquisition import qExpectedImprovement
from botorch.optim import optimize_acqf
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.animation import FuncAnimation
import seaborn as sns
import pandas as pd
from botorch.utils.sampling import draw_sobol_samples
from torch.quasirandom import SobolEngine
# Define the Branin function with input scaling, i.e, takes between [0,1] x [0,1] instead of the typical [0,15] x [-5,10]
def branin(x, negate=False):
a = 1.0
b = 5.1 / (4 * torch.pi**2)
c = 5 / torch.pi
d = 6
e = 10
f = 1 / (8 * torch.pi)
x1 = 15 * x[:, 0] - 5
x2 = 15 * x[:, 1]
result = a * (x2 - b * x1**2 + c * x1 - d)**2 + e * (1 - f) * torch.cos(x1) + e
if negate:
return -result
else:
return result
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
dtype = torch.float64
bounds = torch.tensor([[0., 0.], [1., 1.]], dtype=dtype, device=device)
N = 7
batch_size=5
iteration_number=20
iterations = list(range(1, iteration_number+1))
supra_best=[]
sobol_engine = SobolEngine(dimension=2, scramble=False) # 2 dimensions for your input space
train_x = draw_sobol_samples(bounds=bounds, n=1, q=N).squeeze(0)
print(train_x)
train_y = branin(train_x, negate=True).unsqueeze(-1)
models = []
gp_model = SingleTaskGP(train_x, train_y).to(device=device, dtype=dtype)
mll = ExactMarginalLogLikelihood(gp_model.likelihood, gp_model)
fit_gpytorch_model(mll)
models.append(gp_model)
# To store data for animation
frames_x = [train_x.cpu().numpy()]
frames_y = [train_y.cpu().numpy()]
best_points = []
best_y_values=[]
for iteration in range(iteration_number):
EI = qExpectedImprovement(model=gp_model, best_f=train_y.max())#, maximize=True)
candidate, _ = optimize_acqf(
acq_function=EI,
bounds=bounds,
q=batch_size,
num_restarts=5,
raw_samples=20,
options={"dtype": dtype, "device": device}
)
new_y = branin(candidate, negate=True).unsqueeze(-1)
train_x = torch.cat([train_x, candidate])
train_y = torch.cat([train_y, new_y])
gp_model = SingleTaskGP(train_x, train_y).to(device=device, dtype=dtype)
mll = ExactMarginalLogLikelihood(gp_model.likelihood, gp_model)
fit_gpytorch_model(mll)
frames_x.append(train_x.cpu().numpy())
frames_y.append(train_y.cpu().numpy())
best_points.append(train_x[train_y.argmax(), :].cpu().numpy())
best_y_values.append(train_y.max().cpu().numpy())
models.append(gp_model)
best_point = train_x[train_y.argmax(), :]
best_value = train_y.max().item()
best_y_values = np.array([element for element in best_y_values])
supra_best.append(best_y_values)
print("Best observed point:", best_point.cpu().numpy(), "Best observed value:", best_value)
# Function to create the contour plot of the Branin function
def plot_gp_mean(model, bounds, resolution=100):
x1 = torch.linspace(bounds[0, 0], bounds[1, 0], resolution, dtype=dtype, device=device)
x2 = torch.linspace(bounds[0, 1], bounds[1, 1], resolution, dtype=dtype, device=device)
X1, X2 = torch.meshgrid(x1, x2)
grid = torch.stack([X1.flatten(), X2.flatten()], -1)
with torch.no_grad():
mean = model.posterior(grid).mean.cpu().numpy().reshape(resolution, resolution)
return X1.cpu().numpy(), X2.cpu().numpy(), mean
# Update function for the animation
def update(frame):
plt.clf()
X1, X2, mean = plot_gp_mean(models[frame], bounds)
cp = plt.contourf(X1, X2, mean, levels=50, cmap=cm.viridis)
plt.colorbar(cp)
plt.scatter(frames_x[frame][:, 0], frames_x[frame][:, 1], color="red")
plt.title(f"Iteration {frame+1}")
fig, ax = plt.subplots(figsize=(10, 6))
ani = FuncAnimation(fig, update, frames=range(len(frames_x)), repeat=False)
# Save the animation
ani.save("branin_optimization_ini_"+str(N)+"_q_"+str(batch_size)+"_inum_"+str(iteration_number)+".mp4", writer="ffmpeg", dpi=200)
plt.close()
print("Animation saved")
# Plotting the best Y values vs iterations
plt.figure(figsize=(10, 6))
plt.plot(iterations, best_y_values, marker='o', linestyle='-', color='b')
plt.title('Best Y Values vs Iterations')
plt.xlabel('Iteration')
plt.ylabel('Best Y Value')
plt.grid(True)
plt.tight_layout()
plt.savefig("best_value_vs_iterations_ini_"+str(N)+"_q_"+str(batch_size)+"_inum_"+str(iteration_number)+".png",bbox_inches="tight",dpi=600)
plt.close()