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() ss = get_search_space_for_value() modelbridge = ModelBridge(ss, 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"] == []) self.assertTrue(fit_args["observation_data"] == []) self.assertTrue(mock_observations_from_data.called) # Test prediction on out of design features. modelbridge._predict = mock.MagicMock( "ax.modelbridge.base.ModelBridge._predict", autospec=True, side_effect=ValueError("Out of Design"), ) # This point is in design, and thus failures in predict are legitimate. with mock.patch.object(ModelBridge, "model_space", return_value=get_search_space_for_range_values): with self.assertRaises(ValueError): modelbridge.predict([get_observation2().features]) # This point is out of design, and not in training data. with self.assertRaises(ValueError): modelbridge.predict([get_observation_status_quo0().features]) # Now it's in the training data. with mock.patch.object( ModelBridge, "get_training_data", return_value=[get_observation_status_quo0()], ): # Return raw training value. self.assertEqual( modelbridge.predict([get_observation_status_quo0().features]), unwrap_observation_data([get_observation_status_quo0().data]), ) # 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]) # Check that _single_predict is equivalent here. modelbridge._single_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, [False, False]) with self.assertRaises(ValueError): modelbridge.training_in_design = [True, True, False] with self.assertRaises(ValueError): modelbridge.training_in_design = [True, True, False] # Test feature_importances with self.assertRaises(NotImplementedError): modelbridge.feature_importances("a")
class BaseModelBridgeTest(TestCase): @mock.patch( "ax.modelbridge.base.observations_from_data", autospec=True, return_value=([get_observation1(), get_observation2()]), ) @mock.patch("ax.modelbridge.base.gen_arms", autospec=True, return_value=[Arm(parameters={})]) @mock.patch("ax.modelbridge.base.ModelBridge._fit", autospec=True) 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() ss = get_search_space_for_value() modelbridge = ModelBridge(ss, 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"] == []) self.assertTrue(fit_args["observation_data"] == []) self.assertTrue(mock_observations_from_data.called) # Test prediction on out of design features. modelbridge._predict = mock.MagicMock( "ax.modelbridge.base.ModelBridge._predict", autospec=True, side_effect=ValueError("Out of Design"), ) # This point is in design, and thus failures in predict are legitimate. with mock.patch.object(ModelBridge, "model_space", return_value=get_search_space_for_range_values): with self.assertRaises(ValueError): modelbridge.predict([get_observation2().features]) # This point is out of design, and not in training data. with self.assertRaises(ValueError): modelbridge.predict([get_observation_status_quo0().features]) # Now it's in the training data. with mock.patch.object( ModelBridge, "get_training_data", return_value=[get_observation_status_quo0()], ): # Return raw training value. self.assertEqual( modelbridge.predict([get_observation_status_quo0().features]), unwrap_observation_data([get_observation_status_quo0().data]), ) # 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]) # Check that _single_predict is equivalent here. modelbridge._single_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, [False, False]) with self.assertRaises(ValueError): modelbridge.training_in_design = [True, True, False] with self.assertRaises(ValueError): modelbridge.training_in_design = [True, True, False] # Test feature_importances with self.assertRaises(NotImplementedError): modelbridge.feature_importances("a") @mock.patch( "ax.modelbridge.base.observations_from_data", autospec=True, return_value=([get_observation1()]), ) @mock.patch("ax.modelbridge.base.ModelBridge._fit", autospec=True) 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) @mock.patch( "ax.modelbridge.base.observations_from_data", autospec=True, return_value=([ get_observation_status_quo0(), get_observation_status_quo1(), get_observation1(), get_observation2(), ]), ) @mock.patch("ax.modelbridge.base.ModelBridge._fit", autospec=True) def testSetStatusQuoMultipleObs(self, mock_fit, mock_observations_from_data): exp = get_experiment_with_repeated_arms(2) trial_index = 1 status_quo_features = ObservationFeatures( parameters=exp.trials[trial_index].status_quo.parameters, trial_index=trial_index, ) modelbridge = ModelBridge( get_search_space_for_value(), 0, [], exp, 0, status_quo_features=status_quo_features, ) # Check that for experiments with many trials the status quo is set # to the value of the status quo of the last trial. if len(exp.trials) >= 1: self.assertEqual(modelbridge.status_quo, get_observation_status_quo1()) @mock.patch( "ax.modelbridge.base.observations_from_data", autospec=True, return_value=([get_observation1(), get_observation1()]), ) @mock.patch("ax.modelbridge.base.ModelBridge._fit", autospec=True) 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 testUnwrapObservationData(self): observation_data = [get_observation1().data, get_observation2().data] f, cov = unwrap_observation_data(observation_data) self.assertEqual(f["a"], [2.0, 2.0]) self.assertEqual(f["b"], [4.0, 1.0]) self.assertEqual(cov["a"]["a"], [1.0, 2.0]) self.assertEqual(cov["b"]["b"], [4.0, 5.0]) self.assertEqual(cov["a"]["b"], [2.0, 3.0]) self.assertEqual(cov["b"]["a"], [3.0, 4.0]) # Check that errors if metric mismatch od3 = ObservationData(metric_names=["a"], means=np.array([2.0]), covariance=np.array([[4.0]])) with self.assertRaises(ValueError): unwrap_observation_data(observation_data + [od3]) def testGenArms(self): p1 = {"x": 0, "y": 1} p2 = {"x": 4, "y": 8} observation_features = [ ObservationFeatures(parameters=p1), ObservationFeatures(parameters=p2), ] arms = gen_arms(observation_features=observation_features) self.assertEqual(arms[0].parameters, p1) arm = Arm(name="1_1", parameters=p1) arms_by_signature = {arm.signature: arm} arms = gen_arms( observation_features=observation_features, arms_by_signature=arms_by_signature, ) self.assertEqual(arms[0].name, "1_1") @mock.patch( "ax.modelbridge.base.ModelBridge._gen", autospec=True, return_value=([get_observation1trans().features], [2], None, {}), ) @mock.patch("ax.modelbridge.base.ModelBridge.predict", autospec=True, return_value=None) 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={}, ) @mock.patch( "ax.modelbridge.base.ModelBridge._gen", autospec=True, side_effect=[ ([get_observation1trans().features], [2], None, {}), ([get_observation2trans().features], [2], None, {}), ([get_observation2().features], [2], None, {}), ], ) @mock.patch("ax.modelbridge.base.ModelBridge._update", autospec=True) 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, ), ss, ) 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, )