MNT: refactor rocket drawing plot methods

Author: MateusStano

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.