Update initialization of trials

optimas-org · Oct 21, 2024 · 318d7f0 · 318d7f0
1 parent ac7c618
commit 318d7f0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/optimas/generators/ax/service/base.py b/optimas/generators/ax/service/base.py
@@ -23,7 +23,7 @@
     Trial,
     VaryingParameter,
     Parameter,
-    TrialStatus,
+    TrialParameter,
 )
 from optimas.generators.ax.base import AxGenerator
 from optimas.utils.ax import AxModelManager
@@ -109,6 +109,7 @@ def __init__(
         model_save_period: Optional[int] = 5,
         model_history_dir: Optional[str] = "model_history",
     ) -> None:
+        custom_trial_parameters = [ TrialParameter("trial_index", "ax_trial_index", dtype=int)]
         super().__init__(
             varying_parameters=varying_parameters,
             objectives=objectives,
@@ -119,6 +120,7 @@ def __init__(
             save_model=save_model,
             model_save_period=model_save_period,
             model_history_dir=model_history_dir,
+            custom_trial_parameters=custom_trial_parameters,
             allow_fixed_parameters=True,
             allow_updating_parameters=True,
         )
@@ -153,7 +155,7 @@ def ask(self, num_points: Optional[int]) -> List[dict]:
                 var.name: parameters.get(var.name)
                 for var in self._varying_parameters
             }
-            point["_id"] = trial_id
+            point["ax_trial_id"] = trial_id
             points.append(point)
         return points
 
@@ -163,10 +165,11 @@ def tell(self, results: List[dict]) -> None:
         """
         for result in results:
             trial = Trial.from_dict(
-                result,
-                self._varying_parameters,
-                self._objectives,
-                self._analyzed_parameters
+                trial_dict=result,
+                varying_parameters=self._varying_parameters,
+                objectives=self._objectives,
+                analyzed_parameters=self._analyzed_parameters,
+                custom_parameters=self._custom_trial_parameters,
                 )
             try:
                 trial_id = trial.ax_trial_id

diff --git a/optimas/generators/base.py b/optimas/generators/base.py
@@ -209,11 +209,9 @@ def ask_trials(self, n_trials: int) -> List[Trial]:
             # Convert the points to the Trial format
             gen_trials = []
             for point in gen_points:
-                trial = Trial(
+                trial = Trial.from_dict(
+                    trial_dict=point,
                     varying_parameters=self._varying_parameters,
-                    parameter_values=[
-                        point[var.name] for var in self._varying_parameters
-                    ],
                     objectives=self._objectives,
                     analyzed_parameters=self._analyzed_parameters,
                     custom_parameters=self._custom_trial_parameters,