236 lines
7.7 KiB
Python
236 lines
7.7 KiB
Python
"""
|
|
Visualization function for quadrotor/maglev pod
|
|
Ported from visualizeQuad.m
|
|
"""
|
|
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from mpl_toolkits.mplot3d import Axes3D
|
|
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
|
|
from matplotlib.animation import FuncAnimation, PillowWriter
|
|
from scipy.interpolate import interp1d
|
|
from utils import euler2dcm
|
|
from parameters import QuadParams
|
|
|
|
|
|
def visualize_quad(S):
|
|
"""
|
|
Takes in an input structure S and visualizes the resulting 3D motion
|
|
in approximately real-time. Outputs the data used to form the plot.
|
|
|
|
Parameters
|
|
----------
|
|
S : dict
|
|
Structure with the following elements:
|
|
- rMat : (M, 3) matrix of quad positions, in meters
|
|
- eMat : (M, 3) matrix of quad attitudes, in radians
|
|
- tVec : (M,) vector of times corresponding to each measurement
|
|
- plotFrequency : scalar number of frames per second (Hz)
|
|
- bounds : 6-element list [xmin, xmax, ymin, ymax, zmin, zmax]
|
|
- makeGifFlag : boolean (if True, export plot to .gif)
|
|
- gifFileName : string with the file name of the .gif
|
|
|
|
Returns
|
|
-------
|
|
P : dict
|
|
Structure with the following elements:
|
|
- tPlot : (N,) vector of time points used in the plot
|
|
- rPlot : (N, 3) vector of positions used to generate the plot
|
|
- ePlot : (N, 3) vector of attitudes used to generate the plot
|
|
"""
|
|
# UT colors
|
|
burntOrangeUT = np.array([191, 87, 0]) / 255
|
|
darkGrayUT = np.array([51, 63, 72]) / 255
|
|
|
|
# Get quad parameters
|
|
quad_params = QuadParams()
|
|
|
|
frame_l = 1.2192
|
|
frame_w = 0.711 # in meters
|
|
yh = quad_params.yh
|
|
|
|
# Parameters for the rotors
|
|
rotorLocations = np.array([
|
|
[frame_l/2, frame_l/2, -frame_l/2, -frame_l/2],
|
|
[frame_w/2, -frame_w/2, frame_w/2, -frame_w/2],
|
|
[yh, yh, yh, yh]
|
|
])
|
|
r_rotor = 0.13
|
|
|
|
# Determine the location of the corners of the body box in the body frame
|
|
bpts = np.array([
|
|
[frame_l/2, frame_l/2, -frame_l/2, -frame_l/2, frame_l/2, frame_l/2, -frame_l/2, -frame_l/2],
|
|
[frame_w/2, -frame_w/2, frame_w/2, -frame_w/2, frame_w/2, -frame_w/2, frame_w/2, -frame_w/2],
|
|
[0.03, 0.03, 0.03, 0.03, -0.030, -0.030, -0.030, -0.030]
|
|
])
|
|
|
|
# Rectangles representing each side of the body box
|
|
b1 = bpts[:, [0, 4, 5, 1]]
|
|
b2 = bpts[:, [0, 4, 6, 2]]
|
|
b3 = bpts[:, [2, 6, 7, 3]]
|
|
b4 = bpts[:, [0, 2, 3, 1]]
|
|
b5 = bpts[:, [4, 6, 7, 5]]
|
|
b6 = bpts[:, [1, 5, 7, 3]]
|
|
|
|
# Create a circle for each rotor
|
|
t_circ = np.linspace(0, 2*np.pi, 20)
|
|
circpts = np.zeros((3, 20))
|
|
for i in range(20):
|
|
circpts[:, i] = r_rotor * np.array([np.cos(t_circ[i]), np.sin(t_circ[i]), 0])
|
|
|
|
m = len(S['tVec'])
|
|
|
|
# Single epoch plot
|
|
if m == 1:
|
|
fig = plt.figure(figsize=(10, 8))
|
|
ax = fig.add_subplot(111, projection='3d')
|
|
|
|
# Extract params
|
|
RIB = euler2dcm(S['eMat'][0, :]).T
|
|
r = S['rMat'][0, :]
|
|
|
|
# Plot rotors
|
|
for i in range(4):
|
|
rotor_circle = r.reshape(3, 1) + RIB @ (circpts + rotorLocations[:, i].reshape(3, 1))
|
|
ax.plot(rotor_circle[0, :], rotor_circle[1, :], rotor_circle[2, :],
|
|
color=darkGrayUT, linewidth=2)
|
|
|
|
# Plot body
|
|
plot_body(ax, r, RIB, [b1, b2, b3, b4, b5, b6])
|
|
|
|
# Plot body axes
|
|
plot_axes(ax, r, RIB)
|
|
|
|
ax.set_xlim([S['bounds'][0], S['bounds'][1]])
|
|
ax.set_ylim([S['bounds'][2], S['bounds'][3]])
|
|
ax.set_zlim([S['bounds'][4], S['bounds'][5]])
|
|
ax.set_xlabel('X')
|
|
ax.set_ylabel('Y')
|
|
ax.set_zlabel('Z')
|
|
ax.grid(True)
|
|
|
|
plt.show()
|
|
|
|
P = {
|
|
'tPlot': S['tVec'],
|
|
'rPlot': S['rMat'],
|
|
'ePlot': S['eMat']
|
|
}
|
|
|
|
else: # Animation
|
|
# Create time vectors
|
|
tf = 1 / S['plotFrequency']
|
|
tmax = S['tVec'][-1]
|
|
tmin = S['tVec'][0]
|
|
tPlot = np.arange(tmin, tmax + tf/2, tf)
|
|
tPlotLen = len(tPlot)
|
|
|
|
# Interpolate to regularize times
|
|
t2unique, indUnique = np.unique(S['tVec'], return_index=True)
|
|
|
|
rPlot_interp = interp1d(t2unique, S['rMat'][indUnique, :], axis=0,
|
|
kind='linear', fill_value='extrapolate')
|
|
ePlot_interp = interp1d(t2unique, S['eMat'][indUnique, :], axis=0,
|
|
kind='linear', fill_value='extrapolate')
|
|
|
|
rPlot = rPlot_interp(tPlot)
|
|
ePlot = ePlot_interp(tPlot)
|
|
|
|
# Create figure and axis
|
|
fig = plt.figure(figsize=(10, 8))
|
|
ax = fig.add_subplot(111, projection='3d')
|
|
|
|
# Storage for plot elements
|
|
plot_elements = []
|
|
|
|
def init():
|
|
ax.set_xlim([S['bounds'][0], S['bounds'][1]])
|
|
ax.set_ylim([S['bounds'][2], S['bounds'][3]])
|
|
ax.set_zlim([S['bounds'][4], S['bounds'][5]])
|
|
ax.set_xlabel('X')
|
|
ax.set_ylabel('Y')
|
|
ax.set_zlabel('Z')
|
|
ax.view_init(elev=0, azim=0)
|
|
ax.grid(True)
|
|
return []
|
|
|
|
def update(frame):
|
|
nonlocal plot_elements
|
|
|
|
# Clear previous frame
|
|
for elem in plot_elements:
|
|
elem.remove()
|
|
plot_elements = []
|
|
|
|
# Extract data
|
|
RIB = euler2dcm(ePlot[frame, :]).T
|
|
r = rPlot[frame, :]
|
|
|
|
# Plot rotors
|
|
for i in range(4):
|
|
rotor_circle = r.reshape(3, 1) + RIB @ (circpts + rotorLocations[:, i].reshape(3, 1))
|
|
line, = ax.plot(rotor_circle[0, :], rotor_circle[1, :], rotor_circle[2, :],
|
|
color=darkGrayUT, linewidth=2)
|
|
plot_elements.append(line)
|
|
|
|
# Plot body
|
|
body_elem = plot_body(ax, r, RIB, [b1, b2, b3, b4, b5, b6])
|
|
plot_elements.append(body_elem)
|
|
|
|
# Plot axes
|
|
axis_elems = plot_axes(ax, r, RIB)
|
|
plot_elements.extend(axis_elems)
|
|
|
|
ax.set_title(f'Time: {tPlot[frame]:.3f} s')
|
|
|
|
return plot_elements
|
|
|
|
anim = FuncAnimation(fig, update, init_func=init, frames=tPlotLen,
|
|
interval=tf*1000, blit=False, repeat=True)
|
|
|
|
# Save as GIF if requested
|
|
if S.get('makeGifFlag', False):
|
|
writer = PillowWriter(fps=S['plotFrequency'])
|
|
anim.save(S['gifFileName'], writer=writer)
|
|
print(f"Animation saved to {S['gifFileName']}")
|
|
|
|
plt.show()
|
|
|
|
P = {
|
|
'tPlot': tPlot,
|
|
'rPlot': rPlot,
|
|
'ePlot': ePlot
|
|
}
|
|
|
|
return P
|
|
|
|
|
|
def plot_body(ax, r, RIB, body_faces):
|
|
"""Plot the body box"""
|
|
vertices = []
|
|
for face in body_faces:
|
|
face_r = r.reshape(3, 1) + RIB @ face
|
|
vertices.append(face_r.T)
|
|
|
|
# Create 3D polygon collection
|
|
poly = Poly3DCollection(vertices, alpha=0.5, facecolor=[0.5, 0.5, 0.5],
|
|
edgecolor='black', linewidths=1)
|
|
ax.add_collection3d(poly)
|
|
return poly
|
|
|
|
|
|
def plot_axes(ax, r, RIB):
|
|
"""Plot body axes"""
|
|
bodyX = 0.5 * RIB @ np.array([1, 0, 0])
|
|
bodyY = 0.5 * RIB @ np.array([0, 1, 0])
|
|
bodyZ = 0.5 * RIB @ np.array([0, 0, 1])
|
|
|
|
q1 = ax.quiver(r[0], r[1], r[2], bodyX[0], bodyX[1], bodyX[2],
|
|
color='red', arrow_length_ratio=0.3)
|
|
q2 = ax.quiver(r[0], r[1], r[2], bodyY[0], bodyY[1], bodyY[2],
|
|
color='blue', arrow_length_ratio=0.3)
|
|
q3 = ax.quiver(r[0], r[1], r[2], bodyZ[0], bodyZ[1], bodyZ[2],
|
|
color='green', arrow_length_ratio=0.3)
|
|
|
|
return [q1, q2, q3]
|