Merge pull request #249 from optimas-org/standardization

Attempt to make `optimas` generators compatible with generator standard

Author: RemiLehe

doc/source/examples/ps_grid_sampling.rst

@@ -43,7 +43,7 @@ where :math:`l_b=0` and :math:`u_b=15`, the grid of sample looks like:
 
    all_trials = []
    while True:
-       trial = gen.ask(1)
+       trial = gen.ask_trials(1)
        if trial:
            all_trials.append(trial[0])
        else:

doc/source/examples/ps_line_sampling.rst

@@ -46,7 +46,7 @@ where :math:`x_0` and :math:`x_1` have a default values of :math:`5` and
 
    all_trials = []
    while True:
-       trial = gen.ask(1)
+       trial = gen.ask_trials(1)
        if trial:
            all_trials.append(trial[0])
        else:

doc/source/examples/ps_random_sampling.rst

@@ -45,7 +45,7 @@ distribution such as:
 
    all_trials = []
    while len(all_trials) <= 100:
-       trial = gen.ask(1)
+       trial = gen.ask_trials(1)
        if trial:
            all_trials.append(trial[0])
        else:

optimas/core/trial.py

@@ -63,12 +63,12 @@ def __init__(
         )
         evaluations = [] if evaluations is None else evaluations
         self._index = index
+        self._gen_id = None
         self._custom_parameters = (
             [] if custom_parameters is None else custom_parameters
         )
         self._ignored = False
         self._ignored_reason = None
-
         # Add custom parameters as trial attributes.
         for param in self._custom_parameters:
             setattr(self, param.name, None)
