Spaces:
Running
Running
| 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 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) | |
| # Iterate through all top-level events (this is the corrected part) | |
| 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 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(json_data): | |
| G = nx.DiGraph() | |
| build_graph_from_json(json_data, G) # Build graph from the provided JSON data | |
| draw_global_tree_3d(G, filename='global_tree.png') | |
| 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(G, best_path, 'best_path.png', 'blue') | |
| if worst_path: | |
| draw_path_3d(G, worst_path, 'worst_path.png', 'red') | |
| if longest_path: | |
| draw_path_3d(G, longest_path, 'longest_duration_path.png', 'green') | |
| if shortest_path: | |
| draw_path_3d(G, shortest_path, 'shortest_duration_path.png', 'purple') | |
| return ('global_tree.png','best_path.png','longest_duration_path.png') # Return the filename of the global tree | |