Fix ax_trial_id

Author: shuds13

optimas/core/trial.py

@@ -68,7 +68,6 @@ def __init__(
         )
         self._ignored = False
         self._ignored_reason = None
-
         # Add custom parameters as trial attributes.
         for param in self._custom_parameters:
             setattr(self, param.name, None)
@@ -93,6 +92,10 @@ def to_dict(self) -> Dict:
             "_ignored_reason": self._ignored_reason,
             "_status": self._status,
         }
+
+        if hasattr(self, "_ax_trial_id"):
+            trial_dict["ax_trial_id"] = self._ax_trial_id
+
         return trial_dict
 
     @classmethod
@@ -139,6 +142,8 @@ def from_dict(
         )
         if "_id" in trial_dict:
             trial._index = trial_dict["_id"]
+        if "ax_trial_id" in trial_dict:
+            trial._ax_trial_id = trial_dict["ax_trial_id"]
         if "_ignored" in trial_dict:
             trial._ignored = trial_dict["_ignored"]
         if "_ignored_reason" in trial_dict:

optimas/generators/base.py

@@ -216,6 +216,8 @@ def ask_trials(self, n_trials: int) -> List[Trial]:
                     analyzed_parameters=self._analyzed_parameters,
                     custom_parameters=self._custom_trial_parameters,
                 )
+                if 'ax_trial_id' in point:
+                    trial.ax_trial_id = point['ax_trial_id']
                 gen_trials.append(trial)
             # Keep only trials that have been given data.
             for trial in gen_trials: