Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

MNT: Modularize Rocket Draw #580

Merged
merged 5 commits into from
Apr 25, 2024
Merged

MNT: Modularize Rocket Draw #580

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6805e32
MNT: refactor rocket drawing plot methods
MateusStano Mar 24, 2024
51bdd91
MNT: changelog
MateusStano Mar 24, 2024
e3e1a59
DEV: fix changelog order
MateusStano Apr 4, 2024
822a89e
Merge branch 'develop' into mnt/modularize-rocket-draw
MateusStano Apr 7, 2024
2cbaa77
Merge branch 'develop' into mnt/modularize-rocket-draw
MateusStano Apr 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
### Changed

- DEP: delete deprecated rocketpy.tools.cached_property [#587](https://github.com/RocketPy-Team/RocketPy/pull/587)
- MNT: Modularize Rocket Draw [#580](https://github.com/RocketPy-Team/RocketPy/pull/580)
- DOC: Improvements of Environment docstring phrasing [#565](https://github.com/RocketPy-Team/RocketPy/pull/565)
- MNT: Refactor flight prints module [#579](https://github.com/RocketPy-Team/RocketPy/pull/579)
- DOC: Convert CompareFlights example notebooks to .rst files [#576](https://github.com/RocketPy-Team/RocketPy/pull/576)
Expand Down
292 changes: 152 additions & 140 deletions rocketpy/plots/rocket_plots.py
View file
Open in desktop
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])
)
Comment on lines +275 to +278
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is strange to me. You add the surface to a list, which was passed as an argument to the function. If you don't return the list with the added surfaces, is this code line really necessary? I'm afraid that everytime you are starting with an empty list.

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

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This behavior is also present in other parts of the code.

Copy link
Member Author

@MateusStano MateusStano Apr 7, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The drawn_surfaces list is created and then appended to in the various draw methods I created. That same list is used later to draw the tubes later.

Remember that a list in Python is just a pointer. Changing the name of the list variable does not change the address that it points to.


# 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]

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]

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
Loading