ENH: rocket drawing

rocketpy/plots/__init__.py

@@ -1,18 +0,0 @@
-from .aero_surface_plots import (
-    _EllipticalFinsPlots,
-    _FinsPlots,
-    _NoseConePlots,
-    _TailPlots,
-    _TrapezoidalFinsPlots,
-)
-from .compare import Compare, CompareFlights
-from .environment_analysis_plots import _EnvironmentAnalysisPlots
-from .environment_plots import _EnvironmentPlots
-from .flight_plots import _FlightPlots
-from .hybrid_motor_plots import _HybridMotorPlots
-from .liquid_motor_plots import _LiquidMotorPlots
-from .motor_plots import _MotorPlots
-from .rocket_plots import _RocketPlots
-from .solid_motor_plots import _SolidMotorPlots
-from .tank_geometry_plots import _TankGeometryPlots
-from .tank_plots import _TankPlots

rocketpy/plots/rocket_plots.py

@@ -1,5 +1,6 @@
 import matplotlib.pyplot as plt
-import numpy as np
+
+from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail
 
 
 class _RocketPlots:
@@ -104,6 +105,181 @@ def thrust_to_weight(self):
 
         return None
 
+    def draw(self, vis_args=None):
+        """Draws the rocket in a matplotlib figure.
+
+        Parameters
+        ----------
+        vis_args : dict, optional
+            Determines the visual aspects when drawing the rocket. If None,
+            default values are used. Default values are:
+            {
+                "background": "#EEEEEE",
+                "tail": "black",
+                "nose": "black",
+                "body": "dimgrey",
+                "fins": "black",
+                "motor": "black",
+                "buttons": "black",
+                "line_width": 2.0,
+            }
+            A full list of color names can be found at:
+            https://matplotlib.org/stable/gallery/color/named_colors
+        """
+        if vis_args is None:
+            vis_args = {
+                "background": "#EEEEEE",
+                "tail": "black",
+                "nose": "black",
+                "body": "dimgrey",
+                "fins": "black",
+                "motor": "black",
+                "buttons": "black",
+                "line_width": 2.0,
+            }
+
+        # Create the figure and axis
+        _, ax = plt.subplots(figsize=(10, 5))
+        ax.set_aspect("equal")
+        ax.set_facecolor(vis_args["background"])
+        ax.grid(True, linestyle="--", linewidth=0.5)
+
+        csys = self.rocket._csys
+
+        # Draw rocket body
+        reverse = csys == 1
+        self.rocket.aerodynamic_surfaces.sort_by_position(reverse=reverse)
+        y_tube, x_tube = [], []
+        for surface, position in self.rocket.aerodynamic_surfaces:
+            if isinstance(surface, (NoseCone, Tail)):
+                # Append the x and y coordinates of the surface shape_vec to the respective lists
+                x_tube.extend((-csys) * surface.shape_vec[0] + position)
+                y_tube.extend(surface.shape_vec[1])
+            if isinstance(surface, Fins):
+                pass
+
+        # Negate each element in the y_tube list using a list comprehension
+        y_tube_negated = [-y for y in y_tube]
+        plt.plot(
+            x_tube, y_tube, color=vis_args["body"], linewidth=vis_args["line_width"]
+        )
+        plt.plot(
+            x_tube,
+            y_tube_negated,
+            color=vis_args["body"],
+            linewidth=vis_args["line_width"],
+        )
+
+        # Get nozzle position (kinda not working for EmptyMotor class)
+        x_nozzle = self.rocket.motor_position + self.rocket.motor.nozzle_position
+
+        # Find the last point of the rocket
+        idx = -1 if csys == 1 else 0
+        surface, position = self.rocket.aerodynamic_surfaces[idx]
+        length = surface.shape_vec[0][-1] - surface.shape_vec[0][0]
+        x_last = position + (-1 * csys) * length
+        y_last = surface.shape_vec[1][-1]
+
+        plt.plot(
+            [x_nozzle, x_last],
+            [0, y_last],
+            color=vis_args["body"],
+            linewidth=vis_args["line_width"],
+        )
+        plt.plot(
+            [x_nozzle, x_last],
+            [0, -y_last],
+            color=vis_args["body"],
+            linewidth=vis_args["line_width"],
+        )
+
+        # Draw nosecone
+        nosecones = self.rocket.aerodynamic_surfaces.get_tuple_by_type(NoseCone)
+        for nose, position in nosecones:
+            x_nosecone = -csys * nose.shape_vec[0] + position
+            y_nosecone = nose.shape_vec[1]
+
+            plt.plot(
+                x_nosecone,
+                y_nosecone,
+                color=vis_args["nose"],
+                linewidth=vis_args["line_width"] - 0.05,
+            )
+            plt.plot(
+                x_nosecone,
+                -y_nosecone,
+                color=vis_args["nose"],
+                linewidth=vis_args["line_width"] - 0.05,
+            )
+
+        # Draw transitions
+        tails = self.rocket.aerodynamic_surfaces.get_tuple_by_type(Tail)
+        for tail, position in tails:
+            x_tail = -csys * tail.shape_vec[0] + position
+            y_tail = tail.shape_vec[1]
+
+            plt.plot(
+                x_tail,
+                y_tail,
+                color=vis_args["tail"],
+                linewidth=vis_args["line_width"],
+            )
+            plt.plot(
+                x_tail,
+                -y_tail,
+                color=vis_args["tail"],
+                linewidth=vis_args["line_width"],
+            )
+
+        # Draw fins
+        fins = self.rocket.aerodynamic_surfaces.get_tuple_by_type(Fins)
+        for fin, position in fins:
+            x_fin = -csys * fin.shape_vec[0] + position
+            y_fin = fin.shape_vec[1] + self.rocket.radius
+
+            plt.plot(
+                x_fin,
+                y_fin,
+                color=vis_args["fins"],
+                linewidth=vis_args["line_width"],
+            )
+            plt.plot(
+                x_fin,
+                -y_fin,
+                color=vis_args["fins"],
+                linewidth=vis_args["line_width"],
+            )
+
+        # Draw rail buttons
+        buttons, pos = self.rocket.rail_buttons[0]
+        lower = pos
+        upper = pos + buttons.buttons_distance * csys
+        plt.scatter(
+            lower, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=10
+        )
+        plt.scatter(
+            upper, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=10
+        )
+
+        # Draw center of mass and center of pressure
+        cm = self.rocket.center_of_mass(0)
+        plt.scatter(cm, 0, color="black", label="Center of Mass", s=30)
+        plt.scatter(cm, 0, facecolors="none", edgecolors="black", s=100)
+
+        cp = self.rocket.cp_position
+        plt.scatter(cp, 0, label="Center Of Pressure", color="red", s=30, zorder=10)
+        plt.scatter(cp, 0, facecolors="none", edgecolors="red", s=100, zorder=10)
+
+        # Set plot attributes
+        plt.title(f"Rocket Geometry")
+        plt.ylim([-self.rocket.radius * 4, self.rocket.radius * 6])
+        plt.xlabel("Position (m)")
+        plt.ylabel("Radius (m)")
+        plt.legend(loc="best")
+        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.

