ENH: individual fin in rocket draw

Author: MateusStano

rocketpy/plots/rocket_plots.py

@@ -1,9 +1,9 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
+from rocketpy.mathutils.vector_matrix import Vector
 from rocketpy.motors import EmptyMotor, HybridMotor, LiquidMotor, SolidMotor
 from rocketpy.rocket.aero_surface import Fin, Fins, NoseCone, Tail
-from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail
 from rocketpy.rocket.aero_surface.generic_surface import GenericSurface
 
 from .plot_helpers import show_or_save_plot
@@ -178,7 +178,7 @@ def draw(self, vis_args=None, plane="xz", *, filename=None):
             and webp (these are the formats supported by matplotlib).
         """
 
-        self.__validate_aerodynamic_surfaces()
+        self.__validate_aerodynamic_surfaces(plane)
 
         if vis_args is None:
             vis_args = {
@@ -198,9 +198,9 @@ def draw(self, vis_args=None, plane="xz", *, filename=None):
 
         csys = self.rocket._csys
         reverse = csys == 1
-        self.rocket.aerodynamic_surfaces.sort_by_position(reverse=reverse)
+        surfaces = self.rocket.aerodynamic_surfaces.sort_by_position(reverse=reverse)
 
-        drawn_surfaces = self._draw_aerodynamic_surfaces(ax, vis_args, plane)
+        drawn_surfaces = self._draw_aerodynamic_surfaces(ax, vis_args, plane, surfaces)
         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)
@@ -216,13 +216,15 @@ def draw(self, vis_args=None, plane="xz", *, filename=None):
         plt.tight_layout()
         show_or_save_plot(filename)
 
-    def __validate_aerodynamic_surfaces(self):
+    def __validate_aerodynamic_surfaces(self, plane):
         if not self.rocket.aerodynamic_surfaces:
             raise ValueError(
                 "The rocket must have at least one aerodynamic surface to be drawn."
             )
+        if plane != "xz" and plane != "yz":
+            raise ValueError("The plane must be 'xz' or 'yz'. The default is 'xz'.")
 
-    def _draw_aerodynamic_surfaces(self, ax, vis_args, plane):
+    def _draw_aerodynamic_surfaces(self, ax, vis_args, plane, surfaces):
         """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
@@ -235,15 +237,17 @@ def _draw_aerodynamic_surfaces(self, ax, vis_args, plane):
         # diameter changes. The final point of the last surface is the final
         # point of the last tube
 
-        for surface, position in self.rocket.aerodynamic_surfaces:
+        for surface, position in surfaces:
             if isinstance(surface, NoseCone):
                 self._draw_nose_cone(ax, surface, position.z, drawn_surfaces, vis_args)
             elif isinstance(surface, Tail):
                 self._draw_tail(ax, surface, position.z, drawn_surfaces, vis_args)
             elif isinstance(surface, Fins):
-                self._draw_fins(ax, surface, position.z, drawn_surfaces, vis_args)
+                self._draw_fins(
+                    ax, surface, position.z, drawn_surfaces, vis_args, plane
+                )
             elif isinstance(surface, Fin):
