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 calculate_angles_between_points(path, pos): """Calcula los ángulos en radianes y grados entre puntos consecutivos en los planos XZ, XY y ZY.""" angles = {'xz': [], 'xy': [], 'zy': []} for i in range(1, len(path)): # Obtener las posiciones de los puntos consecutivos x1, y1, z1 = pos[path[i-1]] x2, y2, z2 = pos[path[i]] # Cálculo del ángulo en el plano XZ delta_x = x2 - x1 delta_z = z2 - z1 angle_xz = math.atan2(delta_z, delta_x) # Radianes angles['xz'].append((angle_xz, math.degrees(angle_xz))) # Convertir a grados # Cálculo del ángulo en el plano XY delta_y = y2 - y1 angle_xy = math.atan2(delta_y, delta_x) # Radianes angles['xy'].append((angle_xy, math.degrees(angle_xy))) # Convertir a grados # Cálculo del ángulo en el plano ZY angle_zy = math.atan2(delta_y, delta_z) # Radianes angles['zy'].append((angle_zy, math.degrees(angle_zy))) # Convertir a grados return angles def draw_path_3d(G, path, filename='path_plot_3d.png', highlight_color='blue'): """Dibuja solo un path específico en 3D y guarda la imagen con etiquetas.""" H = G.subgraph(path).copy() pos = nx.get_node_attributes(G, 'pos') labels = nx.get_node_attributes(G, 'label') # Obtener etiquetas 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') 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') ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7) # Dibujar aristas 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) # Agregar etiquetas for node, (x, y, z) in pos.items(): if node in path: label = labels.get(node, str(node)) ax.text(x, y, z, label, fontsize=12, color='black') 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') plt.close() 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) 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}") best_angles = calculate_angles_between_points(best_path, nx.get_node_attributes(G, 'pos')) print("\nAngles for the most probable path:") print(f"XZ plane angles (radians, degrees): {best_angles['xz']}") print(f"XY plane angles (radians, degrees): {best_angles['xy']}") print(f"ZY plane angles (radians, degrees): {best_angles['zy']}") 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))}") longest_angles = calculate_angles_between_points(longest_path, nx.get_node_attributes(G, 'pos')) print("\nAngles for the longest duration path:") print(f"XZ plane angles (radians, degrees): {longest_angles['xz']}") print(f"XY plane angles (radians, degrees): {longest_angles['xy']}") print(f"ZY plane angles (radians, degrees): {longest_angles['zy']}") if shortest_path: print(f"\nPath with the shortest duration: {' -> '.join(map(str, shortest_path))}") # Dibujar paths y añadir etiquetas a los nodos 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',longest_angles,best_angles)