better graph

psychohistory.py

@@ -6,158 +6,181 @@ import json
 import sys
 import random
 
-def generate_tree(current_x, current_y, depth, max_depth, max_nodes, x_range, G, parent=None, node_count_per_depth=None):
-    """Generates a tree of nodes with positions adjusted on the x-axis, y-axis, and number of nodes on the z-axis."""
-    if node_count_per_depth is None:
-        node_count_per_depth = {}
-
-    if depth > max_depth:
-        return node_count_per_depth
-
-    if depth not in node_count_per_depth:
-        node_count_per_depth[depth] = 0
-
-    num_children = random.randint(1, max_nodes)
-    x_positions = [current_x + i * x_range / (num_children + 1) for i in range(num_children)]
-
-    for x in x_positions:
-        node_id = len(G.nodes)
-        node_count_per_depth[depth] += 1
-        prob = random.uniform(0, 1)
-        G.add_node(node_id, pos=(x, prob, depth))
-        if parent is not None:
-            G.add_edge(parent, node_id)
-        generate_tree(x, current_y + 1, depth + 1, max_depth, max_nodes, x_range, G, parent=node_id, node_count_per_depth=node_count_per_depth)
-
-    return node_count_per_depth
-
-
-
-
 def build_graph_from_json(json_data, G):
-    """Builds a graph from JSON data, handling subevents recursively."""
-
+    """Builds a graph from JSON data."""
     def add_event(parent_id, event_data, depth):
+        """Recursively adds events and subevents to the graph."""
+        # Add the current event node
         node_id = len(G.nodes)
-        prob = event_data['probability'] / 100.0
-        # Use event_number as the z-coordinate for better visualization
-        pos = (depth, prob, event_data['event_number'])
-        label = event_data['name']
+        prob = event_data['probability'] / 100.0  # Convert percentage to probability
+        pos = (depth, prob, event_data['event_number'])  # Use event_number for z position
+        label = event_data['name']  # Use event name as label
         G.add_node(node_id, pos=pos, label=label)
         if parent_id is not None:
-            G.add_edge(parent_id, node_id)  # Connect to parent
+            G.add_edge(parent_id, node_id)
 
+        # Add child events
         subevents = event_data.get('subevents', {}).get('event', [])
         if not isinstance(subevents, list):
-            subevents = [subevents]
+            subevents = [subevents]  # Ensure subevents is a list
 
         for subevent in subevents:
-            add_event(node_id, subevent, depth + 1)  # Recursively add subevents
-
-    # Iterate through all top-level events
-    for event_data in json_data.get('events', {}).values():
-        add_event(None, event_data, 0)  # Add each event as a root node
+            add_event(node_id, subevent, depth + 1)
 
+    # Start from the root event (assuming there's only one top-level event)
+    root_event = list(json_data.get('events', {}).values())[0]
+    add_event(None, root_event, 0)
 
 
 def find_paths(G):
