class BoTorchModelTest(TestCase):
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.surrogate_fit_options = {
Keys.REFIT_ON_UPDATE: True,
Keys.STATE_DICT: {"non-empty": "non-empty"},
Keys.WARM_START_REFITTING: False,
}
self.model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=self.acquisition_class,
acquisition_options=self.acquisition_options,
surrogate_fit_options=self.surrogate_fit_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.training_data = TrainingData(X=self.X, Y=self.Y, Yvar=self.Yvar)
self.bounds = [(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)]
self.task_features = []
self.feature_names = ["x1", "x2", "x3"]
self.metric_names = ["y"]
self.metric_names_for_list_surrogate = ["y1", "y2"]
self.fidelity_features = [2]
self.target_fidelities = {1: 1.0}
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.outcome_constraints = None
self.linear_constraints = None
self.fixed_features = None
self.pending_observations = None
self.rounding_func = "func"
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_surrogate_property(self):
self.assertEqual(self.surrogate, self.model.surrogate)
self.model._surrogate = None
with self.assertRaisesRegex(ValueError, "Surrogate has not yet been set."):
self.model.surrogate
def test_botorch_acqf_class_property(self):
self.assertEqual(self.botorch_acqf_class, self.model.botorch_acqf_class)
self.model._botorch_acqf_class = None
with self.assertRaisesRegex(
ValueError, "`AcquisitionFunction` has not yet been set."
):
self.model.botorch_acqf_class
@patch(f"{SURROGATE_PATH}.Surrogate.fit")
@patch(f"{MODEL_PATH}.choose_model_class", return_value=SingleTaskGP)
def test_fit(self, mock_choose_model_class, mock_fit):
# If surrogate is not yet set, initialize it with dispatcher functions.
self.model._surrogate = None
self.model.fit(
Xs=[self.X],
Ys=[self.Y],
Yvars=[self.Yvar],
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
metric_names=self.metric_names,
fidelity_features=self.fidelity_features,
target_fidelities=self.target_fidelities,
candidate_metadata=self.candidate_metadata,
)
# `choose_model_class` is called.
mock_choose_model_class.assert_called_with(
Yvars=[self.Yvar],
task_features=self.task_features,
fidelity_features=self.fidelity_features,
)
# Since we want to refit on updates but not warm start refit, we clear the
# state dict.
mock_fit.assert_called_with(
training_data=self.training_data,
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
fidelity_features=self.fidelity_features,
target_fidelities=self.target_fidelities,
metric_names=self.metric_names,
candidate_metadata=self.candidate_metadata,
state_dict=None,
refit=True,
)
@patch(f"{SURROGATE_PATH}.Surrogate.predict")
def test_predict(self, mock_predict):
self.model.predict(X=self.X)
mock_predict.assert_called_with(X=self.X)
@patch(
f"{MODEL_PATH}.construct_acquisition_and_optimizer_options",
return_value=({"num_fantasies": 64}, {"num_restarts": 40, "raw_samples": 1024}),
)
@patch(f"{CURRENT_PATH}.KnowledgeGradient")
@patch(f"{MODEL_PATH}.get_rounding_func", return_value="func")
@patch(f"{MODEL_PATH}._to_inequality_constraints", return_value=[])
@patch(f"{MODEL_PATH}.choose_botorch_acqf_class", return_value=qKnowledgeGradient)
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_best_point(self):
with self.assertRaises(NotImplementedError):
self.model.best_point(
bounds=self.bounds, objective_weights=self.objective_weights
)
@patch(
f"{MODEL_PATH}.construct_acquisition_and_optimizer_options",
return_value=({"num_fantasies": 64}, {"num_restarts": 40, "raw_samples": 1024}),
)
@patch(f"{CURRENT_PATH}.KnowledgeGradient", autospec=True)
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, fidelity_features=self.fidelity_features
)
model.evaluate_acquisition_function(
X=self.X,
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,
acq_options=self.acquisition_options,
target_fidelities=self.target_fidelities,
)
# `_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()
@patch(
f"{ACQUISITION_PATH}.Acquisition._extract_training_data",
# Mock to register calls, but still execute the function.
side_effect=Acquisition._extract_training_data,
)
@patch(
f"{ACQUISITION_PATH}.Acquisition.optimize",
# Dummy candidates and acquisition function value.
return_value=(torch.tensor([[2.0]]), torch.tensor([1.0])),
)
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
)
class BoTorchModelTest(TestCase):
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.bounds = [(0.0, 10.0), (0.0, 10.0), (0.0, 10.0)]
self.task_features = []
self.feature_names = ["x1", "x2", "x3"]
self.metric_names = ["y"]
self.metric_names_for_list_surrogate = ["y1", "y2"]
self.fidelity_features = [2]
self.target_fidelities = {1: 1.0}
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.outcome_constraints = None
self.linear_constraints = None
self.fixed_features = None
self.pending_observations = None
self.rounding_func = "func"
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_surrogate_property(self):
self.assertEqual(self.surrogate, self.model.surrogate)
self.model._surrogate = None
with self.assertRaisesRegex(ValueError,
"Surrogate has not yet been set."):
self.model.surrogate
def test_botorch_acqf_class_property(self):
self.assertEqual(self.botorch_acqf_class,
self.model.botorch_acqf_class)
self.model._botorch_acqf_class = None
with self.assertRaisesRegex(
ValueError, "`AcquisitionFunction` has not yet been set."):
self.model.botorch_acqf_class
@patch(f"{SURROGATE_PATH}.Surrogate.fit")
@patch(f"{MODEL_PATH}.choose_model_class", return_value=SingleTaskGP)
def test_fit(self, mock_choose_model_class, mock_fit):
# If surrogate is not yet set, initialize it with dispatcher functions.
self.model._surrogate = None
self.model.fit(
Xs=[self.X],
Ys=[self.Y],
Yvars=[self.Yvar],
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
metric_names=self.metric_names,
fidelity_features=self.fidelity_features,
target_fidelities=self.target_fidelities,
candidate_metadata=self.candidate_metadata,
)
# `choose_model_class` is called.
mock_choose_model_class.assert_called_with(
Yvars=[self.Yvar],
task_features=self.task_features,
fidelity_features=self.fidelity_features,
)
# Since we want to refit on updates but not warm start refit, we clear the
# state dict.
mock_fit.assert_called_with(
training_data=self.training_data,
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
fidelity_features=self.fidelity_features,
target_fidelities=self.target_fidelities,
metric_names=self.metric_names,
candidate_metadata=self.candidate_metadata,
state_dict=None,
refit=True,
)
@patch(f"{SURROGATE_PATH}.Surrogate.update")
def test_update(self, mock_update):
fit_update_shared_kwargs = {
"bounds": self.bounds,
"task_features": self.task_features,
"feature_names": self.feature_names,
"metric_names": self.metric_names,
"fidelity_features": self.fidelity_features,
"target_fidelities": self.target_fidelities,
"candidate_metadata": self.candidate_metadata,
}
self.model.fit(Xs=[self.X],
Ys=[self.Y],
Yvars=[self.Yvar],
**fit_update_shared_kwargs)
for refit_on_update, warm_start_refit in [
(True, True),
(True, False),
(False, True),
]:
self.model.refit_on_update = refit_on_update
self.model.warm_start_refit = warm_start_refit
self.model.update(Xs=[self.X],
Ys=[self.Y],
Yvars=[self.Yvar],
**fit_update_shared_kwargs)
expected_state_dict = (None if refit_on_update
and not warm_start_refit else
self.model.surrogate.model.state_dict())
# Check correct training data and `fit_update_shared_kwargs` values (can't
# directly use `assert_called_with` due to tensor comparison ambiguity in
# state dict).
self.assertEqual(
mock_update.call_args_list[-1][1].get("training_data"),
self.training_data,
)
for key in fit_update_shared_kwargs:
self.assertEqual(
mock_update.call_args_list[-1][1].get(key),
fit_update_shared_kwargs.get(key),
)
# Check correct `refit` and `state_dict` values.
self.assertEqual(mock_update.call_args_list[-1][1].get("refit"),
refit_on_update)
if expected_state_dict is None:
self.assertIsNone(
mock_update.call_args_list[-1][1].get("state_dict"),
expected_state_dict,
)
else:
self.assertEqual(
mock_update.call_args_list[-1][1].get("state_dict").keys(),
expected_state_dict.keys(),
)
@patch(f"{SURROGATE_PATH}.Surrogate.predict")
def test_predict(self, mock_predict):
self.model.predict(X=self.X)
mock_predict.assert_called_with(X=self.X)
@patch(f"{MODEL_PATH}.BoTorchModel.fit")
def test_cross_validate(self, mock_fit):
fit_cv_shared_kwargs = {
"bounds": self.bounds,
"task_features": self.task_features,
"feature_names": self.feature_names,
"metric_names": self.metric_names,
"fidelity_features": self.fidelity_features,
}
self.model.fit(
Xs=[self.X],
Ys=[self.Y],
Yvars=[self.Yvar],
candidate_metadata=self.candidate_metadata,
**fit_cv_shared_kwargs,
)
old_surrogate = self.model.surrogate
old_surrogate._model = MagicMock()
old_surrogate._model.state_dict.return_value = {"key": "val"}
for refit_on_cv, warm_start_refit in [
(True, True),
(True, False),
(False, True),
]:
self.model.refit_on_cv = refit_on_cv
self.model.warm_start_refit = warm_start_refit
with patch(
f"{SURROGATE_PATH}.Surrogate.clone_reset",
return_value=MagicMock(spec=Surrogate),
) as mock_clone_reset:
self.model.cross_validate(
Xs_train=[self.X],
Ys_train=[self.Y],
Yvars_train=[self.Yvar],
X_test=self.X2,
**fit_cv_shared_kwargs,
)
# Check that `predict` is called on the cloned surrogate, not
# on the original one.
mock_predict = mock_clone_reset.return_value.predict
mock_predict.assert_called_once()
# Check correct X_test.
self.assertTrue(
torch.equal(
mock_predict.call_args_list[-1][1].get("X"),
self.X2,
), )
# Check that surrogate is reset back to `old_surrogate` at the
# end of cross-validation.
self.model.surrogate is old_surrogate
expected_state_dict = (None
if refit_on_cv and not warm_start_refit else
self.model.surrogate.model.state_dict())
# Check correct `refit` and `state_dict` values.
self.assertEqual(mock_fit.call_args_list[-1][1].get("refit"),
refit_on_cv)
if expected_state_dict is None:
self.assertIsNone(
mock_fit.call_args_list[-1][1].get("state_dict"),
expected_state_dict,
)
else:
self.assertEqual(
mock_fit.call_args_list[-1][1].get("state_dict").keys(),
expected_state_dict.keys(),
)
@patch(
f"{MODEL_PATH}.construct_acquisition_and_optimizer_options",
return_value=({
"num_fantasies": 64
}, {
"num_restarts": 40,
"raw_samples": 1024
}),
)
@patch(f"{CURRENT_PATH}.KnowledgeGradient")
@patch(f"{MODEL_PATH}.get_rounding_func", return_value="func")
@patch(f"{MODEL_PATH}._to_inequality_constraints", return_value=[])
@patch(f"{MODEL_PATH}.choose_botorch_acqf_class",
return_value=qKnowledgeGradient)
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_best_point(self):
with self.assertRaises(NotImplementedError):
self.model.best_point(bounds=self.bounds,
objective_weights=self.objective_weights)
@patch(
f"{MODEL_PATH}.construct_acquisition_and_optimizer_options",
return_value=({
"num_fantasies": 64
}, {
"num_restarts": 40,
"raw_samples": 1024
}),
)
@patch(f"{CURRENT_PATH}.KnowledgeGradient", autospec=True)
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,
fidelity_features=self.fidelity_features)
model.evaluate_acquisition_function(
X=self.X,
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,
acq_options=self.acquisition_options,
target_fidelities=self.target_fidelities,
)
# `_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()
@patch(f"{LIST_SURROGATE_PATH}.ListSurrogate.__init__", return_value=None)
@patch(f"{LIST_SURROGATE_PATH}.ListSurrogate.fit", return_value=None)
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",
)
@patch(
f"{ACQUISITION_PATH}.Acquisition._extract_training_data",
# Mock to register calls, but still execute the function.
side_effect=Acquisition._extract_training_data,
)
@patch(
f"{ACQUISITION_PATH}.Acquisition.optimize",
# Dummy candidates and acquisition function value.
return_value=(torch.tensor([[2.0]]), torch.tensor([1.0])),
)
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,
)
# 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.target_fidelities,
)
mock_extract_training_data.assert_called_once()
self.assertIsInstance(
mock_extract_training_data.call_args[1]["surrogate"],
ListSurrogate)
class BoTorchModelTest(TestCase):
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"]
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)
# 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_surrogate_property(self):
self.assertEqual(self.surrogate, self.model.surrogate)
self.model._surrogate = None
with self.assertRaisesRegex(ValueError,
"Surrogate has not yet been set."):
self.model.surrogate
def test_botorch_acqf_class_property(self):
self.assertEqual(self.botorch_acqf_class,
self.model.botorch_acqf_class)
self.model._botorch_acqf_class = None
with self.assertRaisesRegex(
ValueError,
"BoTorch `AcquisitionFunction` has not yet been set."):
self.model.botorch_acqf_class
@mock.patch(f"{SURROGATE_PATH}.Surrogate.fit")
@mock.patch(f"{MODEL_PATH}.choose_model_class", return_value=SingleTaskGP)
def test_fit(self, mock_choose_model_class, mock_fit):
# If surrogate is not yet set, initialize it with dispatcher functions.
self.model._surrogate = None
self.model.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
search_space_digest=self.mf_search_space_digest,
metric_names=self.metric_names,
candidate_metadata=self.candidate_metadata,
)
# `choose_model_class` is called.
mock_choose_model_class.assert_called_with(
Yvars=self.Yvars,
search_space_digest=self.mf_search_space_digest,
)
# Since we want to refit on updates but not warm start refit, we clear the
# state dict.
mock_fit.assert_called_with(
training_data=self.block_design_training_data,
search_space_digest=self.mf_search_space_digest,
metric_names=self.metric_names,
candidate_metadata=self.candidate_metadata,
state_dict=None,
refit=True,
)
@mock.patch(f"{SURROGATE_PATH}.Surrogate.update")
def test_update(self, mock_update):
self.model.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
search_space_digest=self.mf_search_space_digest,
metric_names=self.metric_names,
candidate_metadata=self.candidate_metadata,
)
for refit_on_update, warm_start_refit in [
(True, True),
(True, False),
(False, True),
]:
self.model.refit_on_update = refit_on_update
self.model.warm_start_refit = warm_start_refit
self.model.update(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
search_space_digest=self.mf_search_space_digest,
metric_names=self.metric_names,
candidate_metadata=self.candidate_metadata,
)
expected_state_dict = (None if refit_on_update
and not warm_start_refit else
self.model.surrogate.model.state_dict())
# Check for correct call args
call_args = mock_update.call_args_list[-1][1]
self.assertEqual(call_args.get("training_data"),
self.block_design_training_data)
self.assertEqual(call_args.get("search_space_digest"),
self.mf_search_space_digest)
self.assertEqual(call_args.get("metric_names"), self.metric_names)
self.assertEqual(call_args.get("candidate_metadata"),
self.candidate_metadata)
# Check correct `refit` and `state_dict` values.
self.assertEqual(mock_update.call_args_list[-1][1].get("refit"),
refit_on_update)
if expected_state_dict is None:
self.assertIsNone(
mock_update.call_args_list[-1][1].get("state_dict"),
expected_state_dict,
)
else:
self.assertEqual(
mock_update.call_args_list[-1][1].get("state_dict").keys(),
expected_state_dict.keys(),
)
@mock.patch(f"{SURROGATE_PATH}.Surrogate.predict")
def test_predict(self, mock_predict):
self.model.predict(X=self.X_test)
mock_predict.assert_called_with(X=self.X_test)
@mock.patch(f"{MODEL_PATH}.BoTorchModel.fit")
def test_cross_validate(self, mock_fit):
self.model.fit(
Xs=self.Xs,
Ys=self.Ys,
Yvars=self.Yvars,
search_space_digest=self.mf_search_space_digest,
candidate_metadata=self.candidate_metadata,
metric_names=self.metric_names,
)
old_surrogate = self.model.surrogate
old_surrogate._model = mock.MagicMock()
old_surrogate._model.state_dict.return_value = {"key": "val"}
for refit_on_cv, warm_start_refit in [
(True, True),
(True, False),
(False, True),
]:
self.model.refit_on_cv = refit_on_cv
self.model.warm_start_refit = warm_start_refit
with mock.patch(
f"{SURROGATE_PATH}.Surrogate.clone_reset",
return_value=mock.MagicMock(spec=Surrogate),
) as mock_clone_reset:
self.model.cross_validate(
Xs_train=self.Xs,
Ys_train=self.Ys,
Yvars_train=self.Yvars,
X_test=self.X_test,
search_space_digest=self.mf_search_space_digest,
metric_names=self.metric_names,
)
# Check that `predict` is called on the cloned surrogate, not
# on the original one.
mock_predict = mock_clone_reset.return_value.predict
mock_predict.assert_called_once()
# Check correct X_test.
self.assertTrue(
torch.equal(
mock_predict.call_args_list[-1][1].get("X"),
self.X_test,
), )
# Check that surrogate is reset back to `old_surrogate` at the
# end of cross-validation.
self.assertTrue(self.model.surrogate is old_surrogate)
expected_state_dict = (None
if refit_on_cv and not warm_start_refit else
self.model.surrogate.model.state_dict())
# Check correct `refit` and `state_dict` values.
self.assertEqual(mock_fit.call_args_list[-1][1].get("refit"),
refit_on_cv)
if expected_state_dict is None:
self.assertIsNone(
mock_fit.call_args_list[-1][1].get("state_dict"),
expected_state_dict,
)
else:
self.assertEqual(
mock_fit.call_args_list[-1][1].get("state_dict").keys(),
expected_state_dict.keys(),
)
@mock.patch(
f"{MODEL_PATH}.construct_acquisition_and_optimizer_options",
return_value=(
ACQ_OPTIONS,
{
"num_restarts": 40,
"raw_samples": 1024
},
),
)
@mock.patch(f"{CURRENT_PATH}.Acquisition")
@mock.patch(f"{MODEL_PATH}.get_rounding_func", return_value="func")
@mock.patch(f"{MODEL_PATH}._to_inequality_constraints", return_value=[])
@mock.patch(f"{MODEL_PATH}.choose_botorch_acqf_class",
return_value=qExpectedImprovement)
def test_gen(
self,
mock_choose_botorch_acqf_class,
mock_inequality_constraints,
mock_rounding,
mock_acquisition,
mock_construct_options,
):
mock_acquisition.return_value.optimize.return_value = (
torch.tensor([1.0]),
torch.tensor([2.0]),
)
model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=Acquisition,
acquisition_options=self.acquisition_options,
)
model.surrogate.construct(
training_data=self.block_design_training_data,
fidelity_features=self.mf_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.mf_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.mf_search_space_digest.
target_fidelities,
)
# Add search space digest reference to make the model think it's been fit
model._search_space_digest = self.mf_search_space_digest
model.gen(
n=1,
bounds=self.mf_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.mf_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, qExpectedImprovement)
# Assert `acquisition_class` called with kwargs
mock_acquisition.assert_called_with(
surrogate=self.surrogate,
botorch_acqf_class=model.botorch_acqf_class,
search_space_digest=self.mf_search_space_digest,
objective_weights=self.objective_weights,
objective_thresholds=None,
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_acquisition.return_value.optimize.assert_called_with(
n=1,
search_space_digest=self.mf_search_space_digest,
inequality_constraints=[],
fixed_features=self.fixed_features,
rounding_func="func",
optimizer_options=self.optimizer_options,
)
def test_best_point(self):
with self.assertRaises(NotImplementedError):
self.model.best_point(bounds=self.bounds,
objective_weights=self.objective_weights)
@mock.patch(
f"{MODEL_PATH}.construct_acquisition_and_optimizer_options",
return_value=({
"num_fantasies": 64
}, {
"num_restarts": 40,
"raw_samples": 1024
}),
)
@mock.patch(f"{CURRENT_PATH}.Acquisition", autospec=True)
def test_evaluate_acquisition_function(self, mock_ei,
_mock_construct_options):
model = BoTorchModel(
surrogate=self.surrogate,
acquisition_class=Acquisition,
acquisition_options=self.acquisition_options,
)
model.surrogate.construct(
training_data=self.block_design_training_data,
search_space_digest=SearchSpaceDigest(
feature_names=[],
bounds=[],
fidelity_features=self.mf_search_space_digest.
fidelity_features,
),
)
model.evaluate_acquisition_function(
X=self.X_test,
search_space_digest=self.mf_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_ei` is a mock of class, so to check the mock `evaluate` on
# instance of that class, we use `mock_ei.return_value.evaluate`
mock_ei.return_value.evaluate.assert_called()
@mock.patch(f"{LIST_SURROGATE_PATH}.ListSurrogate.__init__",
return_value=None)
@mock.patch(f"{LIST_SURROGATE_PATH}.ListSurrogate.fit", return_value=None)
def test_surrogate_options_propagation(self, _, mock_init):
model = BoTorchModel(surrogate_options={"some_option": "some_value"})
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,
)
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",
)
@mock.patch(
f"{ACQUISITION_PATH}.Acquisition.optimize",
# Dummy candidates and acquisition function value.
return_value=(torch.tensor([[2.0]]), torch.tensor([1.0])),
)
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)
@mock.patch(
f"{ACQUISITION_PATH}.Acquisition.optimize",
# Dummy candidates and acquisition function value.
return_value=(torch.tensor([[2.0]]), torch.tensor([1.0])),
)
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,
)