Skip to content

Commit

Permalink
MNT: refactor rocket drawing plot methods
Browse files Browse the repository at this point in the history
  • Loading branch information
@MateusStano
MateusStano committed Mar 24, 2024
1 parent f286413 commit 6805e32
Showing 1 changed file with 152 additions and 140 deletions.
292 changes: 152 additions & 140 deletions rocketpy/plots/rocket_plots.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,6 @@ def draw(self, vis_args=None):
A full list of color names can be found at:
https://matplotlib.org/stable/gallery/color/named_colors
"""
# TODO: we need to modularize this function, it is too big
if vis_args is None:
vis_args = {
"background": "#EEEEEE",
Expand All @@ -212,10 +211,27 @@ def draw(self, vis_args=None):
reverse = csys == 1
self.rocket.aerodynamic_surfaces.sort_by_position(reverse=reverse)

drawn_surfaces = self._draw_aerodynamic_surfaces(ax, vis_args)
last_radius, last_x = self._draw_tubes(ax, drawn_surfaces, vis_args)
self._draw_motor(last_radius, last_x, ax, vis_args)
self._draw_rail_buttons(ax, vis_args)
self._draw_center_of_mass_and_pressure(ax)

plt.title("Rocket Representation")
plt.xlim()
plt.ylim([-self.rocket.radius * 4, self.rocket.radius * 6])
plt.xlabel("Position (m)")
plt.ylabel("Radius (m)")
plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
plt.tight_layout()
plt.show()

def _draw_aerodynamic_surfaces(self, ax, vis_args):
"""Draws the aerodynamic surfaces and saves the position of the points
of interest for the tubes."""
# List of drawn surfaces with the position of points of interest
# and the radius of the rocket at that point
drawn_surfaces = []

# Idea is to get the shape of each aerodynamic surface in their own
# coordinate system and then plot them in the rocket coordinate system
# using the position of each surface
Expand All @@ -225,113 +241,105 @@ def draw(self, vis_args=None):

for surface, position in self.rocket.aerodynamic_surfaces:
if isinstance(surface, NoseCone):
x_nosecone = -csys * surface.shape_vec[0] + position
y_nosecone = surface.shape_vec[1]

ax.plot(
x_nosecone,
y_nosecone,
color=vis_args["nose"],
linewidth=vis_args["line_width"],
)
ax.plot(
x_nosecone,
-y_nosecone,
color=vis_args["nose"],
linewidth=vis_args["line_width"],
)
# close the nosecone
ax.plot(
[x_nosecone[-1], x_nosecone[-1]],
[y_nosecone[-1], -y_nosecone[-1]],
color=vis_args["nose"],
linewidth=vis_args["line_width"],
)

# Add the nosecone to the list of drawn surfaces
drawn_surfaces.append(
(surface, x_nosecone[-1], surface.rocket_radius, x_nosecone[-1])
)

self._draw_nose_cone(ax, surface, position, drawn_surfaces, vis_args)
elif isinstance(surface, Tail):
x_tail = -csys * surface.shape_vec[0] + position
y_tail = surface.shape_vec[1]

ax.plot(
x_tail,
y_tail,
color=vis_args["tail"],
linewidth=vis_args["line_width"],
)
ax.plot(
x_tail,
-y_tail,
color=vis_args["tail"],
linewidth=vis_args["line_width"],
)
# close above and below the tail
ax.plot(
[x_tail[-1], x_tail[-1]],
[y_tail[-1], -y_tail[-1]],
color=vis_args["tail"],
linewidth=vis_args["line_width"],
)
ax.plot(
[x_tail[0], x_tail[0]],
[y_tail[0], -y_tail[0]],
color=vis_args["tail"],
linewidth=vis_args["line_width"],
)
self._draw_tail(ax, surface, position, drawn_surfaces, vis_args)
elif isinstance(surface, Fins):
self._draw_fins(ax, surface, position, drawn_surfaces, vis_args)
return drawn_surfaces

def _draw_nose_cone(self, ax, surface, position, drawn_surfaces, vis_args):
"""Draws the nosecone and saves the position of the points of interest
for the tubes."""
x_nosecone = -self.rocket._csys * surface.shape_vec[0] + position
y_nosecone = surface.shape_vec[1]
ax.plot(
x_nosecone,
y_nosecone,
color=vis_args["nose"],
linewidth=vis_args["line_width"],
)
ax.plot(
x_nosecone,
-y_nosecone,
color=vis_args["nose"],
linewidth=vis_args["line_width"],
)
# close the nosecone
ax.plot(
[x_nosecone[-1], x_nosecone[-1]],
[y_nosecone[-1], -y_nosecone[-1]],
color=vis_args["nose"],
linewidth=vis_args["line_width"],
)
# Add the nosecone to the list of drawn surfaces
drawn_surfaces.append(
(surface, x_nosecone[-1], surface.rocket_radius, x_nosecone[-1])
)

# Add the tail to the list of drawn surfaces
drawn_surfaces.append(
(surface, position, surface.bottom_radius, x_tail[-1])
)
def _draw_tail(self, ax, surface, position, drawn_surfaces, vis_args):
"""Draws the tail and saves the position of the points of interest
for the tubes."""
x_tail = -self.rocket._csys * surface.shape_vec[0] + position
y_tail = surface.shape_vec[1]
ax.plot(
x_tail, y_tail, color=vis_args["tail"], linewidth=vis_args["line_width"]
)
ax.plot(
x_tail, -y_tail, color=vis_args["tail"], linewidth=vis_args["line_width"]
)
# close above and below the tail
ax.plot(
[x_tail[-1], x_tail[-1]],
[y_tail[-1], -y_tail[-1]],
color=vis_args["tail"],
linewidth=vis_args["line_width"],
)
ax.plot(
[x_tail[0], x_tail[0]],
[y_tail[0], -y_tail[0]],
color=vis_args["tail"],
linewidth=vis_args["line_width"],
)
# Add the tail to the list of drawn surfaces
drawn_surfaces.append((surface, position, surface.bottom_radius, x_tail[-1]))

def _draw_fins(self, ax, surface, position, drawn_surfaces, vis_args):
"""Draws the fins and saves the position of the points of interest
for the tubes."""
num_fins = surface.n
x_fin = -self.rocket._csys * surface.shape_vec[0] + position
y_fin = surface.shape_vec[1] + surface.rocket_radius
rotation_angles = [2 * np.pi * i / num_fins for i in range(num_fins)]

for angle in rotation_angles:
# Create a rotation matrix for the current angle around the x-axis
rotation_matrix = np.array([[1, 0], [0, np.cos(angle)]])

# Apply the rotation to the original fin points
rotated_points_2d = np.dot(rotation_matrix, np.vstack((x_fin, y_fin)))

# Extract x and y coordinates of the rotated points
x_rotated, y_rotated = rotated_points_2d

# Project points above the XY plane back into the XY plane (set z-coordinate to 0)
x_rotated = np.where(
rotated_points_2d[1] > 0, rotated_points_2d[0], x_rotated
)
y_rotated = np.where(
rotated_points_2d[1] > 0, rotated_points_2d[1], y_rotated
)
ax.plot(
x_rotated,
y_rotated,
color=vis_args["fins"],
linewidth=vis_args["line_width"],
)

# Draw fins
elif isinstance(surface, Fins):
num_fins = surface.n
x_fin = -csys * surface.shape_vec[0] + position
y_fin = surface.shape_vec[1] + surface.rocket_radius

# Calculate the rotation angles for the other two fins (symmetrically)
rotation_angles = [2 * np.pi * i / num_fins for i in range(num_fins)]

# Apply rotation transformations to get points for the other fins in 2D space
for angle in rotation_angles:
# Create a rotation matrix for the current angle around the x-axis
rotation_matrix = np.array([[1, 0], [0, np.cos(angle)]])

# Apply the rotation to the original fin points
rotated_points_2d = np.dot(
rotation_matrix, np.vstack((x_fin, y_fin))
)

# Extract x and y coordinates of the rotated points
x_rotated, y_rotated = rotated_points_2d

# Project points above the XY plane back into the XY plane (set z-coordinate to 0)
x_rotated = np.where(
rotated_points_2d[1] > 0, rotated_points_2d[0], x_rotated
)
y_rotated = np.where(
rotated_points_2d[1] > 0, rotated_points_2d[1], y_rotated
)

# Plot the fins
ax.plot(
x_rotated,
y_rotated,
color=vis_args["fins"],
linewidth=vis_args["line_width"],
)

# Add the fin to the list of drawn surfaces
drawn_surfaces.append(
(surface, position, surface.rocket_radius, x_rotated[-1])
)
drawn_surfaces.append((surface, position, surface.rocket_radius, x_rotated[-1]))

# Draw tubes
def _draw_tubes(self, ax, drawn_surfaces, vis_args):
"""Draws the tubes between the aerodynamic surfaces."""
for i, d_surface in enumerate(drawn_surfaces):
# Draw the tubes, from the end of the first surface to the beginning
# of the next surface, with the radius of the rocket at that point
Expand Down Expand Up @@ -368,18 +376,41 @@ def draw(self, vis_args=None):
color=vis_args["body"],
linewidth=vis_args["line_width"],
)
return radius, last_x

# Draw motor
def _draw_motor(self, last_radius, last_x, ax, vis_args):
"""Draws the motor from motor patches"""
total_csys = self.rocket._csys * self.rocket.motor._csys
nozzle_position = (
self.rocket.motor_position + self.rocket.motor.nozzle_position * total_csys
)

# List of motor patches
# Get motor patches translated to the correct position
motor_patches = self._generate_motor_patches(total_csys, ax, vis_args)

# Draw patches
if not isinstance(self.rocket.motor, EmptyMotor):
# Add nozzle last so it is in front of the other patches
nozzle = self.rocket.motor.plots._generate_nozzle(
translate=(nozzle_position, 0), csys=self.rocket._csys
)
motor_patches += [nozzle]

outline = self.rocket.motor.plots._generate_motor_region(
list_of_patches=motor_patches
)
# add outline first so it is behind the other patches
ax.add_patch(outline)
for patch in motor_patches:
ax.add_patch(patch)

self._draw_nozzle_tube(last_radius, last_x, nozzle_position, ax, vis_args)

def _generate_motor_patches(self, total_csys, ax, vis_args):
"""Generates motor patches for drawing"""
motor_patches = []

# Get motor patches translated to the correct position
if isinstance(self.rocket.motor, (SolidMotor)):
if isinstance(self.rocket.motor, SolidMotor):
grains_cm_position = (
self.rocket.motor_position
+ self.rocket.motor.grains_center_of_mass_position * total_csys
Expand Down Expand Up @@ -452,27 +483,17 @@ def draw(self, vis_args=None):
)
motor_patches += [tank]

# add nozzle last so it is in front of the other patches
if not isinstance(self.rocket.motor, EmptyMotor):
nozzle = self.rocket.motor.plots._generate_nozzle(
translate=(nozzle_position, 0), csys=self.rocket._csys
)
motor_patches += [nozzle]
outline = self.rocket.motor.plots._generate_motor_region(
list_of_patches=motor_patches
)
# add outline first so it is behind the other patches
ax.add_patch(outline)
for patch in motor_patches:
ax.add_patch(patch)
return motor_patches

def _draw_nozzle_tube(self, last_radius, last_x, nozzle_position, ax, vis_args):
"""Draws the tube from the last surface to the nozzle position."""
# Check if nozzle is beyond the last surface, if so draw a tube
# to it, with the radius of the last surface
if self.rocket._csys == 1:
if nozzle_position < last_x:
x_tube = [last_x, nozzle_position]
y_tube = [radius, radius]
y_tube_negated = [-radius, -radius]
y_tube = [last_radius, last_radius]
y_tube_negated = [-last_radius, -last_radius]

Check warning on line 496 in rocketpy/plots/rocket_plots.py

View check run for this annotation

Codecov / codecov/patch

rocketpy/plots/rocket_plots.py#L495-L496 

Added lines #L495 - L496 were not covered by tests

ax.plot(
x_tube,
Expand All @@ -489,8 +510,8 @@ def draw(self, vis_args=None):
else: # if self.rocket._csys == -1:
if nozzle_position > last_x:
x_tube = [last_x, nozzle_position]
y_tube = [radius, radius]
y_tube_negated = [-radius, -radius]
y_tube = [last_radius, last_radius]
y_tube_negated = [-last_radius, -last_radius]

Check warning on line 514 in rocketpy/plots/rocket_plots.py

View check run for this annotation

Codecov / codecov/patch

rocketpy/plots/rocket_plots.py#L513-L514 

Added lines #L513 - L514 were not covered by tests

ax.plot(
x_tube,
Expand All @@ -505,11 +526,12 @@ def draw(self, vis_args=None):
linewidth=vis_args["line_width"],
)

# Draw rail buttons
def _draw_rail_buttons(self, ax, vis_args):
"""Draws the rail buttons of the rocket."""
try:
buttons, pos = self.rocket.rail_buttons[0]
lower = pos
upper = pos + buttons.buttons_distance * csys
upper = pos + buttons.buttons_distance * self.rocket._csys
ax.scatter(
lower, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=15
)
Expand All @@ -519,6 +541,8 @@ def draw(self, vis_args=None):
except IndexError:
pass

def _draw_center_of_mass_and_pressure(self, ax):
"""Draws the center of mass and center of pressure of the rocket."""
# Draw center of mass and center of pressure
cm = self.rocket.center_of_mass(0)
ax.scatter(cm, 0, color="#1565c0", label="Center of Mass", s=10)
Expand All @@ -528,18 +552,6 @@ def draw(self, vis_args=None):
cp, 0, label="Static Center of Pressure", color="red", s=10, zorder=10
)

# Set plot attributes
plt.title("Rocket Representation")
plt.xlim()
plt.ylim([-self.rocket.radius * 4, self.rocket.radius * 6])
plt.xlabel("Position (m)")
plt.ylabel("Radius (m)")
plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
plt.tight_layout()
plt.show()

return None

def all(self):
"""Prints out all graphs available about the Rocket. It simply calls
all the other plotter methods in this class.
Expand Down

0 comments on commit 6805e32

Please sign in to comment.