Make the data plot into a 3D interactive plotly plot

regularization.py

@@ -3,14 +3,15 @@ from pathlib import Path
 import pickle
 
 import gradio as gr
+import inspect
+import io
+from jinja2 import Template
 import matplotlib.pyplot as plt
 import matplotlib.lines as mlines
-import inspect
 import numpy as np
 import pandas as pd
-import io
-from jinja2 import Template
 from PIL import Image
+import plotly.graph_objects as go
 import sklearn
 from sklearn.linear_model import LogisticRegression
 from sklearn.svm import LinearSVC
@@ -265,24 +266,39 @@ class Regularization:
 
     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
+        _, _, coef = make_regression(n_samples=200, n_features=2, noise=15, random_state=0, coef=True)
+
+        x1 = np.linspace(-1, 1, 50)
+        x2 = np.linspace(-1, 1, 50)
+        mesh_x1, mesh_x2 = np.meshgrid(x1, x2)
+        X = np.stack((mesh_x1.ravel(), mesh_x2.ravel()), axis=-1)
+        y = X @ coef
+        
+        z = y.reshape(mesh_x1.shape)
+
+        fig = go.Figure(data=go.Surface(
+            z=z, 
+            x=mesh_x1, 
+            y=mesh_x2,
+            colorscale='Viridis',
+            opacity=0.8
+        ))
+        
+        fig.update_layout(
+            title="Data",
+            scene={
+                "xaxis": {"title": "X1", "nticks": 6},
+                "yaxis": {"title": "X2", "nticks": 6}, 
+                "zaxis": {"title": "Y", "nticks": 6},
+                "camera": {"eye": {"x": -1.5, "y": -1.5, "z": 1.2}},
+            },
+            width=800,
+            height=600,
+        )
+        return fig
 
     def plot_strength_vs_weight(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)
         alphas = np.concat([np.zeros(1), np.logspace(-2, 2, 100)])
         if self.loss_type == "l2":
@@ -357,7 +373,7 @@ class Regularization:
                     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)
+                        self.data_3d_plot = gr.Plot(value=self.plot_data(), container=True)
                     with gr.Tab("Strength vs weight"):
                         self.strength_vs_weight = gr.Image(value=self.plot_strength_vs_weight(), container=True)