Repository: secondmind-labs/trieste

Issue:
additional tensorboard logs for batches
currently we have histograms of query points when it comes to batch BO that are useful for checking the diversity, that all seems to work fine - there are two more that would be quite useful:

1. histogram of pairwise distances between query points - its easy to compute pairwise euclidean distances with Scipy's pdist and then use `logging.histogram`  
2. pairplot with individual scatterplots for each input dimension and model predicted output - this one is a bit more work and we would need to use image summary  


Discussion:

Resolving patch:
diff --git a/trieste/ask_tell_optimization.py b/trieste/ask_tell_optimization.py
--- a/trieste/ask_tell_optimization.py
+++ b/trieste/ask_tell_optimization.py
@@ -25,6 +25,7 @@
 
 import numpy as np
 import tensorflow as tf
+from scipy.spatial.distance import pdist
 
 from . import logging
 from .acquisition.rule import AcquisitionRule, EfficientGlobalOptimization
@@ -457,6 +458,9 @@ def ask(self) -> TensorType:
                             logging.scalar(f"query_points/[{i}]", float(query_points[0, i]))
                         else:
                             logging.histogram(f"query_points/[{i}]", query_points[:, i])
+                    logging.histogram(
+                        "query_points/euclidean_distances", lambda: pdist(query_points)
+                    )
                 logging.scalar(
                     "wallclock/query_point_generation",
                     query_point_generation_timer.time,
diff --git a/trieste/bayesian_optimizer.py b/trieste/bayesian_optimizer.py
--- a/trieste/bayesian_optimizer.py
+++ b/trieste/bayesian_optimizer.py
@@ -40,6 +40,7 @@
 import dill
 import numpy as np
 import tensorflow as tf
+from scipy.spatial.distance import pdist
 
 from .acquisition.multi_objective import non_dominated
 
@@ -910,15 +911,22 @@ def _write_summary_step(
                     logging.scalar(f"query_points/[{i}]", float(query_points[0, i]))
                 else:
                     logging.histogram(f"query_points/[{i}]", query_points[:, i])
+            logging.histogram("query_points/euclidean_distances", lambda: pdist(query_points))
 
         if pd and sns and logging.include_summary("query_points/_pairplot"):
             columns = [f"x{i}" for i in range(tf.shape(query_points)[1])]
-            query_new_df = pd.DataFrame(query_points, columns=columns).applymap(float)
+            qp_preds = query_points
+            for tag in datasets:
+                qp_preds = tf.concat(
+                    [qp_preds, tf.cast(models[tag].predict(query_points)[0], query_points.dtype)], 1
+                )
+                columns.append(f"{tag} predicted")
+            query_new_df = pd.DataFrame(qp_preds, columns=columns).applymap(float)
             query_new_df["query points"] = "new"
             query_plot_df = pd.concat(
                 (query_plot_dfs.get(0), query_new_df), copy=False, ignore_index=True
             )
-            pairplot = sns.pairplot(query_plot_df, hue="query points")
+            pairplot = sns.pairplot(query_plot_df, hue="query points", height=2.25)
             padding = 0.025 * (self._search_space.upper - self._search_space.lower)
             upper_limits = self._search_space.upper + padding
             lower_limits = self._search_space.lower - padding
diff --git a/trieste/logging.py b/trieste/logging.py
--- a/trieste/logging.py
+++ b/trieste/logging.py
@@ -214,7 +214,7 @@ def pyplot(
         try:
             figure = evaluate_data(figure)
             with io.BytesIO() as buffer:
-                figure.savefig(buffer, format="png")
+                figure.savefig(buffer, dpi=150.0, format="png")
                 buffer.seek(0)
                 image = tf.image.decode_png(buffer.getvalue(), channels=4)
             image = tf.expand_dims(image, 0)


Test patch: