import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D import networkx as nx 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, 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 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 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, 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 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) 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 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 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 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') # 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 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 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') # Save to file with adjusted margins plt.close() # Close the figure to free resources def main(mode, input_file=None): G = nx.DiGraph() if mode == 'random': starting_x = 0 starting_y = 0 max_depth = 5 # Maximum depth of the tree max_nodes = 3 # Maximum number of child nodes x_range = 10 # Maximum range for x position of nodes # 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() build_graph_from_json(json_data, G) else: 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 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 [input_file]") else: mode = sys.argv[1] input_file = sys.argv[2] if len(sys.argv) > 2 else None main(mode, input_file)