From ec0e0f130cace30258f91cc1fe965bbc76db1be5 Mon Sep 17 00:00:00 2001
From: Karlson Pfannschmidt <kiudee@mail.upb.de>
Date: Sun, 6 Feb 2022 18:30:10 +0100
Subject: [PATCH] Add plot_optima to CLI

---
 tune/cli.py   |  2 ++
 tune/local.py | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/tune/cli.py b/tune/cli.py
index c7ac5e4..9fea28d 100644
--- a/tune/cli.py
+++ b/tune/cli.py
@@ -426,6 +426,8 @@ def local(  # noqa: C901
                 plot_results(
                     optimizer=opt,
                     result_object=result_object,
+                    iterations=np.array(performance)[:, 0],
+                    optima=np.array(optima),
                     plot_path=settings.get("plot_path", plot_path),
                     parameter_names=list(param_ranges.keys()),
                     confidence=settings.get("confidence", confidence),
diff --git a/tune/local.py b/tune/local.py
index eab63ac..7100ccb 100644
--- a/tune/local.py
+++ b/tune/local.py
@@ -19,7 +19,7 @@
 from skopt.space import Categorical, Dimension, Integer, Real, Space
 from skopt.utils import normalize_dimensions
 
-from tune.plots import plot_objective, plot_objective_1d
+from tune.plots import plot_objective, plot_objective_1d, plot_optima
 from tune.summary import confidence_intervals
 from tune.utils import TimeControl, confidence_to_mult, expected_ucb
 
@@ -518,6 +518,8 @@ def print_results(
 def plot_results(
     optimizer: Optimizer,
     result_object: OptimizeResult,
+    iterations: np.ndarray,
+    optima: np.ndarray,
     plot_path: str,
     parameter_names: Sequence[str],
     confidence: float = 0.9,
@@ -530,6 +532,10 @@ def plot_results(
         Fitted Optimizer object.
     result_object : scipy.optimize.OptimizeResult
         Result object containing the data and the last fitted model.
+    iterations : np.ndarray
+        Array containing the iterations at which optima were collected.
+    optima : np.ndarray
+        Array containing the predicted optimal parameters.
     plot_path : str
         Path to the directory to which the plots should be saved.
     parameter_names : Sequence of str
@@ -541,6 +547,8 @@ def plot_results(
     logger.debug("Starting to compute the next plot.")
     timestr = time.strftime("%Y%m%d-%H%M%S")
     dark_gray = "#36393f"
+
+    # First save the landscape:
     save_params = dict()
     if optimizer.space.n_dims == 1:
         fig, ax = plot_objective_1d(
@@ -569,6 +577,17 @@ def plot_results(
     logger.info(f"Saving a plot to {full_plotpath}.")
     plt.close(fig)
 
+    # Now plot the history of optima:
+    fig, ax = plot_optima(
+        iterations=iterations,
+        optima=optima,
+        space=optimizer.space,
+        parameter_names=parameter_names,
+    )
+    full_plotpath = plotpath / f"optima-{timestr}-{len(optimizer.Xi)}.png"
+    fig.savefig(full_plotpath, dpi=150, facecolor=dark_gray)
+    plt.close(fig)
+
 
 def run_match(
     rounds=10,