def testSetTrainingDataDupFeatures(self, mock_fit,
mock_observations_from_data):
# Throws an error if repeated features in observations.
with self.assertRaises(ValueError):
ModelBridge(
get_search_space_for_value(),
0,
[],
get_experiment_for_value(),
0,
status_quo_name="1_1",
)
def testGenWithDefaults(self, _, mock_gen):
exp = get_experiment_for_value()
exp.optimization_config = get_optimization_config_no_constraints()
ss = get_search_space_for_range_value()
modelbridge = ModelBridge(ss, None, [], exp)
modelbridge.gen(1)
mock_gen.assert_called_with(
modelbridge,
n=1,
search_space=ss,
fixed_features=ObservationFeatures(parameters={}),
model_gen_options=None,
optimization_config=OptimizationConfig(
objective=Objective(metric=Metric("test_metric"),
minimize=False),
outcome_constraints=[],
),
pending_observations={},
)
def test_update(self, _mock_update, _mock_gen):
exp = get_experiment_for_value()
exp.optimization_config = get_optimization_config_no_constraints()
ss = get_search_space_for_range_values()
exp.search_space = ss
modelbridge = ModelBridge(ss, None, [Log], exp)
exp.new_trial(generator_run=modelbridge.gen(1))
modelbridge._set_training_data(
observations_from_data(
data=Data(
pd.DataFrame([{
"arm_name": "0_0",
"metric_name": "m1",
"mean": 3.0,
"sem": 1.0,
}])),
experiment=exp,
))
exp.new_trial(generator_run=modelbridge.gen(1))
modelbridge.update(
data=Data(
pd.DataFrame([{
"arm_name": "1_0",
"metric_name": "m1",
"mean": 5.0,
"sem": 0.0
}])),
experiment=exp,
)
exp.new_trial(generator_run=modelbridge.gen(1))
# Trying to update with unrecognised metric should error.
with self.assertRaisesRegex(ValueError, "Unrecognised metric"):
modelbridge.update(
data=Data(
pd.DataFrame([{
"arm_name": "1_0",
"metric_name": "m2",
"mean": 5.0,
"sem": 0.0,
}])),
experiment=exp,
)
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 = [transform_1, transform_2]
exp = get_experiment_for_value()
modelbridge = ModelBridge(get_search_space_for_value(), 0, transforms,
exp, 0)
self.assertEqual(list(modelbridge.transforms.keys()),
["transform_1", "transform_2"])
fit_args = mock_fit.mock_calls[0][2]
self.assertTrue(
fit_args["search_space"] == get_search_space_for_value(8.0))
self.assertTrue(fit_args["observation_features"] == [
get_observation1trans().features,
get_observation2trans().features
])
self.assertTrue(
fit_args["observation_data"] ==
[get_observation1trans().data,
get_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=[get_observation2trans().data],
)
modelbridge.predict([get_observation2().features])
# Observation features sent to _predict are un-transformed afterwards
modelbridge._predict.assert_called_with([get_observation2().features])
# Test transforms applied on gen
modelbridge._gen = mock.MagicMock(
"ax.modelbridge.base.ModelBridge._gen",
autospec=True,
return_value=([get_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=get_search_space_for_value(),
optimization_config=oc,
pending_observations={"a": [get_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": [get_observation2trans().features]},
fixed_features=ObservationFeatures({"x": 36}),
model_gen_options=None,
)
mock_gen_arms.assert_called_with(
arms_by_signature={},
observation_features=[get_observation1().features])
# Gen with no pending observations and no fixed features
modelbridge.gen(n=1,
search_space=get_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=get_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=[get_observation1trans().data],
)
cv_training_data = [get_observation2()]
cv_test_points = [get_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=[get_observation2trans().features],
obs_data=[get_observation2trans().data],
cv_test_points=[get_observation1().features
], # untransformed after
)
self.assertTrue(cv_predictions == [get_observation1().data])
# Test stored training data
obs = modelbridge.get_training_data()
self.assertTrue(obs == [get_observation1(), get_observation2()])
self.assertEqual(modelbridge.metric_names, {"a", "b"})
self.assertIsNone(modelbridge.status_quo)
self.assertTrue(
modelbridge.model_space == get_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([get_observation2().data]))
def testSetStatusQuo(self, mock_fit, mock_observations_from_data):
# NOTE: If empty data object is not passed, observations are not
# extracted, even with mock.
modelbridge = ModelBridge(
search_space=get_search_space_for_value(),
model=0,
experiment=get_experiment_for_value(),
data=Data(),
status_quo_name="1_1",
)
self.assertEqual(modelbridge.status_quo, get_observation1())
# Alternatively, we can specify by features
modelbridge = ModelBridge(
get_search_space_for_value(),
0,
[],
get_experiment_for_value(),
0,
status_quo_features=get_observation1().features,
)
self.assertEqual(modelbridge.status_quo, get_observation1())
# Alternatively, we can specify on experiment
# Put a dummy arm with SQ name 1_1 on the dummy experiment.
exp = get_experiment_for_value()
sq = Arm(name="1_1", parameters={"x": 3.0})
exp._status_quo = sq
# Check that we set SQ to arm 1_1
modelbridge = ModelBridge(get_search_space_for_value(), 0, [], exp, 0)
self.assertEqual(modelbridge.status_quo, get_observation1())
# Errors if features and name both specified
with self.assertRaises(ValueError):
modelbridge = ModelBridge(
get_search_space_for_value(),
0,
[],
exp,
0,
status_quo_features=get_observation1().features,
status_quo_name="1_1",
)
# Left as None if features or name don't exist
modelbridge = ModelBridge(get_search_space_for_value(),
0, [],
exp,
0,
status_quo_name="1_0")
self.assertIsNone(modelbridge.status_quo)
modelbridge = ModelBridge(
get_search_space_for_value(),
0,
[],
get_experiment_for_value(),
0,
status_quo_features=ObservationFeatures(parameters={
"x": 3.0,
"y": 10.0
}),
)
self.assertIsNone(modelbridge.status_quo)