def test_gen_validates_one_model_spec(self):
generation_node = GenerationNode(
model_specs=[self.sobol_model_spec, self.sobol_model_spec]
)
# Base generation node can only handle one model spec at the moment
# (this might change in the future), so it should raise a `NotImplemented
# Error` if we attempt to generate from a generation node that has
# more than one model spec. Note that the check is done in `gen` and
# not in the constructor to make `GenerationNode` mode convenient to
# subclass.
with self.assertRaises(NotImplementedError):
generation_node.gen()
class TestGenerationNodeWithBestModelSelector(TestCase):
@fast_botorch_optimize
def setUp(self):
self.branin_experiment = get_branin_experiment()
sobol = Models.SOBOL(search_space=self.branin_experiment.search_space)
sobol_run = sobol.gen(n=20)
self.branin_experiment.new_batch_trial().add_generator_run(
sobol_run
).run().mark_completed()
data = self.branin_experiment.fetch_data()
ms_gpei = ModelSpec(model_enum=Models.GPEI)
ms_gpei.fit(experiment=self.branin_experiment, data=data)
ms_gpkg = ModelSpec(model_enum=Models.GPKG)
ms_gpkg.fit(experiment=self.branin_experiment, data=data)
self.fitted_model_specs = [ms_gpei, ms_gpkg]
self.model_selection_node = GenerationNode(
model_specs=self.fitted_model_specs,
best_model_selector=SingleDiagnosticBestModelSelector(
diagnostic="Fisher exact test p",
criterion=MetricAggregation.MEAN,
metric_aggregation=DiagnosticCriterion.MIN,
),
)
@fast_botorch_optimize
def test_gen(self):
self.model_selection_node.fit(
experiment=self.branin_experiment, data=self.branin_experiment.lookup_data()
)
# Check that with `ModelSelectionNode` generation from a node with
# multiple model specs does not fail.
gr = self.model_selection_node.gen(n=1, pending_observations={"branin": []})
# Currently, `ModelSelectionNode` should just pick the first model
# spec as the one to generate from.
# TODO[adamobeng]: Test correct behavior here when implemented.
self.assertEqual(gr._model_key, "GPEI")
def test_fixed_features_is_from_model_to_gen_from(self) -> None:
self.model_selection_node.model_specs[0].fixed_features = ObservationFeatures(
parameters={"x": 0}
)
self.model_selection_node.model_specs[1].fixed_features = ObservationFeatures(
parameters={"x": 1}
)
self.assertEqual(
self.model_selection_node.fixed_features,
self.model_selection_node.model_spec_to_gen_from.fixed_features,
)
class TestGenerationNode(TestCase):
def setUp(self):
self.sobol_model_spec = ModelSpec(
model_enum=Models.SOBOL,
model_kwargs={"init_position": 3},
model_gen_kwargs={"some_gen_kwarg": "some_value"},
)
self.sobol_generation_node = GenerationNode(model_specs=[self.sobol_model_spec])
self.branin_experiment = get_branin_experiment(with_completed_trial=True)
def test_init(self):
self.assertEqual(
self.sobol_generation_node.model_specs, [self.sobol_model_spec]
)
def test_fit(self):
dat = self.branin_experiment.lookup_data()
with patch.object(
self.sobol_model_spec, "fit", wraps=self.sobol_model_spec.fit
) as mock_model_spec_fit:
self.sobol_generation_node.fit(
experiment=self.branin_experiment,
data=dat,
)
mock_model_spec_fit.assert_called_with(
experiment=self.branin_experiment,
data=dat,
search_space=None,
optimization_config=None,
)
def test_gen(self):
dat = self.branin_experiment.lookup_data()
self.sobol_generation_node.fit(
experiment=self.branin_experiment,
data=dat,
)
with patch.object(
self.sobol_model_spec, "gen", wraps=self.sobol_model_spec.gen
) as mock_model_spec_gen:
gr = self.sobol_generation_node.gen(
n=1, pending_observations={"branin": []}
)
mock_model_spec_gen.assert_called_with(n=1, pending_observations={"branin": []})
self.assertEqual(gr._model_key, self.sobol_model_spec.model_key)
self.assertEqual(gr._model_kwargs.get("init_position"), 3)
def test_gen_validates_one_model_spec(self):
generation_node = GenerationNode(
model_specs=[self.sobol_model_spec, self.sobol_model_spec]
)
# Base generation node can only handle one model spec at the moment
# (this might change in the future), so it should raise a `NotImplemented
# Error` if we attempt to generate from a generation node that has
# more than one model spec. Note that the check is done in `gen` and
# not in the constructor to make `GenerationNode` mode convenient to
# subclass.
with self.assertRaises(NotImplementedError):
generation_node.gen()
def test_properties(self):
node = GenerationNode(
model_specs=[
ModelSpec(
model_enum=Models.GPEI,
model_kwargs={},
model_gen_kwargs={
"n": 1,
"fixed_features": ObservationFeatures(
parameters={}, trial_index=0
),
},
),
],
)
dat = self.branin_experiment.lookup_data()
node.fit(
experiment=self.branin_experiment,
data=dat,
)
self.assertEqual(node.model_enum, node.model_specs[0].model_enum)
self.assertEqual(node.model_kwargs, node.model_specs[0].model_kwargs)
self.assertEqual(node.model_gen_kwargs, node.model_specs[0].model_gen_kwargs)
self.assertEqual(node.model_cv_kwargs, node.model_specs[0].model_cv_kwargs)
self.assertEqual(node.fixed_features, node.model_specs[0].fixed_features)
self.assertEqual(node.cv_results, node.model_specs[0].cv_results)
self.assertEqual(node.diagnostics, node.model_specs[0].diagnostics)
def test_single_fixed_features(self) -> None:
node = GenerationNode(
model_specs=[
ModelSpec(
model_enum=Models.GPEI,
model_kwargs={},
model_gen_kwargs={
"n": 2,
"fixed_features": ObservationFeatures(parameters={"x": 0}),
},
),
],
)
self.assertEqual(node.fixed_features, ObservationFeatures(parameters={"x": 0}))
def test_multiple_same_fixed_features(self) -> None:
node = GenerationNode(
model_specs=[
ModelSpec(
model_enum=Models.GPEI,
model_kwargs={},
model_gen_kwargs={
"n": 2,
"fixed_features": ObservationFeatures(parameters={"x": 0}),
},
),
ModelSpec(
model_enum=Models.GPEI,
model_kwargs={},
model_gen_kwargs={
"n": 3,
"fixed_features": ObservationFeatures(parameters={"x": 0}),
},
),
],
)
self.assertEqual(node.fixed_features, ObservationFeatures(parameters={"x": 0}))