Added 3d Plot

app.py

@@ -98,7 +98,6 @@ X /= X.std(axis=0)
 edge_model.fit(X)
 
 
-
 from sklearn import cluster
 
 _, labels = cluster.affinity_propagation(edge_model.covariance_, random_state=0)
@@ -118,13 +117,10 @@ embedding = node_position_model.fit_transform(X.T).T
 
 import matplotlib.pyplot as plt
 from matplotlib.collections import LineCollection
+import plotly.graph_objs as go
 
-def visualize_stocks():
-    fig  = plt.figure(1, facecolor="w", figsize=(10, 8))
-    plt.clf()
-    ax = plt.axes([0.0, 0.0, 1.0, 1.0])
-    plt.axis("off")
 
+def visualize_stocks():
     # Plot the graph of partial correlations
     partial_correlations = edge_model.precision_.copy()
     d = 1 / np.sqrt(np.diag(partial_correlations))
@@ -133,72 +129,53 @@ def visualize_stocks():
     non_zero = np.abs(np.triu(partial_correlations, k=1)) > 0.02
 
     # Plot the nodes using the coordinates of our embedding
-    plt.scatter(
-        embedding[0], embedding[1], s=100 * d**2, c=labels, cmap=plt.cm.nipy_spectral
+    scatter = go.Scatter3d(
+        x=embedding[0],
+        y=embedding[1],
+        z=embedding[2],
+        mode="markers",
+        marker=dict(size=35 * d**2, color=labels, colorscale="Viridis"),
+        hovertext=names,
+        hovertemplate="%{hovertext}<br>",
     )
 
-    # Plot the edges
+    # # Plot the edges
     start_idx, end_idx = np.where(non_zero)
-    # a sequence of (*line0*, *line1*, *line2*), where::
-    #            linen = (x0, y0), (x1, y1), ... (xm, ym)
+    # print(non_zero, non_zero.shape)
+    # print(start_idx, start_idx.shape)
     segments = [
-        [embedding[:, start], embedding[:, stop]] for start, stop in zip(start_idx, end_idx)
+        dict(
+            x=[embedding[0][start], embedding[0][stop]],
+            y=[embedding[1][start], embedding[1][stop]],
+            z=[embedding[2][start], embedding[2][stop]],
+            colorscale="Hot",
+            color=np.abs(partial_correlations[start, stop]),
+            line=dict(width=10 * np.abs(partial_correlations[start, stop])),
+        )
+        for start, stop in zip(start_idx, end_idx)
     ]
-    values = np.abs(partial_correlations[non_zero])
-    lc = LineCollection(
-        segments, zorder=0, cmap=plt.cm.hot_r, norm=plt.Normalize(0, 0.7 * values.max())
-    )
-    lc.set_array(values)
-    lc.set_linewidths(15 * values)
-    ax.add_collection(lc)
-
-    # Add a label to each node. The challenge here is that we want to
-    # position the labels to avoid overlap with other labels
-    for index, (name, label, (x, y)) in enumerate(zip(names, labels, embedding.T)):
-
-        dx = x - embedding[0]
-        dx[index] = 1
-        dy = y - embedding[1]
-        dy[index] = 1
-        this_dx = dx[np.argmin(np.abs(dy))]
-        this_dy = dy[np.argmin(np.abs(dx))]
-        if this_dx > 0:
-            horizontalalignment = "left"
-            x = x + 0.002
-        else:
-            horizontalalignment = "right"
-            x = x - 0.002
-        if this_dy > 0:
-            verticalalignment = "bottom"
-            y = y + 0.002
-        else:
-            verticalalignment = "top"
-            y = y - 0.002
-        plt.text(
-            x,
-            y,
-            name,
-            size=10,
-            horizontalalignment=horizontalalignment,
-            verticalalignment=verticalalignment,
-            bbox=dict(
-                facecolor="w",
-                edgecolor=plt.cm.nipy_spectral(label / float(n_labels)),
-                alpha=0.6,
+    fig = go.Figure(data=[scatter])
+
+    for idx, segment in enumerate(segments, 1):
+        fig.add_trace(
+            go.Scatter3d(
+                x=segment["x"],  # x-coordinates of the line segment
+                y=segment["y"],  # y-coordinates of the line segment
+                z=segment["z"],  # z-coordinates of the line segment
+                mode="lines",  # type of the plot (line)
+                line=dict(
+                    color=segment["color"],  # color of the line
+                    colorscale=segment["colorscale"],  # color scale of the line
+                    width=segment["line"]["width"] * 2.5,  # width of the line
+                ),
+                hoverinfo="none",  # disable hover for the line segments
             ),
         )
-
-    plt.xlim(
-        embedding[0].min() - 0.15 * embedding[0].ptp(),
-        embedding[0].max() + 0.10 * embedding[0].ptp(),
-    )
-    plt.ylim(
-        embedding[1].min() - 0.03 * embedding[1].ptp(),
-        embedding[1].max() + 0.03 * embedding[1].ptp(),
-    )
+        fig.data[idx].showlegend = False
 
     return fig
-    
+
+
 import gradio as gr
 
 title = " 📈 Visualizing the stock market structure 📈"
@@ -206,14 +183,20 @@ title = " 📈 Visualizing the stock market structure 📈"
 with gr.Blocks(title=title) as demo:
     gr.Markdown(f"# {title}")
     gr.Markdown(" Data is of 56 stocks between the period of 2003 - 2008 <br>")
-    gr.Markdown(" Stocks the move in together with each other are grouped together in a cluster <br>")
+    gr.Markdown(
+        " Stocks the move in together with each other are grouped together in a cluster <br>"
+    )
 
-    gr.Markdown(" **[Demo is based on sklearn docs](https://scikit-learn.org/stable/auto_examples/applications/plot_stock_market.html)**")
+    gr.Markdown(
+        " **[Demo is based on sklearn docs](https://scikit-learn.org/stable/auto_examples/applications/plot_stock_market.html)**"
+    )
 
     for i in range(n_labels + 1):
-        gr.Markdown( f"Cluster {i + 1}: {', '.join(names[labels == i])}")
-        
+        gr.Markdown(f"Cluster {i + 1}: {', '.join(names[labels == i])}")
+
     btn = gr.Button(value="Visualize")
-    btn.click(visualize_stocks, outputs= gr.Plot(label='Visualizing stock into clusters') )
-    gr.Markdown( f"## In progress")
-demo.launch()
+    btn.click(
+        visualize_stocks, outputs=gr.Plot(label="Visualizing stock into clusters")
+    )
+    gr.Markdown(f"## In progress")
+demo.launch()