added Python DigitalTwin code
This commit is contained in:
235
MAGLEV_DIGITALTWIN_PYTHON/visualize.py
Normal file
235
MAGLEV_DIGITALTWIN_PYTHON/visualize.py
Normal file
@@ -0,0 +1,235 @@
|
||||
"""
|
||||
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]
|
||||
Reference in New Issue
Block a user