def testModelBridge(self, mock_fit, mock_gen_arms, mock_observations_from_data):
# Test that on init transforms are stored and applied in the correct order
transforms = [t1, t2]
exp = get_experiment()
modelbridge = ModelBridge(search_space_for_value(), 0, transforms, exp, 0)
self.assertEqual(list(modelbridge.transforms.keys()), ["t1", "t2"])
fit_args = mock_fit.mock_calls[0][2]
self.assertTrue(fit_args["search_space"] == search_space_for_value(8.0))
self.assertTrue(
fit_args["observation_features"]
== [observation1trans().features, observation2trans().features]
)
self.assertTrue(
fit_args["observation_data"]
== [observation1trans().data, observation2trans().data]
)
self.assertTrue(mock_observations_from_data.called)
# Test that transforms are applied correctly on predict
modelbridge._predict = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._predict",
autospec=True,
return_value=[observation2trans().data],
)
modelbridge.predict([observation2().features])
# Observation features sent to _predict are un-transformed afterwards
modelbridge._predict.assert_called_with([observation2().features])
# Test transforms applied on gen
modelbridge._gen = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._gen",
autospec=True,
return_value=([observation1trans().features], [2], None),
)
oc = OptimizationConfig(objective=Objective(metric=Metric(name="test_metric")))
modelbridge._set_kwargs_to_save(
model_key="TestModel", model_kwargs={}, bridge_kwargs={}
)
gr = modelbridge.gen(
n=1,
search_space=search_space_for_value(),
optimization_config=oc,
pending_observations={"a": [observation2().features]},
fixed_features=ObservationFeatures({"x": 5}),
)
self.assertEqual(gr._model_key, "TestModel")
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace([FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=oc,
pending_observations={"a": [observation2trans().features]},
fixed_features=ObservationFeatures({"x": 36}),
model_gen_options=None,
)
mock_gen_arms.assert_called_with(
arms_by_signature={}, observation_features=[observation1().features]
)
# Gen with no pending observations and no fixed features
modelbridge.gen(
n=1, search_space=search_space_for_value(), optimization_config=None
)
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace([FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=None,
pending_observations={},
fixed_features=ObservationFeatures({}),
model_gen_options=None,
)
# Gen with multi-objective optimization config.
oc2 = OptimizationConfig(
objective=ScalarizedObjective(
metrics=[Metric(name="test_metric"), Metric(name="test_metric_2")]
)
)
modelbridge.gen(
n=1, search_space=search_space_for_value(), optimization_config=oc2
)
modelbridge._gen.assert_called_with(
n=1,
search_space=SearchSpace([FixedParameter("x", ParameterType.FLOAT, 8.0)]),
optimization_config=oc2,
pending_observations={},
fixed_features=ObservationFeatures({}),
model_gen_options=None,
)
# Test transforms applied on cross_validate
modelbridge._cross_validate = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._cross_validate",
autospec=True,
return_value=[observation1trans().data],
)
cv_training_data = [observation2()]
cv_test_points = [observation1().features]
cv_predictions = modelbridge.cross_validate(
cv_training_data=cv_training_data, cv_test_points=cv_test_points
)
modelbridge._cross_validate.assert_called_with(
obs_feats=[observation2trans().features],
obs_data=[observation2trans().data],
cv_test_points=[observation1().features], # untransformed after
)
self.assertTrue(cv_predictions == [observation1().data])
# Test stored training data
obs = modelbridge.get_training_data()
self.assertTrue(obs == [observation1(), observation2()])
self.assertEqual(modelbridge.metric_names, {"a", "b"})
self.assertIsNone(modelbridge.status_quo)
self.assertTrue(modelbridge.model_space == search_space_for_value())
self.assertEqual(modelbridge.training_in_design, [True, True])
modelbridge.training_in_design = [True, False]
with self.assertRaises(ValueError):
modelbridge.training_in_design = [True, True, False]
ood_obs = modelbridge.out_of_design_data()
self.assertTrue(ood_obs == unwrap_observation_data([observation2().data]))
def testNoOutOfDesign(self, mock_fit, mock_observations_from_data):
exp = get_experiment_for_value()
modelbridge = ModelBridge(get_search_space_for_value(), 0, [], exp, 0)
self.assertEqual(modelbridge.out_of_design_data(), None)
