Split the regularization and data plots into two tabs

regularization.py

@@ -134,7 +134,7 @@ class Regularization:
 
         self.plot_regularization_path = False
 
-    def plot(self):
+    def plot_regularization_contour(self):
         '''
         '''
 
@@ -192,23 +192,22 @@ class Regularization:
         print(loss_levels)
 
         # plot contour plots
-        # fig = plt.figure(figsize=(5, 5))
-        fig, axs = plt.subplots(1, 2, figsize=(10, 5))
-        axs[0].set_title("")
-        axs[0].set_xlabel("w1")
-        axs[0].set_ylabel("w2")
+        fig, ax = plt.subplots(figsize=(8, 8))
+        ax.set_title("")
+        ax.set_xlabel("w1")
+        ax.set_ylabel("w2")
 
         cmap = plt.get_cmap("viridis")   
         N = len(reg_levels)
         colors = [cmap(i / (N - 1)) for i in range(N)]
 
         # regularizer contours
-        cs1 = axs[0].contour(W1, W2, regs, levels=reg_levels, colors=colors, linestyles="dashed")
-        axs[0].clabel(cs1, inline=True, fontsize=8) # show contour levels
+        cs1 = ax.contour(W1, W2, regs, levels=reg_levels, colors=colors, linestyles="dashed")
+        ax.clabel(cs1, inline=True, fontsize=8) # show contour levels
 
         # loss contours
-        cs2 = axs[0].contour(W1, W2, losses, levels=loss_levels, colors=colors[::-1])
-        axs[0].clabel(cs2, inline=True, fontsize=8)
+        cs2 = ax.contour(W1, W2, losses, levels=loss_levels, colors=colors[::-1])
+        ax.clabel(cs2, inline=True, fontsize=8)
 
         # plot path
         if self.plot_regularization_path:
@@ -223,7 +222,7 @@ class Regularization:
                 path_w.append(stacked[mask][idx])
             
             path_w = np.array(path_w)
-            axs[0].plot(path_w[:, 0], path_w[:, 1], "r-")
+            ax.plot(path_w[:, 0], path_w[:, 1], "r-")
 
         # custom legend
         loss_line = mlines.Line2D([], [], color='black', linestyle='-', label='loss')
@@ -232,13 +231,7 @@ class Regularization:
         if self.plot_regularization_path:
             path_line = mlines.Line2D([], [], color='red', linestyle='-', label='regularization path')
             handles.append(path_line)
-        axs[0].legend(handles=handles)
-
-        # plot data points
-        axs[1].set_xlabel("X1")
-        axs[1].set_ylabel("X2")
-        sc = axs[1].scatter(X[:, 0], X[:, 1], c=y, cmap="viridis")
-        fig.colorbar(sc, ax=axs[1], label="y")
+        ax.legend(handles=handles)
 
         # plot solutions
         #for alpha, w, norm, mse in solutions:
@@ -253,42 +246,61 @@ class Regularization:
 
         return img
 
+    def plot_data(self):
+        # make sure the data is the same as the one used in plot_regularization_contour
+        X, y = make_regression(n_samples=200, n_features=2, noise=15, random_state=0)
+        fig, ax = plt.subplots(figsize=(8, 8))
+
+        # plot data points
+        ax.set_xlabel("X1")
+        ax.set_ylabel("X2")
+        sc = ax.scatter(X[:, 0], X[:, 1], c=y, cmap="viridis")
+        fig.colorbar(sc, ax=ax, label="y")
+
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+        plt.close(fig)
+        buf.seek(0)
+        img = Image.open(buf)
+
+        return img
+
     def update_loss(self, loss_type):
         self.loss_type = loss_type
         self.loss = self.losses[loss_type]
-        return self.plot()
+        return self.plot_regularization_contour()
 
     def update_regularizer(self, reg_type):
         self.reg_type = reg_type
         self.regularizer = self.regularizers[reg_type]
 
-        return self.plot()
+        return self.plot_regularization_contour()
 
     def update_reg_levels(self, reg_levels):
         self.reg_levels = [float(reg_level) for reg_level in reg_levels.split(",")]
 
-        return self.plot()
+        return self.plot_regularization_contour()
 
     def update_w1_range(self, w1_range):
         self.w1_range = [float(w1) for w1 in w1_range.split(",")]
         logger.info("Updated w1 range to " + str(self.w1_range))
 
-        return self.plot()
+        return self.plot_regularization_contour()
 
     def update_w2_range(self, w2_range):
         self.w2_range = [float(w2) for w2 in w2_range.split(",")]
         logger.info("Updated w2 range to " + str(self.w2_range))
 
-        return self.plot()
+        return self.plot_regularization_contour()
 
     def update_resolution(self, num_dots):
         self.num_dots = num_dots
         logger.info("updated resolution to " + str(num_dots))
-        return self.plot()
+        return self.plot_regularization_contour()
 
     def update_plot_path(self, plot_path):
         self.plot_regularization_path = plot_path
-        return self.plot()
+        return self.plot_regularization_contour()
 
     def launch(self):
         # build the Gradio interface
@@ -299,7 +311,10 @@ class Regularization:
             # GUI elements and layout 
             with gr.Row():
                 with gr.Column(scale=2):
-                    self.data_image = gr.Image(value=self.plot(), container=True)
+                    with gr.Tab("Regularization contour"):
+                        self.regularization_contour = gr.Image(value=self.plot_regularization_contour(), container=True)
+                    with gr.Tab("Data"):
+                        self.data_image = gr.Image(value=self.plot_data(), container=True)
 
                 with gr.Column(scale=1): 
                     with gr.Tab("Settings"):
@@ -366,23 +381,23 @@ class Regularization:
             
                     # event handlers for GUI elements
                     loss_type.change(fn=self.update_loss, inputs=loss_type,
-                                     outputs=self.data_image)
+                                     outputs=self.regularization_contour)
                     regularizer_type.change(fn=self.update_regularizer,
-                                            inputs=regularizer_type, outputs=self.data_image)
+                                            inputs=regularizer_type, outputs=self.regularization_contour)
 
                     reg_textbox.submit(self.update_reg_levels, inputs=reg_textbox,
-                                       outputs=self.data_image)
+                                       outputs=self.regularization_contour)
 
                     w1_textbox.submit(self.update_w1_range, inputs=w1_textbox,
-                                       outputs=self.data_image)
+                                       outputs=self.regularization_contour)
 
                     w2_textbox.submit(self.update_w2_range, inputs=w2_textbox,
-                                       outputs=self.data_image)
+                                       outputs=self.regularization_contour)
 
-                    slider.change(self.update_resolution, inputs=slider, outputs=self.data_image)
+                    slider.change(self.update_resolution, inputs=slider, outputs=self.regularization_contour)
 
                     path_checkbox.change(
-                        self.update_plot_path, inputs=path_checkbox, outputs=self.data_image
+                        self.update_plot_path, inputs=path_checkbox, outputs=self.regularization_contour
                     )
 
         demo.launch()