@@ -81,6 +81,76 @@ def __init__(
             self._mapped_evaluations[ev.parameter.name] = ev
         self.mark_as(TrialStatus.CANDIDATE)
 
+    def to_dict(self) -> Dict:
+        """Convert the trial to a dictionary."""
+        trial_dict = {
+            **self.parameters_as_dict(),
+            **self.objectives_as_dict(),
+            **self.analyzed_parameters_as_dict(),
+            **self.custom_parameters_as_dict(),
+            "_id": self._gen_id,
+            "_ignored": self._ignored,
+            "_ignored_reason": self._ignored_reason,
+            "_status": self._status,
+        }
+
+        return trial_dict
+
+    @classmethod
+    def from_dict(
+        cls,
+        trial_dict: Dict,
+        varying_parameters: List[VaryingParameter],
+        objectives: List[Objective],
+        analyzed_parameters: List[Parameter],
+        custom_parameters: Optional[List[TrialParameter]] = None,
+    ) -> "Trial":
+        """Create a trial from a dictionary.
+
+        Parameters
+        ----------
+        trial_dict : dict
+            Dictionary containing the trial information.
+        varying_parameters : list of VaryingParameter
+            The varying parameters of the optimization.
+        objectives : list of Objective
+            The optimization objectives.
+        analyzed_parameters : list of Parameter, optional
+            Additional parameters to be analyzed during the optimization.
+        custom_parameters : list of TrialParameter, optional
+            Additional parameters needed to identify or carry out the trial, and
+            which will be included in the optimization history.
+        """
+        # Prepare values of the input parameters
+        parameter_values = [trial_dict[var.name] for var in varying_parameters]
+        # Prepare evaluations
+        evaluations = [
+            Evaluation(parameter=par, value=trial_dict[par.name])
+            for par in objectives + analyzed_parameters
+            if par.name in trial_dict
+        ]
+        # Create the trial object
+        trial = cls(
+            varying_parameters=varying_parameters,
+            objectives=objectives,
+            analyzed_parameters=analyzed_parameters,
+            parameter_values=parameter_values,
+            evaluations=evaluations,
+            custom_parameters=custom_parameters,
+        )
+        if "_id" in trial_dict:
+            trial._gen_id = trial_dict["_id"]
+        if "_ignored" in trial_dict:
+            trial._ignored = trial_dict["_ignored"]
+        if "_ignored_reason" in trial_dict:
+            trial._ignored_reason = trial_dict["_ignored_reason"]
+        if "_status" in trial_dict:
+            trial._status = trial_dict["_status"]
+        for custom_param in custom_parameters:
+            setattr(trial, custom_param.name, trial_dict[custom_param.name])
+
+        return trial
+
     @property
     def varying_parameters(self) -> List[VaryingParameter]:
         """Get the list of varying parameters."""
@@ -129,6 +199,15 @@ def index(self) -> int:
     def index(self, value):
         self._index = value
 
+    @property
+    def gen_id(self) -> int:
+        """Get the index of the trial."""
+        return self._gen_id
+
+    @gen_id.setter
+    def gen_id(self, value):
+        self._gen_id = value
+
     @property
     def ignored(self) -> bool:
         """Get whether the trial is ignored by the generator."""
@@ -225,7 +304,8 @@ def objectives_as_dict(self) -> Dict:
         params = {}
         for obj in self._objectives:
             ev = self._mapped_evaluations[obj.name]
-            params[obj.name] = (ev.value, ev.sem)
+            if ev is not None:
+                params[obj.name] = ev.value
         return params
 
     def analyzed_parameters_as_dict(self) -> Dict:
@@ -237,5 +317,13 @@ def analyzed_parameters_as_dict(self) -> Dict:
         params = {}
         for par in self._analyzed_parameters:
             ev = self._mapped_evaluations[par.name]
-            params[par.name] = (ev.value, ev.sem)
+            if ev is not None:
+                params[par.name] = ev.value
+        return params
+
+    def custom_parameters_as_dict(self) -> Dict:
+        """Get a mapping between names and values of the custom parameters."""
+        params = {}
+        for param in self._custom_parameters:
+            params[param.name] = getattr(self, param.name)
         return params

optimas/gen_functions.py

@@ -20,7 +20,7 @@ def persistent_generator(H, persis_info, gen_specs, libE_info):
     """Generate and launch evaluations with the optimas generators.
 
     This function gets the generator object and uses it to generate new
-    evaluations via the `ask` method. Once finished, the result of the
+    evaluations via the `ask_trials` method. Once finished, the result of the
     evaluations is communicated back to the generator via the `tell` method.
 
     This is a persistent generator function, i.e., it is called by a dedicated
@@ -68,7 +68,7 @@ def persistent_generator(H, persis_info, gen_specs, libE_info):
         # Ask the optimizer to generate `batch_size` new points
         # Store this information in the format expected by libE
         H_o = np.zeros(number_of_gen_points, dtype=gen_specs["out"])
-        generated_trials = generator.ask(number_of_gen_points)
+        generated_trials = generator.ask_trials(number_of_gen_points)
         for i, trial in enumerate(generated_trials):
             for var, val in zip(
                 trial.varying_parameters, trial.parameter_values
@@ -107,8 +107,8 @@ def persistent_generator(H, persis_info, gen_specs, libE_info):
                         y = calc_in[par.name][i]
                         ev = Evaluation(parameter=par, value=y)
                         trial.complete_evaluation(ev)
-                # Register trial with unknown SEM
-                generator.tell([trial])
+                generator.tell_trials([trial])
+
             # Set the number of points to generate to that number:
             number_of_gen_points = min(n + n_failed_gens, max_evals - n_gens)
             n_failed_gens = 0

optimas/generators/__init__.py

@@ -22,7 +22,7 @@
 from .grid_sampling import GridSamplingGenerator
 from .line_sampling import LineSamplingGenerator
 from .random_sampling import RandomSamplingGenerator
-
+from .external import ExternalGenerator
 
 __all__ = [
     "AxSingleFidelityGenerator",
@@ -32,4 +32,5 @@
     "GridSamplingGenerator",
     "LineSamplingGenerator",
     "RandomSamplingGenerator",
+    "ExternalGenerator",
 ]

optimas/generators/ax/developer/multitask.py

@@ -196,10 +196,13 @@ def __init__(
         model_save_period: Optional[int] = 5,
         model_history_dir: Optional[str] = "model_history",
     ) -> None:
+        # As trial parameters these get written to history array
+        # Ax trial_index and arm toegther locate a point
+        # Multiple points (Optimas trials) can share the same Ax trial_index
         custom_trial_parameters = [
             TrialParameter("arm_name", "ax_arm_name", dtype="U32"),
             TrialParameter("trial_type", "ax_trial_type", dtype="U32"),
-            TrialParameter("trial_index", "ax_trial_index", dtype=int),
+            TrialParameter("ax_trial_id", "ax_trial_index", dtype=int),
         ]
         self._check_inputs(varying_parameters, objectives, lofi_task, hifi_task)
         super().__init__(
@@ -260,23 +263,38 @@ def _check_inputs(
             "to the number of high-fidelity trials"
         )
 
-    def _ask(self, trials: List[Trial]) -> List[Trial]:
-        """Fill in the parameter values of the requested trials."""
-        for trial in trials:
+    def suggest(self, num_points: Optional[int]) -> List[dict]:
+        """Request the next set of points to evaluate."""
+        points = []
+        for _ in range(num_points):
             next_trial = self._get_next_trial_arm()
             if next_trial is not None:
                 arm, trial_type, trial_index = next_trial
-                trial.parameter_values = [
-                    arm.parameters.get(var.name)
+                point = {
+                    var.name: arm.parameters.get(var.name)
                     for var in self._varying_parameters
-                ]
-                trial.trial_type = trial_type
-                trial.arm_name = arm.name
-                trial.trial_index = trial_index
-        return trials
-
-    def _tell(self, trials: List[Trial]) -> None:
+                }
+                # SH for VOCS standard these will need to be 'variables'
+                # For now much match the trial parameter names.
+                point["ax_trial_id"] = trial_index
+                point["arm_name"] = arm.name
+                point["trial_type"] = trial_type
+                points.append(point)
+        return points
+
+    def ingest(self, results: List[dict]) -> None:
         """Incorporate evaluated trials into experiment."""
+        # reconstruct Optimas trials
+        trials = []
+        for result in results:
+            trial = Trial.from_dict(
+                trial_dict=result,
+                varying_parameters=self._varying_parameters,
+                objectives=self._objectives,
+                analyzed_parameters=self._analyzed_parameters,
+                custom_parameters=self._custom_trial_parameters,
+            )
+            trials.append(trial)
         if self.gen_state == NOT_STARTED:
             self._incorporate_external_data(trials)
         else:
@@ -285,17 +303,18 @@ def _tell(self, trials: List[Trial]) -> None:
     def _incorporate_external_data(self, trials: List[Trial]) -> None:
         """Incorporate external data (e.g., from history) into experiment."""
         # Get trial indices.
+        # SH should have handling if ax_trial_ids are None...
         trial_indices = []
         for trial in trials:
-            trial_indices.append(trial.trial_index)
+            trial_indices.append(trial.ax_trial_id)
         trial_indices = np.unique(np.array(trial_indices))
 
         # Group trials by index.
         grouped_trials = {}
         for index in trial_indices:
             grouped_trials[index] = []
         for trial in trials:
-            grouped_trials[trial.trial_index].append(trial)
+            grouped_trials[trial.ax_trial_id].append(trial)
 
         # Add trials to experiment.
         for index in trial_indices: