From d1f9d0a0c216e5c815414b142e7dd400c884e9d4 Mon Sep 17 00:00:00 2001
From: Filippo Airaldi <f.airaldi@tudelft.nl>
Date: Wed, 31 Jul 2024 10:32:32 +0200
Subject: [PATCH] small refinements

---
 benchmarking/plot.py              | 15 +++++++++++----
 examples/nonmyopic_acquisition.py | 24 ++++++++++--------------
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/benchmarking/plot.py b/benchmarking/plot.py
index e3f5eef..0606a91 100644
--- a/benchmarking/plot.py
+++ b/benchmarking/plot.py
@@ -8,6 +8,7 @@
 from functools import partial
 from math import ceil
 from typing import Literal, Optional
+from warnings import warn
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -337,7 +338,8 @@ def _format_row(
     problem = row.name[0]
     best_method_idx = getattr(row, f"arg{order}")()
     strs[best_method_idx] = f"\033[1;34m{strs[best_method_idx]}\033[0m"
-    best_method_data = src_data[(problem, methods[best_method_idx])]
+    best_method = methods[best_method_idx]
+    best_method_data = src_data[(problem, best_method)]
 
     # then, loop over the rest of the methods, and compare them in a pairwise
     # fashion via the Wilcoxon (one-sided) signed-rank test. If the null hypothesis
@@ -345,10 +347,15 @@ def _format_row(
     # highlight the other method in italic violet
     side = "less" if order == "min" else "greater"
     for i in filter(lambda i: i != best_method_idx, range(len(strs))):
-        other_method_data = src_data[(problem, methods[i])]
+        method = methods[i]
+        method_data = src_data[(problem, method)]
         try:
-            _, alpha = wilcoxon(best_method_data, other_method_data, alternative=side)
-        except ValueError:
+            _, alpha = wilcoxon(best_method_data, method_data, alternative=side)
+        except ValueError as e:
+            warn(
+                f"Exception during statistical test of `{best_method}` vs `{method}` "
+                f" for `{problem}`: {e}"
+            )
             alpha = float("nan")
         if alpha > threshold_alpha:
             strs[i] = f"\033[35m{strs[i]}\033[0m"
diff --git a/examples/nonmyopic_acquisition.py b/examples/nonmyopic_acquisition.py
index a42dc9a..493e812 100644
--- a/examples/nonmyopic_acquisition.py
+++ b/examples/nonmyopic_acquisition.py
@@ -65,8 +65,8 @@
 # because now the acquisition is a function of trajectories of query points, not just
 # the poins themselves. In other words, the input to the `forward` method of the
 # acquisition function is a tensor of shape n x horizon x d, where n is the number of
-# query points, horizon is the number of steps in the lookahead, and d is the dimension
-# of the input space.
+# query trajectories (not points), horizon is the number of steps in the lookahead, and
+# d is the dimension of the input space.
 # Therefore, to plot it, we generate all possible combinations of trajectories,
 # evaluate them, and plot the best one.
 acqfun = Ms(
@@ -104,18 +104,14 @@
 ax.plot(train_X.cpu().squeeze(), train_Y.cpu().squeeze(), "o", label=None, color="C0")
 ax.plot(X, y_hat.squeeze(), label=r"Surrogate model $\hat{f}(x) \pm s(x)$", color="C1")
 ax.fill_between(X, (y_hat - s).squeeze(), (y_hat + s).squeeze(), color="C1", alpha=0.2)
-names = [
-    r"Vanilla acquisition $\Lambda$",
-]
-data = [(a, x_opt, a_opt)]
-for i, (name, (a_, x_opt_, a_opt_)) in enumerate(zip(names, data)):
-    c = f"C{i + 3}"
-    a_ = a_.cpu().squeeze()
-    x_opt_ = x_opt_.cpu().squeeze()
-    a_opt_ = a_opt_.cpu().squeeze()
-    a_min = a_.amin()
-    ax.plot(X_subset.cpu(), a_ - a_min, "--", lw=1, label=name, color=c)
-    ax.plot(x_opt_, a_opt_ - a_min, "*", markersize=17, color=c)
+a = a.cpu().squeeze()
+x_opt = x_opt.cpu().squeeze()
+a_opt = a_opt.cpu().squeeze()
+a_min = a.amin()
+ax.plot(
+    X_subset.cpu(), a - a_min, "--", lw=1, label="Nonmyopic acquisition", color="C3"
+)
+ax.plot(x_opt, a_opt - a_min, "*", markersize=17, color="C3")
 ax.set_xlim(lb, ub)
 ax.set_ylim(0, 2.5)
 ax.legend()