def get_botorch_objective(
model: Model,
objective_weights: Tensor,
use_scalarized_objective: bool = True,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
objective_thresholds: Optional[Tensor] = None,
X_observed: Optional[Tensor] = None,
) -> AcquisitionObjective:
"""Constructs a BoTorch `AcquisitionObjective` object.
Args:
model: A BoTorch Model
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
use_scalarized_objective: A boolean parameter that defaults to True,
specifying whether ScalarizedObjective should be used.
NOTE: when using outcome_constraints, use_scalarized_objective
will be ignored.
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. (Not used by single task models)
objective_thresholds: A tensor containing thresholds forming a reference point
from which to calculate pareto frontier hypervolume. Points that do not
dominate the objective_thresholds contribute nothing to hypervolume.
X_observed: Observed points that are feasible and appear in the
objective or the constraints. None if there are no such points.
Returns:
A BoTorch `AcquisitionObjective` object. It will be one of:
`ScalarizedObjective`, `LinearMCOObjective`, `ConstrainedMCObjective`.
"""
if objective_thresholds is not None:
nonzero_idcs = torch.nonzero(objective_weights).view(-1)
objective_weights = objective_weights[nonzero_idcs]
return WeightedMCMultiOutputObjective(weights=objective_weights,
outcomes=nonzero_idcs.tolist())
if X_observed is None:
raise UnsupportedError(
"X_observed is required to construct a BoTorch Objective.")
if outcome_constraints:
if use_scalarized_objective:
logger.warning(
"Currently cannot use ScalarizedObjective when there are outcome "
"constraints. Ignoring (default) kwarg `use_scalarized_objective`"
"= True. Creating ConstrainedMCObjective.")
obj_tf = get_objective_weights_transform(objective_weights)
def objective(samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
return obj_tf(samples)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
return ConstrainedMCObjective(objective=objective,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
elif use_scalarized_objective:
return ScalarizedObjective(weights=objective_weights)
return LinearMCObjective(weights=objective_weights)
def get_NEI(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
**kwargs: Any,
) -> AcquisitionFunction:
r"""Instantiates a qNoisyExpectedImprovement acquisition function.
Args:
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. (Not used by single task models)
X_observed: A tensor containing points observed for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
X_pending: A tensor containing points whose evaluation is pending (i.e.
that have been submitted for evaluation) present for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
mc_samples: The number of MC samples to use (default: 512).
qmc: If True, use qMC instead of MC (default: True).
prune_baseline: If True, prune the baseline points for NEI (default: True).
Returns:
qNoisyExpectedImprovement: The instantiated acquisition function.
"""
if X_observed is None:
raise ValueError("There are no feasible observed points.")
# Parse random_scalarization params
objective_weights = _extract_random_scalarization_settings(
objective_weights, outcome_constraints, **kwargs)
# construct Objective module
if outcome_constraints is None:
objective = LinearMCObjective(weights=objective_weights)
else:
obj_tf = get_objective_weights_transform(objective_weights)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
X_observed = torch.as_tensor(X_observed)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return get_acquisition_function(
acquisition_function_name="qNEI",
model=model,
objective=objective,
X_observed=X_observed,
X_pending=X_pending,
prune_baseline=kwargs.get("prune_baseline", True),
mc_samples=kwargs.get("mc_samples", 512),
qmc=kwargs.get("qmc", True),
seed=torch.randint(1, 10000, (1, )).item(),
)
def get_PosteriorMean(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
**kwargs: Any,
) -> AcquisitionFunction:
r"""Instantiates a PosteriorMean acquisition function.
Note: If no OutcomeConstraints given, return an analytic acquisition
function. This requires {optimizer_kwargs: {joint_optimization: True}} or an
optimizer that does not assume pending point support.
Args:
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. (Not used by single task models)
X_observed: A tensor containing points observed for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
X_pending: A tensor containing points whose evaluation is pending (i.e.
that have been submitted for evaluation) present for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
Returns:
PosteriorMean: The instantiated acquisition function.
"""
if X_observed is None:
raise ValueError("There are no feasible observed points.")
# construct Objective module
if kwargs.get("chebyshev_scalarization", False):
obj_tf = get_chebyshev_scalarization(
weights=objective_weights,
Y=torch.stack(kwargs.get("Ys")).transpose(0, 1).squeeze(-1),
)
else:
obj_tf = get_objective_weights_transform(objective_weights)
def obj_fn(samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
return obj_tf(samples)
if outcome_constraints is None:
objective = GenericMCObjective(objective=obj_fn)
else:
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_fn)
objective = ConstrainedMCObjective(objective=obj_fn,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
# Use qSimpleRegret, not analytic posterior, to handle arbitrary objective fns.
acq_func = qSimpleRegret(model, objective=objective)
return acq_func
def test_get_infeasible_cost(self):
for dtype in (torch.float, torch.double):
X = torch.zeros(5, 1, device=self.device, dtype=dtype)
means = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0],
device=self.device,
dtype=dtype)
variances = torch.tensor([0.09, 0.25, 0.36, 0.25, 0.09],
device=self.device,
dtype=dtype)
mm = MockModel(MockPosterior(mean=means, variance=variances))
# means - 6 * std = [-0.8, -1, -0.6, 1, 3.2]. After applying the
# objective, the minimum becomes -6.0, so 6.0 should be returned.
M = get_infeasible_cost(X=X,
model=mm,
objective=lambda Y: Y.squeeze(-1) - 5.0)
self.assertEqual(M, 6.0)
# test default objective (squeeze last dim)
M2 = get_infeasible_cost(X=X, model=mm)
self.assertEqual(M2, 1.0)
def get_botorch_objective_and_transform(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
objective_thresholds: Optional[Tensor] = None,
X_observed: Optional[Tensor] = None,
) -> Tuple[Optional[MCAcquisitionObjective], Optional[PosteriorTransform]]:
"""Constructs a BoTorch `AcquisitionObjective` object.
Args:
model: A BoTorch Model
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. (Not used by single task models)
objective_thresholds: A tensor containing thresholds forming a reference point
from which to calculate pareto frontier hypervolume. Points that do not
dominate the objective_thresholds contribute nothing to hypervolume.
X_observed: Observed points that are feasible and appear in the
objective or the constraints. None if there are no such points.
Returns:
A two-tuple containing (optioally) an `MCAcquisitionObjective` and
(optionally) a `PosteriorTransform`.
"""
if objective_thresholds is not None:
# we are doing multi-objective optimization
nonzero_idcs = torch.nonzero(objective_weights).view(-1)
objective_weights = objective_weights[nonzero_idcs]
objective = WeightedMCMultiOutputObjective(
weights=objective_weights, outcomes=nonzero_idcs.tolist())
return objective, None
if X_observed is None:
raise UnsupportedError(
"X_observed is required to construct a BoTorch objective.")
if outcome_constraints:
# If there are outcome constraints, we use MC Acquistion functions
obj_tf = get_objective_weights_transform(objective_weights)
def objective(samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
return obj_tf(samples)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=objective,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return objective, None
# Case of linear weights - use ScalarizedPosteriorTransform
transform = ScalarizedPosteriorTransform(weights=objective_weights)
return None, transform
def get_botorch_objective(
model: Model,
objective_weights: Tensor,
use_scalarized_objective: bool = True,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
) -> AcquisitionObjective:
"""Constructs a BoTorch `AcquisitionObjective` object.
Args:
model: A BoTorch Model
objective_weights: The objective is to maximize a weighted sum of
the columns of f(x). These are the weights.
use_scalarized_objective: A boolean parameter that defaults to True,
specifying whether ScalarizedObjective should be used.
NOTE: when using outcome_constraints, use_scalarized_objective
will be ignored.
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. (Not used by single task models)
X_observed: Observed points that are feasible and appear in the
objective or the constraints. None if there are no such points.
Returns:
A BoTorch `AcquisitionObjective` object. It will be one of:
`ScalarizedObjective`, `LinearMCOObjective`, `ConstrainedMCObjective`.
"""
if X_observed is None:
raise UnsupportedError(
"X_observed is required to construct a BoTorch Objective.")
if outcome_constraints:
if use_scalarized_objective:
logger.warning(
"Currently cannot use ScalarizedObjective when there are outcome "
"constraints. Ignoring (default) kwarg `use_scalarized_objective`"
"= True. Creating ConstrainedMCObjective.")
obj_tf = get_objective_weights_transform(objective_weights)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
return ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
if use_scalarized_objective:
return ScalarizedObjective(weights=objective_weights)
return LinearMCObjective(weights=objective_weights)
def _get_objective(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
) -> AcquisitionObjective:
if outcome_constraints is None:
objective = ScalarizedObjective(weights=objective_weights)
else:
X_observed = torch.as_tensor(X_observed)
obj_tf = get_objective_weights_transform(objective_weights)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return objective
def get_botorch_objective(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
) -> AcquisitionObjective:
"""Constructs a BoTorch `Objective`."""
if X_observed is None:
raise UnsupportedError(
"X_observed is required to construct a BoTorch Objective.")
if outcome_constraints is None:
objective = ScalarizedObjective(weights=objective_weights)
else:
obj_tf = get_objective_weights_transform(objective_weights)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return objective
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 get_NEI(
model: Model,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
**kwargs: Any,
) -> AcquisitionFunction:
r"""Instantiates a qNoisyExpectedImprovement acquisition function.
Args:
model: The underlying model which the acqusition function uses
to estimate acquisition values of candidates.
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. (Not used by single task models)
X_observed: A tensor containing points observed for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
X_pending: A tensor containing points whose evaluation is pending (i.e.
that have been submitted for evaluation) present for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
mc_samples: The number of MC samples to use (default: 512).
qmc: If True, use qMC instead of MC (default: True).
prune_baseline: If True, prune the baseline points for NEI (default: True).
chebyshev_scalarization: Use augmented Chebyshev scalarization.
Returns:
qNoisyExpectedImprovement: The instantiated acquisition function.
"""
if X_observed is None:
raise ValueError("There are no feasible observed points.")
# construct Objective module
if kwargs.get("chebyshev_scalarization", False):
if "Ys" not in kwargs:
raise ValueError("Chebyshev Scalarization requires Ys argument")
Y_tensor = torch.cat(kwargs.get("Ys"), dim=-1)
obj_tf = get_chebyshev_scalarization(weights=objective_weights,
Y=Y_tensor)
else:
obj_tf = get_objective_weights_transform(objective_weights)
if outcome_constraints is None:
objective = GenericMCObjective(objective=obj_tf)
else:
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed,
model=model,
objective=obj_tf)
objective = ConstrainedMCObjective(objective=obj_tf,
constraints=con_tfs or [],
infeasible_cost=inf_cost)
return get_acquisition_function(
acquisition_function_name="qNEI",
model=model,
objective=objective,
X_observed=X_observed,
X_pending=X_pending,
prune_baseline=kwargs.get("prune_baseline", True),
mc_samples=kwargs.get("mc_samples", 512),
qmc=kwargs.get("qmc", True),
# pyre-fixme[6]: Expected `Optional[int]` for 9th param but got
# `Union[float, int]`.
seed=torch.randint(1, 10000, (1, )).item(),
)
def _get_acquisition_func(
model: Model,
acquisition_function_name: str,
objective_weights: Tensor,
outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
X_observed: Optional[Tensor] = None,
X_pending: Optional[Tensor] = None,
mc_objective: Type[GenericMCObjective] = GenericMCObjective,
constrained_mc_objective: Optional[
Type[ConstrainedMCObjective]
] = ConstrainedMCObjective,
mc_objective_kwargs: Optional[Dict] = None,
**kwargs: Any,
) -> AcquisitionFunction:
r"""Instantiates a acquisition function.
Args:
model: The underlying model which the acqusition function uses
to estimate acquisition values of candidates.
acquisition_function_name: Name of the acquisition function.
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. (Not used by single task models)
X_observed: A tensor containing points observed for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
X_pending: A tensor containing points whose evaluation is pending (i.e.
that have been submitted for evaluation) present for all objective
outcomes and outcomes that appear in the outcome constraints (if
there are any).
mc_objective: GenericMCObjective class, used for constructing a
MC-objective. If constructing a penalized MC-objective, pass in
PenalizedMCObjective together with mc_objective_kwargs .
constrained_mc_objective: ConstrainedMCObjective class, used when
applying constraints on the outcomes.
mc_objective_kwargs: kwargs for constructing MC-objective.
For GenericMCObjective, leave it as None. For PenalizedMCObjective,
it needs to be specified in the format of kwargs.
mc_samples: The number of MC samples to use (default: 512).
qmc: If True, use qMC instead of MC (default: True).
prune_baseline: If True, prune the baseline points for NEI (default: True).
chebyshev_scalarization: Use augmented Chebyshev scalarization.
Returns:
The instantiated acquisition function.
"""
if X_observed is None:
raise ValueError("There are no feasible observed points.")
# construct Objective module
if kwargs.get("chebyshev_scalarization", False):
with torch.no_grad():
Y = model.posterior(X_observed).mean
obj_tf = get_chebyshev_scalarization(weights=objective_weights, Y=Y)
else:
obj_tf = get_objective_weights_transform(objective_weights)
def objective(samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
return obj_tf(samples)
if outcome_constraints is None:
mc_objective_kwargs = {} if mc_objective_kwargs is None else mc_objective_kwargs
objective = mc_objective(objective=objective, **mc_objective_kwargs)
else:
if constrained_mc_objective is None:
raise ValueError(
"constrained_mc_objective cannot be set to None "
"when applying outcome constraints."
)
if issubclass(mc_objective, PenalizedMCObjective):
raise RuntimeError(
"Outcome constraints are not supported for PenalizedMCObjective."
)
con_tfs = get_outcome_constraint_transforms(outcome_constraints)
inf_cost = get_infeasible_cost(X=X_observed, model=model, objective=objective)
objective = constrained_mc_objective(
objective=objective, constraints=con_tfs or [], infeasible_cost=inf_cost
)
return get_acquisition_function(
acquisition_function_name=acquisition_function_name,
model=model,
objective=objective,
X_observed=X_observed,
X_pending=X_pending,
prune_baseline=kwargs.get("prune_baseline", True),
mc_samples=kwargs.get("mc_samples", 512),
qmc=kwargs.get("qmc", True),
# pyre-fixme[6]: Expected `Optional[int]` for 9th param but got
# `Union[float, int]`.
seed=torch.randint(1, 10000, (1,)).item(),
marginalize_dim=kwargs.get("marginalize_dim"),
)
def test_FullyBayesianBotorchModel(self,
dtype=torch.float,
cuda=False):
Xs1, Ys1, Yvars1, bounds, tfs, fns, mns = 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)
Yvars_inferred_noise = [
torch.full_like(Yvars1[0], float("nan")),
torch.full_like(Yvars2[0], float("nan")),
]
# make input different for each output
Xs2_diff = [Xs2[0] + 0.1]
Xs = Xs1 + Xs2_diff
Ys = Ys1 + Ys2
for inferred_noise, use_input_warping, use_saas in product(
(True, False), repeat=3):
Yvars = Yvars_inferred_noise if inferred_noise else Yvars1 + Yvars2
model = self.model_cls(
use_input_warping=use_input_warping,
thinning=1,
num_samples=4,
use_saas=use_saas,
disable_progbar=True,
max_tree_depth=1,
)
if use_input_warping:
self.assertTrue(model.use_input_warping)
# Test ModelListGP
# make training data different for each output
tkwargs = {"dtype": dtype, "device": Xs1[0].device}
dummy_samples_list = _get_dummy_mcmc_samples(num_samples=4,
num_outputs=2,
**tkwargs)
for dummy_samples in dummy_samples_list:
if use_input_warping:
dummy_samples["c0"] = (
torch.rand(4, 1, Xs1[0].shape[-1], **tkwargs) * 0.5
+ 0.1)
dummy_samples["c1"] = (
torch.rand(4, 1, Xs1[0].shape[-1], **tkwargs) * 0.5
+ 0.1)
if inferred_noise:
dummy_samples["noise"] = torch.rand(
4, 1,
**tkwargs).clamp_min(MIN_INFERRED_NOISE_LEVEL)
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model:
model.fit(
Xs=Xs,
Ys=Ys,
Yvars=Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=tfs,
),
metric_names=mns,
)
self.assertEqual(_mock_fit_model.call_count, 2)
for i, call in enumerate(_mock_fit_model.call_args_list):
_, ckwargs = call
X = Xs[i]
Y = Ys[i]
Yvar = Yvars[i]
self.assertIs(ckwargs["pyro_model"], pyro_model)
self.assertTrue(torch.equal(ckwargs["X"], X))
self.assertTrue(torch.equal(ckwargs["Y"], Y))
if inferred_noise:
self.assertTrue(torch.isnan(ckwargs["Yvar"]).all())
else:
self.assertTrue(torch.equal(ckwargs["Yvar"], Yvar))
self.assertEqual(ckwargs["num_samples"], 4)
self.assertEqual(ckwargs["warmup_steps"], 1024)
self.assertEqual(ckwargs["max_tree_depth"], 1)
self.assertTrue(ckwargs["disable_progbar"])
self.assertEqual(ckwargs["use_input_warping"],
use_input_warping)
self.assertEqual(ckwargs["use_saas"], use_saas)
# Check attributes
self.assertTrue(torch.equal(model.Xs[i], Xs[i]))
self.assertEqual(model.dtype, Xs[i].dtype)
self.assertEqual(model.device, Xs[i].device)
self.assertIsInstance(model.model, ModelListGP)
# Check fitting
# Note each model in the model list is a batched model, where
# the batch dim corresponds to the MCMC samples
model_list = model.model.models
# Put model in `eval` mode to transform the train inputs.
m = model_list[i].eval()
# check mcmc samples
dummy_samples = dummy_samples_list[i]
expected_train_inputs = Xs[i].expand(4, *Xs[i].shape)
if use_input_warping:
# train inputs should be warped inputs
expected_train_inputs = m.input_transform(
expected_train_inputs)
self.assertTrue(
torch.equal(
m.train_inputs[0],
expected_train_inputs,
))
self.assertTrue(
torch.equal(
m.train_targets,
Ys[i].view(1, -1).expand(4, Ys[i].numel()),
))
expected_noise = (dummy_samples["noise"].view(
m.likelihood.noise.shape) if inferred_noise else
Yvars[i].view(1, -1).expand(
4, Yvars[i].numel()))
self.assertTrue(
torch.allclose(
m.likelihood.noise.detach(),
expected_noise,
))
self.assertIsInstance(m.likelihood,
_GaussianLikelihoodBase)
self.assertTrue(
torch.allclose(
m.covar_module.base_kernel.lengthscale.detach(
),
dummy_samples["lengthscale"].view(
m.covar_module.base_kernel.lengthscale.
shape),
))
self.assertTrue(
torch.allclose(
m.covar_module.outputscale.detach(),
dummy_samples["outputscale"].view(
m.covar_module.outputscale.shape),
))
self.assertTrue(
torch.allclose(
m.mean_module.constant.detach(),
dummy_samples["mean"].view(
m.mean_module.constant.shape),
))
if use_input_warping:
self.assertTrue(hasattr(m, "input_transform"))
self.assertIsInstance(m.input_transform, Warp)
self.assertTrue(
torch.equal(
m.input_transform.concentration0,
dummy_samples_list[i]["c0"],
))
self.assertTrue(
torch.equal(
m.input_transform.concentration1,
dummy_samples_list[i]["c1"],
))
else:
self.assertFalse(hasattr(m, "input_transform"))
# test that multi-task is not implemented
(
Xs_mt,
Ys_mt,
Yvars_mt,
bounds_mt,
tfs_mt,
fns_mt,
mns_mt,
) = get_torch_test_data(dtype=dtype,
cuda=cuda,
constant_noise=True,
task_features=[2])
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model, self.assertRaises(NotImplementedError):
model.fit(
Xs=Xs_mt,
Ys=Ys_mt,
Yvars=Yvars_mt,
search_space_digest=SearchSpaceDigest(
feature_names=fns_mt,
bounds=bounds_mt,
task_features=tfs_mt,
),
metric_names=mns_mt,
)
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model, self.assertRaises(NotImplementedError):
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
fidelity_features=[0],
),
metric_names=mns,
)
# fit model with same inputs (otherwise X_observed will be None)
model = self.model_cls(
use_input_warping=use_input_warping,
thinning=1,
num_samples=4,
use_saas=use_saas,
disable_progbar=True,
max_tree_depth=1,
)
Yvars = Yvars1 + Yvars2
dummy_samples_list = _get_dummy_mcmc_samples(num_samples=4,
num_outputs=2,
**tkwargs)
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model:
model.fit(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars,
search_space_digest=SearchSpaceDigest(
feature_names=fns,
bounds=bounds,
task_features=tfs,
),
metric_names=mns,
)
# 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]], **tkwargs)
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] if self.model_cls is FullyBayesianBotorchModel
else [1.0, 1.0],
**tkwargs,
)
outcome_constraints = (
torch.tensor([[0.0, 1.0]], **tkwargs),
torch.tensor([[5.0]], **tkwargs),
)
gen_kwargs = ({
"objective_thresholds": torch.zeros(2, **tkwargs)
} if self.model_cls is FullyBayesianMOOBotorchModel else {})
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]], **tkwargs),
torch.tensor([[2.0, 6.0, 8.0]], **tkwargs),
]
n = 3
X_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
acqfv_dummy = torch.tensor([[[1.0, 2.0, 3.0]]], **tkwargs)
model_gen_options = {
Keys.OPTIMIZER_KWARGS: {
"maxiter": 1
},
Keys.ACQF_KWARGS: {
"mc_samples": 3
},
}
# test sequential optimize with constraints
with mock.patch(
"ax.models.torch.botorch_defaults.optimize_acqf",
return_value=(X_dummy, acqfv_dummy),
) as _:
Xgen, wgen, gen_metadata, cand_metadata = 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,
**gen_kwargs,
)
# note: gen() always returns CPU tensors
self.assertTrue(torch.equal(Xgen, X_dummy.cpu()))
self.assertTrue(
torch.equal(wgen, torch.ones(n, dtype=dtype)))
# actually test optimization for 1 step without constraints
with mock.patch(
"ax.models.torch.botorch_defaults.optimize_acqf",
wraps=optimize_acqf,
return_value=(X_dummy, acqfv_dummy),
) as _:
Xgen, wgen, gen_metadata, cand_metadata = model.gen(
n=n,
bounds=bounds,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
fixed_features=fixed_features,
pending_observations=pending_observations,
model_gen_options=model_gen_options,
**gen_kwargs,
)
# note: gen() always returns CPU tensors
self.assertTrue(torch.equal(Xgen, X_dummy.cpu()))
self.assertTrue(
torch.equal(wgen, torch.ones(n, dtype=dtype)))
# 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=Ys,
Yvars_train=Yvars,
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
with mock.patch(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list,
) as _mock_fit_model:
mean, variance = model.cross_validate(
Xs_train=Xs1 + Xs2,
Ys_train=Ys,
Yvars_train=Yvars,
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(
RUN_INFERENCE_PATH,
side_effect=dummy_samples_list) 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 = self.model_cls()
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()