Reduce number of lines

optimas-org · Mar 6, 2024 · 5b1880d · 5b1880d
1 parent dc2d9aa
commit 5b1880d
Showing 1 changed file with 3 additions and 13 deletions.
diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
@@ -125,19 +125,9 @@ def _build_ax_client_from_dataframe(
         # a random sampling step will be set up by Ax, and this step does not
         # allow calling `model.predict`. Any strategy that uses a GP surrogate
         # should work.
-        ax_client = AxClient(
-            generation_strategy=GenerationStrategy(
-                [
-                    GenerationStep(
-                        model=(
-                            Models.GPEI if len(objectives) == 1 else Models.MOO
-                        ),
-                        num_trials=-1,
-                    )
-                ]
-            ),
-            verbose_logging=False,
-        )
+        model = Models.GPEI if len(objectives) == 1 else Models.MOO
+        gs = GenerationStrategy([GenerationStep(model=model, num_trials=-1)])
+        ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
         ax_client.create_experiment(
             parameters=axparameters, objectives=axobjectives
         )