def test_init(self):
# Default model with no specifications.
model = BoTorchModel()
self.assertEqual(model.acquisition_class, Acquisition)
# Model that specifies `botorch_acqf_class`.
model = BoTorchModel(botorch_acqf_class=qExpectedImprovement)
self.assertEqual(model.acquisition_class, Acquisition)
self.assertEqual(model.botorch_acqf_class, qExpectedImprovement)
# Model with `Acquisition` that specifies a `default_botorch_acqf_class`.
model = BoTorchModel(acquisition_class=KnowledgeGradient)
self.assertEqual(model.acquisition_class, KnowledgeGradient)
self.assertEqual(model.botorch_acqf_class, qKnowledgeGradient)
# Check defaults for refitting settings.
self.assertTrue(model.refit_on_update)
self.assertFalse(model.refit_on_cv)
self.assertTrue(model.warm_start_refit)
# Check setting non-default refitting settings
mdl2 = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=self.acquisition_class,
acquisition_options=self.acquisition_options,
refit_on_update=False,
refit_on_cv=True,
warm_start_refit=False,
)
self.assertFalse(mdl2.refit_on_update)
self.assertTrue(mdl2.refit_on_cv)
self.assertFalse(mdl2.warm_start_refit)
def test_init(self):
# Default model with no specifications.
model = BoTorchModel()
self.assertEqual(model.acquisition_class, Acquisition)
# Model that specifies `botorch_acqf_class`.
model = BoTorchModel(botorch_acqf_class=qExpectedImprovement)
self.assertEqual(model.acquisition_class, Acquisition)
self.assertEqual(model.botorch_acqf_class, qExpectedImprovement)
# Model with `Acquisition` that specifies a `default_botorch_acqf_class`.
model = BoTorchModel(acquisition_class=KnowledgeGradient)
self.assertEqual(model.acquisition_class, KnowledgeGradient)
self.assertEqual(model.botorch_acqf_class, qKnowledgeGradient)
def test_list_surrogate_choice(self, _, mock_extract_training_data):
model = BoTorchModel()
model.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
search_space_digest=self.search_space_digest,
metric_names=self.metric_names_for_list_surrogate,
candidate_metadata=self.candidate_metadata,
)
# A list surrogate should be chosen, since Xs are not all the same.
self.assertIsInstance(model.surrogate.model, ModelListGP)
for submodel in model.surrogate.model.models:
# There are fidelity features and nonempty Yvars, so
# fixed noise MFGP should be chosen.
self.assertIsInstance(submodel, FixedNoiseMultiFidelityGP)
model.gen(
n=1,
bounds=self.bounds,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
model_gen_options=self.model_gen_options,
rounding_func=self.rounding_func,
target_fidelities=self.search_space_digest.target_fidelities,
)
mock_extract_training_data.assert_called_once()
self.assertIsInstance(
mock_extract_training_data.call_args[1]["surrogate"],
ListSurrogate)
def test_evaluate_acquisition_function(self, _mock_kg,
_mock_construct_options):
model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=KnowledgeGradient,
acquisition_options=self.acquisition_options,
)
model.surrogate.construct(
training_data=self.training_data,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
fidelity_features=self.search_space_digest.fidelity_features,
),
)
model.evaluate_acquisition_function(
X=self.X,
search_space_digest=self.search_space_digest,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
acq_options=self.acquisition_options,
)
# `_mock_kg` is a mock of class, so to check the mock `evaluate` on
# instance of that class, we use `_mock_kg.return_value.evaluate`
_mock_kg.return_value.evaluate.assert_called()
def test_list_surrogate_choice(self, _, mock_extract_training_data):
model = BoTorchModel()
model.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
metric_names=self.metric_names_for_list_surrogate,
fidelity_features=self.fidelity_features,
target_fidelities=self.target_fidelities,
candidate_metadata=self.candidate_metadata,
)
model.gen(
n=1,
bounds=self.bounds,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
model_gen_options=self.model_gen_options,
rounding_func=self.rounding_func,
target_fidelities=self.target_fidelities,
)
mock_extract_training_data.assert_called_once()
self.assertIsInstance(
mock_extract_training_data.call_args[1]["surrogate"], ListSurrogate
)
def setUp(self):
self.botorch_model_class = SingleTaskGP
self.surrogate = Surrogate(
botorch_model_class=self.botorch_model_class)
self.acquisition_class = KnowledgeGradient
self.botorch_acqf_class = qKnowledgeGradient
self.acquisition_options = {Keys.NUM_FANTASIES: 64}
self.model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=self.acquisition_class,
acquisition_options=self.acquisition_options,
)
self.dtype = torch.float
Xs1, Ys1, Yvars1, self.bounds, _, _, _ = get_torch_test_data(
dtype=self.dtype)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(dtype=self.dtype,
offset=1.0)
self.Xs = Xs1 + Xs2
self.Ys = Ys1 + Ys2
self.Yvars = Yvars1 + Yvars2
self.X = Xs1[0]
self.Y = Ys1[0]
self.Yvar = Yvars1[0]
self.X2 = Xs2[0]
self.training_data = TrainingData(X=self.X, Y=self.Y, Yvar=self.Yvar)
self.search_space_digest = SearchSpaceDigest(
feature_names=["x1", "x2", "x3"],
bounds=[(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)],
task_features=[],
fidelity_features=[2],
target_fidelities={1: 1.0},
)
self.metric_names = ["y"]
self.metric_names_for_list_surrogate = ["y1", "y2"]
self.candidate_metadata = []
self.optimizer_options = {
Keys.NUM_RESTARTS: 40,
Keys.RAW_SAMPLES: 1024
}
self.model_gen_options = {
Keys.OPTIMIZER_KWARGS: self.optimizer_options
}
self.objective_weights = torch.tensor([1.0])
self.objective_thresholds = None
self.outcome_constraints = None
self.linear_constraints = None
self.fixed_features = None
self.pending_observations = None
self.rounding_func = "func"
def test_surrogate_options_propagation(self, _, mock_init):
model = BoTorchModel(surrogate_options={"some_option": "some_value"})
model.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
search_space_digest=self.search_space_digest,
metric_names=self.metric_names_for_list_surrogate,
candidate_metadata=self.candidate_metadata,
)
mock_init.assert_called_with(
botorch_submodel_class_per_outcome={
outcome: FixedNoiseMultiFidelityGP
for outcome in self.metric_names_for_list_surrogate
},
some_option="some_value",
)
def test_list_surrogate_choice(self, _): # , mock_extract_training_data):
model = BoTorchModel()
model.fit(
Xs=self.non_block_design_training_data.Xs,
Ys=self.non_block_design_training_data.Ys,
Yvars=self.non_block_design_training_data.Yvars,
search_space_digest=self.mf_search_space_digest,
metric_names=self.metric_names_for_list_surrogate,
candidate_metadata=self.candidate_metadata,
)
# A list surrogate should be chosen, since Xs are not all the same.
self.assertIsInstance(model.surrogate, ListSurrogate)
self.assertIsInstance(model.surrogate.model, ModelListGP)
for submodel in model.surrogate.model.models:
# There are fidelity features and nonempty Yvars, so
# fixed noise MFGP should be chosen.
self.assertIsInstance(submodel, FixedNoiseMultiFidelityGP)
def test_surrogate_options_propagation(self, _, mock_init):
model = BoTorchModel(surrogate_options={"some_option": "some_value"})
model.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
metric_names=self.metric_names_for_list_surrogate,
fidelity_features=self.fidelity_features,
target_fidelities=self.target_fidelities,
candidate_metadata=self.candidate_metadata,
)
mock_init.assert_called_with(
botorch_submodel_class_per_outcome={
outcome: FixedNoiseMultiFidelityGP
for outcome in self.metric_names_for_list_surrogate
},
some_option="some_value",
)
def get_botorch_model_with_default_acquisition_class() -> BoTorchModel:
return BoTorchModel(
surrogate=get_surrogate(),
acquisition_class=Acquisition,
botorch_acqf_class=get_acquisition_function_type(),
)
def get_botorch_model() -> BoTorchModel:
return BoTorchModel(
surrogate=get_surrogate(), acquisition_class=get_acquisition_type()
)
def test_gen(
self,
mock_choose_botorch_acqf_class,
mock_inequality_constraints,
mock_rounding,
mock_kg,
mock_construct_options,
):
mock_kg.return_value.optimize.return_value = (
torch.tensor([1.0]),
torch.tensor([2.0]),
)
model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=KnowledgeGradient,
acquisition_options=self.acquisition_options,
)
model.surrogate.construct(training_data=self.training_data,
fidelity_features=self.fidelity_features)
model._botorch_acqf_class = None
model.gen(
n=1,
bounds=self.bounds,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
model_gen_options=self.model_gen_options,
rounding_func=self.rounding_func,
target_fidelities=self.target_fidelities,
)
# Assert `construct_acquisition_and_optimizer_options` called with kwargs
mock_construct_options.assert_called_with(
acqf_options=self.acquisition_options,
model_gen_options=self.model_gen_options,
)
# Assert `choose_botorch_acqf_class` is called
mock_choose_botorch_acqf_class.assert_called_once()
self.assertEqual(model._botorch_acqf_class, qKnowledgeGradient)
# Assert `acquisition_class` called with kwargs
mock_kg.assert_called_with(
surrogate=self.surrogate,
botorch_acqf_class=model.botorch_acqf_class,
bounds=self.bounds,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
target_fidelities=self.target_fidelities,
options=self.acquisition_options,
)
# Assert `optimize` called with kwargs
mock_kg.return_value.optimize.assert_called_with(
bounds=ANY,
n=1,
inequality_constraints=[],
fixed_features=self.fixed_features,
rounding_func="func",
optimizer_options=self.optimizer_options,
)
def test_gen(
self,
mock_choose_botorch_acqf_class,
mock_inequality_constraints,
mock_rounding,
mock_kg,
mock_construct_options,
):
mock_kg.return_value.optimize.return_value = (
torch.tensor([1.0]),
torch.tensor([2.0]),
)
model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=KnowledgeGradient,
acquisition_options=self.acquisition_options,
)
model.surrogate.construct(
training_data=self.training_data,
fidelity_features=self.search_space_digest.fidelity_features,
)
model._botorch_acqf_class = None
# Assert that error is raised if we haven't fit the model
with self.assertRaises(RuntimeError):
model.gen(
n=1,
bounds=self.search_space_digest.bounds,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
model_gen_options=self.model_gen_options,
rounding_func=self.rounding_func,
target_fidelities=self.search_space_digest.target_fidelities,
)
# Add search space digest reference to make the model think it's been fit
model._search_space_digest = self.search_space_digest
model.gen(
n=1,
bounds=self.search_space_digest.bounds,
objective_weights=self.objective_weights,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
model_gen_options=self.model_gen_options,
rounding_func=self.rounding_func,
target_fidelities=self.search_space_digest.target_fidelities,
)
# Assert `construct_acquisition_and_optimizer_options` called with kwargs
mock_construct_options.assert_called_with(
acqf_options=self.acquisition_options,
model_gen_options=self.model_gen_options,
)
# Assert `choose_botorch_acqf_class` is called
mock_choose_botorch_acqf_class.assert_called_once()
self.assertEqual(model._botorch_acqf_class, qKnowledgeGradient)
# Assert `acquisition_class` called with kwargs
mock_kg.assert_called_with(
surrogate=self.surrogate,
botorch_acqf_class=model.botorch_acqf_class,
search_space_digest=self.search_space_digest,
objective_weights=self.objective_weights,
objective_thresholds=self.objective_thresholds,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
options=self.acquisition_options,
)
# Assert `optimize` called with kwargs
mock_kg.return_value.optimize.assert_called_with(
n=1,
search_space_digest=self.search_space_digest,
inequality_constraints=[],
fixed_features=self.fixed_features,
rounding_func="func",
optimizer_options=self.optimizer_options,
)
def test_MOO(self, _):
# Add mock for qNEHVI input constructor to catch arguments passed to it.
qNEHVI_input_constructor = get_acqf_input_constructor(
qNoisyExpectedHypervolumeImprovement)
mock_input_constructor = mock.MagicMock(
qNEHVI_input_constructor, side_effect=qNEHVI_input_constructor)
_register_acqf_input_constructor(
acqf_cls=qNoisyExpectedHypervolumeImprovement,
input_constructor=mock_input_constructor,
)
model = BoTorchModel()
model.fit(
Xs=self.moo_training_data.Xs,
Ys=self.moo_training_data.Ys,
Yvars=self.moo_training_data.Yvars,
search_space_digest=self.search_space_digest,
metric_names=self.moo_metric_names,
candidate_metadata=self.candidate_metadata,
)
self.assertIsInstance(model.surrogate.model, FixedNoiseGP)
_, _, gen_metadata, _ = model.gen(
n=1,
bounds=self.search_space_digest.bounds,
objective_weights=self.moo_objective_weights,
objective_thresholds=self.moo_objective_thresholds,
outcome_constraints=self.outcome_constraints,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
model_gen_options=self.model_gen_options,
rounding_func=self.rounding_func,
target_fidelities=self.mf_search_space_digest.target_fidelities,
)
ckwargs = mock_input_constructor.call_args[1]
self.assertIs(model.botorch_acqf_class,
qNoisyExpectedHypervolumeImprovement)
mock_input_constructor.assert_called_once()
m = ckwargs["model"]
self.assertIsInstance(m, FixedNoiseGP)
self.assertEqual(m.num_outputs, 2)
training_data = ckwargs["training_data"]
for attr in ("Xs", "Ys", "Yvars"):
self.assertTrue(
all(
torch.equal(x1, x2) for x1, x2 in zip(
getattr(training_data, attr),
getattr(self.moo_training_data, attr),
)))
self.assertTrue(
torch.equal(ckwargs["objective_thresholds"],
self.moo_objective_thresholds[:2]))
self.assertIsNone(ckwargs["outcome_constraints"], )
self.assertIsNone(ckwargs["X_pending"], )
obj_t = gen_metadata["objective_thresholds"]
self.assertTrue(
torch.equal(obj_t[:2], self.moo_objective_thresholds[:2]))
self.assertTrue(np.isnan(obj_t[2].item()))
self.assertIsInstance(
ckwargs.get("objective"),
WeightedMCMultiOutputObjective,
)
self.assertTrue(
torch.equal(
mock_input_constructor.call_args[1].get("objective").weights,
self.moo_objective_weights[:2],
))
expected_X_baseline = _filter_X_observed(
Xs=self.moo_training_data.Xs,
objective_weights=self.moo_objective_weights,
outcome_constraints=self.outcome_constraints,
bounds=self.search_space_digest.bounds,
linear_constraints=self.linear_constraints,
fixed_features=self.fixed_features,
)
self.assertTrue(
torch.equal(
mock_input_constructor.call_args[1].get("X_baseline"),
expected_X_baseline,
))
# test inferred objective_thresholds
with ExitStack() as es:
_mock_model_infer_objective_thresholds = es.enter_context(
mock.patch(
"ax.models.torch.botorch_modular.acquisition."
"infer_objective_thresholds",
return_value=torch.tensor([9.9, 3.3,
float("nan")]),
))
objective_weights = torch.tensor([-1.0, -1.0, 0.0])
outcome_constraints = (
torch.tensor([[1.0, 0.0, 0.0]]),
torch.tensor([[10.0]]),
)
linear_constraints = (
torch.tensor([[1.0, 0.0, 0.0]]),
torch.tensor([[2.0]]),
)
_, _, gen_metadata, _ = model.gen(
n=1,
bounds=self.search_space_digest.bounds,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
linear_constraints=linear_constraints,
fixed_features=self.fixed_features,
pending_observations=self.pending_observations,
model_gen_options=self.model_gen_options,
rounding_func=self.rounding_func,
target_fidelities=self.mf_search_space_digest.
target_fidelities,
)
expected_X_baseline = _filter_X_observed(
Xs=self.moo_training_data.Xs,
objective_weights=objective_weights,
outcome_constraints=outcome_constraints,
bounds=self.search_space_digest.bounds,
linear_constraints=linear_constraints,
fixed_features=self.fixed_features,
)
ckwargs = _mock_model_infer_objective_thresholds.call_args[1]
self.assertTrue(
torch.equal(
ckwargs["objective_weights"],
objective_weights,
))
oc = ckwargs["outcome_constraints"]
self.assertTrue(torch.equal(oc[0], outcome_constraints[0]))
self.assertTrue(torch.equal(oc[1], outcome_constraints[1]))
m = ckwargs["model"]
self.assertIsInstance(m, FixedNoiseGP)
self.assertEqual(m.num_outputs, 2)
self.assertIn("objective_thresholds", gen_metadata)
obj_t = gen_metadata["objective_thresholds"]
self.assertTrue(torch.equal(obj_t[:2], torch.tensor([9.9, 3.3])))
self.assertTrue(np.isnan(obj_t[2].item()))
# Avoid polluting the registry for other tests; re-register correct input
# contructor for qNEHVI.
_register_acqf_input_constructor(
acqf_cls=qNoisyExpectedHypervolumeImprovement,
input_constructor=qNEHVI_input_constructor,
)
def setUp(self):
self.botorch_model_class = SingleTaskGP
self.surrogate = Surrogate(
botorch_model_class=self.botorch_model_class)
self.acquisition_class = Acquisition
self.botorch_acqf_class = qExpectedImprovement
self.acquisition_options = ACQ_OPTIONS
self.model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=self.acquisition_class,
botorch_acqf_class=self.botorch_acqf_class,
acquisition_options=self.acquisition_options,
)
self.dtype = torch.float
self.device = torch.device("cpu")
tkwargs = {"dtype": self.dtype, "device": self.device}
Xs1, Ys1, Yvars1, self.bounds, _, _, _ = get_torch_test_data(
dtype=self.dtype)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(dtype=self.dtype,
offset=1.0)
self.Xs = Xs1
self.Ys = Ys1
self.Yvars = Yvars1
self.X_test = Xs2[0]
self.block_design_training_data = TrainingData(Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars)
self.non_block_design_training_data = TrainingData(
Xs=Xs1 + Xs2,
Ys=Ys1 + Ys2,
Yvars=Yvars1 + Yvars2,
)
self.search_space_digest = SearchSpaceDigest(
feature_names=["x1", "x2", "x3"],
bounds=[(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)],
)
self.mf_search_space_digest = SearchSpaceDigest(
feature_names=["x1", "x2", "x3"],
bounds=[(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)],
task_features=[],
fidelity_features=[2],
target_fidelities={1: 1.0},
)
self.metric_names = ["y"]
self.metric_names_for_list_surrogate = ["y1", "y2"]
self.candidate_metadata = []
self.optimizer_options = {
Keys.NUM_RESTARTS: 40,
Keys.RAW_SAMPLES: 1024
}
self.model_gen_options = {
Keys.OPTIMIZER_KWARGS: self.optimizer_options
}
self.objective_weights = torch.tensor([1.0], **tkwargs)
self.moo_objective_weights = torch.tensor([1.0, 1.5, 0.0], **tkwargs)
self.moo_objective_thresholds = torch.tensor(
[0.5, 1.5, float("nan")], **tkwargs)
self.outcome_constraints = None
self.linear_constraints = None
self.fixed_features = None
self.pending_observations = None
self.rounding_func = "func"
self.moo_training_data = TrainingData(
Xs=self.Xs * 3,
Ys=self.non_block_design_training_data.Ys + self.Ys,
Yvars=self.Yvars * 3,
)
self.moo_metric_names = ["y1", "y2", "y3"]