247 lines
9.8 KiB
Python
247 lines
9.8 KiB
Python
import numpy as np
|
|
from scipy.spatial.transform import Rotation
|
|
import matplotlib.pyplot as plt
|
|
|
|
from rotorpy.utils.axes3ds import Axes3Ds
|
|
from rotorpy.utils.animate import animate
|
|
|
|
import os
|
|
|
|
"""
|
|
Functions for showing the results from the simulator.
|
|
|
|
"""
|
|
|
|
class Plotter():
|
|
|
|
def __init__(self, results, world):
|
|
|
|
(self.time, self.x, self.x_des, self.v,
|
|
self.v_des, self.q, self.q_des, self.w,
|
|
self.s, self.s_des, self.M, self.T, self.wind,
|
|
self.accel, self.gyro, self.accel_gt,
|
|
self.x_mc, self.v_mc, self.q_mc, self.w_mc,
|
|
self.filter_state, self.covariance, self.sd) = self.unpack_results(results)
|
|
|
|
self.R = Rotation.from_quat(self.q).as_matrix()
|
|
self.R_mc = Rotation.from_quat(self.q_mc).as_matrix() # Rotation as measured by motion capture.
|
|
|
|
self.world = world
|
|
|
|
return
|
|
|
|
def plot_results(self, plot_mocap, plot_estimator, plot_imu):
|
|
"""
|
|
Plot the results
|
|
|
|
"""
|
|
|
|
# 3D Paths
|
|
fig = plt.figure('3D Path')
|
|
ax = Axes3Ds(fig)
|
|
self.world.draw(ax)
|
|
ax.plot3D(self.x[:,0], self.x[:,1], self.x[:,2], 'b.')
|
|
ax.plot3D(self.x_des[:,0], self.x_des[:,1], self.x_des[:,2], 'k')
|
|
|
|
# Position and Velocity vs. Time
|
|
(fig, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num='Pos/Vel vs Time')
|
|
ax = axes[0]
|
|
ax.plot(self.time, self.x_des[:,0], 'r', self.time, self.x_des[:,1], 'g', self.time, self.x_des[:,2], 'b')
|
|
ax.plot(self.time, self.x[:,0], 'r.', self.time, self.x[:,1], 'g.', self.time, self.x[:,2], 'b.')
|
|
ax.legend(('x', 'y', 'z'))
|
|
ax.set_ylabel('position, m')
|
|
ax.grid('major')
|
|
ax.set_title('Position')
|
|
ax = axes[1]
|
|
ax.plot(self.time, self.v_des[:,0], 'r', self.time, self.v_des[:,1], 'g', self.time, self.v_des[:,2], 'b')
|
|
ax.plot(self.time, self.v[:,0], 'r.', self.time, self.v[:,1], 'g.', self.time, self.v[:,2], 'b.')
|
|
ax.legend(('x', 'y', 'z'))
|
|
ax.set_ylabel('velocity, m/s')
|
|
ax.set_xlabel('time, s')
|
|
ax.grid('major')
|
|
|
|
# Orientation and Angular Velocity vs. Time
|
|
(fig, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num='Attitude/Rate vs Time')
|
|
ax = axes[0]
|
|
ax.plot(self.time, self.q_des[:,0], 'r', self.time, self.q_des[:,1], 'g', self.time, self.q_des[:,2], 'b', self.time, self.q_des[:,3], 'm')
|
|
ax.plot(self.time, self.q[:,0], 'r.', self.time, self.q[:,1], 'g.', self.time, self.q[:,2], 'b.', self.time, self.q[:,3], 'm.')
|
|
ax.legend(('i', 'j', 'k', 'w'))
|
|
ax.set_ylabel('quaternion')
|
|
ax.set_xlabel('time, s')
|
|
ax.grid('major')
|
|
ax = axes[1]
|
|
ax.plot(self.time, self.w[:,0], 'r.', self.time, self.w[:,1], 'g.', self.time, self.w[:,2], 'b.')
|
|
ax.legend(('x', 'y', 'z'))
|
|
ax.set_ylabel('angular velocity, rad/s')
|
|
ax.set_xlabel('time, s')
|
|
ax.grid('major')
|
|
|
|
if plot_mocap: # if mocap should be plotted.
|
|
# Motion capture position and velocity vs time
|
|
(fig, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num='Motion Capture Pos/Vel vs Time')
|
|
ax = axes[0]
|
|
ax.plot(self.time, self.x_mc[:,0], 'r.', self.time, self.x_mc[:,1], 'g.', self.time, self.x_mc[:,2], 'b.')
|
|
ax.legend(('x', 'y', 'z'))
|
|
ax.set_ylabel('position, m')
|
|
ax.grid('major')
|
|
ax.set_title('MOTION CAPTURE Position/Velocity')
|
|
ax = axes[1]
|
|
ax.plot(self.time, self.v_mc[:,0], 'r.', self.time, self.v_mc[:,1], 'g.', self.time, self.v_mc[:,2], 'b.')
|
|
ax.legend(('x', 'y', 'z'))
|
|
ax.set_ylabel('velocity, m/s')
|
|
ax.set_xlabel('time, s')
|
|
ax.grid('major')
|
|
# Motion Capture Orientation and Angular Velocity vs. Time
|
|
(fig, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num='Motion Capture Attitude/Rate vs Time')
|
|
ax = axes[0]
|
|
ax.plot(self.time, self.q_mc[:,0], 'r.', self.time, self.q_mc[:,1], 'g.', self.time, self.q_mc[:,2], 'b.', self.time, self.q_mc[:,3], 'm.')
|
|
ax.legend(('i', 'j', 'k', 'w'))
|
|
ax.set_ylabel('quaternion')
|
|
ax.set_xlabel('time, s')
|
|
ax.grid('major')
|
|
ax.set_title("MOTION CAPTURE Attitude/Rate")
|
|
ax = axes[1]
|
|
ax.plot(self.time, self.w_mc[:,0], 'r.', self.time, self.w_mc[:,1], 'g.', self.time, self.w_mc[:,2], 'b.')
|
|
ax.legend(('x', 'y', 'z'))
|
|
ax.set_ylabel('angular velocity, rad/s')
|
|
ax.set_xlabel('time, s')
|
|
ax.grid('major')
|
|
|
|
# Commands vs. Time
|
|
(fig, axes) = plt.subplots(nrows=3, ncols=1, sharex=True, num='Commands vs Time')
|
|
ax = axes[0]
|
|
ax.plot(self.time, self.s_des[:,0], 'r', self.time, self.s_des[:,1], 'g', self.time, self.s_des[:,2], 'b', self.time, self.s_des[:,3], 'k')
|
|
ax.plot(self.time, self.s[:,0], 'r.', self.time, self.s[:,1], 'g.', self.time, self.s[:,2], 'b.', self.time, self.s[:,3], 'k.')
|
|
ax.legend(('1', '2', '3', '4'))
|
|
ax.set_ylabel('motor speeds, rad/s')
|
|
ax.grid('major')
|
|
ax.set_title('Commands')
|
|
ax = axes[1]
|
|
ax.plot(self.time, self.M[:,0], 'r.', self.time, self.M[:,1], 'g.', self.time, self.M[:,2], 'b.')
|
|
ax.legend(('x', 'y', 'z'))
|
|
ax.set_ylabel('moment, N*m')
|
|
ax.grid('major')
|
|
ax = axes[2]
|
|
ax.plot(self.time, self.T, 'k.')
|
|
ax.set_ylabel('thrust, N')
|
|
ax.set_xlabel('time, s')
|
|
ax.grid('major')
|
|
|
|
# Winds
|
|
(fig, axes) = plt.subplots(nrows=3, ncols=1, sharex=True, num='Winds vs Time')
|
|
ax = axes[0]
|
|
ax.plot(self.time, self.wind[:,0], 'r')
|
|
ax.set_ylabel("wind X, m/s")
|
|
ax.grid('major')
|
|
ax.set_title('Winds')
|
|
ax = axes[1]
|
|
ax.plot(self.time, self.wind[:,1], 'g')
|
|
ax.set_ylabel("wind Y, m/s")
|
|
ax.grid('major')
|
|
ax = axes[2]
|
|
ax.plot(self.time, self.wind[:,2], 'b')
|
|
ax.set_ylabel("wind Z, m/s")
|
|
ax.set_xlabel("time, s")
|
|
ax.grid('major')
|
|
|
|
# IMU sensor
|
|
if plot_imu:
|
|
(fig, axes) = plt.subplots(nrows=2, ncols=1, sharex=True, num="IMU Measurements vs Time")
|
|
ax = axes[0]
|
|
ax.plot(self.time, self.accel[:,0], 'r.', self.time, self.accel[:,1], 'g.', self.time, self.accel[:,2], 'b.')
|
|
ax.plot(self.time, self.accel_gt[:,0], 'k', self.time, self.accel_gt[:,1], 'c', self.time, self.accel_gt[:,2], 'm')
|
|
ax.set_ylabel("linear acceleration, m/s/s")
|
|
ax.grid()
|
|
ax = axes[1]
|
|
ax.plot(self.time, self.gyro[:,0], 'r.', self.time, self.gyro[:,1], 'g.', self.time, self.gyro[:,2], 'b.')
|
|
ax.set_ylabel("angular velocity, rad/s")
|
|
ax.grid()
|
|
ax.legend(('x','y','z'))
|
|
ax.set_xlabel("Time, s")
|
|
|
|
if plot_estimator:
|
|
if self.estimator_exists:
|
|
N_filter = self.filter_state.shape[1]
|
|
(fig, axes) = plt.subplots(nrows=N_filter, ncols=1, sharex=True, num="Filter States vs Time")
|
|
fig.set_size_inches(11, 8.5)
|
|
for i in range(N_filter):
|
|
ax = axes[i]
|
|
ax.plot(self.time, self.filter_state[:,i], 'k', )
|
|
ax.fill_between(self.time, self.filter_state[:,i]-self.sd[:,i], self.filter_state[:,i]+self.sd[:,i], alpha=0.3, color='k')
|
|
ax.set_ylabel("x"+str(i))
|
|
ax.set_xlabel("Time, s")
|
|
|
|
(fig, axes) = plt.subplots(nrows=N_filter, ncols=1, sharex=True, num="Filter Covariance vs Time")
|
|
fig.set_size_inches(11, 8.5)
|
|
for i in range(N_filter):
|
|
ax = axes[i]
|
|
ax.plot(self.time, self.sd[:,i]**2, 'k', )
|
|
ax.set_ylabel("cov(x"+str(i)+")")
|
|
ax.set_xlabel("Time, s")
|
|
|
|
plt.show()
|
|
|
|
return
|
|
|
|
def animate_results(self, fname=None):
|
|
"""
|
|
Animate the results
|
|
|
|
"""
|
|
|
|
# Animation (Slow)
|
|
# Instead of viewing the animation live, you may provide a .mp4 filename to save.
|
|
ani = animate(self.time, self.x, self.R, world=self.world, filename=fname)
|
|
plt.show()
|
|
|
|
return
|
|
|
|
def unpack_results(self, result):
|
|
|
|
# Unpack the dictionary of results
|
|
time = result['time']
|
|
state = result['state']
|
|
control = result['control']
|
|
flat = result['flat']
|
|
imu_measurements = result['imu_measurements']
|
|
imu_gt = result['imu_gt']
|
|
mocap = result['mocap_measurements']
|
|
state_estimate = result['state_estimate']
|
|
|
|
# Unpack each result into NumPy arrays
|
|
x = state['x']
|
|
x_des = flat['x']
|
|
v = state['v']
|
|
v_des = flat['x_dot']
|
|
|
|
q = state['q']
|
|
q_des = control['cmd_q']
|
|
w = state['w']
|
|
|
|
s_des = control['cmd_motor_speeds']
|
|
s = state['rotor_speeds']
|
|
M = control['cmd_moment']
|
|
T = control['cmd_thrust']
|
|
|
|
wind = state['wind']
|
|
|
|
accel = imu_measurements['accel']
|
|
gyro = imu_measurements['gyro']
|
|
|
|
accel_gt = imu_gt['accel']
|
|
|
|
x_mc = mocap['x']
|
|
v_mc = mocap['v']
|
|
q_mc = mocap['q']
|
|
w_mc = mocap['w']
|
|
|
|
filter_state = state_estimate['filter_state']
|
|
covariance = state_estimate['covariance']
|
|
if filter_state.shape[1] > 0:
|
|
sd = 3*np.sqrt(np.diagonal(covariance, axis1=1, axis2=2))
|
|
self.estimator_exists = True
|
|
else:
|
|
sd = []
|
|
self.estimator_exists = False
|
|
|
|
return (time, x, x_des, v, v_des, q, q_des, w, s, s_des, M, T, wind, accel, gyro, accel_gt, x_mc, v_mc, q_mc, w_mc, filter_state, covariance, sd) |