DroneSimulator/plot_trajectory.py at main · Salmon490/DroneSimulator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
#!/usr/bin/env python3
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def plot_trajectory():
    # Read CSV data
    try:
        df = pd.read_csv('build/data/trajectory.csv')
    except FileNotFoundError:
        print("Error: 'build/data/trajectory.csv' not found.")
        return

    # --- WINDOW 1: 3D Trajectory ---
    fig1 = plt.figure(figsize=(10, 8))
    fig1.canvas.manager.set_window_title('3D Flight Path') # Sets the window title

    ax1 = fig1.add_subplot(111, projection='3d')
    ax1.plot(df['x'], df['y'], df['z'], 'b-', linewidth=1, label='Actual Path', alpha=0.7)

    # Check if target columns exist before plotting them
    if {'target_x', 'target_y', 'target_z'}.issubset(df.columns):
        ax1.plot(df['target_x'], df['target_y'], df['target_z'], 'r--', linewidth=2, label='Target Waypoints', marker='o', markersize=8)

    ax1.scatter(df['x'].iloc[0], df['y'].iloc[0], df['z'].iloc[0], c='green', marker='o', s=100, label='Start')
    ax1.scatter(df['x'].iloc[-1], df['y'].iloc[-1], df['z'].iloc[-1], c='red', marker='s', s=100, label='End')
    ax1.set_xlabel('X Position (m)')
    ax1.set_ylabel('Y Position (m)')
    ax1.set_zlabel('Z Position (m)')
    ax1.set_title('3D Drone Trajectory')
    ax1.legend()
    ax1.grid(True)

    # Save the 3D plot specifically
    fig1.savefig('data/trajectory_3d.png', dpi=300)

    # --- WINDOW 2: 2D Analysis Graphs ---
    fig2 = plt.figure(figsize=(10, 12))
    fig2.canvas.manager.set_window_title('Flight Data Analysis') # Sets the window title

    # Position vs time (Top)
    ax2 = fig2.add_subplot(311)
    ax2.plot(df['time'], df['x'], label='X')
    ax2.plot(df['time'], df['y'], label='Y')
    ax2.plot(df['time'], df['z'], label='Z')
    ax2.set_ylabel('Position (m)')
    ax2.set_title('Position vs Time')
    ax2.legend()
    ax2.grid(True)

    # Velocity vs time (Middle)
    ax3 = fig2.add_subplot(312)
    # Check if velocity columns exist
    if {'vx', 'vy', 'vz'}.issubset(df.columns):
        ax3.plot(df['time'], df['vx'], label='Vx')
        ax3.plot(df['time'], df['vy'], label='Vy')
        ax3.plot(df['time'], df['vz'], label='Vz')
        ax3.set_ylabel('Velocity (m/s)')
        ax3.set_title('Velocity vs Time')
        ax3.legend()
        ax3.grid(True)

    # Attitude vs time (Bottom)
    ax4 = fig2.add_subplot(313)
    ax4.plot(df['time'], np.degrees(df['roll']), label='Roll')
    ax4.plot(df['time'], np.degrees(df['pitch']), label='Pitch')
    ax4.plot(df['time'], np.degrees(df['yaw']), label='Yaw')
    ax4.set_xlabel('Time (s)')
    ax4.set_ylabel('Angle (degrees)')
    ax4.set_title('Attitude vs Time')
    ax4.legend()
    ax4.grid(True)

    plt.tight_layout()
    fig2.savefig('data/trajectory_analysis.png', dpi=300)

    print("Plots saved to data/trajectory_3d.png and data/trajectory_analysis.png")

    # This will open both windows at once and keep them open until you close them
    plt.show()

if __name__ == "__main__":
    plot_trajectory()