diff --git a/ax/modelbridge/registry.py b/ax/modelbridge/registry.py index 8158df7f88b..df9caf45722 100644 --- a/ax/modelbridge/registry.py +++ b/ax/modelbridge/registry.py @@ -424,7 +424,7 @@ class Models(ModelRegistryBase): To instantiate a model in this enum, simply call an enum member like so: `Models.SOBOL(search_space=search_space)` or - `Models.GPEI(experiment=experiment, data=data)`. Keyword arguments + `Models.BOTORCH(experiment=experiment, data=data)`. Keyword arguments specified to the call will be passed into the model or the model bridge constructors according to their keyword. diff --git a/ax/modelbridge/tests/test_aepsych_criterion.py b/ax/modelbridge/tests/test_aepsych_criterion.py index 23d094b295c..eda6a4c6632 100644 --- a/ax/modelbridge/tests/test_aepsych_criterion.py +++ b/ax/modelbridge/tests/test_aepsych_criterion.py @@ -31,13 +31,13 @@ def test_single_criterion(self) -> None: experiment = get_experiment() generation_strategy = GenerationStrategy( - name="SOBOL+GPEI::default", + name="SOBOL+MBM::default", steps=[ GenerationStep( model=Models.SOBOL, num_trials=-1, completion_criteria=[criterion] ), GenerationStep( - model=Models.GPEI, + model=Models.BOTORCH_MODULAR, num_trials=-1, max_parallelism=1, ), @@ -77,7 +77,8 @@ def test_single_criterion(self) -> None: ) self.assertEqual( - generation_strategy._curr.model_spec_to_gen_from.model_enum, Models.GPEI + generation_strategy._curr.model_spec_to_gen_from.model_enum, + Models.BOTORCH_MODULAR, ) def test_many_criteria(self) -> None: @@ -89,13 +90,13 @@ def test_many_criteria(self) -> None: experiment = get_experiment() generation_strategy = GenerationStrategy( - name="SOBOL+GPEI::default", + name="SOBOL+MBM::default", steps=[ GenerationStep( model=Models.SOBOL, num_trials=-1, completion_criteria=criteria ), GenerationStep( - model=Models.GPEI, + model=Models.BOTORCH_MODULAR, num_trials=-1, max_parallelism=1, ), @@ -153,5 +154,6 @@ def test_many_criteria(self) -> None: ) self.assertEqual( - generation_strategy._curr.model_spec_to_gen_from.model_enum, Models.GPEI + generation_strategy._curr.model_spec_to_gen_from.model_enum, + Models.BOTORCH_MODULAR, ) diff --git a/ax/modelbridge/tests/test_generation_node.py b/ax/modelbridge/tests/test_generation_node.py index 7df2ae88059..c2d3d8d4b4d 100644 --- a/ax/modelbridge/tests/test_generation_node.py +++ b/ax/modelbridge/tests/test_generation_node.py @@ -116,7 +116,7 @@ def test_properties(self) -> None: node_name="test", model_specs=[ ModelSpec( - model_enum=Models.GPEI, + model_enum=Models.BOTORCH_MODULAR, model_kwargs={}, model_gen_kwargs={ "n": 1, @@ -140,7 +140,7 @@ def test_properties(self) -> None: self.assertEqual( node.model_spec_to_gen_from.model_kwargs, node.model_specs[0].model_kwargs ) - self.assertEqual(node.model_to_gen_from_name, "GPEI") + self.assertEqual(node.model_to_gen_from_name, "BoTorch") self.assertEqual( node.model_spec_to_gen_from.model_gen_kwargs, node.model_specs[0].model_gen_kwargs, @@ -167,7 +167,7 @@ def test_node_string_representation(self) -> None: node_name="test", model_specs=[ ModelSpec( - model_enum=Models.GPEI, + model_enum=Models.BOTORCH_MODULAR, model_kwargs={}, model_gen_kwargs={}, ), @@ -180,7 +180,7 @@ def test_node_string_representation(self) -> None: self.assertEqual( string_rep, ( - "GenerationNode(model_specs=[ModelSpec(model_enum=GPEI," + "GenerationNode(model_specs=[ModelSpec(model_enum=BoTorch," " model_kwargs={}, model_gen_kwargs={}, model_cv_kwargs={}," " )], node_name=test, " "transition_criteria=[MaxTrials({'threshold': 5, " @@ -197,7 +197,7 @@ def test_single_fixed_features(self) -> None: node_name="test", model_specs=[ ModelSpec( - model_enum=Models.GPEI, + model_enum=Models.BOTORCH_MODULAR, model_kwargs={}, model_gen_kwargs={ "n": 2, diff --git a/ax/modelbridge/tests/test_model_fit_metrics.py b/ax/modelbridge/tests/test_model_fit_metrics.py index f9c77eefeff..41e35be047d 100644 --- a/ax/modelbridge/tests/test_model_fit_metrics.py +++ b/ax/modelbridge/tests/test_model_fit_metrics.py @@ -58,7 +58,7 @@ def setUp(self) -> None: GenerationStep( model=Models.SOBOL, num_trials=NUM_SOBOL, max_parallelism=NUM_SOBOL ), - GenerationStep(model=Models.GPEI, num_trials=-1), + GenerationStep(model=Models.BOTORCH_MODULAR, num_trials=-1), ] ) diff --git a/ax/modelbridge/tests/test_model_spec.py b/ax/modelbridge/tests/test_model_spec.py index ee24532bcf6..f5207093d62 100644 --- a/ax/modelbridge/tests/test_model_spec.py +++ b/ax/modelbridge/tests/test_model_spec.py @@ -37,7 +37,7 @@ def setUp(self) -> None: class ModelSpecTest(BaseModelSpecTest): @fast_botorch_optimize def test_construct(self) -> None: - ms = ModelSpec(model_enum=Models.GPEI) + ms = ModelSpec(model_enum=Models.BOTORCH_MODULAR) with self.assertRaises(UserInputError): ms.gen(n=1) ms.fit(experiment=self.experiment, data=self.data) @@ -51,7 +51,7 @@ def test_construct(self) -> None: wraps=extract_search_space_digest, ) def test_fit(self, wrapped_extract_ssd: Mock) -> None: - ms = ModelSpec(model_enum=Models.GPEI) + ms = ModelSpec(model_enum=Models.BOTORCH_MODULAR) # This should fit the model as usual. ms.fit(experiment=self.experiment, data=self.data) wrapped_extract_ssd.assert_called_once() @@ -159,7 +159,7 @@ def test_cross_validate_with_non_GP_model( mock_diagnostics.assert_not_called() def test_fixed_features(self) -> None: - ms = ModelSpec(model_enum=Models.GPEI) + ms = ModelSpec(model_enum=Models.BOTORCH_MODULAR) self.assertIsNone(ms.fixed_features) new_features = ObservationFeatures(parameters={"a": 1.0}) ms.fixed_features = new_features @@ -177,7 +177,7 @@ def test_gen_attaches_empty_model_fit_metadata_if_fit_not_applicable(self) -> No self.assertEqual(gen_metadata["model_std_generalization"], None) def test_gen_attaches_model_fit_metadata_if_applicable(self) -> None: - ms = ModelSpec(model_enum=Models.GPEI) + ms = ModelSpec(model_enum=Models.BOTORCH_MODULAR) ms.fit(experiment=self.experiment, data=self.data) gr = ms.gen(n=1) gen_metadata = not_none(gr.gen_metadata) diff --git a/ax/modelbridge/tests/test_transition_criterion.py b/ax/modelbridge/tests/test_transition_criterion.py index d00d7e5e400..a2b73af9a69 100644 --- a/ax/modelbridge/tests/test_transition_criterion.py +++ b/ax/modelbridge/tests/test_transition_criterion.py @@ -58,7 +58,7 @@ def test_minimum_preference_criterion(self) -> None: completion_criteria=[criterion], ), GenerationStep( - model=Models.GPEI, + model=Models.BOTORCH_MODULAR, num_trials=-1, max_parallelism=1, ), @@ -92,7 +92,8 @@ def test_minimum_preference_criterion(self) -> None: ) ) self.assertEqual( - generation_strategy._curr.model_spec_to_gen_from.model_enum, Models.GPEI + generation_strategy._curr.model_spec_to_gen_from.model_enum, + Models.BOTORCH_MODULAR, ) def test_default_step_criterion_setup(self) -> None: @@ -106,21 +107,21 @@ def test_default_step_criterion_setup(self) -> None: """ experiment = get_experiment() gs = GenerationStrategy( - name="SOBOL+GPEI::default", + name="SOBOL+MBM::default", steps=[ GenerationStep( model=Models.SOBOL, num_trials=3, ), GenerationStep( - model=Models.GPEI, + model=Models.BOTORCH_MODULAR, num_trials=4, max_parallelism=1, min_trials_observed=2, enforce_num_trials=False, ), GenerationStep( - model=Models.GPEI, + model=Models.BOTORCH_MODULAR, num_trials=-1, ), ], diff --git a/ax/plot/tests/test_traces.py b/ax/plot/tests/test_traces.py index 5679aa6ba53..bd7465834d3 100644 --- a/ax/plot/tests/test_traces.py +++ b/ax/plot/tests/test_traces.py @@ -77,7 +77,7 @@ def test_plot_objective_value_vs_trial_index(self) -> None: for _ in range(2): t = exp.new_trial(sobol.gen(1)).run() t.mark_completed() - model = Models.GPEI( + model = Models.BOTORCH_MODULAR( experiment=exp, data=exp.fetch_data(), ) diff --git a/ax/service/tests/test_report_utils.py b/ax/service/tests/test_report_utils.py index a31b8809c92..ae35e82dc34 100644 --- a/ax/service/tests/test_report_utils.py +++ b/ax/service/tests/test_report_utils.py @@ -522,7 +522,7 @@ def test_skip_contour_high_dimensional(self) -> None: sobol = Models.SOBOL(search_space=exp.search_space) for _ in range(1): exp.new_trial(sobol.gen(1)).run() - model = Models.GPEI( + model = Models.BOTORCH_MODULAR( experiment=exp, data=exp.fetch_data(), ) @@ -1206,7 +1206,7 @@ def test_compare_to_baseline_select_baseline_arm(self) -> None: experiment=experiment, arms_df=arms_df, baseline_arm_name=wrong_baseline_name, - ), + ) # status quo baseline arm experiment_with_status_quo = copy.deepcopy(experiment) @@ -1272,7 +1272,9 @@ def test_warn_if_unpredictable_metrics(self) -> None: min_trials_observed=3, max_parallelism=3, ), - GenerationStep(model=Models.GPEI, num_trials=-1, max_parallelism=3), + GenerationStep( + model=Models.BOTORCH_MODULAR, num_trials=-1, max_parallelism=3 + ), ] ) gs.experiment = exp diff --git a/ax/telemetry/tests/test_scheduler.py b/ax/telemetry/tests/test_scheduler.py index faa67ef29be..49e7865179a 100644 --- a/ax/telemetry/tests/test_scheduler.py +++ b/ax/telemetry/tests/test_scheduler.py @@ -167,7 +167,7 @@ def test_scheduler_model_fit_metrics_logging(self) -> None: GenerationStep( model=Models.SOBOL, num_trials=NUM_SOBOL, max_parallelism=NUM_SOBOL ), - GenerationStep(model=Models.GPEI, num_trials=-1), + GenerationStep(model=Models.BOTORCH_MODULAR, num_trials=-1), ] ) diff --git a/docs/api.md b/docs/api.md index c517537c3d2..c3dace88022 100644 --- a/docs/api.md +++ b/docs/api.md @@ -155,7 +155,7 @@ for i in range(5): best_arm = None for i in range(15): - gpei = Models.GPEI(experiment=exp, data=exp.fetch_data()) + gpei = Models.BOTORCH_MODULAR(experiment=exp, data=exp.fetch_data()) generator_run = gpei.gen(1) best_arm, _ = generator_run.best_arm_predictions trial = exp.new_trial(generator_run=generator_run) diff --git a/docs/trial-evaluation.md b/docs/trial-evaluation.md index 7b675058948..a4d81e50ae2 100644 --- a/docs/trial-evaluation.md +++ b/docs/trial-evaluation.md @@ -122,7 +122,7 @@ for i in range(5): trial.mark_completed() for i in range(15): - gpei = Models.GPEI(experiment=exp, data=exp.fetch_data()) + gpei = Models.BOTORCH_MODULAR(experiment=exp, data=exp.fetch_data()) generator_run = gpei.gen(1) trial = exp.new_trial(generator_run=generator_run) trial.run()