Update genesis/visualization.py

genesis/visualization.py

@@ -1,62 +1,155 @@
 # genesis/visualization.py
-import networkx as nx
-import matplotlib.pyplot as plt
-import io
-import base64
-
-def _graph_to_image_bytes(G, layout_func=nx.spring_layout):
-    """Helper: Convert NetworkX graph to PNG bytes."""
-    plt.figure(figsize=(8, 6))
-    pos = layout_func(G)
-    nx.draw_networkx_nodes(G, pos, node_size=1500, node_color="#34d399", alpha=0.9)
-    nx.draw_networkx_edges(G, pos, width=2, alpha=0.5, edge_color="#6b7280")
-    nx.draw_networkx_labels(G, pos, font_size=10, font_family="sans-serif", font_color="white")
-    plt.axis("off")
-
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", bbox_inches="tight")
-    plt.close()
-    buf.seek(0)
-    return buf
-
-def generate_pathway_graph(entities, relationships):
+import os
+from neo4j import GraphDatabase
+import plotly.graph_objects as go
+
+# =========================
+# CONFIGURATION
+# =========================
+NEO4J_URI = os.getenv("NEO4J_URI")
+NEO4J_USER = os.getenv("NEO4J_USER")
+NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")
+
+COLOR_PATHWAY_NODE = "#00FF99"  # Biotech green
+COLOR_FUNDING_NODE = "#FF9900"  # Orange
+COLOR_EDGE = "#AAAAAA"
+
+# =========================
+# NEO4J CONNECTION
+# =========================
+driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
+
+def run_neo4j_query(query, params=None):
+    """Run a Cypher query and return results."""
+    with driver.session() as session:
+        return list(session.run(query, params or {}))
+
+# =========================
+# GRAPH UTILS
+# =========================
+def create_plotly_graph(nodes, edges, title):
+    """Creates an interactive Plotly network graph."""
+    node_x, node_y, node_text, node_color = [], [], [], []
+    edge_x, edge_y = [], []
+
+    # Layout algorithm (spring layout for better spacing)
+    import networkx as nx
+    G = nx.Graph()
+    for node_id, label, color in nodes:
+        G.add_node(node_id, label=label, color=color)
+    for src, dst in edges:
+        G.add_edge(src, dst)
+
+    pos = nx.spring_layout(G, seed=42, k=0.5)
+
+    for src, dst in edges:
+        x0, y0 = pos[src]
+        x1, y1 = pos[dst]
+        edge_x.extend([x0, x1, None])
+        edge_y.extend([y0, y1, None])
+
+    for node_id, label, color in nodes:
+        x, y = pos[node_id]
+        node_x.append(x)
+        node_y.append(y)
+        node_text.append(label)
+        node_color.append(color)
+
+    edge_trace = go.Scatter(
+        x=edge_x, y=edge_y,
+        line=dict(width=1, color=COLOR_EDGE),
+        hoverinfo='none',
+        mode='lines'
+    )
+
+    node_trace = go.Scatter(
+        x=node_x, y=node_y,
+        mode='markers+text',
+        text=node_text,
+        textposition="top center",
+        hoverinfo="text",
+        marker=dict(
+            color=node_color,
+            size=18,
+            line=dict(width=2, color='#000000')
+        )
+    )
+
+    fig = go.Figure(data=[edge_trace, node_trace])
+    fig.update_layout(
+        title=title,
+        title_x=0.5,
+        plot_bgcolor="#111111",
+        paper_bgcolor="#111111",
+        font=dict(color="#FFFFFF"),
+        showlegend=False,
+        margin=dict(l=10, r=10, t=40, b=10)
+    )
+    return fig
+
+# =========================
+# PATHWAY GRAPH
+# =========================
+def generate_pathway_graph(pathway_name):
+    """
+    Generate an interactive graph for a given metabolic pathway.
     """
-    Create a pathway relationship graph for synthetic biology concepts.
-    entities: list of strings
-    relationships: list of dicts {source, target, type}
-    Returns: file path to PNG
+    query = """
+    MATCH (p:Pathway {name: $name})-[r:INVOLVES]->(m:Molecule)
+    RETURN p.name AS pathway, m.name AS molecule
     """
-    G = nx.DiGraph()
-    for entity in entities:
-        G.add_node(entity)
+    results = run_neo4j_query(query, {"name": pathway_name})
+
+    if not results:
+        return None
 
-    for rel in relationships:
-        G.add_edge(rel["source"], rel["target"], label=rel["type"])
+    nodes = [(pathway_name, pathway_name, COLOR_PATHWAY_NODE)]
+    edges = []
+    for record in results:
+        mol_name = record["molecule"]
+        nodes.append((mol_name, mol_name, "#00BFFF"))  # Blue for molecules
+        edges.append((pathway_name, mol_name))
 
-    buf = _graph_to_image_bytes(G)
-    file_path = "/tmp/pathway_graph.png"
-    with open(file_path, "wb") as f:
-        f.write(buf.read())
-    return file_path
+    return create_plotly_graph(nodes, edges, f"Metabolic Pathway: {pathway_name}")
 
-def generate_funding_network(companies):
+# =========================
+# FUNDING NETWORK
+# =========================
+def generate_funding_network(industry_keyword):
     """
-    Create a funding network graph.
-    companies: list of dicts {name, investors}
-    Returns: file path to PNG
+    Generate an interactive funding network graph for companies in a given biotech domain.
     """
-    G = nx.Graph()
-    for comp in companies:
-        company_node = f"🏢 {comp['name']}"
-        G.add_node(company_node)
-        investors = [i.strip() for i in comp["investors"].split(",")]
-        for inv in investors:
-            inv_node = f"💰 {inv}"
-            G.add_node(inv_node)
-            G.add_edge(company_node, inv_node)
-
-    buf = _graph_to_image_bytes(G, layout_func=nx.kamada_kawai_layout)
-    file_path = "/tmp/funding_network.png"
-    with open(file_path, "wb") as f:
-        f.write(buf.read())
-    return file_path
+    query = """
+    MATCH (c:Company)-[f:FUNDED_BY]->(i:Investor)
+    WHERE toLower(c.industry) CONTAINS toLower($keyword)
+    RETURN c.name AS company, i.name AS investor
+    """
+    results = run_neo4j_query(query, {"keyword": industry_keyword})
+
+    if not results:
+        return None
+
+    nodes = []
+    edges = []
+    seen_nodes = set()
+
+    for record in results:
+        comp = record["company"]
+        inv = record["investor"]
+
+        if comp not in seen_nodes:
+            nodes.append((comp, comp, COLOR_FUNDING_NODE))
+            seen_nodes.add(comp)
+        if inv not in seen_nodes:
+            nodes.append((inv, inv, "#FFD700"))  # Gold for investors
+            seen_nodes.add(inv)
+
+        edges.append((comp, inv))
+
+    return create_plotly_graph(nodes, edges, f"Funding Network: {industry_keyword}")
+
+# =========================
+# CLEANUP
+# =========================
+def close_driver():
+    driver.close()