Spaces:
Runtime error
Runtime error
| 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 | |