ESSOptimizer: Include results for local searches in OptimizeResult

pypesto/optimize/ess/ess.py

@@ -6,7 +6,7 @@
 import enum
 import logging
 import time
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, Optional, Union
 from warnings import warn
 
 import numpy as np
@@ -65,6 +65,7 @@ def __init__(
         n_procs=None,
         n_threads=None,
         max_walltime_s=None,
+        result_includes_refset: bool = False,
     ):
         """Construct new ESS instance.
 
@@ -113,6 +114,10 @@ def __init__(
         history:
             History of the best values/parameters found so far.
             (Monotonously decreasing objective values.)
+        result_includes_refset:
+            Whether the :meth:`minimize` result should include the final
+            RefSet, or just the local search results and the overall best
+            parameters.
         """
         if max_eval is None and max_walltime_s is None and max_iter is None:
             # in this case, we'd run forever
@@ -151,6 +156,7 @@ def __init__(
         self.logger = logging.getLogger(
             f"{self.__class__.__name__}-{id(self)}"
         )
+        self._result_includes_refset = result_includes_refset
 
     def _initialize(self):
         """(Re-)Initialize."""
@@ -160,8 +166,8 @@ def _initialize(self):
         self.x_best: Optional[np.array] = None
         # Overall best function value found so far
         self.fx_best: float = np.inf
-        # Final parameters from local searches
-        self.local_solutions: List[np.array] = []
+        # Results from local searches
+        self.local_solutions: list[OptimizerResult] = []
         # Index of current iteration
         self.n_iter: int = 0
         # ESS iteration at which the last local search took place
@@ -313,25 +319,29 @@ def _create_result(self) -> pypesto.Result:
             **common_result_fields,
         )
         optimizer_result.update_to_full(result.problem)
-        # TODO DW: Create a single History with the global best?
         result.optimize_result.append(optimizer_result)
 
-        # save refset
-        for i in range(self.refset.dim):
+        # save local solutions
+        for i, optimizer_result in enumerate(self.local_solutions):
             i_result += 1
-            result.optimize_result.append(
-                pypesto.OptimizerResult(
-                    id=str(i_result),
-                    x=self.refset.x[i],
-                    fval=self.refset.fx[i],
-                    message=f"RefSet[{i}]",
-                    **common_result_fields,
+            optimizer_result.id = f"Local solution {i}"
+            optimizer_result.optimizer = str(self.local_optimizer)
+            result.optimize_result.append(optimizer_result)
+
+        if self._result_includes_refset:
+            # save refset
+            for i in range(self.refset.dim):
+                i_result += 1
+                result.optimize_result.append(
+                    pypesto.OptimizerResult(
+                        id=str(i_result),
+                        x=self.refset.x[i],
+                        fval=self.refset.fx[i],
+                        message=f"RefSet[{i}]",
+                        **common_result_fields,
+                    )
                 )
-            )
-            result.optimize_result[-1].update_to_full(result.problem)
-
-        # TODO DW: also save local solutions?
-        #  (need to track fvals or re-evaluate)
+                result.optimize_result[-1].update_to_full(result.problem)
 
         return result
 
@@ -370,20 +380,20 @@ def _get_remaining_eval(self):
             return np.inf
         return self.max_eval - self.evaluator.n_eval
 
-    def _combine_solutions(self) -> Tuple[np.array, np.array]:
+    def _combine_solutions(self) -> tuple[np.array, np.array]:
         """Combine solutions and evaluate.
 
-        Creates the next generation from the RefSet by pair-wise combinations
+        Creates the next generation from the RefSet by pair-wise combination
         of all RefSet members. Creates ``RefSet.dim ** 2 - RefSet.dim`` new
         parameter vectors, tests them, and keeps the best child of each parent.
 
         Returns
         -------
         y:
-            Contains the best of all pairwise combinations of all RefSet
-            members, for each RefSet members.
+            The next generation of parameter vectors
+            (`dim_refset` x `dim_problem`).
         fy:
-            The corresponding objective values.
+            The objective values corresponding to the parameters in `y`.
         """
         y = np.zeros(shape=(self.refset.dim, self.evaluator.problem.dim))
         fy = np.full(shape=self.refset.dim, fill_value=np.inf)
@@ -471,8 +481,10 @@ def _do_local_search(
             #  optima found so far
             min_distances = np.array(
                 np.min(
-                    np.linalg.norm(y_i - local_solution)
-                    for local_solution in self.local_solutions
+                    np.linalg.norm(
+                        y_i - optimizer_result.x[optimizer_result.free_indices]
+                    )
+                    for optimizer_result in self.local_solutions
                 )
                 for y_i in x_best_children
             )
@@ -521,15 +533,11 @@ def _do_local_search(
                 f"{optimizer_result.exitflag}: {optimizer_result.message}"
             )
             if np.isfinite(optimizer_result.fval):
-                local_solution_x = optimizer_result.x[
-                    optimizer_result.free_indices
-                ]
-                local_solution_fx = optimizer_result.fval
-
-                self.local_solutions.append(local_solution_x)
+                self.local_solutions.append(optimizer_result)
 
                 self._maybe_update_global_best(
-                    local_solution_x, local_solution_fx
+                    optimizer_result.x[optimizer_result.free_indices],
+                    optimizer_result.fval,
                 )
                 break
 

pypesto/optimize/ess/sacess.py

@@ -543,25 +543,16 @@ def run(
             )
         )
 
-        ess_results = pypesto.Result(problem=problem)
         ess = self._setup_ess(startpoint_method)
 
         # run ESS until exit criteria are met, but start at least one iteration
         while self._keep_going() or ess.n_iter == 0:
             # perform one ESS iteration
             ess._do_iteration()
 
-            # TODO maybe not in every iteration?
-            # drop all but the 50 best results
-            cur_ess_results = ess._create_result()
-            ess_results.optimize_result.append(
-                cur_ess_results.optimize_result,
-                prefix=f"{self._worker_idx}_{ess.n_iter}_",
-            )
-            ess_results.optimize_result.list = (
-                ess_results.optimize_result.list[:50]
-            )
             if self._tmp_result_file:
+                # TODO maybe not in every iteration?
+                ess_results = ess._create_result()
                 write_result(
                     ess_results,
                     self._tmp_result_file,
@@ -580,6 +571,7 @@ def run(
                 f"sacess worker {self._worker_idx} iteration {ess.n_iter} "
                 f"(best: {self._best_known_fx})."
             )
+
         ess.history.finalize(exitflag=ess.exit_flag.name)
         self._manager._result_queue.put(ess.history)
         ess._report_final()