[WIP] Attempt to make optimas generators compatible with generator standard

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

@@ -81,6 +81,69 @@ 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(),
+            "_id": self._index,
+            "_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]],
+    ) -> "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.
+        """
+        # 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._index = 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"]
+        # TODO: Handle custom parameters
+        return trial
+
     @property
     def varying_parameters(self) -> List[VaryingParameter]:
         """Get the list of varying parameters."""
@@ -225,7 +288,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, ev.sem)
         return params
 
     def analyzed_parameters_as_dict(self) -> Dict:

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
@@ -108,7 +108,7 @@ def persistent_generator(H, persis_info, gen_specs, libE_info):
                         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/ax/developer/multitask.py

@@ -162,7 +162,7 @@ def _check_inputs(
             "to the number of high-fidelity trials"
         )
 
-    def _ask(self, trials: List[Trial]) -> List[Trial]:
+    def ask(self, trials: List[Trial]) -> List[Trial]:
         """Fill in the parameter values of the requested trials."""
         for trial in trials:
             next_trial = self._get_next_trial_arm()
@@ -177,7 +177,7 @@ def _ask(self, trials: List[Trial]) -> List[Trial]:
                 trial.trial_index = trial_index
         return trials
 
-    def _tell(self, trials: List[Trial]) -> None:
+    def tell(self, trials: List[Trial]) -> None:
         """Incorporate evaluated trials into experiment."""
         if self.gen_state == NOT_STARTED:
             self._incorporate_external_data(trials)

optimas/generators/ax/service/base.py

@@ -142,19 +142,22 @@ def model(self) -> AxModelManager:
         """Get access to the underlying model using an `AxModelManager`."""
         return self._model
 
-    def _ask(self, trials: List[Trial]) -> List[Trial]:
-        """Fill in the parameter values of the requested trials."""
-        for trial in trials:
+    def ask(self, num_points: Optional[int]) -> List[dict]:
+        """Request the next set of points to evaluate."""
+        points = []
+        for _ in range(num_points):
             parameters, trial_id = self._ax_client.get_next_trial(
                 fixed_features=self._fixed_features
             )
-            trial.parameter_values = [
-                parameters.get(var.name) for var in self._varying_parameters
-            ]
-            trial.ax_trial_id = trial_id
-        return trials
+            point = {
+                var.name: parameters.get(var.name)
+                for var in self._varying_parameters
+            }
+            point["_id"] = trial_id
+            points.append(point)
+        return points
 
-    def _tell(self, trials: List[Trial]) -> None:
+    def tell(self, trials: List[Trial]) -> None:
         """Incorporate evaluated trials into Ax client."""
         for trial in trials:
             try:

optimas/generators/base.py

@@ -187,7 +187,7 @@ def n_evaluated_trials(self) -> int:
                 n_evaluated += 1
         return n_evaluated
 
-    def ask(self, n_trials: int) -> List[Trial]:
+    def ask_trials(self, n_trials: int) -> List[Trial]:
         """Ask the generator to suggest the next ``n_trials`` to evaluate.
 
         Parameters
@@ -204,18 +204,21 @@ def ask(self, n_trials: int) -> List[Trial]:
         # Generate as many trials as needed and add them to the queue.
         if n_trials > self.n_queued_trials:
             n_gen = n_trials - self.n_queued_trials
+            # Ask generator for n_gen points, using the standardized API
+            gen_points = self.ask(n_gen)
+            # Convert the points to the Trial format
             gen_trials = []
