Create psychohistory.py

psychohistory.py

@@ -0,0 +1,260 @@
+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
+
+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."""
+    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
+
+    if depth > max_depth:
+        return node_count_per_depth
+
+    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
+        if parent is not None:
+            G.add_edge(parent, node_id)
+        # Recursively add child nodes
+        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."""
+    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'
+            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
+    data = json.loads(json_data)
+    add_event(root_id, data, 'medium_probability')
+
+def find_paths(G):
+    """Finds the paths with the highest and lowest average probability, and the maximum and minimum duration 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 average probability of the path
+                    probabilities = [G.nodes[node]['pos'][1] for node in path]  # Get probabilities of the nodes in the path
+                    mean_prob = np.mean(probabilities)
+                    
+                    # Evaluate the path with the highest average probability
+                    if mean_prob > best_mean_prob:
+                        best_mean_prob = mean_prob
+                        best_path = path
+                    
+                    # Evaluate the path with the lowest average probability
+                    if mean_prob < worst_mean_prob:
+                        worst_mean_prob = mean_prob
+                        worst_path = path
+                    
+                    # Calculate the duration of the path
+                    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
+                    if duration > max_duration:
+                        max_duration = duration
+                        longest_duration_path = path
+                    
+                    # Evaluate the path with the minimum 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 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')
+
+    # Assign colors to the 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 the 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
+    ax.set_xlabel('Time (weeks)')
+    ax.set_ylabel('Event Probability')
+    ax.set_zlabel('Event Number')
+    ax.set_title('Event Tree in 3D - Path')
+
+    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')
+    
+    # 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
+    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 the nodes
+    for node, (x, y, z) in pos.items():
+        ax.text(x, y, z, str(node), fontsize=12, color='black')
+
+    # Adjust 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')
+
+    plt.savefig(filename, bbox_inches='tight')  # Save to a file with adjusted margins
+    plt.close()  # Close the figure to free up resources
+
+def main(mode, input_file=None):
+    G = nx.DiGraph()
+
+    if mode == 'random':
+        starting_x = 0
+        starting_y = 0
+        max_depth = 5          # Maximum tree depth
+        max_nodes = 3          # Maximum number of child nodes
+        x_range = 10           # Maximum range for node x positions
+
+        # Generate the tree and get the 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()
+        build_graph_from_json(json_data, G)
+    else:
+        print("Invalid mode or input file not provided.")
+        return
+
+    # 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
+    if best_path:
+        print(f"\nPath with the highest average probability:")
+        print(" -> ".join(map(str, best_path)))
+        print(f"Average probability: {best_mean_prob:.2f}")
+
+    if worst_path:
+        print(f"\nPath with the lowest average probability:")
+        print(" -> ".join(map(str, worst_path)))
+        print(f"Average probability: {worst_mean_prob:.2f}")
+
+    if longest_duration_path:
+        print(f"\nPath with the longest duration:")
+        print(" -> ".join(map(str, longest_duration_path)))
+        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in longest_duration_path) - min(G.nodes[node]['pos'][0] for node in longest_duration_path):.2f}")
+
+    if shortest_duration_path:
+        print(f"\nPath with the shortest duration:")
+        print(" -> ".join(map(str, shortest_duration_path)))
+        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in shortest_duration_path) - min(G.nodes[node]['pos'][0] for node in shortest_duration_path):.2f}")
+
+    # Save the global visualization
+    draw_global_tree_3d(G, filename='global_tree.png')
+
+    # Draw and save the 3D figure for each relevant path
+    if best_path:
+        draw_path_3d(G, path=best_path, filename='best_path.png', highlight_color='blue')
+
+    if worst_path:
+        draw_path_3d(G, path=worst_path, filename='worst_path.png', highlight_color='red')
+
+    if longest_duration_path:
+        draw_path_3d(G, path=longest_duration_path, filename='longest_duration_path.png', highlight_color='green')
+
+    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]")
+    else:
+        mode = sys.argv[1]
+        input_file = sys.argv[2] if len(sys.argv) > 2 else None
+        main(mode, input_file)