WIP: nicer plots

mpc_demo/main.py

@@ -10,13 +10,17 @@ from cvxpy_mpc import optimize
 import sys
 import time
 
-# classes
+SIM_START_X=0
+SIM_START_Y=2
+SIM_START_H=0
+
+# Classes
 class MPC():
 
     def __init__(self):
 
         # State for the robot mathematical model [x,y,heading]
-        self.state =  np.zeros(3)
+        self.state =  [SIM_START_X, SIM_START_Y, SIM_START_H]
         # Sim step
         self.dt = 0.25
 
@@ -31,8 +35,14 @@ class MPC():
         # Interpolated Path to follow given waypoints
         self.path = compute_path_from_wp([0,20,30,30],[0,0,10,20])
 
+        # Sim help vars
+        self.sim_time=0
+        self.x_history=[]
+        self.y_history=[]
+        self.v_history=[]
+        self.w_history=[]
+
         #Initialise plot
-        # First set up the figure, the axis, and the plot element we want to animate
         plt.style.use("ggplot")
         self.fig = plt.figure()
         plt.ion()
@@ -70,13 +80,39 @@ class MPC():
         self.state[2] = self.state[2] +ang_v*self.dt
 
     def plot_sim(self):
-        plt.clf()
-        self.ax = plt.axes(xlim=(np.min(self.path[0,:])-1, np.max(self.path[0,:])+1),
-                      ylim=(np.min(self.path[1,:])-1, np.max(self.path[1,:])+1))
 
-        self.track, = self.ax.plot(self.path[0,:],self.path[1,:], "g-", label="reference track")
+        self.sim_time = self.sim_time+self.dt
-        self.vehicle, = self.ax.plot([self.state[0]], [self.state[1]], "r*", label="vehicle path")
+        self.x_history.append(self.state[0])
+        self.y_history.append(self.state[1])
+        self.v_history.append(self.opt_u[0,1])
+        self.w_history.append(self.opt_u[1,1])
+
+        plt.clf()
+        plt.title("MPC Simulation \n" + "Simulation elapsed time {}s".format(self.sim_time))
+
+        grid = plt.GridSpec(2, 3)
+        plt.subplot(grid[0:2, 0:2])
+
+        # self.ax = plt.axes(xlim=(np.min(self.path[0,:])-1, np.max(self.path[0,:])+1),
+        #               ylim=(np.min(self.path[1,:])-1, np.max(self.path[1,:])+1))
+
+        plt.plot(self.path[0,:],self.path[1,:], c='tab:orange', marker=".", label="reference track")
+        plt.plot(self.x_history, self.y_history, c='tab:blue', marker=".", alpha=0.5, label="vehicle trajectory")
+        plt.plot(self.x_history[-1], self.y_history[-1], c='tab:blue', marker=".", markersize=12, label="vehicle position")
+        plt.ylabel('map y')
+        plt.xlabel('map x')
         plt.legend()
+
+        plt.subplot(grid[0, 2])
+        plt.plot(self.v_history,c='tab:orange')
+        plt.ylabel('v(t) [m/s]')
+
+        plt.subplot(grid[1, 2])
+        plt.plot(np.degrees(self.w_history),c='tab:orange')
+        plt.ylabel('w(t) [deg/s]')
+
+        plt.tight_layout()
+
         plt.draw()
         plt.pause(0.1)