-    """Finds paths with highest/lowest probability and longest/shortest durations."""
-    best_path, worst_path = None, None
-    longest_path, shortest_path = None, None
-    best_mean_prob, worst_mean_prob = -1, float('inf')
-    max_duration, min_duration = -1, float('inf')
-
-    # Use nx.all_pairs_shortest_path for efficiency
-    all_paths_dict = dict(nx.all_pairs_shortest_path(G))
-
-    for source, paths_from_source in all_paths_dict.items():
-        for target, path in paths_from_source.items():
-            if source != target and all('pos' in G.nodes[node] for node in path):
-                probabilities = [G.nodes[node]['pos'][1] for node in path]
-                mean_prob = np.mean(probabilities)
-
-                if mean_prob > best_mean_prob:
-                    best_mean_prob = mean_prob
-                    best_path = path
-                if mean_prob < worst_mean_prob:
-                    worst_mean_prob = mean_prob
-                    worst_path = path
-
-                x_positions = [G.nodes[node]['pos'][0] for node in path]
-                duration = max(x_positions) - min(x_positions)
-
-                if duration > max_duration:
-                    max_duration = duration
-                    longest_path = path
-                if duration < min_duration and duration > 0:  # Avoid paths with 0 duration
-                    min_duration = duration
-                    shortest_path = path
-
-    return best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path
+    """Finds the paths with the highest and lowest average probability,
+       and the longest and shortest durations in graph G."""
+    best_path = None
+    worst_path = None
+    longest_duration_path = None
+    shortest_duration_path = None
+    best_mean_prob = -1
+    worst_mean_prob = float('inf')
+    max_duration = -1
+    min_duration = float('inf')
+
+    for source in G.nodes:
+        for target in G.nodes:
+            if source != target:
+                all_paths = list(nx.all_simple_paths(G, source=source, target=target))
+                for path in all_paths:
+                    # Check if all nodes in the path have the 'pos' attribute
+                    if not all('pos' in G.nodes[node] for node in path):
+                        continue  # Skip paths with nodes missing the 'pos' attribute
+
+                    # Calculate the mean probability of the path
+                    probabilities = [G.nodes[node]['pos'][1] for node in path]  # Get node probabilities
+                    mean_prob = np.mean(probabilities)
+
+                    # Evaluate path with the highest mean probability
+                    if mean_prob > best_mean_prob:
+                        best_mean_prob = mean_prob
+                        best_path = path
+
+                    # Evaluate path with the lowest mean probability
+                    if mean_prob < worst_mean_prob:
+                        worst_mean_prob = mean_prob
+                        worst_path = path
+
+                    # Calculate path duration
+                    x_positions = [G.nodes[node]['pos'][0] for node in path]
+                    duration = max(x_positions) - min(x_positions)
+
+                    # Evaluate path with the longest duration
+                    if duration > max_duration:
+                        max_duration = duration
+                        longest_duration_path = path
+
+                    # Evaluate path with the shortest duration
+                    if duration < min_duration:
+                        min_duration = duration
+                        shortest_duration_path = path
+
+    return best_path, best_mean_prob, worst_path, worst_mean_prob, longest_duration_path, shortest_duration_path
 
 def draw_path_3d(G, path, filename='path_plot_3d.png', highlight_color='blue'):
-    """Draws a specific path in 3D."""
+    """Draws only the specific path in 3D using networkx and matplotlib
+       and saves the figure to a file."""
+    # Create a subgraph containing only the nodes and edges of the path
     H = G.subgraph(path).copy()
+
     pos = nx.get_node_attributes(G, 'pos')
+
+    # Get data for 3D visualization
     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]]
+    # Assign colors to nodes based on probability
+    node_colors = []
+    for node in path:
+        prob = G.nodes[node]['pos'][1]
+        if prob < 0.33:
+            node_colors.append('red')
+        elif prob < 0.67:
+            node_colors.append('blue')
+        else:
+            node_colors.append('green')
+
+    # Draw nodes
     ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7)
 
+    # Draw edges
     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)
 
+    # Add labels to nodes
     for node, (x, y, z) in pos.items():
         if node in path:
             ax.text(x, y, z, str(node), fontsize=12, color='black')
 
+    # Set labels and title
     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()
+    plt.savefig(filename, bbox_inches='tight')  # Save to file with adjusted margins
+    plt.close()  # Close the figure to free resources
 
 
 def draw_global_tree_3d(G, filename='global_tree.png'):
-    """Draws the entire graph in 3D."""
+    """Draws the entire graph in 3D using networkx and matplotlib
+       and saves the figure to a file."""
     pos = nx.get_node_attributes(G, 'pos')
     labels = nx.get_node_attributes(G, 'label')
 
+    # Check if the graph is empty
     if not pos:
         print("Graph is empty. No nodes to visualize.")
         return
 
+    # Get data for 3D visualization
     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()]
+    # Assign colors to nodes based on probability
+    node_colors = []
+    for node, (x, prob, z) in pos.items():
+        if prob < 0.33:
+            node_colors.append('red')
+        elif prob < 0.67:
+            node_colors.append('blue')
+        else:
+            node_colors.append('green')
+
+    # Draw nodes
     ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7)
 
+    # Draw edges
     for edge in G.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='gray', lw=2)
 
+    # Add labels to nodes
     for node, (x, y, z) in pos.items():
         label = labels.get(node, f"{node}")
         ax.text(x, y, z, label, fontsize=12, color='black')
 
+    # Set labels and title
     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()
-
+    plt.savefig(filename, bbox_inches='tight')  # Save to file with adjusted margins
+    plt.close()  # Close the figure to free resources
 
 def main(json_data):
     G = nx.DiGraph()
@@ -189,4 +212,4 @@ def main(json_data):
     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
+    return 'global_tree.png'  # Return the filename of the global tree