Update app.py

app.py

@@ -1,6 +1,6 @@
 import pandas as pd
 import networkx as nx
-import matplotlib.pyplot as plt
+import plotly.graph_objects as go
 from io import BytesIO
 from PIL import Image
 import gradio as gr
@@ -108,12 +108,12 @@ def filter_investors_by_country_and_industry(selected_country, selected_industry
     
     return list(investor_valuations.keys()), filtered_data
 
-# Function to generate the graph and return node information
-def generate_graph_and_get_node_info(selected_investors, filtered_data, clicked_node=None):
+# Function to generate the Plotly graph
+def generate_graph(selected_investors, filtered_data):
     if not selected_investors:
-        logger.warning("No investors selected. Returning None for graph.")
-        return None, "No investors selected."
-
+        logger.warning("No investors selected. Returning empty figure.")
+        return go.Figure()
+    
     investor_company_mapping_filtered = build_investor_company_mapping(filtered_data)
     filtered_mapping = {inv: investor_company_mapping_filtered[inv] for inv in selected_investors if inv in investor_company_mapping_filtered}
     
@@ -125,66 +125,94 @@ def generate_graph_and_get_node_info(selected_investors, filtered_data, clicked_
         for company in companies:
             G.add_edge(investor, company)
     
-    # Node size based on valuation
-    max_valuation = filtered_data["Valuation_Billions"].max()
-    node_sizes = []
-    for node in G.nodes:
+    # Generate positions using spring layout
+    pos = nx.spring_layout(G, k=0.2, seed=42)
+    
+    # Prepare Plotly traces
+    edge_x = []
+    edge_y = []
+    for edge in G.edges():
+        x0, y0 = pos[edge[0]]
+        x1, y1 = pos[edge[1]]
+        edge_x += [x0, x1, None]
+        edge_y += [y0, y1, None]
+    
+    edge_trace = go.Scatter(
+        x=edge_x, y=edge_y,
+        line=dict(width=0.5, color='#888'),
+        hoverinfo='none',
+        mode='lines'
+    )
+    
+    node_x = []
+    node_y = []
+    node_text = []
+    node_size = []
+    node_color = []
+    customdata = []
+    for node in G.nodes():
+        x, y = pos[node]
+        node_x.append(x)
+        node_y.append(y)
         if node in filtered_mapping:
-            node_sizes.append(1500)  # Fixed size for investors
+            node_text.append(f"Investor: {node}")
+            node_size.append(20)  # Investors have larger size
+            node_color.append('orange')
+            customdata.append(None)  # Investors do not have a single valuation
         else:
             valuation = filtered_data.loc[filtered_data["Company"] == node, "Valuation_Billions"].sum()
-            size = (valuation / max_valuation) * 1500 if max_valuation else 100
-            node_sizes.append(size)
-    
-    # Node color: Investors (orange), Companies (green)
-    node_colors = ["#FF8C00" if node in filtered_mapping else "#32CD32" for node in G.nodes]
+            node_text.append(f"Company: {node}<br>Valuation: ${valuation}B")
+            node_size.append(10 + (valuation / filtered_data["Valuation_Billions"].max()) * 30 if filtered_data["Valuation_Billions"].max() else 10)
+            node_color.append('green')
+            customdata.append(f"${valuation}B")
     
-    # Draw the graph
-    plt.figure(figsize=(15, 15))
-    pos = nx.spring_layout(G, k=0.2, seed=42)
-    nx.draw(
-        G, pos,
-        with_labels=True,
-        node_size=node_sizes,
-        node_color=node_colors,
-        font_size=10,
-        edge_color="#A9A9A9",  # Light gray edges
-        alpha=0.9
+    node_trace = go.Scatter(
+        x=node_x, y=node_y,
+        mode='markers',
+        hoverinfo='text',
+        text=node_text,
+        customdata=customdata,
+        marker=dict(
+            showscale=False,
+            colorscale='YlGnBu',
+            color=node_color,
+            size=node_size,
+            line_width=2
+        )
     )
     
-    # Legend
-    from matplotlib.lines import Line2D
-    legend_elements = [
-        Line2D([0], [0], marker='o', color='w', label='Investor', markersize=10, markerfacecolor='#FF8C00'),
-        Line2D([0], [0], marker='o', color='w', label='Company', markersize=10, markerfacecolor='#32CD32')
-    ]
-    plt.legend(handles=legend_elements, loc='upper left')
-    
-    plt.title("Venture Network Visualization", fontsize=20)
-    plt.axis("off")
-    
-    # Save plot to BytesIO
-    buf = BytesIO()
-    plt.savefig(buf, format="png", bbox_inches="tight")
-    plt.close()
-    buf.seek(0)
-    
-    # Get node information if clicked
-    node_info = "No node clicked."
-    if clicked_node:
-        if clicked_node in filtered_data["Company"].values:
-            valuation = filtered_data.loc[filtered_data["Company"] == clicked_node, "Valuation_Billions"].iloc[0]
-            node_info = f"Company: {clicked_node}, Valuation: ${valuation}B"
-        elif clicked_node in filtered_mapping:
-            node_info = f"Investor: {clicked_node}"
-    
-    return Image.open(buf), node_info
-
-# Gradio app function
-def app(selected_country, selected_industry, selected_investors, clicked_node=None):
+    fig = go.Figure(data=[edge_trace, node_trace],
+             layout=go.Layout(
+                title='Venture Network Visualization',
+                titlefont_size=16,
+                showlegend=False,
+                hovermode='closest',
+                margin=dict(b=20,l=5,r=5,t=40),
+                annotations=[ dict(
+                    text="",
+                    showarrow=False,
+                    xref="paper", yref="paper") ],
+                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False))
+                )
+    
+    fig.update_traces(marker=dict(line=dict(width=0.5, color='white')), selector=dict(mode='markers'))
+    
+    logger.info("Plotly graph generated successfully.")
+    
+    return fig
+
+# Gradio app function to update CheckboxGroup and filtered data
+def app(selected_country, selected_industry):
     investor_list, filtered_data = filter_investors_by_country_and_industry(selected_country, selected_industry)
-    graph_image, node_info = generate_graph_and_get_node_info(selected_investors, filtered_data, clicked_node)
-    return graph_image, node_info, gr.update(choices=investor_list, value=selected_investors)
+    logger.info("Updating CheckboxGroup and filtered data holder.")
+    
+    # Use gr.update() to create an update dictionary for the CheckboxGroup
+    return gr.update(
+        choices=investor_list,
+        value=investor_list,
+        visible=True
+    ), filtered_data
 
 # Gradio Interface
 def main():
@@ -196,33 +224,57 @@ def main():
     
     with gr.Blocks() as demo:
         with gr.Row():
+            # Set default value to "US" for country and "Enterprise Tech" for industry
             country_filter = gr.Dropdown(choices=country_list, label="Filter by Country", value="US")
             industry_filter = gr.Dropdown(choices=industry_list, label="Filter by Industry", value="Enterprise Tech")
         
-        filtered_investor_list = gr.CheckboxGroup(choices=[], label="Select Investors")
-        clicked_node_input = gr.Textbox(label="Clicked Node (Enter name to simulate a click)", placeholder="Enter a node name")
-        graph_output = gr.Image(type="pil", label="Venture Network Graph")
-        node_info_output = gr.Textbox(label="Node Information", interactive=False)
+        filtered_investor_list = gr.CheckboxGroup(choices=[], label="Select Investors", visible=False)
+        graph_output = gr.Plot(label="Venture Network Graph")
+        valuation_display = gr.Markdown(value="Click on a company node to see its valuation.", label="Company Valuation")
+        
+        filtered_data_holder = gr.State()
         
+        # Event handlers for filters
         country_filter.change(
             app,
-            inputs=[country_filter, industry_filter, filtered_investor_list, clicked_node_input],
-            outputs=[graph_output, node_info_output, filtered_investor_list]
+            inputs=[country_filter, industry_filter],
+            outputs=[filtered_investor_list, filtered_data_holder]
         )
         industry_filter.change(
             app,
-            inputs=[country_filter, industry_filter, filtered_investor_list, clicked_node_input],
-            outputs=[graph_output, node_info_output, filtered_investor_list]
+            inputs=[country_filter, industry_filter],
+            outputs=[filtered_investor_list, filtered_data_holder]
         )
+        
+        # Generate graph when investors are selected
         filtered_investor_list.change(
-            app,
-            inputs=[country_filter, industry_filter, filtered_investor_list, clicked_node_input],
-            outputs=[graph_output, node_info_output, filtered_investor_list]
+            generate_graph,
+            inputs=[filtered_investor_list, filtered_data_holder],
+            outputs=graph_output
         )
-        clicked_node_input.change(
-            app,
-            inputs=[country_filter, industry_filter, filtered_investor_list, clicked_node_input],
-            outputs=[graph_output, node_info_output, filtered_investor_list]
+        
+        # Handle plot click to display valuation
+        def display_valuation(plotly_click):
+            if plotly_click is None:
+                return "Click on a company node to see its valuation."
+            point = plotly_click
+            if 'text' in point and point['text']:
+                text = point['text']
+                if "Company:" in text:
+                    # Extract valuation
+                    parts = text.split("<br>")
+                    company_part = parts[0]
+                    valuation_part = parts[1]
+                    company = company_part.replace("Company: ", "")
+                    valuation = valuation_part.replace("Valuation: ", "")
+                    return f"**{company}** has a valuation of **{valuation}**."
+            return "Click on a company node to see its valuation."
+        
+        graph_output.event(
+            "plotly_click",
+            fn=display_valuation,
+            inputs=graph_output,
+            outputs=valuation_display
         )
     
     demo.launch()