diff --git a/rocketpy/plots/rocket_plots.py b/rocketpy/plots/rocket_plots.py index 7074704e6..2be7a4a73 100644 --- a/rocketpy/plots/rocket_plots.py +++ b/rocketpy/plots/rocket_plots.py @@ -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", @@ -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 @@ -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 @@ -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 @@ -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, @@ -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, @@ -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 ) @@ -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) @@ -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.