import plotly.graph_objects as go # Import Plotly for interactive plots 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 > 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.""" def add_event(parent_id, event_data, depth): 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'] 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 subevents = event_data.get('subevents', {}).get('event', []) if not isinstance(subevents, list): subevents = [subevents] 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 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 def draw_path_3d_interactive(G, path, highlight_color='blue'): """Draws a specific path in 3D using Plotly for interactivity.""" 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]) node_colors = ['red' if prob < 0.33 else 'blue' if prob < 0.67 else 'green' for _, prob, _ in [pos[node] for node in path]] 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(map(str, path)), textposition='top center', hoverinfo='text') edge_traces = [] for edge in H.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) layout = go.Layout(scene=dict(xaxis_title='Time (weeks)', yaxis_title='Event Probability', zaxis_title='Event Number'), title='3D Event Tree - Path') fig = go.Figure(data=[node_trace] + edge_traces, layout=layout) fig.show() def draw_global_tree_3d_interactive(G): """Draws the entire graph in 3D using Plotly for interactivity.""" 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 x_vals, y_vals, z_vals = zip(*pos.values()) 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') 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='gray'), hoverinfo='none') edge_traces.append(edge_trace) layout = go.Layout(scene=dict(xaxis_title='Time', yaxis_title='Probability', zaxis_title='Event Number'), title='3D Event Tree') fig = go.Figure(data=[node_trace] + edge_traces, layout=layout) fig.show() def main(json_data): G = nx.DiGraph() build_graph_from_json(json_data, G) # Build graph from the provided JSON data # Draw the interactive graph using Plotly draw_global_tree_3d_interactive(G) best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path = find_paths(G) if best_path: print(f"\nPath with the highest average probability: {' -> '.join(map(str, best_path))}") print(f"Average probability: {best_mean_prob:.2f}") if worst_path: print(f"\nPath with the lowest average probability: {' -> '.join(map(str, worst_path))}") print(f"Average probability: {worst_mean_prob:.2f}") if longest_path: 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: 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_interactive(G, best_path, 'blue') if worst_path: draw_path_3d_interactive(G, worst_path, 'red') if longest_path: draw_path_3d_interactive(G, longest_path, 'green') if shortest_path: draw_path_3d_interactive(G, shortest_path, 'purple') return None # No need to return a filename for interactive plot