Update psychohistory.py

psychohistory.py

@@ -84,43 +84,21 @@ def find_paths(G):
     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."""
+    """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')
-    labels = nx.get_node_attributes(G, 'label')  # Obtener etiquetas
 
+    # 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]
@@ -131,27 +109,28 @@ def draw_path_3d(G, path, filename='path_plot_3d.png', highlight_color='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)
 
-    # Dibujar aristas
+    # 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)
 
-    # Agregar etiquetas
+    # Add labels to nodes
     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.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')
-    plt.close()
+    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'):
@@ -204,12 +183,9 @@ def draw_global_tree_3d(G, filename='global_tree.png'):
     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_json(json_data, G)  # Build graph from the provided JSON data
 
     draw_global_tree_3d(G, filename='global_tree.png')
 
@@ -218,28 +194,16 @@ def main(json_data):
     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']}")
-
+        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))}")
-    
-    # Dibujar paths y añadir etiquetas a los nodos
+        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:
@@ -249,4 +213,4 @@ def main(json_data):
     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 ('global_tree.png','best_path.png','longest_duration_path.png')  # Return the filename of the global tree