def optimize(self, acqf: MCAcquisitionFunction) -> Tuple[Tensor, Tensor]:
"""
Optimizes the acquisition function
:param acqf: The acquisition function being optimized
:return: Best solution and value
"""
initial_conditions = self.generate_restart_points(acqf)
# shape = num_restarts x *acqf.batch_shape x 1 x dim_X
if self.inequality_constraints is not None:
org_shape = initial_conditions.shape
initial_conditions = initial_conditions.reshape(
self.num_restarts, -1, self.dim_x)
options = {"maxiter": int(self.maxiter / 25)}
with settings.propagate_grads(True):
solutions, values = gen_candidates_scipy(
initial_conditions=initial_conditions,
acquisition_function=acqf,
lower_bounds=self.bounds[0],
upper_bounds=self.bounds[1],
options=options,
inequality_constraints=self.inequality_constraints,
)
self.add_solutions(solutions.view(-1, 1, self.dim_x).detach())
best_ind = torch.argmax(values, dim=0)
if self.inequality_constraints is not None:
solutions = solutions.reshape(org_shape)
solution = solutions.gather(
dim=0,
index=best_ind.view(1, *best_ind.shape, 1,
1).repeat(*[1] * (best_ind.dim() + 2),
self.dim_x),
)
if self.inequality_constraints is not None:
org_shape = solution.shape
solution = solution.reshape(1, -1, self.dim_x)
options = {"maxiter": self.maxiter}
with settings.propagate_grads(True):
solution, value = gen_candidates_scipy(
initial_conditions=solution,
acquisition_function=acqf,
lower_bounds=self.bounds[0],
upper_bounds=self.bounds[1],
options=options,
inequality_constraints=self.inequality_constraints,
)
# This is needed due to nested optimization
value = acqf(solution)
if self.inequality_constraints is not None:
solution = solution.reshape(org_shape)
return solution, value.reshape(*acqf.batch_shape)
def optimize_inner(self,
acqf: InnerRho,
return_best_only: bool = True) -> Tuple[Tensor, Tensor]:
"""
Optimizes the acquisition function
:param acqf: The acquisition function being optimized
:param return_best_only: If True, returns only the best solution. Otherwise,
returns all solutions returned by `gen_candidates_scipy`.
:return: Best solution and value
"""
X = self.generate_inner_raw_samples()
initial_conditions = self.generate_restart_points_from_samples(X, acqf)
solutions, values = gen_candidates_scipy(
initial_conditions=initial_conditions,
acquisition_function=acqf,
lower_bounds=self.inner_bounds[0],
upper_bounds=self.inner_bounds[1],
options={"maxiter": self.maxiter},
inequality_constraints=self.inequality_constraints,
)
solutions = solutions.detach()
values = values.detach()
self.add_inner_solutions(solutions, values)
best = torch.argmax(values.view(-1), dim=0)
if return_best_only:
solutions = solutions[best]
values = values[best]
self.current_best = -values
else:
self.current_best = -values[best]
return solutions, -values
def evaluate(self, X_actual: Tensor, bounds: Tensor, **kwargs: Any) -> Tensor:
r"""Evaluate qKnowledgeGradient on the candidate set `X_actual` by
solving the inner optimization problem.
Args:
X_actual: A `b x q x d` Tensor with `b` t-batches of `q` design points
each. Unlike `forward()`, this does not include solutions of the
inner optimization problem.
bounds: A `2 x d` tensor of lower and upper bounds for each column of
the solutions to the inner problem.
kwargs: Additional keyword arguments. This includes the options for
optimization of the inner problem, i.e. `num_restarts`, `raw_samples`,
an `options` dictionary to be passed on to the optimization helpers, and
a `scipy_options` dictionary to be passed to `scipy.minimize`.
Returns:
A Tensor of shape `b`. For t-batch b, the q-KG value of the design
`X_actual[b]` is averaged across the fantasy models.
NOTE: If `current_value` is not provided, then this is not the
true KG value of `X_actual[b]`.
"""
# construct the fantasy model of shape `num_fantasies x b`
fantasy_model = self.model.fantasize(
X=X_actual, sampler=self.sampler, observation_noise=True
)
# get the value function
value_function = _get_value_function(
model=fantasy_model, objective=self.objective, sampler=self.inner_sampler
)
# optimize the inner problem
from botorch.optim.initializers import gen_value_function_initial_conditions
from botorch.generation.gen import gen_candidates_scipy
initial_conditions = gen_value_function_initial_conditions(
acq_function=value_function,
bounds=bounds,
num_restarts=kwargs.get("num_restarts", 20),
raw_samples=kwargs.get("raw_samples", 1024),
current_model=self.model,
options={**kwargs.get("options", {}), **kwargs.get("scipy_options", {})},
)
_, values = gen_candidates_scipy(
initial_conditions=initial_conditions,
acquisition_function=value_function,
lower_bounds=bounds[0],
upper_bounds=bounds[1],
options=kwargs.get("scipy_options"),
)
# get the maximizer for each batch
values, _ = torch.max(values, dim=0)
if self.current_value is not None:
values = values - self.current_value
# return average over the fantasy samples
return values.mean(dim=0)
def test_gen_candidates_scipy_nan_handling(self):
for dtype, expected_regex in [
(torch.float, "Consider using"),
(torch.double, "gradient array"),
]:
ckwargs = {"dtype": dtype, "device": self.device}
test_ics = torch.rand(3, 1, **ckwargs)
test_grad = torch.tensor([0.5, 0.2, float("nan")], **ckwargs)
# test NaN in grad
with mock.patch("torch.autograd.grad", return_value=[test_grad]):
with self.assertRaisesRegex(RuntimeError, expected_regex):
gen_candidates_scipy(
initial_conditions=test_ics,
acquisition_function=mock.Mock(return_value=test_ics),
)
# test NaN in `x`
test_ics = torch.tensor([0.0, 0.0, float("nan")], **ckwargs)
with self.assertRaisesRegex(RuntimeError, "array `x` are NaN."):
gen_candidates_scipy(
initial_conditions=test_ics,
acquisition_function=mock.Mock(),
)
def test_random_restart_optimization(self):
for double in (True, False):
self._setUp(double=double)
with gpt_settings.debug(False):
best_f = self.model(self.train_x).mean.max().item()
qEI = qExpectedImprovement(self.model, best_f=best_f)
bounds = torch.tensor([[0.0], [1.0]]).type_as(self.train_x)
batch_ics = torch.rand(2, 1).type_as(self.train_x)
batch_candidates, batch_acq_values = gen_candidates_scipy(
initial_conditions=batch_ics,
acquisition_function=qEI,
lower_bounds=bounds[0],
upper_bounds=bounds[1],
options={"maxiter": 3},
)
candidates = get_best_candidates(batch_candidates=batch_candidates,
batch_values=batch_acq_values)
self.assertTrue(-EPS <= candidates <= 1 + EPS)
def test_gen_candidates_scipy_with_fixed_features_inequality_constraints(self):
options = {"maxiter": 5}
for double in (True, False):
self._setUp(double=double, expand=True)
qEI = qExpectedImprovement(self.model, best_f=self.f_best)
candidates, _ = gen_candidates_scipy(
initial_conditions=self.initial_conditions.reshape(1, 1, -1),
acquisition_function=qEI,
inequality_constraints=[
(torch.tensor([0]), torch.tensor([1]), 0),
(torch.tensor([1]), torch.tensor([-1]), -1),
],
fixed_features={1: 0.25},
options=options,
)
# candidates is of dimension 1 x 1 x 2
# so we are squeezing all the singleton dimensions
candidates = candidates.squeeze()
self.assertTrue(-EPS <= candidates[0] <= 1 + EPS)
self.assertTrue(candidates[1].item() == 0.25)
def optimize_acqf(
acq_function: AcquisitionFunction,
bounds: Tensor,
q: int,
num_restarts: int,
raw_samples: int,
options: Optional[Dict[str, Union[bool, float, int, str]]] = None,
inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None,
equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None,
fixed_features: Optional[Dict[int, float]] = None,
post_processing_func: Optional[Callable[[Tensor], Tensor]] = None,
batch_initial_conditions: Optional[Tensor] = None,
return_best_only: bool = True,
sequential: bool = False,
) -> Tuple[Tensor, Tensor]:
r"""Generate a set of candidates via multi-start optimization.
Args:
acq_function: An AcquisitionFunction.
bounds: A `2 x d` tensor of lower and upper bounds for each column of `X`.
q: The number of candidates.
num_restarts: The number of starting points for multistart acquisition
function optimization.
raw_samples: The number of samples for initialization.
options: Options for candidate generation.
inequality constraints: A list of tuples (indices, coefficients, rhs),
with each tuple encoding an inequality constraint of the form
`\sum_i (X[indices[i]] * coefficients[i]) >= rhs`
equality constraints: A list of tuples (indices, coefficients, rhs),
with each tuple encoding an inequality constraint of the form
`\sum_i (X[indices[i]] * coefficients[i]) = rhs`
fixed_features: A map `{feature_index: value}` for features that
should be fixed to a particular value during generation.
post_processing_func: A function that post-processes an optimization
result appropriately (i.e., according to `round-trip`
transformations).
batch_initial_conditions: A tensor to specify the initial conditions. Set
this if you do not want to use default initialization strategy.
return_best_only: If False, outputs the solutions corresponding to all
random restart initializations of the optimization.
sequential: If False, uses joint optimization, otherwise uses sequential
optimization.
Returns:
A two-element tuple containing
- a `(num_restarts) x q x d`-dim tensor of generated candidates.
- a tensor of associated acquisiton values. If `sequential=False`,
this is a `(num_restarts)`-dim tensor of joint acquisition values
(with explicit restart dimension if `return_best_only=False`). If
`sequential=True`, this is a `q`-dim tensor of expected acquisition
values conditional on having observed canidates `0,1,...,i-1`.
Example:
>>> # generate `q=2` candidates jointly using 20 random restarts
>>> # and 512 raw samples
>>> candidates, acq_value = optimize_acqf(qEI, bounds, 2, 20, 512)
>>> generate `q=3` candidates sequentially using 15 random restarts
>>> # and 256 raw samples
>>> qEI = qExpectedImprovement(model, best_f=0.2)
>>> bounds = torch.tensor([[0.], [1.]])
>>> candidates, acq_value_list = optimize_acqf(
>>> qEI, bounds, 3, 15, 256, sequential=True
>>> )
"""
if sequential and q > 1:
if not return_best_only:
raise NotImplementedError(
"return_best_only=False only supported for joint optimization"
)
if isinstance(acq_function, OneShotAcquisitionFunction):
raise NotImplementedError(
"sequential optimization currently not supported for one-shot "
"acquisition functions. Must have `sequential=False`."
)
candidate_list, acq_value_list = [], []
candidates = torch.tensor([])
base_X_pending = acq_function.X_pending
for i in range(q):
candidate, acq_value = optimize_acqf(
acq_function=acq_function,
bounds=bounds,
q=1,
num_restarts=num_restarts,
raw_samples=raw_samples,
options=options or {},
inequality_constraints=inequality_constraints,
equality_constraints=equality_constraints,
fixed_features=fixed_features,
post_processing_func=post_processing_func,
batch_initial_conditions=None,
return_best_only=True,
sequential=False,
)
candidate_list.append(candidate)
acq_value_list.append(acq_value)
candidates = torch.cat(candidate_list, dim=-2)
acq_function.set_X_pending(
torch.cat([base_X_pending, candidates], dim=-2)
if base_X_pending is not None
else candidates
)
logger.info(f"Generated sequential candidate {i+1} of {q}")
# Reset acq_func to previous X_pending state
acq_function.set_X_pending(base_X_pending)
return candidates, torch.stack(acq_value_list)
options = options or {}
if batch_initial_conditions is None:
ic_gen = (
gen_one_shot_kg_initial_conditions
if isinstance(acq_function, qKnowledgeGradient)
else gen_batch_initial_conditions
)
# TODO: Generating initial candidates should use parameter constraints.
batch_initial_conditions = ic_gen(
acq_function=acq_function,
bounds=bounds,
q=q,
num_restarts=num_restarts,
raw_samples=raw_samples,
options=options,
)
batch_limit: int = options.get("batch_limit", num_restarts)
batch_candidates_list: List[Tensor] = []
batch_acq_values_list: List[Tensor] = []
start_idcs = list(range(0, num_restarts, batch_limit))
for start_idx in start_idcs:
end_idx = min(start_idx + batch_limit, num_restarts)
# optimize using random restart optimization
batch_candidates_curr, batch_acq_values_curr = gen_candidates_scipy(
initial_conditions=batch_initial_conditions[start_idx:end_idx],
acquisition_function=acq_function,
lower_bounds=bounds[0],
upper_bounds=bounds[1],
options={
k: v
for k, v in options.items()
if k not in ("batch_limit", "nonnegative")
},
inequality_constraints=inequality_constraints,
equality_constraints=equality_constraints,
fixed_features=fixed_features,
)
batch_candidates_list.append(batch_candidates_curr)
batch_acq_values_list.append(batch_acq_values_curr)
logger.info(f"Generated candidate batch {start_idx+1} of {len(start_idcs)}.")
batch_candidates = torch.cat(batch_candidates_list)
batch_acq_values = torch.cat(batch_acq_values_list)
if post_processing_func is not None:
batch_candidates = post_processing_func(batch_candidates)
if return_best_only:
best = torch.argmax(batch_acq_values.view(-1), dim=0)
batch_candidates = batch_candidates[best]
batch_acq_values = batch_acq_values[best]
if isinstance(acq_function, OneShotAcquisitionFunction):
batch_candidates = acq_function.extract_candidates(X_full=batch_candidates)
return batch_candidates, batch_acq_values
def optimize_acqf(
acq_function: AcquisitionFunction,
bounds: Tensor,
q: int,
num_restarts: int,
raw_samples: Optional[int] = None,
options: Optional[Dict[str, Union[bool, float, int, str]]] = None,
inequality_constraints: Optional[List[Tuple[Tensor, Tensor,
float]]] = None,
equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None,
nonlinear_inequality_constraints: Optional[List[Callable]] = None,
fixed_features: Optional[Dict[int, float]] = None,
post_processing_func: Optional[Callable[[Tensor], Tensor]] = None,
batch_initial_conditions: Optional[Tensor] = None,
return_best_only: bool = True,
sequential: bool = False,
**kwargs: Any,
) -> Tuple[Tensor, Tensor]:
r"""Generate a set of candidates via multi-start optimization.
Args:
acq_function: An AcquisitionFunction.
bounds: A `2 x d` tensor of lower and upper bounds for each column of `X`.
q: The number of candidates.
num_restarts: The number of starting points for multistart acquisition
function optimization.
raw_samples: The number of samples for initialization. This is required
if `batch_initial_conditions` is not specified.
options: Options for candidate generation.
inequality constraints: A list of tuples (indices, coefficients, rhs),
with each tuple encoding an inequality constraint of the form
`\sum_i (X[indices[i]] * coefficients[i]) >= rhs`
equality constraints: A list of tuples (indices, coefficients, rhs),
with each tuple encoding an inequality constraint of the form
`\sum_i (X[indices[i]] * coefficients[i]) = rhs`
nonlinear_inequality_constraints: A list of callables with that represent
non-linear inequality constraints of the form `callable(x) >= 0`. Each
callable is expected to take a `(num_restarts) x q x d`-dim tensor as an
input and return a `(num_restarts) x q`-dim tensor with the constraint
values. The constraints will later be passed to SLSQP. You need to pass in
`batch_initial_conditions` in this case. Using non-linear inequality
constraints also requires that `batch_limit` is set to 1, which will be
done automatically if not specified in `options`.
fixed_features: A map `{feature_index: value}` for features that
should be fixed to a particular value during generation.
post_processing_func: A function that post-processes an optimization
result appropriately (i.e., according to `round-trip`
transformations).
batch_initial_conditions: A tensor to specify the initial conditions. Set
this if you do not want to use default initialization strategy.
return_best_only: If False, outputs the solutions corresponding to all
random restart initializations of the optimization.
sequential: If False, uses joint optimization, otherwise uses sequential
optimization.
kwargs: Additonal keyword arguments.
Returns:
A two-element tuple containing
- a `(num_restarts) x q x d`-dim tensor of generated candidates.
- a tensor of associated acquisition values. If `sequential=False`,
this is a `(num_restarts)`-dim tensor of joint acquisition values
(with explicit restart dimension if `return_best_only=False`). If
`sequential=True`, this is a `q`-dim tensor of expected acquisition
values conditional on having observed canidates `0,1,...,i-1`.
Example:
>>> # generate `q=2` candidates jointly using 20 random restarts
>>> # and 512 raw samples
>>> candidates, acq_value = optimize_acqf(qEI, bounds, 2, 20, 512)
>>> generate `q=3` candidates sequentially using 15 random restarts
>>> # and 256 raw samples
>>> qEI = qExpectedImprovement(model, best_f=0.2)
>>> bounds = torch.tensor([[0.], [1.]])
>>> candidates, acq_value_list = optimize_acqf(
>>> qEI, bounds, 3, 15, 256, sequential=True
>>> )
"""
if sequential and q > 1:
if not return_best_only:
raise NotImplementedError(
"`return_best_only=False` only supported for joint optimization."
)
if isinstance(acq_function, OneShotAcquisitionFunction):
raise NotImplementedError(
"sequential optimization currently not supported for one-shot "
"acquisition functions. Must have `sequential=False`.")
candidate_list, acq_value_list = [], []
base_X_pending = acq_function.X_pending
for i in range(q):
candidate, acq_value = optimize_acqf(
acq_function=acq_function,
bounds=bounds,
q=1,
num_restarts=num_restarts,
raw_samples=raw_samples,
options=options or {},
inequality_constraints=inequality_constraints,
equality_constraints=equality_constraints,
nonlinear_inequality_constraints=
nonlinear_inequality_constraints,
fixed_features=fixed_features,
post_processing_func=post_processing_func,
batch_initial_conditions=None,
return_best_only=True,
sequential=False,
)
candidate_list.append(candidate)
acq_value_list.append(acq_value)
candidates = torch.cat(candidate_list, dim=-2)
acq_function.set_X_pending(
torch.cat([base_X_pending, candidates], dim=-2
) if base_X_pending is not None else candidates)
logger.info(f"Generated sequential candidate {i+1} of {q}")
# Reset acq_func to previous X_pending state
acq_function.set_X_pending(base_X_pending)
return candidates, torch.stack(acq_value_list)
options = options or {}
# Handle the trivial case when all features are fixed
if fixed_features is not None and len(fixed_features) == bounds.shape[-1]:
X = torch.tensor(
[fixed_features[i] for i in range(bounds.shape[-1])],
device=bounds.device,
dtype=bounds.dtype,
)
X = X.expand(q, *X.shape)
with torch.no_grad():
acq_value = acq_function(X)
return X, acq_value
if batch_initial_conditions is None:
if nonlinear_inequality_constraints:
raise NotImplementedError(
"`batch_initial_conditions` must be given if there are non-linear "
"inequality constraints.")
if raw_samples is None:
raise ValueError(
"Must specify `raw_samples` when `batch_initial_conditions` is `None`."
)
ic_gen = (gen_one_shot_kg_initial_conditions if isinstance(
acq_function, qKnowledgeGradient) else
gen_batch_initial_conditions)
batch_initial_conditions = ic_gen(
acq_function=acq_function,
bounds=bounds,
q=q,
num_restarts=num_restarts,
raw_samples=raw_samples,
fixed_features=fixed_features,
options=options,
inequality_constraints=inequality_constraints,
equality_constraints=equality_constraints,
)
batch_limit: int = options.get(
"batch_limit",
num_restarts if not nonlinear_inequality_constraints else 1)
batch_candidates_list: List[Tensor] = []
batch_acq_values_list: List[Tensor] = []
batched_ics = batch_initial_conditions.split(batch_limit)
for i, batched_ics_ in enumerate(batched_ics):
# optimize using random restart optimization
batch_candidates_curr, batch_acq_values_curr = gen_candidates_scipy(
initial_conditions=batched_ics_,
acquisition_function=acq_function,
lower_bounds=bounds[0],
upper_bounds=bounds[1],
options={
k: v
for k, v in options.items() if k not in INIT_OPTION_KEYS
},
inequality_constraints=inequality_constraints,
equality_constraints=equality_constraints,
nonlinear_inequality_constraints=nonlinear_inequality_constraints,
fixed_features=fixed_features,
)
batch_candidates_list.append(batch_candidates_curr)
batch_acq_values_list.append(batch_acq_values_curr)
logger.info(f"Generated candidate batch {i+1} of {len(batched_ics)}.")
batch_candidates = torch.cat(batch_candidates_list)
batch_acq_values = torch.cat(batch_acq_values_list)
if post_processing_func is not None:
batch_candidates = post_processing_func(batch_candidates)
if return_best_only:
best = torch.argmax(batch_acq_values.view(-1), dim=0)
batch_candidates = batch_candidates[best]
batch_acq_values = batch_acq_values[best]
if isinstance(acq_function, OneShotAcquisitionFunction):
if not kwargs.get("return_full_tree", False):
batch_candidates = acq_function.extract_candidates(
X_full=batch_candidates)
return batch_candidates, batch_acq_values
def optimize_outer(
self,
acqf: MCAcquisitionFunction,
w_samples: Tensor = None,
batch_size: int = 5,
random_w: bool = False,
) -> Tuple[Tensor, Tensor]:
"""
rhoKGapx, Nested or Tts optimizer with w component restricted to w_samples
:param acqf: rhoKGapx or rhoKG object
:param w_samples: the set W to consider. If None, assumes continuous optimization.
:param batch_size: We will do the optimization in mini batches to save on memory
:param random_w: If this is True, the w component of the candidate is fixed to
a random realization instead of being optimized. This is only for
presenting a comparison in the paper, and should not be used.
:return: Optimal solution and value
"""
if self.low_fantasies is not None:
# set the low num_fantasies
acqf.change_num_fantasies(num_fantasies=self.low_fantasies)
if random_w:
if w_samples is not None:
w_samples = w_samples[int(
torch.randint(w_samples.shape[0], (1, )))].reshape(1, -1)
else:
w_samples = torch.rand(1,
self.dim - self.dim_x,
device=self.device,
dtype=self.dtype)
initial_conditions = self.generate_outer_restart_points(
acqf, w_samples)
if self.low_fantasies is not None:
# recover the original num_fantasies
acqf.change_num_fantasies()
if w_samples is not None:
fixed_features = dict()
# If q-batch is used, then we only need to fix dim_x:dim
if self.solution_shape[0] == self.q:
for i in range(self.dim_x, self.dim):
fixed_features[i] = None
else:
# If solutions are one-shot, then we need to fix dim_x:dim for each q
for j in range(self.q):
for i in range(self.dim_x, self.dim):
fixed_features[j * self.dim + i] = None
else:
fixed_features = None
acqf.tts_reset()
init_size = initial_conditions.shape[0]
num_batches = ceil(init_size / batch_size)
solutions = torch.empty(init_size,
*self.solution_shape,
dtype=self.dtype,
device=self.device)
values = torch.empty(init_size, dtype=self.dtype, device=self.device)
options = {"maxiter": int(self.maxiter / 25)}
for i in range(num_batches):
l_idx = i * batch_size
if i == num_batches - 1:
r_idx = init_size
else:
r_idx = (i + 1) * batch_size
solutions[l_idx:r_idx], values[l_idx:r_idx] = gen_candidates_scipy(
initial_conditions=initial_conditions[l_idx:r_idx],
acquisition_function=acqf,
lower_bounds=self.outer_bounds[0],
upper_bounds=self.outer_bounds[1],
options=options,
fixed_features=fixed_features,
inequality_constraints=self.inequality_constraints,
)
_, idx = torch.sort(values)
acqf.tts_reset()
options = {"maxiter": self.maxiter}
solutions, values = gen_candidates_scipy(
initial_conditions=solutions[idx[:self.num_refine_restarts]],
acquisition_function=acqf,
lower_bounds=self.outer_bounds[0],
upper_bounds=self.outer_bounds[1],
options=options,
fixed_features=fixed_features,
inequality_constraints=self.inequality_constraints,
)
best = torch.argmax(values)
return solutions[best].detach(), values[best].detach()
