revert to matplotlib

psychohistory.py

@@ -1,5 +1,5 @@
-import plotly.graph_objects as go  # Import Plotly for interactive plots
-from mpl_toolkits.mplot3d import Axes3D  # Not needed anymore, but you can keep it if you use it elsewhere
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
 import networkx as nx
 import numpy as np
 import json
@@ -32,6 +32,8 @@ def generate_tree(current_x, current_y, depth, max_depth, max_nodes, x_range, G,
     return node_count_per_depth
 
 
+
+
 def build_graph_from_json(json_data, G):
     """Builds a graph from JSON data, handling subevents recursively."""
 
@@ -57,6 +59,7 @@ def build_graph_from_json(json_data, G):
         add_event(None, event_data, 0)  # Add each event as a root node
 
 
+
 def find_paths(G):
     """Finds paths with highest/lowest probability and longest/shortest durations."""
     best_path, worst_path = None, None
@@ -92,63 +95,98 @@ def find_paths(G):
 
     return best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path
 
-def draw_graph_plotly(G, title="3D Event Tree", highlight_color='gray'):
-    """Draws the graph in 3D using Plotly and returns the HTML string."""
+def draw_path_3d(G, path, filename='path_plot_3d.png', highlight_color='blue'):
+    """Draws a specific path in 3D."""
+    H = G.subgraph(path).copy()
+    pos = nx.get_node_attributes(G, 'pos')
+    x_vals, y_vals, z_vals = zip(*[pos[node] for node in path])
+
+    fig = plt.figure(figsize=(16, 12))
+    ax = fig.add_subplot(111, projection='3d')
+
+    node_colors = ['red' if prob < 0.33 else 'blue' if prob < 0.67 else 'green' for _, prob, _ in [pos[node] for node in path]]
+    ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7)
+
+    for edge in H.edges():
+        x_start, y_start, z_start = pos[edge[0]]
+        x_end, y_end, z_end = pos[edge[1]]
+        ax.plot([x_start, x_end], [y_start, y_end], [z_start, z_end], color=highlight_color, lw=2)
+
+    for node, (x, y, z) in pos.items():
+        if node in path:
+            ax.text(x, y, z, str(node), fontsize=12, color='black')
+
+    ax.set_xlabel('Time (weeks)')
+    ax.set_ylabel('Event Probability')
+    ax.set_zlabel('Event Number')
+    ax.set_title('3D Event Tree - Path')
+
+    plt.savefig(filename, bbox_inches='tight')
+    plt.close()
+
+
+def draw_global_tree_3d(G, filename='global_tree.png'):
+    """Draws the entire graph in 3D."""
     pos = nx.get_node_attributes(G, 'pos')
     labels = nx.get_node_attributes(G, 'label')
 
     if not pos:
         print("Graph is empty. No nodes to visualize.")
-        return ""
+        return
 
     x_vals, y_vals, z_vals = zip(*pos.values())
+    fig = plt.figure(figsize=(16, 12))
+    ax = fig.add_subplot(111, projection='3d')
 
     node_colors = ['red' if prob < 0.33 else 'blue' if prob < 0.67 else 'green' for _, prob, _ in pos.values()]
-    node_trace = go.Scatter3d(x=x_vals, y=y_vals, z=z_vals, mode='markers+text',
-                             marker=dict(size=10, color=node_colors, line=dict(width=1, color='black')),
-                             text=list(labels.values()), textposition='top center', hoverinfo='text')
+    ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7)
 
-    edge_traces = []
     for edge in G.edges():
         x_start, y_start, z_start = pos[edge[0]]
         x_end, y_end, z_end = pos[edge[1]]
-        edge_trace = go.Scatter3d(x=[x_start, x_end], y=[y_start, y_end], z=[z_start, z_end],
-                                 mode='lines', line=dict(width=2, color=highlight_color), hoverinfo='none')
-        edge_traces.append(edge_trace)
+        ax.plot([x_start, x_end], [y_start, y_end], [z_start, z_end], color='gray', lw=2)
 
-    layout = go.Layout(scene=dict(xaxis_title='Time', yaxis_title='Probability', zaxis_title='Event Number'),
-                       title=title)
-    fig = go.Figure(data=[node_trace] + edge_traces, layout=layout)
+    for node, (x, y, z) in pos.items():
+        label = labels.get(node, f"{node}")
+        ax.text(x, y, z, label, fontsize=12, color='black')
 
-    # Convert Plotly figure to HTML string
-    html_str = fig.to_html(full_html=False, include_plotlyjs='cdn')
-    return html_str
+    ax.set_xlabel('Time')
+    ax.set_ylabel('Probability')
+    ax.set_zlabel('Event Number')
+    ax.set_title('3D Event Tree')
+
+    plt.savefig(filename, bbox_inches='tight')
+    plt.close()
 
 
 def main(json_data):
     G = nx.DiGraph()
-    build_graph_from_json(json_data, G)
+    build_graph_from_json(json_data, G)  # Build graph from the provided JSON data
 
-    # Find the best, worst, longest, and shortest paths FIRST
-    best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path = find_paths(G)
+    draw_global_tree_3d(G, filename='global_tree.png')
 
-    # Generate the HTML string for the Plotly graph
-    html_graph = draw_graph_plotly(G)
+    best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path = find_paths(G)
 
-    # Now you can use the path variables
     if best_path:
-        best_path_graph = draw_graph_plotly(G.subgraph(best_path), title="Best Path", highlight_color='blue')
-        html_graph += best_path_graph
+        print(f"\nPath with the highest average probability: {' -> '.join(map(str, best_path))}")
+        print(f"Average probability: {best_mean_prob:.2f}")
     if worst_path:
-        worst_path_graph = draw_graph_plotly(G.subgraph(worst_path), title="Worst Path", highlight_color='red')
-        html_graph += worst_path_graph
+        print(f"\nPath with the lowest average probability: {' -> '.join(map(str, worst_path))}")
+        print(f"Average probability: {worst_mean_prob:.2f}")
     if longest_path:
-        longest_path_graph = draw_graph_plotly(G.subgraph(longest_path), title="Longest Path", highlight_color='green')
-        html_graph += longest_path_graph
+        print(f"\nPath with the longest duration: {' -> '.join(map(str, longest_path))}")
+        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in longest_path) - min(G.nodes[node]['pos'][0] for node in longest_path):.2f}")
     if shortest_path:
-        shortest_path_graph = draw_graph_plotly(G.subgraph(shortest_path), title="Shortest Path", highlight_color='purple')
-        html_graph += shortest_path_graph
-
-    return html_graph  # Return the HTML string
+        print(f"\nPath with the shortest duration: {' -> '.join(map(str, shortest_path))}")
+        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in shortest_path) - min(G.nodes[node]['pos'][0] for node in shortest_path):.2f}")
 
+    if best_path:
+        draw_path_3d(G, best_path, 'best_path.png', 'blue')
+    if worst_path:
+        draw_path_3d(G, worst_path, 'worst_path.png', 'red')
+    if longest_path:
+        draw_path_3d(G, longest_path, 'longest_duration_path.png', 'green')
+    if shortest_path:
+        draw_path_3d(G, shortest_path, 'shortest_duration_path.png', 'purple')
 
+    return 'global_tree.png'  # Return the filename of the global tree