def test_extract_training_data(self):
self.assertEqual( # Base `Surrogate` case.
self.acquisition._extract_training_data(surrogate=self.surrogate),
self.training_data,
)
# `ListSurrogate` case.
list_surrogate = ListSurrogate(
botorch_submodel_class=self.botorch_model_class)
list_surrogate._training_data_per_outcome = {"a": self.training_data}
self.assertEqual(
self.acquisition._extract_training_data(surrogate=list_surrogate),
list_surrogate._training_data_per_outcome,
)
def setUp(self):
self.outcomes = ["outcome_1", "outcome_2"]
self.mll_class = SumMarginalLogLikelihood
self.dtype = torch.float
self.search_space_digest = SearchSpaceDigest(feature_names=[],
bounds=[],
task_features=[0])
self.task_features = [0]
Xs1, Ys1, Yvars1, bounds, _, _, _ = get_torch_test_data(
dtype=self.dtype,
task_features=self.search_space_digest.task_features)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(
dtype=self.dtype,
task_features=self.search_space_digest.task_features)
self.botorch_submodel_class_per_outcome = {
self.outcomes[0]:
choose_model_class(Yvars=Yvars1,
search_space_digest=self.search_space_digest),
self.outcomes[1]:
choose_model_class(Yvars=Yvars2,
search_space_digest=self.search_space_digest),
}
self.expected_submodel_type = FixedNoiseMultiTaskGP
for submodel_cls in self.botorch_submodel_class_per_outcome.values():
self.assertEqual(submodel_cls, FixedNoiseMultiTaskGP)
self.Xs = Xs1 + Xs2
self.Ys = Ys1 + Ys2
self.Yvars = Yvars1 + Yvars2
self.training_data = [
TrainingData(X=X, Y=Y, Yvar=Yvar)
for X, Y, Yvar in zip(self.Xs, self.Ys, self.Yvars)
]
self.submodel_options_per_outcome = {
self.outcomes[0]: {
RANK: 1
},
self.outcomes[1]: {
RANK: 2
},
}
self.surrogate = ListSurrogate(
botorch_submodel_class_per_outcome=self.
botorch_submodel_class_per_outcome,
mll_class=self.mll_class,
submodel_options_per_outcome=self.submodel_options_per_outcome,
)
self.bounds = [(0.0, 1.0), (1.0, 4.0)]
self.feature_names = ["x1", "x2"]
def get_list_surrogate() -> Surrogate:
return ListSurrogate(
botorch_submodel_class_per_outcome={"m": get_model_type()},
submodel_options_per_outcome={"m": {"some_option": "some_value"}},
submodel_options={"shared_option": "shared_option_value"},
mll_class=get_mll_type(),
)
def test_construct_per_outcome_options_no_Yvar(self, _):
surrogate = ListSurrogate(
botorch_submodel_class=MultiTaskGP,
mll_class=self.mll_class,
submodel_options_per_outcome=self.submodel_options_per_outcome,
)
# Test that splitting the training data works correctly when Yvar is None.
training_data_no_Yvar = TrainingData(Xs=self.Xs, Ys=self.Ys)
surrogate.construct(
training_data=training_data_no_Yvar,
task_features=self.task_features,
metric_names=self.outcomes,
)
self.assertTrue(
all(
trd.Yvar is None for trd in surrogate.training_data_per_outcome.values()
)
)
self.assertEqual(len(surrogate.training_data_per_outcome), 2)
def _autoset_surrogate(
self,
Xs: List[Tensor],
Ys: List[Tensor],
Yvars: List[Tensor],
task_features: List[int],
fidelity_features: List[int],
metric_names: List[str],
) -> None:
"""Sets a default surrogate on this model if one was not explicitly
provided.
"""
# To determine whether to use `ListSurrogate`, we need to check for
# the batched multi-output case, so we first see which model would
# be chosen given the Yvars and the properties of data.
botorch_model_class = choose_model_class(
Yvars=Yvars,
task_features=task_features,
fidelity_features=fidelity_features,
)
if use_model_list(Xs=Xs, botorch_model_class=botorch_model_class):
# If using `ListSurrogate` / `ModelListGP`, pick submodels for each
# outcome.
botorch_submodel_class_per_outcome = {
metric_name: choose_model_class(
Yvars=[Yvar],
task_features=task_features,
fidelity_features=fidelity_features,
)
for Yvar, metric_name in zip(Yvars, metric_names)
}
self._surrogate = ListSurrogate(
botorch_submodel_class_per_outcome=
botorch_submodel_class_per_outcome,
**self.surrogate_options,
)
else:
# Using regular `Surrogate`, so botorch model picked at the beginning
# of the function is the one we should use.
self._surrogate = Surrogate(
botorch_model_class=botorch_model_class,
**self.surrogate_options)
def test_fit(self, mock_fit_gpytorch, mock_MLL, mock_state_dict):
surrogate = ListSurrogate(
botorch_submodel_class_per_outcome=self.
botorch_submodel_class_per_outcome,
mll_class=SumMarginalLogLikelihood,
)
# Checking that model is None before `fit` (and `construct`) calls.
self.assertIsNone(surrogate._model)
# Should instantiate mll and `fit_gpytorch_model` when `state_dict`
# is `None`.
surrogate.fit(
training_data=self.training_data,
search_space_digest=SearchSpaceDigest(
feature_names=self.feature_names,
bounds=self.bounds,
task_features=self.task_features,
),
metric_names=self.outcomes,
)
mock_state_dict.assert_not_called()
mock_MLL.assert_called_once()
mock_fit_gpytorch.assert_called_once()
mock_state_dict.reset_mock()
mock_MLL.reset_mock()
mock_fit_gpytorch.reset_mock()
# Should `load_state_dict` when `state_dict` is not `None`
# and `refit` is `False`.
state_dict = {"state_attribute": "value"}
surrogate.fit(
training_data=self.training_data,
search_space_digest=SearchSpaceDigest(
feature_names=self.feature_names,
bounds=self.bounds,
task_features=self.task_features,
),
metric_names=self.outcomes,
refit=False,
state_dict=state_dict,
)
mock_state_dict.assert_called_once()
mock_MLL.assert_not_called()
mock_fit_gpytorch.assert_not_called()
class ListSurrogateTest(TestCase):
def setUp(self):
self.outcomes = ["outcome_1", "outcome_2"]
self.mll_class = SumMarginalLogLikelihood
self.dtype = torch.float
self.task_features = [0]
Xs1, Ys1, Yvars1, bounds, _, _, _ = get_torch_test_data(
dtype=self.dtype, task_features=self.task_features)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(
dtype=self.dtype, task_features=self.task_features)
self.botorch_submodel_class_per_outcome = {
self.outcomes[0]:
choose_model_class(Yvars=Yvars1,
task_features=self.task_features,
fidelity_features=[]),
self.outcomes[1]:
choose_model_class(Yvars=Yvars2,
task_features=self.task_features,
fidelity_features=[]),
}
self.expected_submodel_type = FixedNoiseMultiTaskGP
for submodel_cls in self.botorch_submodel_class_per_outcome.values():
self.assertEqual(submodel_cls, FixedNoiseMultiTaskGP)
self.Xs = Xs1 + Xs2
self.Ys = Ys1 + Ys2
self.Yvars = Yvars1 + Yvars2
self.training_data = [
TrainingData(X=X, Y=Y, Yvar=Yvar)
for X, Y, Yvar in zip(self.Xs, self.Ys, self.Yvars)
]
self.submodel_options_per_outcome = {
self.outcomes[0]: {
RANK: 1
},
self.outcomes[1]: {
RANK: 2
},
}
self.surrogate = ListSurrogate(
botorch_submodel_class_per_outcome=self.
botorch_submodel_class_per_outcome,
mll_class=self.mll_class,
submodel_options_per_outcome=self.submodel_options_per_outcome,
)
self.bounds = [(0.0, 1.0), (1.0, 4.0)]
self.feature_names = ["x1", "x2"]
def check_ranks(self, c: ListSurrogate) -> type(None):
self.assertIsInstance(c, ListSurrogate)
self.assertIsInstance(c.model, ModelListGP)
for idx, submodel in enumerate(c.model.models):
self.assertIsInstance(submodel, self.expected_submodel_type)
self.assertEqual(
submodel._rank,
self.submodel_options_per_outcome[self.outcomes[idx]][RANK],
)
def test_init(self):
self.assertEqual(
self.surrogate.botorch_submodel_class_per_outcome,
self.botorch_submodel_class_per_outcome,
)
self.assertEqual(self.surrogate.mll_class, self.mll_class)
with self.assertRaises(NotImplementedError):
self.surrogate.training_data
with self.assertRaisesRegex(ValueError, NOT_YET_FIT_MSG):
self.surrogate.training_data_per_outcome
with self.assertRaisesRegex(
ValueError, "BoTorch `Model` has not yet been constructed"):
self.surrogate.model
@patch(
f"{CURRENT_PATH}.FixedNoiseMultiTaskGP.construct_inputs",
# Mock to register calls, but still execute the function.
side_effect=FixedNoiseMultiTaskGP.construct_inputs,
)
def test_construct_per_outcome_options(self, mock_MTGP_construct_inputs):
with self.assertRaisesRegex(ValueError, ".* are required"):
self.surrogate.construct(training_data=self.training_data)
with self.assertRaisesRegex(ValueError,
"No model class specified for"):
self.surrogate.construct(training_data=self.training_data,
metric_names=["new_metric"])
self.surrogate.construct(
training_data=self.training_data,
task_features=self.task_features,
metric_names=self.outcomes,
)
self.check_ranks(self.surrogate)
# Should construct inputs for MTGP twice.
self.assertEqual(len(mock_MTGP_construct_inputs.call_args_list), 2)
# First construct inputs should be called for MTGP with training data #0.
for idx in range(len(mock_MTGP_construct_inputs.call_args_list)):
self.assertEqual(
# `call_args` is a tuple of (args, kwargs), and we check kwargs.
mock_MTGP_construct_inputs.call_args_list[idx][1],
{
"fidelity_features": [],
"task_features":
self.task_features,
"training_data":
self.training_data[idx],
"rank":
self.submodel_options_per_outcome[self.outcomes[idx]]
["rank"],
},
)
@patch(
f"{CURRENT_PATH}.FixedNoiseMultiTaskGP.construct_inputs",
# Mock to register calls, but still execute the function.
side_effect=FixedNoiseMultiTaskGP.construct_inputs,
)
def test_construct_shared_shortcut_options(self, mock_construct_inputs):
surrogate = ListSurrogate(
botorch_submodel_class=self.botorch_submodel_class_per_outcome[
self.outcomes[0]],
submodel_options={"shared_option": True},
submodel_options_per_outcome={
outcome: {
"individual_option": f"val_{idx}"
}
for idx, outcome in enumerate(self.outcomes)
},
)
surrogate.construct(
training_data=self.training_data,
task_features=self.task_features,
metric_names=self.outcomes,
)
# 2 submodels should've been constructed, both of type `botorch_submodel_class`.
self.assertEqual(len(mock_construct_inputs.call_args_list), 2)
first_call_args, second_call_args = mock_construct_inputs.call_args_list
for idx in range(len(mock_construct_inputs.call_args_list)):
self.assertEqual(
mock_construct_inputs.call_args_list[idx][1],
{
"fidelity_features": [],
"individual_option": f"val_{idx}",
"shared_option": True,
"task_features": [0],
"training_data": self.training_data[idx],
},
)
@patch(f"{CURRENT_PATH}.ModelListGP.load_state_dict", return_value=None)
@patch(f"{CURRENT_PATH}.SumMarginalLogLikelihood")
@patch(f"{SURROGATE_PATH}.fit_gpytorch_model")
def test_fit(self, mock_fit_gpytorch, mock_MLL, mock_state_dict):
surrogate = ListSurrogate(
botorch_submodel_class_per_outcome=self.
botorch_submodel_class_per_outcome,
mll_class=SumMarginalLogLikelihood,
)
# Checking that model is None before `fit` (and `construct`) calls.
self.assertIsNone(surrogate._model)
# Should instantiate mll and `fit_gpytorch_model` when `state_dict`
# is `None`.
surrogate.fit(
training_data=self.training_data,
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
metric_names=self.outcomes,
fidelity_features=[],
)
mock_state_dict.assert_not_called()
mock_MLL.assert_called_once()
mock_fit_gpytorch.assert_called_once()
mock_state_dict.reset_mock()
mock_MLL.reset_mock()
mock_fit_gpytorch.reset_mock()
# Should `load_state_dict` when `state_dict` is not `None`
# and `refit` is `False`.
state_dict = {}
surrogate.fit(
training_data=self.training_data,
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
metric_names=self.outcomes,
fidelity_features=[],
refit=False,
state_dict=state_dict,
)
mock_state_dict.assert_called_once()
mock_MLL.assert_not_called()
mock_fit_gpytorch.assert_not_called()
def test_serialize_attributes_as_kwargs(self):
expected = self.surrogate.__dict__
# The two attributes below don't need to be saved as part of state,
# so we remove them from the expected dict.
expected.pop("botorch_model_class")
expected.pop("model_options")
self.assertEqual(self.surrogate._serialize_attributes_as_kwargs(),
expected)
class ListSurrogateTest(TestCase):
def setUp(self):
self.outcomes = ["outcome_1", "outcome_2"]
self.mll_class = SumMarginalLogLikelihood
self.dtype = torch.float
self.task_features = [0]
Xs1, Ys1, Yvars1, bounds, _, _, _ = get_torch_test_data(
dtype=self.dtype, task_features=self.task_features
)
Xs2, Ys2, Yvars2, _, _, _, _ = get_torch_test_data(
dtype=self.dtype, task_features=self.task_features
)
self.botorch_model_class_per_outcome = {
self.outcomes[0]: choose_model_class(
Yvars=Yvars1, task_features=self.task_features, fidelity_features=[]
),
self.outcomes[1]: choose_model_class(
Yvars=Yvars2, task_features=self.task_features, fidelity_features=[]
),
}
self.expected_submodel_type = FixedNoiseMultiTaskGP
for submodel_cls in self.botorch_model_class_per_outcome.values():
self.assertEqual(submodel_cls, FixedNoiseMultiTaskGP)
self.Xs = Xs1 + Xs2
self.Ys = Ys1 + Ys2
self.Yvars = Yvars1 + Yvars2
self.training_data = [
TrainingData(X=X, Y=Y, Yvar=Yvar)
for X, Y, Yvar in zip(self.Xs, self.Ys, self.Yvars)
]
self.submodel_options = {
self.outcomes[0]: {RANK: 1},
self.outcomes[1]: {RANK: 2},
}
self.surrogate = ListSurrogate(
botorch_model_class_per_outcome=self.botorch_model_class_per_outcome,
mll_class=self.mll_class,
submodel_options_per_outcome=self.submodel_options,
)
self.bounds = [(0.0, 1.0), (1.0, 4.0)]
self.feature_names = ["x1", "x2"]
def check_ranks(self, c: ListSurrogate) -> type(None):
self.assertIsInstance(c, ListSurrogate)
self.assertIsInstance(c.model, ModelListGP)
for idx, submodel in enumerate(c.model.models):
self.assertIsInstance(submodel, self.expected_submodel_type)
self.assertEqual(
submodel._rank, self.submodel_options[self.outcomes[idx]][RANK]
)
def test_init(self):
self.assertEqual(
self.surrogate.botorch_model_class_per_outcome,
self.botorch_model_class_per_outcome,
)
self.assertEqual(self.surrogate.mll_class, self.mll_class)
with self.assertRaises(NotImplementedError):
self.surrogate.training_data
with self.assertRaisesRegex(ValueError, NOT_YET_FIT_MSG):
self.surrogate.training_data_per_outcome
with self.assertRaisesRegex(
ValueError, "BoTorch `Model` has not yet been constructed"
):
self.surrogate.model
@patch(
f"{CURRENT_PATH}.FixedNoiseMultiTaskGP.construct_inputs",
# Mock to register calls, but still execute the function.
side_effect=FixedNoiseMultiTaskGP.construct_inputs,
)
def test_construct(self, mock_MTGP_construct_inputs):
with self.assertRaisesRegex(ValueError, ".* are required"):
self.surrogate.construct(training_data=self.training_data)
self.surrogate.construct(
training_data=self.training_data,
fidelity_features=[],
task_features=self.task_features,
metric_names=self.outcomes,
)
self.check_ranks(self.surrogate)
# Should construct inputs for MTGP twice.
self.assertEqual(len(mock_MTGP_construct_inputs.call_args_list), 2)
# First construct inputs should be called for MTGP with training data #0.
self.assertEqual(
# `call_args` is a tuple of (args, kwargs), and we are interested in kwargs.
mock_MTGP_construct_inputs.call_args_list[0][1],
{
"fidelity_features": [],
"task_features": self.task_features,
"training_data": self.training_data[0],
},
)
# Then, with training data #1.
self.assertEqual(
# `call_args` is a tuple of (args, kwargs), and we are interested in kwargs.
mock_MTGP_construct_inputs.call_args_list[1][1],
{
"fidelity_features": [],
"task_features": self.task_features,
"training_data": self.training_data[1],
},
)
@patch(f"{CURRENT_PATH}.ModelListGP.load_state_dict", return_value=None)
@patch(f"{CURRENT_PATH}.SumMarginalLogLikelihood")
@patch(f"{SURROGATE_PATH}.fit_gpytorch_model")
def test_fit(self, mock_fit_gpytorch, mock_MLL, mock_state_dict):
surrogate = ListSurrogate(
botorch_model_class_per_outcome=self.botorch_model_class_per_outcome,
mll_class=SumMarginalLogLikelihood,
)
# Checking that model is None before `fit` (and `construct`) calls.
self.assertIsNone(surrogate._model)
# Should instantiate mll and `fit_gpytorch_model` when `state_dict`
# is `None`.
surrogate.fit(
training_data=self.training_data,
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
metric_names=self.outcomes,
fidelity_features=[],
)
mock_state_dict.assert_not_called()
mock_MLL.assert_called_once()
mock_fit_gpytorch.assert_called_once()
mock_state_dict.reset_mock()
mock_MLL.reset_mock()
mock_fit_gpytorch.reset_mock()
# Should `load_state_dict` when `state_dict` is not `None`
# and `refit` is `False`.
state_dict = {}
surrogate.fit(
training_data=self.training_data,
bounds=self.bounds,
task_features=self.task_features,
feature_names=self.feature_names,
metric_names=self.outcomes,
fidelity_features=[],
refit=False,
state_dict=state_dict,
)
mock_state_dict.assert_called_once()
mock_MLL.assert_not_called()
mock_fit_gpytorch.assert_not_called()