Add plot options

gp_visualizer.py

@@ -88,6 +88,13 @@ class GPVisualizer:
 
         self.kernel = eval_kernel(self.DEFAULT_KERNEL)
 
+        self.plot_options = {
+            "show_training_data": True,
+            "show_confidence_interval": True,
+            "show_true_function": True,
+            "show_predictions": True,
+        }
+
         self.css = """
 .hidden-button {
     display: none;
@@ -127,16 +134,21 @@ class GPVisualizer:
 
         R2 = gpr.score(X_tr, y_tr)
 
-        if len(X_tr) > 1:
-            plt.scatter(X_tr.flatten(), y_tr, label='training data (R2=%.2f)' % (R2))
-        else:
-            plt.scatter(X_tr.flatten(), y_tr, label='training data')
+        if self.plot_options["show_training_data"]:
+            if len(X_tr) > 1:
+                plt.scatter(X_tr.flatten(), y_tr, label='training data (R2=%.2f)' % (R2))
+            else:
+                plt.scatter(X_tr.flatten(), y_tr, label='training data')
+
+        if self.plot_options["show_true_function"]:
+            plt.plot(X_ts.flatten(), np.sin(2*np.pi*X_ts.flatten()), color='red', label='true function')
 
-        plt.plot(X_ts.flatten(), np.sin(2*np.pi*X_ts.flatten()), color='red', label='true function')
-        plt.scatter(X_ts.flatten(), y_pred, marker='+', label='predictions')
+        if self.plot_options["show_predictions"]:
+            plt.scatter(X_ts.flatten(), y_pred, marker='+', label='predictions')
 
-        plt.fill_between(X_ts.flatten(), y_pred - 1.96*y_std, y_pred + 1.96*y_std, alpha=0.5,
-                         label='95% confidence interval')
+        if self.plot_options["show_confidence_interval"]:
+            plt.fill_between(X_ts.flatten(), y_pred - 1.96*y_std, y_pred + 1.96*y_std, alpha=0.5,
+                             label='95% confidence interval')
 
         plt.legend()
 
@@ -190,6 +202,12 @@ class GPVisualizer:
         self.kernel = eval_kernel(kernel_spec)
         return self.plot()
 
+    def update_plot_options(self, **kwargs):
+        for key, value in kwargs.items():
+            if key in self.plot_options:
+                self.plot_options[key] = value
+        return self.plot()
+
     def launch(self):
         # build the Gradio interface
         with gr.Blocks(css=self.css) as demo:
@@ -220,6 +238,13 @@ class GPVisualizer:
                             interactive=True,
                         )
 
+                        # plot show options
+                        with gr.Group():
+                            show_training_data = gr.Checkbox(label="Show training data", value=True)
+                            show_confidence_interval = gr.Checkbox(label="Show confidence interval", value=True)
+                            show_true_function = gr.Checkbox(label="Show true function", value=True)
+                            show_predictions = gr.Checkbox(label="Show predictions", value=True)
+
                         #gr.Markdown(''.join(open('kernel_examples.md', 'r').readlines()))
 
                     with gr.Tab("Export"):
@@ -250,10 +275,33 @@ class GPVisualizer:
                     kernel_spec.submit(
                         fn=self.update_kernel_spec,
                         inputs=kernel_spec,
-                        outputs=(self.canvas)
+                        outputs=[self.canvas],
+                    )
+
+                    # plot options
+                    show_training_data.change(
+                        fn=lambda show: self.update_plot_options(show_training_data=show),
+                        inputs=show_training_data,
+                        outputs=[self.canvas],
+                    )
+                    show_confidence_interval.change(
+                        fn=lambda show: self.update_plot_options(show_confidence_interval=show),
+                        inputs=show_confidence_interval,
+                        outputs=[self.canvas],
+                    )
+                    show_true_function.change(
+                        fn=lambda show: self.update_plot_options(show_true_function=show),
+                        inputs=show_true_function,
+                        outputs=[self.canvas],
+                    )
+                    show_predictions.change(
+                        fn=lambda show: self.update_plot_options(show_predictions=show),
+                        inputs=show_predictions,
+                        outputs=[self.canvas],
                     )
 
         demo.launch()
 
 visualizer = GPVisualizer(width=1200, height=900)
 visualizer.launch()
+