import torch
from botorch.acquisition import LogExpectedImprovement
from botorch.fit import fit_gpytorch_mll
from botorch.models import SingleTaskGP
from botorch.models.transforms import Normalize, Standardize
from botorch.optim import optimize_acqf
from gpytorch.mlls import ExactMarginalLogLikelihood

def run(fixed_features):

    train_X = torch.rand(10, 2, dtype=torch.double) * 2
    Y = 1 - torch.linalg.norm(train_X - 0.5, dim=-1, keepdim=True)
    Y = Y + 0.1 * torch.randn_like(Y)

    gp = SingleTaskGP(
        train_X=train_X,
        train_Y=Y,
        input_transform=Normalize(d=2),
        outcome_transform=Standardize(m=1),
    )
    mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
    fit_gpytorch_mll(mll)
    logEI = LogExpectedImprovement(model=gp, best_f=Y.max())
    bounds = torch.stack([torch.zeros(2), torch.ones(2)]).to(torch.double)
    candidate, acq_value = optimize_acqf(
        logEI,
        bounds=bounds,
        q=1,
        num_restarts=5,
        raw_samples=20,
        fixed_features=fixed_features,
    )
    print(acq_value)

run({0: 0})  
run({0: 0, 1: 0})