-            for _ in range(n_gen):
-                gen_trials.append(
-                    Trial(
-                        varying_parameters=self._varying_parameters,
-                        objectives=self._objectives,
-                        analyzed_parameters=self._analyzed_parameters,
-                        custom_parameters=self._custom_trial_parameters,
-                    )
+            for point in gen_points:
+                trial = Trial(
+                    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,
                 )
-            # Ask the generator to fill them.
-            gen_trials = self._ask(gen_trials)
+                gen_trials.append(trial)
             # Keep only trials that have been given data.
             for trial in gen_trials:
                 if len(trial.parameter_values) > 0:
@@ -236,7 +239,7 @@ def ask(self, n_trials: int) -> List[Trial]:
                 trials.append(trial)
         return trials
 
-    def tell(
+    def tell_trials(
         self, trials: List[Trial], allow_saving_model: Optional[bool] = True
     ) -> None:
         """Give trials back to generator once they have been evaluated.
@@ -250,7 +253,10 @@ def tell(
             incorporating the evaluated trials. By default ``True``.
 
         """
-        self._tell(trials)
+        # Feed data to the generator, using the standardized API
+        self.tell([trial.to_dict() for trial in trials])
+
+        # Perform additional checks that rely on the trial format
         for trial in trials:
             if trial not in self._given_trials:
                 self._add_external_evaluated_trial(trial)
@@ -290,7 +296,7 @@ def incorporate_history(self, history: np.ndarray) -> None:
         trials = self._create_trials_from_external_data(
             history_ended, ignore_unrecognized_parameters=True
         )
-        self.tell(trials, allow_saving_model=False)
+        self.tell_trials(trials, allow_saving_model=False)
         # Communicate to history array whether the trial has been ignored.
         for trial in trials:
             idxs = np.where(history["trial_index"] == trial.index)[0]
@@ -306,7 +312,7 @@ def attach_trials(
 
         The given trials are placed at the top of the queue of trials that
         will be proposed by the generator (that is, they will be the first
-        ones to be proposed the next time that `ask` is called).
+        ones to be proposed the next time that `ask_trials` is called).
 
         Parameters
         ----------
@@ -409,7 +415,7 @@ def _add_trial_to_queue(
 
         By default, the trial will be appended to the end of the queue, unless
         a `queue_index` is given. Trials at the top of the queue will be the
-        first ones to be given for evaluation when `ask` is called.
+        first ones to be given for evaluation when `ask_trials` is called.
 
         Parameters
         ----------
@@ -579,28 +585,27 @@ def get_libe_specs(self) -> Dict:
         libE_specs = {}
         return libE_specs
 
-    def _ask(self, trials: List[Trial]) -> List[Trial]:
-        """Ask method to be implemented by the Generator subclasses.
+    def ask(self, num_points: Optional[int]) -> List[dict]:
+        """Request the next set of points to evaluate.
 
         Parameters
         ----------
-        trials : list of Trial
-            A list with as many trials as requested to the generator. The
-            trials do not yet contain the values of the varying parameters.
-            These values should instead be supplied in this method.
+        num_points : int
+            Number of points to generate.
 
         """
-        return trials
+        return []
 
-    def _tell(self, trials: List[Trial]) -> None:
-        """Tell method to be implemented by the Generator subclasses.
+    def tell(self, results: List[dict]) -> None:
+        """
+        Send the results of evaluations to the generator.
 
         Parameters
         ----------
-        trials : list of Trial
-            A list with all evaluated trials. All evaluations included in the
-            trials should be incorporated to the generator model in this
-            method.
+        results : list of dict
+            List with the results of the evaluations.
+            All evaluations included in the results should be incorporated
+            to the generator model in this method.
 
         """
         pass

optimas/generators/grid_sampling.py

@@ -67,15 +67,14 @@ def _create_configurations(self) -> None:
             all_configs.append(config)
         self._all_configs = all_configs
 
-    def _ask(self, trials: List[Trial]) -> List[Trial]:
-        """Fill in the parameter values of the requested trials."""
-        for trial in trials:
+    def ask(self, num_points: Optional[int]) -> List[dict]:
+        """Request the next set of points to evaluate."""
+        points = []
+        for _ in range(num_points):
             if self._all_configs:
                 config = self._all_configs.pop(0)
-                trial.parameter_values = [
-                    config[var.name] for var in trial.varying_parameters
-                ]
-        return trials
+                points.append(config)
+        return points
 
     def _mark_trial_as_failed(self, trial: Trial):
         """No need to do anything, since there is no surrogate model."""

optimas/generators/line_sampling.py

@@ -103,15 +103,14 @@ def _create_configurations(self) -> None:
         # Store configurations.
         self._all_configs = all_configs
 
-    def _ask(self, trials: List[Trial]) -> List[Trial]:
-        """Fill in the parameter values of the requested trials."""
-        for trial in trials:
+    def ask(self, num_points: Optional[int]) -> List[dict]:
+        """Request the next set of points to evaluate."""
+        points = []
+        for _ in range(num_points):
             if self._all_configs:
                 config = self._all_configs.pop(0)
-                trial.parameter_values = [
-                    config[var.name] for var in trial.varying_parameters
-                ]
-        return trials
+                points.append(config)
+        return points
 
     def _mark_trial_as_failed(self, trial: Trial):
         """No need to do anything, since there is no surrogate model."""

optimas/generators/random_sampling.py

@@ -57,13 +57,16 @@ def __init__(
         self._rng = np.random.default_rng(seed)
         self._define_generator_parameters()
 
-    def _ask(self, trials: List[Trial]) -> List[Trial]:
-        """Fill in the parameter values of the requested trials."""
-        n_trials = len(trials)
-        configs = self._generate_sampling[self._distribution](n_trials)
-        for trial, config in zip(trials, configs):
-            trial.parameter_values = config
-        return trials
+    def ask(self, num_points: Optional[int]) -> List[dict]:
+        """Request the next set of points to evaluate."""
+        configs = self._generate_sampling[self._distribution](num_points)
+        points = []
+        for config in configs:
+            point = {}
+            for var, value in zip(self._varying_parameters, config):
+                point[var.name] = value
+            points.append(point)
+        return points
 
     def _check_inputs(
         self,