-                self._draw_fin(ax, surface, position.z, drawn_surfaces, vis_args)
+                self._draw_fin(ax, surface, position, drawn_surfaces, vis_args, plane)
             elif isinstance(surface, GenericSurface):
                 self._draw_generic_surface(
                     ax, surface, position, drawn_surfaces, vis_args, plane
@@ -306,13 +310,15 @@ def _draw_tail(self, ax, surface, position, drawn_surfaces, vis_args):
         # 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):
+    def _draw_fins(self, ax, surface, position, drawn_surfaces, vis_args, plane):
         """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)]
+        rotation_angles = np.array([2 * np.pi * i / num_fins for i in range(num_fins)])
+        if plane == "xz":
+            rotation_angles -= np.pi / 2
 
         for angle in rotation_angles:
             # Create a rotation matrix for the current angle around the x-axis
@@ -324,13 +330,6 @@ def _draw_fins(self, ax, surface, position, drawn_surfaces, vis_args):
             # 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,
@@ -340,34 +339,55 @@ def _draw_fins(self, ax, surface, position, drawn_surfaces, vis_args):
 
         drawn_surfaces.append((surface, position, surface.rocket_radius, x_rotated[-1]))
 
-    def _draw_fin(self, ax, surface, position, drawn_surfaces, vis_args):
-        """Draws the fins and saves the position of the points of interest
-        for the tubes."""
+    def _draw_fin(self, ax, surface, position, drawn_surfaces, vis_args, plane):
+        """Draws individual fins."""
 
-        x_fin = -self.rocket._csys * surface.shape_vec[0] + position
-        y_fin = surface.shape_vec[1] + surface.rocket_radius
-        angle = surface.angular_position
+        # Get shape vec
+        xs = surface.shape_vec[0]
+        ys = surface.shape_vec[1]
+        zs = np.zeros_like(xs)
+
+        # Define shape in fin coordinate system
+        x_fin = -zs
+        y_fin = ys
+        z_fin = xs
+        points = np.column_stack((x_fin, y_fin, z_fin))
+
+        # Move drawing coordinates to center of fin for cant angle rotation
+        xd = np.array([0, 0, max(xs) / 2])
+        points -= xd
 
-        # Create a rotation matrix for the angle around the x-axis
-        rotation_matrix = np.array([[1, 0], [0, np.cos(angle)]])
+        # Rotate to body coordinate system
+        for i, p in enumerate(points):
+            points[i] = surface._rotation_fin_to_body @ Vector(p)
 
-        # Apply the rotation to the original fin points
-        rotated_points_2d = np.dot(rotation_matrix, np.vstack((x_fin, y_fin)))
+        rotated_xd = surface._rotation_fin_to_body @ Vector(xd)
+        points += np.array(rotated_xd)
 
-        # Extract x and y coordinates of the rotated points
-        x_rotated, y_rotated = rotated_points_2d
+        # Back to the drawing system
+        x_fin_rotated = points[:, 0]
+        y_fin_rotated = points[:, 1]
+        z_fin_rotated = points[:, 2]
+
+        if plane == "xz":
+            x_rotated = self.rocket._csys * z_fin_rotated + position.z
+            y_rotated = x_fin_rotated + position.x
+        elif plane == "yz":
+            x_rotated = self.rocket._csys * z_fin_rotated + position.z
+            y_rotated = y_fin_rotated + position.y
+        else:  # pragma: no cover
+            raise ValueError("Plane must be 'xz' or 'yz'.")
 
-        # 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"],
         )
 
-        drawn_surfaces.append((surface, position, surface.rocket_radius, x_rotated[-1]))
+        drawn_surfaces.append(
+            (surface, position.z, surface.rocket_radius, x_rotated[-1])
+        )
 
     def _draw_generic_surface(
         self,
@@ -390,8 +410,6 @@ def _draw_generic_surface(
             x_pos = position[2]
             # y position of the surface is the y position in the plot
             y_pos = position[1]
-        else:  # pragma: no cover
-            raise ValueError("Plane must be 'xz' or 'yz'.")
 
         ax.scatter(
             x_pos,
@@ -470,9 +488,7 @@ def _draw_motor(self, last_radius, last_x, ax, vis_args):
 
         self._draw_nozzle_tube(last_radius, last_x, nozzle_position, ax, vis_args)
 
-    def _generate_motor_patches(
-        self, total_csys, ax
-    ):  # pylint: disable=unused-argument
+    def _generate_motor_patches(self, total_csys, ax):
         """Generates motor patches for drawing"""
         motor_patches = []
 
@@ -680,7 +696,7 @@ def all(self):
 
         # Rocket draw
         if len(self.rocket.aerodynamic_surfaces) > 0:
-            print("\nRocket Draw")
+            print("\nRocket Drawing")
             print("-" * 40)
             self.draw()