def test_prune_baseline(self):
no = "botorch.utils.testing.MockModel.num_outputs"
prune = "botorch.acquisition.monte_carlo.prune_inferior_points"
for dtype in (torch.float, torch.double):
X_baseline = torch.zeros(1, 1, device=self.device, dtype=dtype)
X_pruned = torch.rand(1, 1, device=self.device, dtype=dtype)
with mock.patch(
no, new_callable=mock.PropertyMock) as mock_num_outputs:
mock_num_outputs.return_value = 1
mm = MockModel(mock.Mock())
with mock.patch(prune, return_value=X_pruned) as mock_prune:
acqf = qNoisyExpectedImprovement(
model=mm,
X_baseline=X_baseline,
prune_baseline=True,
cache_root=False,
)
mock_prune.assert_called_once()
self.assertTrue(torch.equal(acqf.X_baseline, X_pruned))
with mock.patch(prune, return_value=X_pruned) as mock_prune:
acqf = qNoisyExpectedImprovement(
model=mm,
X_baseline=X_baseline,
prune_baseline=True,
marginalize_dim=-3,
cache_root=False,
)
_, kwargs = mock_prune.call_args
self.assertEqual(kwargs["marginalize_dim"], -3)
def test_q_noisy_expected_improvement(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
for dtype in (torch.float, torch.double):
# the event shape is `b x q x t` = 1 x 2 x 1
samples_noisy = torch.tensor([1.0, 0.0], device=device, dtype=dtype)
samples_noisy = samples_noisy.view(1, 2, 1)
# X_baseline is `q' x d` = 1 x 1
X_baseline = torch.zeros(1, 1, device=device, dtype=dtype)
mm_noisy = MockModel(MockPosterior(samples=samples_noisy))
# X is `q x d` = 1 x 1
X = torch.zeros(1, 1, device=device, dtype=dtype)
# basic test
sampler = IIDNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
# basic test, no resample
sampler = IIDNormalSampler(num_samples=2, seed=12345)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# basic test, qmc, no resample
sampler = SobolQMCNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# basic test, qmc, resample
sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
def test_prune_baseline(self):
prune = "botorch.acquisition.monte_carlo.prune_inferior_points"
for dtype in (torch.float, torch.double):
X_baseline = torch.zeros(1, 1, device=self.device, dtype=dtype)
X_pruned = torch.rand(1, 1, device=self.device, dtype=dtype)
mm = MockModel(mock.Mock())
with mock.patch(prune, return_value=X_pruned) as mock_prune:
acqf = qNoisyExpectedImprovement(
model=mm, X_baseline=X_baseline, prune_baseline=True
)
mock_prune.assert_called_once()
self.assertTrue(torch.equal(acqf.X_baseline, X_pruned))
def test_acquisition_functions(self):
tkwargs = {"device": self.device, "dtype": torch.double}
train_X, train_Y, train_Yvar, model = self._get_data_and_model(
infer_noise=True, **tkwargs
)
fit_fully_bayesian_model_nuts(
model, warmup_steps=8, num_samples=5, thinning=2, disable_progbar=True
)
sampler = IIDNormalSampler(num_samples=2)
acquisition_functions = [
ExpectedImprovement(model=model, best_f=train_Y.max()),
ProbabilityOfImprovement(model=model, best_f=train_Y.max()),
PosteriorMean(model=model),
UpperConfidenceBound(model=model, beta=4),
qExpectedImprovement(model=model, best_f=train_Y.max(), sampler=sampler),
qNoisyExpectedImprovement(model=model, X_baseline=train_X, sampler=sampler),
qProbabilityOfImprovement(
model=model, best_f=train_Y.max(), sampler=sampler
),
qSimpleRegret(model=model, sampler=sampler),
qUpperConfidenceBound(model=model, beta=4, sampler=sampler),
qNoisyExpectedHypervolumeImprovement(
model=ModelListGP(model, model),
X_baseline=train_X,
ref_point=torch.zeros(2, **tkwargs),
sampler=sampler,
),
qExpectedHypervolumeImprovement(
model=ModelListGP(model, model),
ref_point=torch.zeros(2, **tkwargs),
sampler=sampler,
partitioning=NondominatedPartitioning(
ref_point=torch.zeros(2, **tkwargs), Y=train_Y.repeat([1, 2])
),
),
]
for acqf in acquisition_functions:
for batch_shape in [[5], [6, 5, 2]]:
test_X = torch.rand(*batch_shape, 1, 4, **tkwargs)
self.assertEqual(acqf(test_X).shape, torch.Size(batch_shape))
def main(argv):
dataset = 1
try:
opts, args = getopt.getopt(argv, "hd:", ["dataset="])
except getopt.GetoptError:
print('random parallel with input dataset')
sys.exit(2)
for opt, arg in opts:
if opt == '-h':
print('random parallel with input dataset')
sys.exit()
elif opt in ("-d", "--dataset"):
dataset = int(arg)
# average over multiple trials
for trial in range(1, N_TRIALS + 1):
print(f"\nTrial {trial:>2} of {N_TRIALS} ", end="")
best_observed_ei, best_observed_nei = [], []
# call helper functions to generate initial training data and initialize model
train_x_ei, train_obj_ei, best_observed_value_ei, current_best_config = generate_initial_data(
dataset)
train_x_nei, train_obj_nei = train_x_ei, train_obj_ei
best_observed_value_nei = best_observed_value_ei
mll_nei, model_nei = initialize_model(train_x_nei, train_obj_nei)
best_observed_nei.append(best_observed_value_nei)
# run N_BATCH rounds of BayesOpt after the initial random batch
for iteration in range(1, N_BATCH + 1):
# fit the models
fit_gpytorch_model(mll_nei)
# define the qEI and qNEI acquisition modules using a QMC sampler
qmc_sampler = SobolQMCNormalSampler(num_samples=MC_SAMPLES)
# for best_f, we use the best observed noisy values as an approximation
qNEI = qNoisyExpectedImprovement(
model=model_nei,
X_baseline=train_x_nei,
sampler=qmc_sampler,
)
# optimize and get new observation
new_x_nei, new_obj_nei = optimize_acqf_and_get_observation(
qNEI, dataset)
# update training points
train_x_nei = torch.cat([train_x_nei, new_x_nei])
train_obj_nei = torch.cat([train_obj_nei, new_obj_nei])
# update progress
best_value_nei = train_obj_nei.max().item()
best_observed_nei.append(best_value_nei)
# reinitialize the models so they are ready for fitting on next iteration
# use the current state dict to speed up fitting
mll_nei, model_nei = initialize_model(
train_x_nei,
train_obj_nei,
model_nei.state_dict(),
)
# return the best configuration
best_tensor_nei, indices_nei = torch.max(train_obj_nei, 0)
train_best_x_nei = train_x_nei[indices_nei].cpu().numpy()
from botorch.acquisition import PosteriorMean
argmax_pmean_nei, max_pmean_nei = optimize_acqf(
acq_function=PosteriorMean(model_nei),
bounds=bounds,
q=1,
num_restarts=20,
raw_samples=2048,
)
csv_file_name = '/home/junjie/modes/botorch/' + folder_name + '/modes-i/hp-ngp-qnei-dataset-' + str(
dataset) + '-trail' + str(trial) + '.csv'
with open(csv_file_name, 'w') as csvFile:
writer = csv.writer(csvFile)
writer.writerow([
str(argmax_pmean_nei.cpu().numpy()),
str(max_pmean_nei.cpu().numpy())
]) # nei prediction
writer.writerow(
[str(train_best_x_nei),
str(best_tensor_nei.cpu().numpy())]) # nei observation
csvFile.close()
def get_acquisition_function(
acquisition_function_name: str,
model: Model,
objective: MCAcquisitionObjective,
X_observed: Tensor,
X_pending: Optional[Tensor] = None,
constraints: Optional[List[Callable[[Tensor], Tensor]]] = None,
mc_samples: int = 500,
qmc: bool = True,
seed: Optional[int] = None,
**kwargs,
) -> monte_carlo.MCAcquisitionFunction:
r"""Convenience function for initializing botorch acquisition functions.
Args:
acquisition_function_name: Name of the acquisition function.
model: A fitted model.
objective: A MCAcquisitionObjective.
X_observed: A `m1 x d`-dim Tensor of `m1` design points that have
already been observed.
X_pending: A `m2 x d`-dim Tensor of `m2` design points whose evaluation
is pending.
constraints: A list of callables, each mapping a Tensor of dimension
`sample_shape x batch-shape x q x m` to a Tensor of dimension
`sample_shape x batch-shape x q`, where negative values imply
feasibility. Used when constraint_transforms are not passed
as part of the objective.
mc_samples: The number of samples to use for (q)MC evaluation of the
acquisition function.
qmc: If True, use quasi-Monte-Carlo sampling (instead of iid).
seed: If provided, perform deterministic optimization (i.e. the
function to optimize is fixed and not stochastic).
Returns:
The requested acquisition function.
Example:
>>> model = SingleTaskGP(train_X, train_Y)
>>> obj = LinearMCObjective(weights=torch.tensor([1.0, 2.0]))
>>> acqf = get_acquisition_function("qEI", model, obj, train_X)
"""
# initialize the sampler
if qmc:
sampler = SobolQMCNormalSampler(num_samples=mc_samples, seed=seed)
else:
sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed)
# instantiate and return the requested acquisition function
if acquisition_function_name == "qEI":
best_f = objective(model.posterior(X_observed).mean).max().item()
return monte_carlo.qExpectedImprovement(
model=model,
best_f=best_f,
sampler=sampler,
objective=objective,
X_pending=X_pending,
)
elif acquisition_function_name == "qPI":
best_f = objective(model.posterior(X_observed).mean).max().item()
return monte_carlo.qProbabilityOfImprovement(
model=model,
best_f=best_f,
sampler=sampler,
objective=objective,
X_pending=X_pending,
tau=kwargs.get("tau", 1e-3),
)
elif acquisition_function_name == "qNEI":
return monte_carlo.qNoisyExpectedImprovement(
model=model,
X_baseline=X_observed,
sampler=sampler,
objective=objective,
X_pending=X_pending,
prune_baseline=kwargs.get("prune_baseline", False),
)
elif acquisition_function_name == "qSR":
return monte_carlo.qSimpleRegret(model=model,
sampler=sampler,
objective=objective,
X_pending=X_pending)
elif acquisition_function_name == "qUCB":
if "beta" not in kwargs:
raise ValueError("`beta` must be specified in kwargs for qUCB.")
return monte_carlo.qUpperConfidenceBound(
model=model,
beta=kwargs["beta"],
sampler=sampler,
objective=objective,
X_pending=X_pending,
)
elif acquisition_function_name == "qEHVI":
# pyre-fixme [16]: `Model` has no attribute `train_targets`
try:
ref_point = kwargs["ref_point"]
except KeyError:
raise ValueError(
"`ref_point` must be specified in kwargs for qEHVI")
try:
Y = kwargs["Y"]
except KeyError:
raise ValueError("`Y` must be specified in kwargs for qEHVI")
# get feasible points
if constraints is not None:
feas = torch.stack([c(Y) <= 0 for c in constraints],
dim=-1).all(dim=-1)
Y = Y[feas]
obj = objective(Y)
partitioning = NondominatedPartitioning(
ref_point=torch.as_tensor(ref_point,
dtype=Y.dtype,
device=Y.device),
Y=obj,
alpha=kwargs.get("alpha", 0.0),
)
return moo_monte_carlo.qExpectedHypervolumeImprovement(
model=model,
ref_point=ref_point,
partitioning=partitioning,
sampler=sampler,
objective=objective,
constraints=constraints,
X_pending=X_pending,
)
raise NotImplementedError(
f"Unknown acquisition function {acquisition_function_name}")
def test_q_noisy_expected_improvement_batch(self):
for dtype in (torch.float, torch.double):
# the event shape is `b x q x t` = 2 x 3 x 1
samples_noisy = torch.zeros(2, 3, 1, device=self.device, dtype=dtype)
samples_noisy[0, 0, 0] = 1.0
mm_noisy = MockModel(MockPosterior(samples=samples_noisy))
# X is `q x d` = 1 x 1
X = torch.zeros(1, 1, 1, device=self.device, dtype=dtype)
X_baseline = torch.zeros(1, 1, device=self.device, dtype=dtype)
# test batch mode
sampler = IIDNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
# test batch mode, no resample
sampler = IIDNormalSampler(num_samples=2, seed=12345)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X) # 1-dim batch
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
res = acqf(X.expand(2, 1, 1)) # 2-dim batch
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
# the base samples should have the batch dim collapsed
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X.expand(2, 1, 1))
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# test batch mode, qmc, no resample
sampler = SobolQMCNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# test X_pending w/ batch mode, qmc, resample
sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X) # 1-dim batch
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
res = acqf(X.expand(2, 1, 1)) # 2-dim batch
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
# the base samples should have the batch dim collapsed
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X.expand(2, 1, 1))
self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
def test_q_noisy_expected_improvement(self):
for dtype in (torch.float, torch.double):
# the event shape is `b x q x t` = 1 x 2 x 1
samples_noisy = torch.tensor([1.0, 0.0], device=self.device, dtype=dtype)
samples_noisy = samples_noisy.view(1, 2, 1)
# X_baseline is `q' x d` = 1 x 1
X_baseline = torch.zeros(1, 1, device=self.device, dtype=dtype)
mm_noisy = MockModel(MockPosterior(samples=samples_noisy))
# X is `q x d` = 1 x 1
X = torch.zeros(1, 1, device=self.device, dtype=dtype)
# basic test
sampler = IIDNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
# basic test, no resample
sampler = IIDNormalSampler(num_samples=2, seed=12345)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# basic test, qmc, no resample
sampler = SobolQMCNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# basic test, qmc, resample
sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
# basic test for X_pending and warning
sampler = SobolQMCNormalSampler(num_samples=2)
samples_noisy_pending = torch.tensor(
[1.0, 0.0, 0.0], device=self.device, dtype=dtype
)
samples_noisy_pending = samples_noisy_pending.view(1, 3, 1)
mm_noisy_pending = MockModel(MockPosterior(samples=samples_noisy_pending))
acqf = qNoisyExpectedImprovement(
model=mm_noisy_pending, X_baseline=X_baseline, sampler=sampler
)
acqf.set_X_pending()
self.assertIsNone(acqf.X_pending)
acqf.set_X_pending(None)
self.assertIsNone(acqf.X_pending)
acqf.set_X_pending(X)
self.assertEqual(acqf.X_pending, X)
res = acqf(X)
X2 = torch.zeros(
1, 1, 1, device=self.device, dtype=dtype, requires_grad=True
)
with warnings.catch_warnings(record=True) as ws, settings.debug(True):
acqf.set_X_pending(X2)
self.assertEqual(acqf.X_pending, X2)
self.assertEqual(len(ws), 1)
self.assertTrue(issubclass(ws[-1].category, BotorchWarning))
def optimize():
verbose = True
best_observed_all_ei, best_observed_all_nei, best_random_all = [], [], []
train_x_all_ei, train_x_all_nei, train_x_all_random = [], [], []
train_y_all_ei, train_y_all_nei, train_y_all_random = [], [], []
# statistics over multiple trials
for trial in range(1, N_TRIALS + 1):
print('\nTrial {} of {}'.format(trial, N_TRIALS))
best_observed_ei, best_observed_nei = [], []
best_random = []
# generate initial training data and initialize model
print('\nGenerating {} random samples'.format(N_INITIAL_SAMPLES))
train_x_ei, train_y_ei, best_y_ei, mean_y, std_y = generate_initial_data(
n_samples=N_INITIAL_SAMPLES)
denormalize = lambda x: -(x * std_y + mean_y)
mll_ei, model_ei = initialize_model(train_x_ei, train_y_ei)
train_x_nei, train_y_nei, best_y_nei = train_x_ei, train_y_ei, best_y_ei
mll_nei, model_nei = initialize_model(train_x_nei, train_y_nei)
train_x_random, train_y_random, best_y_random = train_x_ei, train_y_ei, best_y_ei
best_observed_ei.append(denormalize(best_y_ei))
best_observed_nei.append(denormalize(best_y_nei))
best_random.append(denormalize(best_y_random))
# run N_BATCH rounds of BayesOpt after the initial random batch
for iteration in range(1, N_BATCH + 1):
print('\nBatch {} of {}\n'.format(iteration, N_BATCH))
t0 = time.time()
# fit the models
fit_gpytorch_model(mll_ei)
fit_gpytorch_model(mll_nei)
# update acquisition functions
qEI = qExpectedImprovement(
model=model_ei,
best_f=train_y_ei.max(),
sampler=qmc_sampler,
)
qNEI = qNoisyExpectedImprovement(
model=model_nei,
X_baseline=train_x_nei,
sampler=qmc_sampler,
)
# optimize acquisition function and evaluate new sample
new_x_ei, new_y_ei = optimize_acqf_and_get_observation(
qEI, mean_y=mean_y, std_y=std_y)
print('EI: time to traverse is {:.4f}s'.format(
-(new_y_ei.numpy().ravel()[0] * std_y + mean_y)))
new_x_nei, new_y_nei = optimize_acqf_and_get_observation(
qNEI, mean_y=mean_y, std_y=std_y)
print('NEI: time to traverse is {:.4f}s'.format(
-(new_y_nei.numpy().ravel()[0] * std_y + mean_y)))
new_x_random, new_y_random = sample_random_observations(
mean_y=mean_y, std_y=std_y)
print('Random: time to traverse is {:.4f}s'.format(
-(new_y_random.numpy().ravel()[0] * std_y + mean_y)))
# update training points
train_x_ei = torch.cat([train_x_ei, new_x_ei])
train_y_ei = torch.cat([train_y_ei, new_y_ei])
train_x_nei = torch.cat([train_x_nei, new_x_nei])
train_y_nei = torch.cat([train_y_nei, new_y_nei])
train_x_random = torch.cat([train_x_random, new_x_random])
train_y_random = torch.cat([train_y_random, new_y_random])
# update progress
best_value_ei = denormalize(train_y_ei.max().item())
best_value_nei = denormalize(train_y_nei.max().item())
best_value_random = denormalize(train_y_random.max().item())
best_observed_ei.append(best_value_ei)
best_observed_nei.append(best_value_nei)
best_random.append(best_value_random)
# reinitialize the models so they are ready for fitting on next iteration
# use the current state dict to speed up fitting
mll_ei, model_ei = initialize_model(
train_x_ei,
train_y_ei,
model_ei.state_dict(),
)
mll_nei, model_nei = initialize_model(
train_x_nei,
train_y_nei,
model_nei.state_dict(),
)
t1 = time.time()
if verbose:
print(
'best lap time (random, qEI, qNEI) = {:.2f}, {:.2f}, {:.2f}, time to compute = {:.2f}s'
.format(best_value_random, best_value_ei, best_value_nei,
t1 - t0))
else:
print(".")
best_observed_all_ei.append(best_observed_ei)
best_observed_all_nei.append(best_observed_nei)
best_random_all.append(best_random)
train_x_all_ei.append(train_x_ei.cpu().numpy())
train_x_all_nei.append(train_x_nei.cpu().numpy())
train_x_all_random.append(train_x_random.cpu().numpy())
train_y_all_ei.append(denormalize(train_y_ei.cpu().numpy()))
train_y_all_nei.append(denormalize(train_y_nei.cpu().numpy()))
train_y_all_random.append(denormalize(train_y_random.cpu().numpy()))
iters = np.arange(N_BATCH + 1) * BATCH_SIZE
y_ei = np.asarray(best_observed_all_ei)
y_nei = np.asarray(best_observed_all_nei)
y_rnd = np.asarray(best_random_all)
savestr = time.strftime('%Y%m%d%H%M%S')
#####################################################################
# save results
if SAVE_RESULTS:
np.savez(
'results/{}_raceline_data-{}.npz'.format('UCB', savestr),
y_ei=y_ei,
y_nei=y_nei,
y_rnd=y_rnd,
iters=iters,
train_x_all_ei=np.asarray(train_x_all_ei),
train_x_all_nei=np.asarray(train_x_all_nei),
train_x_all_random=np.asarray(train_x_all_random),
train_y_all_ei=np.asarray(train_y_all_ei),
train_y_all_nei=np.asarray(train_y_all_nei),
train_y_all_random=np.asarray(train_y_all_random),
SEED=SEED,
)
#####################################################################
# plot results
if PLOT_RESULTS:
def ci(y):
return 1.96 * y.std(axis=0) / np.sqrt(N_TRIALS)
plt.figure()
plt.gca().set_prop_cycle(None)
plt.plot(iters, y_rnd.mean(axis=0), linewidth=1.5)
plt.plot(iters, y_ei.mean(axis=0), linewidth=1.5)
plt.plot(iters, y_nei.mean(axis=0), linewidth=1.5)
plt.gca().set_prop_cycle(None)
plt.fill_between(iters,
y_rnd.mean(axis=0) - ci(y_rnd),
y_rnd.mean(axis=0) + ci(y_rnd),
label='random',
alpha=0.2)
plt.fill_between(iters,
y_ei.mean(axis=0) - ci(y_ei),
y_ei.mean(axis=0) + ci(y_ei),
label='qEI',
alpha=0.2)
plt.fill_between(iters,
y_nei.mean(axis=0) - ci(y_nei),
y_nei.mean(axis=0) + ci(y_nei),
label='qNEI',
alpha=0.2)
plt.xlabel('number of observations (beyond initial points)')
plt.ylabel('best lap times')
plt.grid(True)
plt.legend(loc=0)
plt.savefig('results/{}_laptimes-{}.png'.format('UCB', savestr),
dpi=600)
plt.show()
def test_q_noisy_expected_improvement(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
for dtype in (torch.float, torch.double):
# the event shape is `b x q x t` = 1 x 2 x 1
samples_noisy = torch.tensor([1.0, 0.0],
device=device,
dtype=dtype)
samples_noisy = samples_noisy.view(1, 2, 1)
# X_baseline is `q' x d` = 1 x 1
X_baseline = torch.zeros(1, 1, device=device, dtype=dtype)
mm_noisy = MockModel(MockPosterior(samples=samples_noisy))
# X is `q x d` = 1 x 1
X = torch.zeros(1, 1, device=device, dtype=dtype)
# basic test
sampler = IIDNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(model=mm_noisy,
X_baseline=X_baseline,
sampler=sampler)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
# basic test, no resample
sampler = IIDNormalSampler(num_samples=2, seed=12345)
acqf = qNoisyExpectedImprovement(model=mm_noisy,
X_baseline=X_baseline,
sampler=sampler)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape,
torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# basic test, qmc, no resample
sampler = SobolQMCNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(model=mm_noisy,
X_baseline=X_baseline,
sampler=sampler)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape,
torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# basic test, qmc, resample
sampler = SobolQMCNormalSampler(num_samples=2,
resample=True,
seed=12345)
acqf = qNoisyExpectedImprovement(model=mm_noisy,
X_baseline=X_baseline,
sampler=sampler)
res = acqf(X)
self.assertEqual(res.item(), 1.0)
self.assertEqual(acqf.sampler.base_samples.shape,
torch.Size([2, 1, 2, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
def run_RGPE(test_task: int, objective, bounds, base_model_list):
input_dim = bounds.shape[1]
best_rgpe_all = []
best_argmax_rgpe_all = []
# Average over multiple trials
for trial in range(N_TRIALS):
print(f"Trial {trial + 1} of {N_TRIALS}")
best_BMs = []
best_rgpe = []
# Initial random observations
raw_x = draw_sobol_samples(bounds=bounds,
n=RANDOM_INITIALIZATION_SIZE,
q=1,
seed=trial).squeeze(1)
train_x = normalize(raw_x, bounds=bounds)
train_y_noiseless = objective(raw_x, shift=test_task)
train_y = train_y_noiseless + noise_std * torch.randn_like(
train_y_noiseless)
train_yvar = torch.full_like(train_y, noise_std**2)
# keep track of the best observed point at each iteration
best_value = train_y.max().item()
best_rgpe.append(best_value)
# Run N_BATCH rounds of BayesOpt after the initial random batch
for iteration in range(N_BATCH):
target_model = get_fitted_model(train_x, train_y, train_yvar)
model_list = base_model_list + [target_model]
rank_weights = compute_rank_weights(
train_x,
train_y,
base_model_list,
target_model,
NUM_POSTERIOR_SAMPLES,
)
# create model and acquisition function
rgpe_model = RGPE(model_list, rank_weights)
sampler_qnei = SobolQMCNormalSampler(num_samples=MC_SAMPLES)
qNEI = qNoisyExpectedImprovement(
model=rgpe_model,
X_baseline=train_x,
sampler=sampler_qnei,
)
# optimize
candidate, _ = optimize_acqf(
acq_function=qNEI,
bounds=bounds,
q=Q_BATCH_SIZE,
num_restarts=N_RESTARTS,
raw_samples=N_RESTART_CANDIDATES,
)
# fetch the new values
new_x = candidate.detach()
new_y_noiseless = objective(unnormalize(new_x, bounds=bounds),
shift=test_task)
new_y = new_y_noiseless + noise_std * torch.randn_like(
new_y_noiseless)
new_yvar = torch.full_like(new_y, noise_std**2)
# update training points
train_x = torch.cat((train_x, new_x))
train_y = torch.cat((train_y, new_y))
train_yvar = torch.cat((train_yvar, new_yvar))
# get the new best observed value
best_value = train_y.max().item()
best_idx = torch.argmax(train_y).item()
best_candidate = train_x[best_idx].view(1, -1)
_, best_BM = objective(unnormalize(best_candidate, bounds=bounds),
shift=test_task,
include_BMs=True)
best_rgpe.append(best_value)
best_BMs.append(best_BM)
best_rgpe_all.append(best_rgpe)
best_argmax_rgpe_all.append(best_BMs)
BM_winner_idx = np.argmax(np.array(best_rgpe_all)[:, -1], axis=0)
BM_winner = np.reshape(np.array(best_argmax_rgpe_all[BM_winner_idx][-1]),
(2, input_dim))
return BM_winner
def test_sample_points_around_best(self):
tkwargs = {"device": self.device}
_bounds = torch.ones(2, 2)
_bounds[1] = 2
for dtype in (torch.float, torch.double):
tkwargs["dtype"] = dtype
bounds = _bounds.to(**tkwargs)
X_train = 1 + torch.rand(20, 2, **tkwargs)
model = MockModel(
MockPosterior(mean=(2 * X_train + 1).sum(dim=-1, keepdim=True))
)
# test NEI with X_baseline
acqf = qNoisyExpectedImprovement(model, X_baseline=X_train)
with mock.patch(
"botorch.optim.initializers.sample_perturbed_subset_dims"
) as mock_subset_dims:
X_rnd = sample_points_around_best(
acq_function=acqf,
n_discrete_points=4,
sigma=1e-3,
bounds=bounds,
)
mock_subset_dims.assert_not_called()
self.assertTrue(X_rnd.shape, torch.Size([4, 2]))
self.assertTrue((X_rnd >= 1).all())
self.assertTrue((X_rnd <= 2).all())
# test model that returns a batched mean
model = MockModel(
MockPosterior(
mean=(2 * X_train + 1).sum(dim=-1, keepdim=True).unsqueeze(0)
)
)
acqf = qNoisyExpectedImprovement(model, X_baseline=X_train)
X_rnd = sample_points_around_best(
acq_function=acqf,
n_discrete_points=4,
sigma=1e-3,
bounds=bounds,
)
self.assertTrue(X_rnd.shape, torch.Size([4, 2]))
self.assertTrue((X_rnd >= 1).all())
self.assertTrue((X_rnd <= 2).all())
# test EI without X_baseline
acqf = qExpectedImprovement(model, best_f=0.0)
with warnings.catch_warnings(record=True) as w, settings.debug(True):
X_rnd = sample_points_around_best(
acq_function=acqf,
n_discrete_points=4,
sigma=1e-3,
bounds=bounds,
)
self.assertEqual(len(w), 1)
self.assertTrue(issubclass(w[-1].category, BotorchWarning))
self.assertIsNone(X_rnd)
# set train inputs
model.train_inputs = (X_train,)
X_rnd = sample_points_around_best(
acq_function=acqf,
n_discrete_points=4,
sigma=1e-3,
bounds=bounds,
)
self.assertTrue(X_rnd.shape, torch.Size([4, 2]))
self.assertTrue((X_rnd >= 1).all())
self.assertTrue((X_rnd <= 2).all())
# test an acquisition function that has objective=None
# and maximize=False
pm = PosteriorMean(model, maximize=False)
self.assertIsNone(pm.objective)
self.assertFalse(pm.maximize)
X_rnd = sample_points_around_best(
acq_function=pm,
n_discrete_points=4,
sigma=0,
bounds=bounds,
best_pct=1e-8, # ensures that we only use best value
)
idx = (-model.posterior(X_train).mean).argmax()
self.assertTrue((X_rnd == X_train[idx : idx + 1]).all(dim=-1).all())
# test acquisition function that has no model
ff = FixedFeatureAcquisitionFunction(pm, d=2, columns=[0], values=[0])
# set X_baseline for testing purposes
ff.X_baseline = X_train
with warnings.catch_warnings(record=True) as w, settings.debug(True):
X_rnd = sample_points_around_best(
acq_function=ff,
n_discrete_points=4,
sigma=1e-3,
bounds=bounds,
)
self.assertEqual(len(w), 1)
self.assertTrue(issubclass(w[-1].category, BotorchWarning))
self.assertIsNone(X_rnd)
# test constraints with NEHVI
constraints = [lambda Y: Y[..., 0]]
ref_point = torch.zeros(2, **tkwargs)
# test cases when there are and are not any feasible points
for any_feas in (True, False):
Y_train = torch.stack(
[
torch.linspace(-0.5, 0.5, X_train.shape[0], **tkwargs)
if any_feas
else torch.ones(X_train.shape[0], **tkwargs),
X_train.sum(dim=-1),
],
dim=-1,
)
moo_model = MockModel(MockPosterior(mean=Y_train, samples=Y_train))
acqf = qNoisyExpectedHypervolumeImprovement(
moo_model,
ref_point=ref_point,
X_baseline=X_train,
constraints=constraints,
cache_root=False,
)
X_rnd = sample_points_around_best(
acq_function=acqf,
n_discrete_points=4,
sigma=0.0,
bounds=bounds,
)
self.assertTrue(X_rnd.shape, torch.Size([4, 2]))
# this should be true since sigma=0
# and we should only be returning feasible points
violation = constraints[0](Y_train)
neg_violation = -violation.clamp_min(0.0)
feas = neg_violation == 0
eq_mask = (X_train.unsqueeze(1) == X_rnd.unsqueeze(0)).all(dim=-1)
if feas.any():
# determine
# create n_train x n_rnd tensor of booleans
eq_mask = (X_train.unsqueeze(1) == X_rnd.unsqueeze(0)).all(dim=-1)
# check that all X_rnd correspond to feasible points
self.assertEqual(eq_mask[feas].sum(), 4)
else:
idcs = torch.topk(neg_violation, k=2).indices
self.assertEqual(eq_mask[idcs].sum(), 4)
self.assertTrue((X_rnd >= 1).all())
self.assertTrue((X_rnd <= 2).all())
# test that subset_dims is called if d>=21
X_train = 1 + torch.rand(20, 21, **tkwargs)
model = MockModel(
MockPosterior(mean=(2 * X_train + 1).sum(dim=-1, keepdim=True))
)
bounds = torch.ones(2, 21, **tkwargs)
bounds[1] = 2
# test NEI with X_baseline
acqf = qNoisyExpectedImprovement(model, X_baseline=X_train)
with mock.patch(
"botorch.optim.initializers.sample_perturbed_subset_dims",
wraps=sample_perturbed_subset_dims,
) as mock_subset_dims:
X_rnd = sample_points_around_best(
acq_function=acqf, n_discrete_points=5, sigma=1e-3, bounds=bounds
)
self.assertTrue(X_rnd.shape, torch.Size([5, 2]))
self.assertTrue((X_rnd >= 1).all())
self.assertTrue((X_rnd <= 2).all())
mock_subset_dims.assert_called_once()
# test tiny prob_perturb to make sure we perturb at least one dimension
X_rnd = sample_points_around_best(
acq_function=acqf,
n_discrete_points=5,
sigma=1e-3,
bounds=bounds,
prob_perturb=1e-8,
)
self.assertTrue(
((X_rnd.unsqueeze(0) == X_train.unsqueeze(1)).all(dim=-1)).sum() == 0
)
def test_get_X_baseline(self):
tkwargs = {"device": self.device}
for dtype in (torch.float, torch.double):
tkwargs["dtype"] = dtype
X_train = torch.rand(20, 2, **tkwargs)
model = MockModel(
MockPosterior(mean=(2 * X_train +
1).sum(dim=-1, keepdim=True)))
# test NEI with X_baseline
acqf = qNoisyExpectedImprovement(model, X_baseline=X_train[:2])
X = get_X_baseline(acq_function=acqf)
self.assertTrue(torch.equal(X, acqf.X_baseline))
# test EI without X_baseline
acqf = qExpectedImprovement(model, best_f=0.0)
with warnings.catch_warnings(
record=True) as w, settings.debug(True):
X_rnd = get_X_baseline(acq_function=acqf, )
self.assertEqual(len(w), 1)
self.assertTrue(issubclass(w[-1].category, BotorchWarning))
self.assertIsNone(X_rnd)
# set train inputs
model.train_inputs = (X_train, )
X = get_X_baseline(acq_function=acqf, )
self.assertTrue(torch.equal(X, X_train))
# test that we fail back to train_inputs if X_baseline is an empty tensor
acqf.register_buffer("X_baseline", X_train[:0])
X = get_X_baseline(acq_function=acqf, )
self.assertTrue(torch.equal(X, X_train))
# test acquisitipon function without X_baseline or model
acqf = FixedFeatureAcquisitionFunction(acqf,
d=2,
columns=[0],
values=[0])
with warnings.catch_warnings(
record=True) as w, settings.debug(True):
X_rnd = get_X_baseline(acq_function=acqf, )
self.assertEqual(len(w), 1)
self.assertTrue(issubclass(w[-1].category, BotorchWarning))
self.assertIsNone(X_rnd)
Y_train = 2 * X_train[:2] + 1
moo_model = MockModel(MockPosterior(mean=Y_train, samples=Y_train))
ref_point = torch.zeros(2, **tkwargs)
# test NEHVI with X_baseline
acqf = qNoisyExpectedHypervolumeImprovement(
moo_model,
ref_point=ref_point,
X_baseline=X_train[:2],
cache_root=False,
)
X = get_X_baseline(acq_function=acqf, )
self.assertTrue(torch.equal(X, acqf.X_baseline))
# test qEHVI without train_inputs
acqf = qExpectedHypervolumeImprovement(
moo_model,
ref_point=ref_point,
partitioning=FastNondominatedPartitioning(
ref_point=ref_point,
Y=Y_train,
),
)
# test extracting train_inputs from model list GP
model_list = ModelListGP(
SingleTaskGP(X_train, Y_train[:, :1]),
SingleTaskGP(X_train, Y_train[:, 1:]),
)
acqf = qExpectedHypervolumeImprovement(
model_list,
ref_point=ref_point,
partitioning=FastNondominatedPartitioning(
ref_point=ref_point,
Y=Y_train,
),
)
X = get_X_baseline(acq_function=acqf, )
self.assertTrue(torch.equal(X, X_train))
# test MESMO for which we need to use
# `acqf.mo_model`
batched_mo_model = SingleTaskGP(X_train, Y_train)
acqf = qMultiObjectiveMaxValueEntropy(
batched_mo_model,
sample_pareto_frontiers=lambda model: torch.rand(
10, 2, **tkwargs),
)
X = get_X_baseline(acq_function=acqf, )
self.assertTrue(torch.equal(X, X_train))
# test that if there is an input transform that is applied
# to the train_inputs when the model is in eval mode, we
# extract the untransformed train_inputs
model = SingleTaskGP(X_train,
Y_train[:, :1],
input_transform=Warp(indices=[0, 1]))
model.eval()
self.assertFalse(torch.equal(model.train_inputs[0], X_train))
acqf = qExpectedImprovement(model, best_f=0.0)
X = get_X_baseline(acq_function=acqf, )
self.assertTrue(torch.equal(X, X_train))
def test_q_noisy_expected_improvement_batch(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
for dtype in (torch.float, torch.double):
# the event shape is `b x q x t` = 2 x 3 x 1
samples_noisy = torch.zeros(2, 3, 1, device=device, dtype=dtype)
samples_noisy[0, 0, 0] = 1.0
mm_noisy = MockModel(MockPosterior(samples=samples_noisy))
# X is `q x d` = 1 x 1
X = torch.zeros(1, 1, 1, device=device, dtype=dtype)
X_baseline = torch.zeros(1, 1, device=device, dtype=dtype)
# test batch mode
sampler = IIDNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
# test batch mode, no resample
sampler = IIDNormalSampler(num_samples=2, seed=12345)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X) # 1-dim batch
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
res = acqf(X.expand(2, 1, 1)) # 2-dim batch
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
# the base samples should have the batch dim collapsed
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X.expand(2, 1, 1))
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# test batch mode, qmc, no resample
sampler = SobolQMCNormalSampler(num_samples=2)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X)
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertTrue(torch.equal(acqf.sampler.base_samples, bs))
# test X_pending w/ batch mode, qmc, resample
sampler = SobolQMCNormalSampler(num_samples=2, resample=True, seed=12345)
acqf = qNoisyExpectedImprovement(
model=mm_noisy, X_baseline=X_baseline, sampler=sampler
)
res = acqf(X) # 1-dim batch
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X)
self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
res = acqf(X.expand(2, 1, 1)) # 2-dim batch
self.assertEqual(res[0].item(), 1.0)
self.assertEqual(res[1].item(), 0.0)
# the base samples should have the batch dim collapsed
self.assertEqual(acqf.sampler.base_samples.shape, torch.Size([2, 1, 3, 1]))
bs = acqf.sampler.base_samples.clone()
acqf(X.expand(2, 1, 1))
self.assertFalse(torch.equal(acqf.sampler.base_samples, bs))
fit_gpytorch_model(mll_nei)
# define the qEI and qNEI acquisition modules using a QMC sampler
qmc_sampler = SobolQMCNormalSampler(num_samples=MC_SAMPLES)
# for best_f, we use the best observed noisy values as an approximation
qEI = qExpectedImprovement(
model=model_ei,
best_f=(train_obj_ei * (train_con_ei <= 0).to(train_obj_ei)).max(),
sampler=qmc_sampler,
objective=constrained_obj,
)
qNEI = qNoisyExpectedImprovement(
model=model_nei,
X_baseline=train_x_nei,
sampler=qmc_sampler,
objective=constrained_obj,
)
# optimize and get new observation
new_x_ei, new_obj_ei, new_con_ei = optimize_acqf_and_get_observation(
qEI)
new_x_nei, new_obj_nei, new_con_nei = optimize_acqf_and_get_observation(
qNEI)
# update training points
train_x_ei = torch.cat([train_x_ei, new_x_ei])
train_obj_ei = torch.cat([train_obj_ei, new_obj_ei])
train_con_ei = torch.cat([train_con_ei, new_con_ei])
train_x_nei = torch.cat([train_x_nei, new_x_nei])
def get_acquisition_function(
acquisition_function_name: str,
model: Model,
objective: MCAcquisitionObjective,
X_observed: Tensor,
X_pending: Optional[Tensor] = None,
mc_samples: int = 500,
qmc: bool = True,
seed: Optional[int] = None,
**kwargs,
) -> monte_carlo.MCAcquisitionFunction:
r"""Convenience function for initializing botorch acquisition functions.
Args:
acquisition_function_name: Name of the acquisition function.
model: A fitted model.
objective: A MCAcquisitionObjective.
X_observed: A `m1 x d`-dim Tensor of `m1` design points that have
already been observed.
X_pending: A `m2 x d`-dim Tensor of `m2` design points whose evaluation
is pending.
mc_samples: The number of samples to use for (q)MC evaluation of the
acquisition function.
qmc: If True, use quasi-Monte-Carlo sampling (instead of iid).
seed: If provided, perform deterministic optimization (i.e. the
function to optimize is fixed and not stochastic).
Returns:
The requested acquisition function.
Example:
>>> model = SingleTaskGP(train_X, train_Y)
>>> obj = LinearMCObjective(weights=torch.tensor([1.0, 2.0]))
>>> acqf = get_acquisition_function("qEI", model, obj, train_X)
"""
# initialize the sampler
if qmc:
sampler = SobolQMCNormalSampler(num_samples=mc_samples, seed=seed)
else:
sampler = IIDNormalSampler(num_samples=mc_samples, seed=seed)
# instantiate and return the requested acquisition function
if acquisition_function_name == "qEI":
best_f = objective(model.posterior(X_observed).mean).max().item()
return monte_carlo.qExpectedImprovement(
model=model,
best_f=best_f,
sampler=sampler,
objective=objective,
X_pending=X_pending,
)
elif acquisition_function_name == "qPI":
best_f = objective(model.posterior(X_observed).mean).max().item()
return monte_carlo.qProbabilityOfImprovement(
model=model,
best_f=best_f,
sampler=sampler,
objective=objective,
X_pending=X_pending,
tau=kwargs.get("tau", 1e-3),
)
elif acquisition_function_name == "qNEI":
return monte_carlo.qNoisyExpectedImprovement(
model=model,
X_baseline=X_observed,
sampler=sampler,
objective=objective,
X_pending=X_pending,
prune_baseline=kwargs.get("prune_baseline", False),
)
elif acquisition_function_name == "qSR":
return monte_carlo.qSimpleRegret(model=model,
sampler=sampler,
objective=objective,
X_pending=X_pending)
elif acquisition_function_name == "qUCB":
if "beta" not in kwargs:
raise ValueError("`beta` must be specified in kwargs for qUCB.")
return monte_carlo.qUpperConfidenceBound(
model=model,
beta=kwargs["beta"],
sampler=sampler,
objective=objective,
X_pending=X_pending,
)
raise NotImplementedError(
f"Unknown acquisition function {acquisition_function_name}")
def test_cache_root(self):
sample_cached_path = (
"botorch.acquisition.cached_cholesky.sample_cached_cholesky")
raw_state_dict = {
"likelihood.noise_covar.raw_noise":
torch.tensor([[0.0895], [0.2594]], dtype=torch.float64),
"mean_module.constant":
torch.tensor([[-0.4545], [-0.1285]], dtype=torch.float64),
"covar_module.raw_outputscale":
torch.tensor([1.4876, 1.4897], dtype=torch.float64),
"covar_module.base_kernel.raw_lengthscale":
torch.tensor([[[-0.7202, -0.2868]], [[-0.8794, -1.2877]]],
dtype=torch.float64),
}
# test batched models (e.g. for MCMC)
for train_batch_shape, m, dtype in product(
(torch.Size([]), torch.Size([3])), (1, 2),
(torch.float, torch.double)):
state_dict = deepcopy(raw_state_dict)
for k, v in state_dict.items():
if m == 1:
v = v[0]
if len(train_batch_shape) > 0:
v = v.unsqueeze(0).expand(*train_batch_shape, *v.shape)
state_dict[k] = v
tkwargs = {"device": self.device, "dtype": dtype}
if m == 2:
objective = GenericMCObjective(lambda Y, X: Y.sum(dim=-1))
else:
objective = None
for k, v in state_dict.items():
state_dict[k] = v.to(**tkwargs)
all_close_kwargs = ({
"atol": 1e-1,
"rtol": 0.0,
} if dtype == torch.float else {
"atol": 1e-4,
"rtol": 0.0
})
torch.manual_seed(1234)
train_X = torch.rand(*train_batch_shape, 3, 2, **tkwargs)
train_Y = (
torch.sin(train_X * 2 * pi) +
torch.randn(*train_batch_shape, 3, 2, **tkwargs))[..., :m]
train_Y = standardize(train_Y)
model = SingleTaskGP(
train_X,
train_Y,
)
if len(train_batch_shape) > 0:
X_baseline = train_X[0]
else:
X_baseline = train_X
model.load_state_dict(state_dict, strict=False)
# test sampler with collapse_batch_dims=False
sampler = IIDNormalSampler(5, seed=0, collapse_batch_dims=False)
with self.assertRaises(UnsupportedError):
qNoisyExpectedImprovement(
model=model,
X_baseline=X_baseline,
sampler=sampler,
objective=objective,
prune_baseline=False,
cache_root=True,
)
sampler = IIDNormalSampler(5, seed=0)
torch.manual_seed(0)
acqf = qNoisyExpectedImprovement(
model=model,
X_baseline=X_baseline,
sampler=sampler,
objective=objective,
prune_baseline=False,
cache_root=True,
)
orig_base_samples = acqf.base_sampler.base_samples.detach().clone()
sampler2 = IIDNormalSampler(5, seed=0)
sampler2.base_samples = orig_base_samples
torch.manual_seed(0)
acqf_no_cache = qNoisyExpectedImprovement(
model=model,
X_baseline=X_baseline,
sampler=sampler2,
objective=objective,
prune_baseline=False,
cache_root=False,
)
for q, batch_shape in product(
(1, 3), (torch.Size([]), torch.Size([3]), torch.Size([4, 3]))):
test_X = (0.3 +
0.05 * torch.randn(*batch_shape, q, 2, **tkwargs)
).requires_grad_(True)
with mock.patch(
sample_cached_path,
wraps=sample_cached_cholesky) as mock_sample_cached:
torch.manual_seed(0)
val = acqf(test_X)
mock_sample_cached.assert_called_once()
val.sum().backward()
base_samples = acqf.sampler.base_samples.detach().clone()
X_grad = test_X.grad.clone()
test_X2 = test_X.detach().clone().requires_grad_(True)
acqf_no_cache.sampler.base_samples = base_samples
with mock.patch(
sample_cached_path,
wraps=sample_cached_cholesky) as mock_sample_cached:
torch.manual_seed(0)
val2 = acqf_no_cache(test_X2)
mock_sample_cached.assert_not_called()
self.assertTrue(torch.allclose(val, val2, **all_close_kwargs))
val2.sum().backward()
self.assertTrue(
torch.allclose(X_grad, test_X2.grad, **all_close_kwargs))
# test we fall back to standard sampling for
# ill-conditioned covariances
acqf._baseline_L = torch.zeros_like(acqf._baseline_L)
with warnings.catch_warnings(
record=True) as ws, settings.debug(True):
with torch.no_grad():
acqf(test_X)
self.assertEqual(len(ws), 1)
self.assertTrue(issubclass(ws[-1].category, BotorchWarning))
