Update psychohistory.py

psychohistory.py

@@ -1,16 +1,16 @@
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
 import networkx as nx
-import random
 import numpy as np
 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, and the number of nodes on the z-axis."""
+    """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 not in node_count_per_depth:
         node_count_per_depth[depth] = 0
 
@@ -19,13 +19,13 @@ def generate_tree(current_x, current_y, depth, max_depth, max_nodes, x_range, G,
 
     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:
         # Add node to the graph
         node_id = len(G.nodes)
         node_count_per_depth[depth] += 1
         prob = random.uniform(0, 1)  # Assign random probability
-        G.add_node(node_id, pos=(x, prob, depth))  # Use `depth` for the z position
+        G.add_node(node_id, pos=(x, prob, depth))  # Use `depth` for z position
         if parent is not None:
             G.add_edge(parent, node_id)
         # Recursively add child nodes
@@ -33,27 +33,39 @@ 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."""
-    def add_event(parent_id, event_data, prob_level):
-        for key, value in event_data.get('events', {}).items():
-            # Add node
-            node_id = len(G.nodes)
-            prob = {'high_probability': 0.9, 'medium_probability': 0.5, 'low_probability': 0.1}[prob_level]
-            G.add_node(node_id, pos=(len(G.nodes), prob, len(G.nodes)))  # Ensure each node has 'pos'
+    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  # 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)
-            # Add child events
-            add_event(node_id, {'events': value}, key)
 
-    root_id = len(G.nodes)
-    G.add_node(root_id, pos=(0, 0.5, 0))  # Root node with default medium probability
-    if len(G.nodes) > 1:
-        G.add_edge(-1, root_id)  # Root node without a parent
+        # Add child events
+        subevents = event_data.get('subevents', {}).get('event', [])
+        if not isinstance(subevents, list):
+            subevents = [subevents]  # Ensure subevents is a list
+
+        for subevent in subevents:
+            add_event(node_id, subevent, depth + 1)
+
     data = json.loads(json_data)
-    add_event(root_id, data, 'medium_probability')
+    root_id = len(G.nodes)
+    root_event = list(data.get('events', {}).values())[0]
+    G.add_node(root_id, pos=(0, root_event['probability'] / 100.0, root_event['event_number']), label=root_event['name'])
+    add_event(None, root_event, 0)  # Start from the root
+
+
 
 def find_paths(G):
-    """Finds the paths with the highest and lowest average probability, and the maximum and minimum duration in graph G."""
+    """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
@@ -72,30 +84,30 @@ def find_paths(G):
                     if not all('pos' in G.nodes[node] for node in path):
                         continue  # Skip paths with nodes missing the 'pos' attribute
                     
-                    # Calculate the average probability of the path
-                    probabilities = [G.nodes[node]['pos'][1] for node in path]  # Get probabilities of the nodes in the path
+                    # 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 the path with the highest average probability
+                    # Evaluate path with the highest mean probability
                     if mean_prob > best_mean_prob:
                         best_mean_prob = mean_prob
                         best_path = path
                     
-                    # Evaluate the path with the lowest average probability
+                    # Evaluate path with the lowest mean probability
                     if mean_prob < worst_mean_prob:
                         worst_mean_prob = mean_prob
                         worst_path = path
                     
-                    # Calculate the duration of the path
+                    # Calculate path duration
                     x_positions = [G.nodes[node]['pos'][0] for node in path]
                     duration = max(x_positions) - min(x_positions)
                     
-                    # Evaluate the path with the maximum duration
+                    # Evaluate path with the longest duration
                     if duration > max_duration:
                         max_duration = duration
                         longest_duration_path = path
                     
-                    # Evaluate the path with the minimum duration
+                    # Evaluate path with the shortest duration
                     if duration < min_duration:
                         min_duration = duration
                         shortest_duration_path = path
@@ -115,7 +127,7 @@ def draw_path_3d(G, path, filename='path_plot_3d.png', highlight_color='blue'):
     fig = plt.figure(figsize=(16, 12))
     ax = fig.add_subplot(111, projection='3d')
 
-    # Assign colors to the nodes based on probability
+    # Assign colors to nodes based on probability
     node_colors = []
     for node in path:
         prob = G.nodes[node]['pos'][1]
@@ -135,31 +147,38 @@ def draw_path_3d(G, path, filename='path_plot_3d.png', highlight_color='blue'):
         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 the nodes
+    # 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')
 
-    # Adjust labels and title
+    # Set labels and title
     ax.set_xlabel('Time (weeks)')
     ax.set_ylabel('Event Probability')
     ax.set_zlabel('Event Number')
-    ax.set_title('Event Tree in 3D - Path')
+    ax.set_title('3D Event Tree - Path')
+
+    plt.savefig(filename, bbox_inches='tight')  # Save to file with adjusted margins
+    plt.close()  # Close the figure to free resources
 
-    plt.savefig(filename, bbox_inches='tight')  # Save to a file with adjusted margins
-    plt.close()  # Close the figure to free up resources
 
 def draw_global_tree_3d(G, filename='global_tree.png'):
     """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')
 
-    # Assign colors to the nodes based on probability
+    # Assign colors to nodes based on probability
     node_colors = []
     for node, (x, prob, z) in pos.items():
         if prob < 0.33:
@@ -178,18 +197,19 @@ def draw_global_tree_3d(G, filename='global_tree.png'):
         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 the nodes
+    # Add labels to nodes
     for node, (x, y, z) in pos.items():
-        ax.text(x, y, z, str(node), fontsize=12, color='black')
+        label = labels.get(node, f"{node}")
+        ax.text(x, y, z, label, fontsize=12, color='black')
 
-    # Adjust labels and title
-    ax.set_xlabel('Time (weeks)')
-    ax.set_ylabel('Event Probability')
+    # Set labels and title
+    ax.set_xlabel('Time')
+    ax.set_ylabel('Probability')
     ax.set_zlabel('Event Number')
-    ax.set_title('Event Tree in 3D')
+    ax.set_title('3D Event Tree')
 
-    plt.savefig(filename, bbox_inches='tight')  # Save to a file with adjusted margins
-    plt.close()  # Close the figure to free up resources
+    plt.savefig(filename, bbox_inches='tight')  # Save to file with adjusted margins
+    plt.close()  # Close the figure to free resources
 
 def main(mode, input_file=None):
     G = nx.DiGraph()
@@ -197,12 +217,14 @@ def main(mode, input_file=None):
     if mode == 'random':
         starting_x = 0
         starting_y = 0
-        max_depth = 5          # Maximum tree depth
+        max_depth = 5          # Maximum depth of the tree
         max_nodes = 3          # Maximum number of child nodes
-        x_range = 10           # Maximum range for node x positions
+        x_range = 10           # Maximum range for x position of nodes
 
-        # Generate the tree and get the node count per depth
+        # Generate the tree and get node count per depth
         generate_tree(starting_x, starting_y, 0, max_depth, max_nodes, x_range, G)
+
+
     elif mode == 'json' and input_file:
         with open(input_file, 'r') as file:
             json_data = file.read()
@@ -211,10 +233,14 @@ def main(mode, input_file=None):
         print("Invalid mode or input file not provided.")
         return
 
+    # Save the global visualization
+    draw_global_tree_3d(G, filename='global_tree.png')
+
+
     # Find relevant paths
     best_path, best_mean_prob, worst_path, worst_mean_prob, longest_duration_path, shortest_duration_path = find_paths(G)
 
-    # Print the results
+    # Print results
     if best_path:
         print(f"\nPath with the highest average probability:")
         print(" -> ".join(map(str, best_path)))
@@ -251,10 +277,14 @@ def main(mode, input_file=None):
     if shortest_duration_path:
         draw_path_3d(G, path=shortest_duration_path, filename='shortest_duration_path.png', highlight_color='purple')
 
+
+
 if __name__ == "__main__":
     if len(sys.argv) < 2:
-        print("Usage: python script.py <mode> [json_file]")
+        print("Usage: python script.py <mode> [input_file]")
     else:
         mode = sys.argv[1]
         input_file = sys.argv[2] if len(sys.argv) > 2 else None
         main(mode, input_file)
+
+