rocketpy/rocket/__init__.py

@@ -1,4 +1,4 @@
-from .aero_surface import (
+from rocketpy.rocket.aero_surface import (
     AeroSurface,
     EllipticalFins,
     Fins,
@@ -7,6 +7,6 @@
     Tail,
     TrapezoidalFins,
 )
-from .components import Components
-from .parachute import Parachute
-from .rocket import Rocket
+from rocketpy.rocket.components import Components
+from rocketpy.rocket.parachute import Parachute
+from rocketpy.rocket.rocket import Rocket

rocketpy/rocket/aero_surface.py

@@ -1,7 +1,6 @@
 from abc import ABC, abstractmethod
 import warnings
 
-import matplotlib.pyplot as plt
 import numpy as np
 from scipy.optimize import fsolve
 
@@ -20,6 +19,8 @@
     _TrapezoidalFinsPrints,
 )
 
+# TODO: all the evaluate_shape() methods need tests and documentation
+
 
 class AeroSurface(ABC):
     """Abstract class used to define aerodynamic surfaces."""
@@ -1215,6 +1216,30 @@ def evaluate_geometrical_parameters(self):
         self.roll_damping_interference_factor = roll_damping_interference_factor
         self.roll_forcing_interference_factor = roll_forcing_interference_factor
 
+        self.evaluate_shape()
+        return None
+
+    def evaluate_shape(self):
+        if self.sweep_length:
+            points = [
+                (0, 0),
+                (self.sweep_length, self.span),
+                (self.sweep_length + self.tip_chord, self.span),
+                (self.root_chord, 0),
+            ]
+        else:
+            points = [
+                (0, 0),
+                (self.root_chord - self.tip_chord, self.span),
+                (self.root_chord, self.span),
+                (self.root_chord, 0),
+            ]
+
+        x_array, y_array = zip(*points)
+        self.shape_vec = [np.array(x_array), np.array(y_array)]
+
+        return None
+
     def info(self):
         self.prints.geometry()
         self.prints.lift()
@@ -1521,6 +1546,16 @@ def evaluate_geometrical_parameters(self):
         self.roll_damping_interference_factor = roll_damping_interference_factor
         self.roll_forcing_interference_factor = roll_forcing_interference_factor
 
+        self.evaluate_shape()
+        return None
+
+    def evaluate_shape(self):
+        angles = np.arange(0, 360, 5)
+        x_array = self.root_chord / 2 + self.root_chord / 2 * np.cos(np.radians(angles))
+        y_array = self.span * np.sin(np.radians(angles))
+        self.shape_vec = [x_array, y_array]
+        return None
+
     def info(self):
         self.prints.geometry()
         self.prints.lift()
@@ -1675,6 +1710,16 @@ def evaluate_geometrical_parameters(self):
         self.surface_area = (
             np.pi * self.slant_length * (self.top_radius + self.bottom_radius)
         )
+        self.evaluate_shape()
+        return None
+
+    def evaluate_shape(self):
+        # Assuming the tail is a cone, calculate the shape vector
+        self.shape_vec = [
+            np.array([0, self.length]),
+            np.array([self.top_radius, self.bottom_radius]),
+        ]
+        return None
 
     def evaluate_lift_coefficient(self):
         """Calculates and returns tail's lift coefficient.

rocketpy/rocket/components.py

@@ -25,7 +25,13 @@ def __init__(self):
 
     def __repr__(self):
         """Return a string representation of the Components instance."""
-        return repr(self._components)
+        components_str = "\n".join(
+            [
+                f"\tComponent: {str(c.component):80} Position: {c.position:>6.3f}"
+                for c in self._components
+            ]
+        )
+        return f"Components:\n{components_str}"
 
     def __len__(self):
         """Return the number of components in the list of components."""
@@ -155,3 +161,19 @@ def clear(self):
         None
         """
         self._components.clear()
+
+    def sort_by_position(self, reverse=False):
+        """Sort the list of components by position.
+
+        Parameters
+        ----------
+        reverse : bool
+            If True, sort in descending order. If False, sort in ascending
+            order.
+
+        Returns
+        -------
+        None
+        """
+        self._components.sort(key=lambda x: x.position, reverse=reverse)
+        return None