Skip to content

Commit

Permalink
Add acquisition function
Browse files Browse the repository at this point in the history
  • Loading branch information
vladislavalerievich committed Dec 7, 2024
1 parent f69ddbe commit 1c4cc83
Showing 1 changed file with 64 additions and 8 deletions.
72 changes: 64 additions & 8 deletions grakel_replace/single_task_gp_usage_example.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,20 @@
from __future__ import annotations

from itertools import product
from typing import TYPE_CHECKING

import torch
from botorch import fit_gpytorch_mll
from botorch.acquisition import LinearMCObjective, qLogNoisyExpectedImprovement
from botorch.models import SingleTaskGP
from botorch.models.gp_regression_mixed import CategoricalKernel, ScaleKernel
from gpytorch.distributions.multivariate_normal import MultivariateNormal
from botorch.optim import optimize_acqf_mixed
from gpytorch import ExactMarginalLogLikelihood
from gpytorch.kernels import AdditiveKernel, MaternKernel

if TYPE_CHECKING:
from gpytorch.distributions.multivariate_normal import MultivariateNormal

TRAIN_CONFIGS = 10
TEST_CONFIGS = 10
TOTAL_CONFIGS = TRAIN_CONFIGS + TEST_CONFIGS
Expand Down Expand Up @@ -51,7 +62,6 @@
)
kernels.append(hamming)


combined_num_cat_kernel = AdditiveKernel(*kernels)

train_x = X[:TRAIN_CONFIGS]
Expand All @@ -61,9 +71,6 @@
test_y = y[TRAIN_CONFIGS:]

K_matrix = combined_num_cat_kernel.forward(train_x, train_x)
print(
"K_matrix: ", K_matrix.to_dense()
)

train_y = train_y.unsqueeze(-1)
test_y = test_y.unsqueeze(-1)
Expand All @@ -75,6 +82,55 @@
)

multivariate_normal: MultivariateNormal = gp.forward(train_x)
print("Mean:", multivariate_normal.mean)
print("Variance:", multivariate_normal.variance)
print("Covariance matrix:", multivariate_normal.covariance_matrix)

# =============== Fitting the GP using botorch ===============

print("\nFitting the GP model using botorch...")

mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
fit_gpytorch_mll(mll)

acq_function = qLogNoisyExpectedImprovement(
model=gp,
X_baseline=train_x,
objective=LinearMCObjective(weights=torch.tensor([-1.0])),
prune_baseline=True,
)

# Define bounds
bounds = torch.tensor(
[
[0.0] * N_NUMERICAL + [0.0] * N_CATEGORICAL,
[1.0] * N_NUMERICAL + [
float(N_CATEGORICAL_VALUES_PER_CATEGORY - 1)] * N_CATEGORICAL
]
)

# Setup categorical feature optimization
cats_per_column: dict[int, list[float]] = {
column_ix: [float(i) for i in range(N_CATEGORICAL_VALUES_PER_CATEGORY)]
for column_ix in range(N_NUMERICAL, N_NUMERICAL + N_CATEGORICAL)
}

# Generate fixed categorical features
fixed_cats: list[dict[int, float]]
if len(cats_per_column) == 1:
col, choice_indices = next(iter(cats_per_column.items()))
fixed_cats = [{col: i} for i in choice_indices]
else:
fixed_cats = [
dict(zip(cats_per_column.keys(), combo))
for combo in product(*cats_per_column.values())
]

best_candidate, best_score = optimize_acqf_mixed(
acq_function=acq_function,
bounds=bounds,
fixed_features_list=fixed_cats,
num_restarts=10,
raw_samples=10,
q=1,
)

print("Best candidate:", best_candidate)
print("Acquisition score:", best_score)

0 comments on commit 1c4cc83

Please sign in to comment.