def gen(
self,
n: int,
bounds: List[Tuple[float, float]],
objective_weights: Tensor,
objective_thresholds: Optional[Tensor] = None,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
pending_observations: Optional[List[Tensor]] = None,
model_gen_options: Optional[TConfig] = None,
rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
) -> Tuple[Tensor, Tensor, TGenMetadata,
Optional[List[TCandidateMetadata]]]:
if self._search_space_digest is None:
raise RuntimeError("Must `fit` the model before calling `gen`.")
acq_options, opt_options = construct_acquisition_and_optimizer_options(
acqf_options=self.acquisition_options,
model_gen_options=model_gen_options)
# update bounds / target fidelities
new_ssd_args = {
**dataclasses.asdict(self._search_space_digest),
"bounds": bounds,
"target_fidelities": target_fidelities or {},
}
search_space_digest = SearchSpaceDigest(**new_ssd_args)
acqf = self._instantiate_acquisition(
search_space_digest=search_space_digest,
objective_weights=objective_weights,
objective_thresholds=objective_thresholds,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
pending_observations=pending_observations,
acq_options=acq_options,
)
botorch_rounding_func = get_rounding_func(rounding_func)
candidates, expected_acquisition_value = acqf.optimize(
n=n,
search_space_digest=search_space_digest,
inequality_constraints=_to_inequality_constraints(
linear_constraints=linear_constraints),
fixed_features=fixed_features,
rounding_func=botorch_rounding_func,
optimizer_options=checked_cast(dict, opt_options),
)
gen_metadata: TGenMetadata = {
Keys.EXPECTED_ACQF_VAL: expected_acquisition_value.tolist()
}
if objective_weights.nonzero().numel() > 1:
gen_metadata["objective_thresholds"] = acqf.objective_thresholds
gen_metadata["objective_weights"] = acqf.objective_weights
return (
candidates.detach().cpu(),
torch.ones(n, dtype=self.surrogate.dtype),
gen_metadata,
None,
)
def _optimize_and_get_candidates(
acq_function: qKnowledgeGradient,
bounds_: Tensor,
n: int,
num_restarts: int,
raw_samples: int,
optimizer_options: Dict,
rounding_func: Optional[Callable[[Tensor], Tensor]],
inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]],
fixed_features: Optional[Dict[int, float]],
) -> Tensor:
r"""Generates initial conditions for optimization, optimize the acquisition
function, and return the candidates.
"""
batch_initial_conditions = gen_one_shot_kg_initial_conditions(
acq_function=acq_function,
bounds=bounds_,
q=n,
num_restarts=num_restarts,
raw_samples=raw_samples,
options={
"frac_random": optimizer_options.get("frac_random", 0.1),
"num_inner_restarts": num_restarts,
"raw_inner_samples": raw_samples,
},
)
botorch_rounding_func = get_rounding_func(rounding_func)
candidates, _ = optimize_acqf(
acq_function=acq_function,
bounds=bounds_,
q=n,
inequality_constraints=inequality_constraints,
fixed_features=fixed_features,
post_processing_func=botorch_rounding_func,
num_restarts=num_restarts,
raw_samples=raw_samples,
options={
"batch_limit": optimizer_options.get("batch_limit", 8),
"maxiter": optimizer_options.get("maxiter", 200),
"method": "L-BFGS-B",
"nonnegative": optimizer_options.get("nonnegative", False),
},
batch_initial_conditions=batch_initial_conditions,
)
new_x = candidates.detach().cpu()
return new_x
def gen(
self,
n: int,
bounds: List[Tuple[float, float]],
objective_weights: Tensor,
objective_thresholds: Optional[Tensor] = None,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
pending_observations: Optional[List[Tensor]] = None,
model_gen_options: Optional[TConfig] = None,
rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
) -> Tuple[Tensor, Tensor, TGenMetadata,
Optional[List[TCandidateMetadata]]]:
acq_options, opt_options = construct_acquisition_and_optimizer_options(
acqf_options=self.acquisition_options,
model_gen_options=model_gen_options)
acqf = self._instantiate_acquisition(
bounds=bounds,
objective_weights=objective_weights,
objective_thresholds=objective_thresholds,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
pending_observations=pending_observations,
target_fidelities=target_fidelities,
acq_options=acq_options,
)
botorch_rounding_func = get_rounding_func(rounding_func)
candidates, expected_acquisition_value = acqf.optimize(
bounds=self._bounds_as_tensor(bounds=bounds),
n=n,
inequality_constraints=_to_inequality_constraints(
linear_constraints=linear_constraints),
fixed_features=fixed_features,
rounding_func=botorch_rounding_func,
optimizer_options=checked_cast(dict, opt_options),
)
return (
candidates.detach().cpu(),
torch.ones(n, dtype=self.surrogate.dtype),
{
Keys.EXPECTED_ACQF_VAL: expected_acquisition_value.tolist()
},
None,
)
def gen(
self,
n: int,
bounds: List[Tuple[float, float]],
objective_weights: Tensor, # objective_directions
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
objective_thresholds: Optional[Tensor] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
pending_observations: Optional[List[Tensor]] = None,
model_gen_options: Optional[TConfig] = None,
rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
) -> Tuple[Tensor, Tensor, TGenMetadata,
Optional[List[TCandidateMetadata]]]:
options = model_gen_options or {}
acf_options = options.get("acquisition_function_kwargs", {})
optimizer_options = options.get("optimizer_kwargs", {})
if target_fidelities:
raise NotImplementedError(
"target_fidelities not implemented for base BotorchModel")
X_pending, X_observed = _get_X_pending_and_observed(
Xs=self.Xs,
pending_observations=pending_observations,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
bounds=bounds,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
)
model = self.model
# subset model only to the outcomes we need for the optimization
if options.get(Keys.SUBSET_MODEL, True):
model, objective_weights, outcome_constraints, Ys = subset_model(
model=model, # pyre-ignore [6]
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
Ys=self.Ys,
)
else:
Ys = self.Ys
bounds_ = torch.tensor(bounds, dtype=self.dtype, device=self.device)
bounds_ = bounds_.transpose(0, 1)
botorch_rounding_func = get_rounding_func(rounding_func)
if acf_options.get("random_scalarization", False) or acf_options.get(
"chebyshev_scalarization", False):
# If using a list of acquisition functions, the algorithm to generate
# that list is configured by acquisition_function_kwargs.
objective_weights_list = [
randomize_objective_weights(objective_weights, **acf_options)
for _ in range(n)
]
acquisition_function_list = [
self.acqf_constructor( # pyre-ignore: [28]
model=model,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
X_observed=X_observed,
X_pending=X_pending,
Ys=Ys, # Required for chebyshev scalarization calculations.
**acf_options,
) for objective_weights in objective_weights_list
]
acquisition_function_list = [
checked_cast(AcquisitionFunction, acq_function)
for acq_function in acquisition_function_list
]
# Multiple acquisition functions require a sequential optimizer
# always use scipy_optimizer_list.
# TODO(jej): Allow any optimizer.
candidates, expected_acquisition_value = scipy_optimizer_list(
acq_function_list=acquisition_function_list,
bounds=bounds_,
inequality_constraints=_to_inequality_constraints(
linear_constraints=linear_constraints),
fixed_features=fixed_features,
rounding_func=botorch_rounding_func,
**optimizer_options,
)
else:
acquisition_function = self.acqf_constructor( # pyre-ignore: [28]
model=model,
objective_weights=objective_weights,
objective_thresholds=objective_thresholds,
outcome_constraints=outcome_constraints,
X_observed=X_observed,
X_pending=X_pending,
Ys=self.Ys, # Required for qEHVI calculations.
**acf_options,
)
acquisition_function = checked_cast(AcquisitionFunction,
acquisition_function)
# pyre-ignore: [28]
candidates, expected_acquisition_value = self.acqf_optimizer(
acq_function=checked_cast(AcquisitionFunction,
acquisition_function),
bounds=bounds_,
n=n,
inequality_constraints=_to_inequality_constraints(
linear_constraints=linear_constraints),
fixed_features=fixed_features,
rounding_func=botorch_rounding_func,
**optimizer_options,
)
return (
candidates.detach().cpu(),
torch.ones(n, dtype=self.dtype),
{
"expected_acquisition_value":
expected_acquisition_value.tolist()
},
None,
)
def test_BotorchModel(self, dtype=torch.float, cuda=False):
Xs1, Ys1, Yvars1, bounds, task_features, feature_names = _get_torch_test_data(
dtype=dtype, cuda=cuda, constant_noise=True)
Xs2, Ys2, Yvars2, _, _, _ = _get_torch_test_data(dtype=dtype,
cuda=cuda,
constant_noise=True)
model = BotorchModel()
# Test ModelListGP
# make training data different for each output
Xs2_diff = [Xs2[0] + 0.1]
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2_diff,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
bounds=bounds,
task_features=task_features,
feature_names=feature_names,
fidelity_features=[],
)
_mock_fit_model.assert_called_once()
# Check attributes
self.assertTrue(torch.equal(model.Xs[0], Xs1[0]))
self.assertTrue(torch.equal(model.Xs[1], Xs2_diff[0]))
self.assertEqual(model.dtype, Xs1[0].dtype)
self.assertEqual(model.device, Xs1[0].device)
self.assertIsInstance(model.model, ModelListGP)
# Check fitting
model_list = model.model.models
self.assertTrue(torch.equal(model_list[0].train_inputs[0], Xs1[0]))
self.assertTrue(torch.equal(model_list[1].train_inputs[0],
Xs2_diff[0]))
self.assertTrue(
torch.equal(model_list[0].train_targets, Ys1[0].view(-1)))
self.assertTrue(
torch.equal(model_list[1].train_targets, Ys2[0].view(-1)))
self.assertIsInstance(model_list[0].likelihood,
_GaussianLikelihoodBase)
self.assertIsInstance(model_list[1].likelihood,
_GaussianLikelihoodBase)
# Test batched multi-output FixedNoiseGP
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
bounds=bounds,
task_features=task_features,
feature_names=feature_names,
fidelity_features=[],
)
_mock_fit_model.assert_called_once()
# Check attributes
self.assertTrue(torch.equal(model.Xs[0], Xs1[0]))
self.assertTrue(torch.equal(model.Xs[1], Xs2[0]))
self.assertEqual(model.dtype, Xs1[0].dtype)
self.assertEqual(model.device, Xs1[0].device)
self.assertIsInstance(model.model, FixedNoiseGP)
# Check fitting
# train inputs should be `o x n x 1`
self.assertTrue(
torch.equal(
model.model.train_inputs[0],
Xs1[0].unsqueeze(0).expand(torch.Size([2]) + Xs1[0].shape),
))
# train targets should be `o x n`
self.assertTrue(
torch.equal(model.model.train_targets,
torch.cat(Ys1 + Ys2, dim=-1).permute(1, 0)))
self.assertIsInstance(model.model.likelihood, _GaussianLikelihoodBase)
# Check infeasible cost can be computed on the model
device = torch.device("cuda") if cuda else torch.device("cpu")
objective_weights = torch.tensor([1.0, 0.0],
dtype=dtype,
device=device)
objective_transform = get_objective_weights_transform(
objective_weights)
infeasible_cost = torch.tensor(
get_infeasible_cost(X=Xs1[0],
model=model.model,
objective=objective_transform))
expected_infeasible_cost = -1 * torch.min(
objective_transform(
model.model.posterior(Xs1[0]).mean -
6 * model.model.posterior(Xs1[0]).variance.sqrt()).min(),
torch.tensor(0.0, dtype=dtype, device=device),
)
self.assertTrue(
torch.abs(infeasible_cost - expected_infeasible_cost) < 1e-5)
# Check prediction
X = torch.tensor([[6.0, 7.0, 8.0]], dtype=dtype, device=device)
f_mean, f_cov = model.predict(X)
self.assertTrue(f_mean.shape == torch.Size([1, 2]))
self.assertTrue(f_cov.shape == torch.Size([1, 2, 2]))
# Check generation
objective_weights = torch.tensor([1.0, 0.0],
dtype=dtype,
device=device)
outcome_constraints = (
torch.tensor([[0.0, 1.0]], dtype=dtype, device=device),
torch.tensor([[5.0]], dtype=dtype, device=device),
)
linear_constraints = (torch.tensor([[0.0, 1.0,
1.0]]), torch.tensor([[100.0]]))
fixed_features = None
pending_observations = [
torch.tensor([[1.0, 3.0, 4.0]], dtype=dtype, device=device),
torch.tensor([[2.0, 6.0, 8.0]], dtype=dtype, device=device),
]
n = 3
X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], dtype=dtype, device=device)
model_gen_options = {}
# test sequential optimize
with mock.patch("ax.models.torch.botorch_defaults.sequential_optimize",
return_value=X_dummy) as mock_optimize_acqf:
Xgen, wgen = model.gen(
n=n,
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
pending_observations=pending_observations,
model_gen_options=model_gen_options,
rounding_func=dummy_func,
)
# note: gen() always returns CPU tensors
self.assertTrue(torch.equal(Xgen, X_dummy.cpu()))
self.assertTrue(torch.equal(wgen, torch.ones(n, dtype=dtype)))
# test joint optimize
with mock.patch("ax.models.torch.botorch_defaults.joint_optimize",
return_value=X_dummy) as mock_optimize_acqf:
Xgen, wgen = model.gen(
n=n,
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=None,
linear_constraints=None,
fixed_features=fixed_features,
pending_observations=pending_observations,
model_gen_options={
"optimizer_kwargs": {
"joint_optimization": True
}
},
)
# note: gen() always returns CPU tensors
self.assertTrue(torch.equal(Xgen, X_dummy.cpu()))
self.assertTrue(torch.equal(wgen, torch.ones(n, dtype=dtype)))
mock_optimize_acqf.assert_called_once()
# test get_rounding_func
dummy_rounding = get_rounding_func(rounding_func=dummy_func)
X_temp = torch.rand(1, 2, 3, 4)
self.assertTrue(torch.equal(X_temp, dummy_rounding(X_temp)))
# Check best point selection
xbest = model.best_point(bounds=bounds,
objective_weights=objective_weights)
xbest = model.best_point(
bounds=bounds,
objective_weights=objective_weights,
fixed_features={0: 100.0},
)
self.assertIsNone(xbest)
# Test cross-validation
mean, variance = model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys1 + Ys2,
Yvars_train=Yvars1 + Yvars2,
X_test=torch.tensor([[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]],
dtype=dtype,
device=device),
)
self.assertTrue(mean.shape == torch.Size([2, 2]))
self.assertTrue(variance.shape == torch.Size([2, 2, 2]))
# Test cross-validation with refit_on_cv
model.refit_on_cv = True
mean, variance = model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys1 + Ys2,
Yvars_train=Yvars1 + Yvars2,
X_test=torch.tensor([[1.2, 3.2, 4.2], [2.4, 5.2, 3.2]],
dtype=dtype,
device=device),
)
self.assertTrue(mean.shape == torch.Size([2, 2]))
self.assertTrue(variance.shape == torch.Size([2, 2, 2]))
# Test update
model.refit_on_update = False
model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2)
# Test feature_importances
importances = model.feature_importances()
self.assertEqual(importances.shape, torch.Size([2, 1, 3]))
# When calling update directly, the data is completely overwritten.
self.assertTrue(torch.equal(model.Xs[0], Xs2[0]))
self.assertTrue(torch.equal(model.Xs[1], Xs2[0]))
self.assertTrue(torch.equal(model.Ys[0], Ys2[0]))
self.assertTrue(torch.equal(model.Yvars[0], Yvars2[0]))
model.refit_on_update = True
with mock.patch(FIT_MODEL_MO_PATH) as _mock_fit_model:
model.update(Xs=Xs2 + Xs2, Ys=Ys2 + Ys2, Yvars=Yvars2 + Yvars2)
# test unfit model CV, update, and feature_importances
unfit_model = BotorchModel()
with self.assertRaises(RuntimeError):
unfit_model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys1 + Ys2,
Yvars_train=Yvars1 + Yvars2,
X_test=Xs1[0],
)
with self.assertRaises(RuntimeError):
unfit_model.update(Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2)
with self.assertRaises(RuntimeError):
unfit_model.feature_importances()
# Test loading state dict
tkwargs = {"device": device, "dtype": dtype}
true_state_dict = {
"mean_module.constant": [3.5004],
"covar_module.raw_outputscale":
2.2438,
"covar_module.base_kernel.raw_lengthscale":
[[-0.9274, -0.9274, -0.9274]],
"covar_module.base_kernel.lengthscale_prior.concentration":
3.0,
"covar_module.base_kernel.lengthscale_prior.rate":
6.0,
"covar_module.outputscale_prior.concentration":
2.0,
"covar_module.outputscale_prior.rate":
0.15,
}
true_state_dict = {
key: torch.tensor(val, **tkwargs)
for key, val in true_state_dict.items()
}
model = get_and_fit_model(
Xs=Xs1,
Ys=Ys1,
Yvars=Yvars1,
task_features=[],
fidelity_features=[],
state_dict=true_state_dict,
refit_model=False,
)
for k, v in chain(model.named_parameters(), model.named_buffers()):
self.assertTrue(torch.equal(true_state_dict[k], v))
# Test for some change in model parameters & buffer for refit_model=True
true_state_dict["mean_module.constant"] += 0.1
true_state_dict["covar_module.raw_outputscale"] += 0.1
true_state_dict["covar_module.base_kernel.raw_lengthscale"] += 0.1
true_state_dict = {
key: torch.tensor(val, **tkwargs)
for key, val in true_state_dict.items()
}
model = get_and_fit_model(
Xs=Xs1,
Ys=Ys1,
Yvars=Yvars1,
task_features=[],
fidelity_features=[],
state_dict=true_state_dict,
refit_model=True,
)
self.assertTrue(
any(not torch.equal(true_state_dict[k], v) for k, v in chain(
model.named_parameters(), model.named_buffers())))
def gen(
self,
n: int,
bounds: List,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
pending_observations: Optional[List[Tensor]] = None,
model_gen_options: Optional[TConfig] = None,
rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
) -> Tuple[Tensor, Tensor, TGenMetadata]:
"""
Generate new candidates.
Args:
n: Number of candidates to generate.
bounds: A list of (lower, upper) tuples for each column of X.
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
outcome_constraints: A tuple of (A, b). For k outcome constraints
and m outputs at f(x), A is (k x m) and b is (k x 1) such that
A f(x) <= b.
linear_constraints: A tuple of (A, b). For k linear constraints on
d-dimensional x, A is (k x d) and b is (k x 1) such that
A x <= b.
fixed_features: A map {feature_index: value} for features that
should be fixed to a particular value during generation.
pending_observations: A list of m (k_i x d) feature tensors X
for m outcomes and k_i pending observations for outcome i.
model_gen_options: A config dictionary that can contain
model-specific options.
rounding_func: A function that rounds an optimization result
appropriately (i.e., according to `round-trip` transformations).
target_fidelities: A map {feature_index: value} of fidelity feature
column indices to their respective target fidelities. Used for
multi-fidelity optimization.
Returns:
2-element tuple containing
- (n x d) tensor of generated points.
- n-tensor of weights for each point.
"""
options = model_gen_options or {}
acf_options = options.get("acquisition_function_kwargs", {})
optimizer_options = options.get("optimizer_kwargs", {})
X_pending, X_observed = _get_X_pending_and_observed(
Xs=self.Xs,
pending_observations=pending_observations,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
bounds=bounds,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
)
objective = _get_objective(
model=self.model, # pyre-ignore: [6]
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
X_observed=X_observed,
)
# get the acquisition function
n_fantasies = acf_options.get("num_fantasies", 64)
qmc = acf_options.get("qmc", True)
seed_inner = acf_options.get("seed_inner", None)
num_restarts = optimizer_options.get("num_restarts", 40)
raw_samples = optimizer_options.get("raw_samples", 1024)
inequality_constraints = _to_inequality_constraints(linear_constraints)
# TODO: update optimizers to handle inequality_constraints
if inequality_constraints is not None:
raise UnsupportedError(
"Inequality constraints are not yet supported for KnowledgeGradient!"
)
# get current value
best_point_acqf, non_fixed_idcs = self._get_best_point_acqf(
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
X_observed=X_observed, # pyre-ignore: [6]
seed_inner=seed_inner,
fixed_features=fixed_features,
target_fidelities=target_fidelities,
qmc=qmc,
)
# solution from previous iteration
recommended_point = self.best_point(
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
model_gen_options=model_gen_options,
target_fidelities=target_fidelities,
)
recommended_point = recommended_point.detach().unsqueeze(
0) # pyre-ignore: [16]
# Extract acquisition value (TODO: Make this less painful and repetitive)
if non_fixed_idcs is not None:
recommended_point = recommended_point[..., non_fixed_idcs]
current_value = best_point_acqf(recommended_point).max()
acq_function = _instantiate_KG(
model=self.model, # pyre-ignore: [6]
objective=objective,
qmc=qmc,
n_fantasies=n_fantasies,
num_trace_observations=options.get("num_trace_observations", 0),
mc_samples=acf_options.get("mc_samples", 256),
seed_inner=seed_inner,
seed_outer=acf_options.get("seed_outer", None),
X_pending=X_pending,
target_fidelities=target_fidelities,
fidelity_weights=options.get("fidelity_weights"),
current_value=current_value,
cost_intercept=self.cost_intercept,
)
# optimize and get new points
bounds_ = torch.tensor(bounds, dtype=self.dtype, device=self.device)
bounds_ = bounds_.transpose(0, 1)
batch_initial_conditions = gen_one_shot_kg_initial_conditions(
acq_function=acq_function,
bounds=bounds_,
q=n,
num_restarts=num_restarts,
raw_samples=raw_samples,
options={
"frac_random": optimizer_options.get("frac_random", 0.1),
"num_inner_restarts": num_restarts,
"raw_inner_samples": raw_samples,
},
)
botorch_rounding_func = get_rounding_func(rounding_func)
candidates, _ = optimize_acqf(
acq_function=acq_function,
bounds=bounds_,
q=n,
inequality_constraints=inequality_constraints,
fixed_features=fixed_features,
post_processing_func=botorch_rounding_func,
num_restarts=num_restarts,
raw_samples=raw_samples,
options={
"batch_limit": optimizer_options.get("batch_limit", 8),
"maxiter": optimizer_options.get("maxiter", 200),
"method": "L-BFGS-B",
"nonnegative": optimizer_options.get("nonnegative", False),
},
batch_initial_conditions=batch_initial_conditions,
)
new_x = candidates.detach().cpu()
return new_x, torch.ones(n, dtype=self.dtype), {}
def gen(
self,
n: int,
bounds: List,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
pending_observations: Optional[List[Tensor]] = None,
model_gen_options: Optional[TConfig] = None,
rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
) -> Tuple[Tensor, Tensor, TGenMetadata]:
if linear_constraints is not None or outcome_constraints is not None:
raise UnsupportedError(
"Constraints are not yet supported by max-value entropy search!"
)
if len(objective_weights) > 1:
raise UnsupportedError(
"Models with multiple outcomes are not yet supported by MES!")
options = model_gen_options or {}
acf_options = options.get("acquisition_function_kwargs", {})
optimizer_options = options.get("optimizer_kwargs", {})
X_pending, X_observed = _get_X_pending_and_observed(
Xs=self.Xs,
pending_observations=pending_observations,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
bounds=bounds,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
)
# get the acquisition function
num_fantasies = acf_options.get("num_fantasies", 16)
num_mv_samples = acf_options.get("num_mv_samples", 10)
num_y_samples = acf_options.get("num_y_samples", 128)
candidate_size = acf_options.get("candidate_size", 1000)
num_restarts = optimizer_options.get("num_restarts", 40)
raw_samples = optimizer_options.get("raw_samples", 1024)
# generate the discrete points in the design space to sample max values
bounds_ = torch.tensor(bounds, dtype=self.dtype, device=self.device)
bounds_ = bounds_.transpose(0, 1)
candidate_set = torch.rand(candidate_size, bounds_.size(1))
candidate_set = bounds_[0] + (bounds_[1] - bounds_[0]) * candidate_set
acq_function = _instantiate_MES(
model=self.model, # pyre-ignore: [6]
candidate_set=candidate_set,
num_fantasies=num_fantasies,
num_trace_observations=options.get("num_trace_observations", 0),
num_mv_samples=num_mv_samples,
num_y_samples=num_y_samples,
X_pending=X_pending,
maximize=True if objective_weights[0] == 1 else False,
target_fidelities=target_fidelities,
fidelity_weights=options.get("fidelity_weights"),
cost_intercept=self.cost_intercept,
)
# optimize and get new points
botorch_rounding_func = get_rounding_func(rounding_func)
candidates, _ = optimize_acqf(
acq_function=acq_function,
bounds=bounds_,
q=n,
inequality_constraints=None,
fixed_features=fixed_features,
post_processing_func=botorch_rounding_func,
num_restarts=num_restarts,
raw_samples=raw_samples,
options={
"batch_limit": optimizer_options.get("batch_limit", 8),
"maxiter": optimizer_options.get("maxiter", 200),
"method": "L-BFGS-B",
"nonnegative": optimizer_options.get("nonnegative", False),
},
sequential=True,
)
new_x = candidates.detach().cpu()
return new_x, torch.ones(n, dtype=self.dtype), {}
def gen(
self,
n: int,
bounds: List[Tuple[float, float]],
objective_weights: Tensor, # objective_directions
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
objective_thresholds: Optional[Tensor] = None,
linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
fixed_features: Optional[Dict[int, float]] = None,
pending_observations: Optional[List[Tensor]] = None,
model_gen_options: Optional[TConfig] = None,
rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
target_fidelities: Optional[Dict[int, float]] = None,
) -> Tuple[Tensor, Tensor, TGenMetadata,
Optional[List[TCandidateMetadata]]]:
options = model_gen_options or {}
acf_options = options.get("acquisition_function_kwargs", {})
optimizer_options = options.get("optimizer_kwargs", {})
if target_fidelities:
raise NotImplementedError(
"target_fidelities not implemented for base BotorchModel")
if (objective_thresholds is not None and
objective_weights.shape[0] != objective_thresholds.shape[0]):
raise AxError(
"Objective weights and thresholds most both contain an element for"
" each modeled metric.")
X_pending, X_observed = _get_X_pending_and_observed(
Xs=self.Xs,
pending_observations=pending_observations,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
bounds=bounds,
linear_constraints=linear_constraints,
fixed_features=fixed_features,
)
model = not_none(self.model)
full_objective_thresholds = objective_thresholds
full_objective_weights = objective_weights
full_outcome_constraints = outcome_constraints
# subset model only to the outcomes we need for the optimization
if options.get(Keys.SUBSET_MODEL, True):
full_objective_weights
subset_model_results = subset_model(
model=model,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
objective_thresholds=objective_thresholds,
)
model = subset_model_results.model
objective_weights = subset_model_results.objective_weights
outcome_constraints = subset_model_results.outcome_constraints
objective_thresholds = subset_model_results.objective_thresholds
idcs = subset_model_results.indices
else:
idcs = None
if objective_thresholds is None:
full_objective_thresholds = infer_objective_thresholds(
model=model,
X_observed=not_none(X_observed),
objective_weights=full_objective_weights,
outcome_constraints=full_outcome_constraints,
subset_idcs=idcs,
)
# subset the objective thresholds
objective_thresholds = (full_objective_thresholds
if idcs is None else
full_objective_thresholds[idcs].clone())
bounds_ = torch.tensor(bounds, dtype=self.dtype, device=self.device)
bounds_ = bounds_.transpose(0, 1)
botorch_rounding_func = get_rounding_func(rounding_func)
if acf_options.get("random_scalarization", False) or acf_options.get(
"chebyshev_scalarization", False):
# If using a list of acquisition functions, the algorithm to generate
# that list is configured by acquisition_function_kwargs.
objective_weights_list = [
randomize_objective_weights(objective_weights, **acf_options)
for _ in range(n)
]
acquisition_function_list = [
self.acqf_constructor( # pyre-ignore: [28]
model=model,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
X_observed=X_observed,
X_pending=X_pending,
**acf_options,
) for objective_weights in objective_weights_list
]
acquisition_function_list = [
checked_cast(AcquisitionFunction, acq_function)
for acq_function in acquisition_function_list
]
# Multiple acquisition functions require a sequential optimizer
# always use scipy_optimizer_list.
# TODO(jej): Allow any optimizer.
candidates, expected_acquisition_value = scipy_optimizer_list(
acq_function_list=acquisition_function_list,
bounds=bounds_,
inequality_constraints=_to_inequality_constraints(
linear_constraints=linear_constraints),
fixed_features=fixed_features,
rounding_func=botorch_rounding_func,
**optimizer_options,
)
else:
acquisition_function = self.acqf_constructor( # pyre-ignore: [28]
model=model,
objective_weights=objective_weights,
objective_thresholds=objective_thresholds,
outcome_constraints=outcome_constraints,
X_observed=X_observed,
X_pending=X_pending,
**acf_options,
)
acquisition_function = checked_cast(AcquisitionFunction,
acquisition_function)
# pyre-ignore: [28]
candidates, expected_acquisition_value = self.acqf_optimizer(
acq_function=checked_cast(AcquisitionFunction,
acquisition_function),
bounds=bounds_,
n=n,
inequality_constraints=_to_inequality_constraints(
linear_constraints=linear_constraints),
fixed_features=fixed_features,
rounding_func=botorch_rounding_func,
**optimizer_options,
)
gen_metadata = {
"expected_acquisition_value": expected_acquisition_value.tolist(),
"objective_thresholds": not_none(full_objective_thresholds).cpu(),
}
return (
candidates.detach().cpu(),
torch.ones(n, dtype=self.dtype),
gen_metadata,
None,
)
