Merge pull request #435 from RocketPy-Team/enh/solid-motors-draw

ENH: solid motors draw

rocketpy/motors/solid_motor.py

@@ -701,6 +701,10 @@ def propellant_I_13(self):
     def propellant_I_23(self):
         return 0
 
+    def draw(self):
+        self.plots.draw()
+        return None
+
     def info(self):
         """Prints out basic data about the SolidMotor."""
         self.prints.all()

rocketpy/plots/rocket_plots.py

@@ -1,6 +1,8 @@
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
 
+from rocketpy.motors import HybridMotor, LiquidMotor, SolidMotor
+from rocketpy.plots import _generate_nozzle
 from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail
 
 
@@ -159,7 +161,7 @@ def draw(self, vis_args=None):
             }
 
         # Create the figure and axis
-        _, ax = plt.subplots(figsize=(8, 5))
+        _, ax = plt.subplots(figsize=(8, 6))
         ax.set_aspect("equal")
         ax.set_facecolor(vis_args["background"])
         ax.grid(True, linestyle="--", linewidth=0.5)
@@ -335,6 +337,30 @@ def draw(self, vis_args=None):
         ax.scatter(
             nozzle_position, 0, label="Nozzle Outlet", color="brown", s=10, zorder=10
         )
+
+        if isinstance(self.rocket.motor, (SolidMotor, HybridMotor)):
+            chamber = self.rocket.motor.plots._generate_combustion_chamber(
+                translate=(nozzle_position, 0), csys=self.rocket._csys
+            )
+            ax.add_patch(chamber)
+            grains = self.rocket.motor.plots._generate_grains(
+                translate=(nozzle_position, 0), csys=self.rocket._csys
+            )
+            for grain in grains:
+                ax.add_patch(grain)
+
+        elif isinstance(self.rocket.motor, LiquidMotor):
+            patches = self.rocket.motor.plots._generate_positioned_tanks(
+                translate=(self.rocket.motor_position, 0), csys=self.rocket._csys
+            )
+            for patch in patches:
+                ax.add_patch(patch)
+
+        nozzle = _generate_nozzle(
+            self.rocket.motor, translate=(nozzle_position, 0), csys=self.rocket._csys
+        )
+        ax.add_patch(nozzle)
+
         # 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:
@@ -393,16 +419,16 @@ def draw(self, vis_args=None):
         ax.scatter(cm, 0, color="black", label="Center of Mass", s=30)
         ax.scatter(cm, 0, facecolors="none", edgecolors="black", s=100)
 
-        cp = self.rocket.cp_position
-        ax.scatter(cp, 0, label="Center Of Pressure", color="red", s=30, zorder=10)
+        cp = self.rocket.cp_position(0)
+        ax.scatter(cp, 0, label="CP (M=0)", color="red", s=30, zorder=10)
         ax.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.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
         plt.tight_layout()
         plt.show()
 

rocketpy/plots/solid_motor_plots.py

@@ -1,3 +1,10 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.patches import Polygon
+
+from rocketpy.plots import _generate_nozzle
+
+
 class _SolidMotorPlots:
     """Class that holds plot methods for SolidMotor class.
 
@@ -361,6 +368,144 @@ def I_23(self, lower_limit=None, upper_limit=None):
 
         return None
 
+    def draw(self):
+        """Draws a simple 2D representation of the SolidMotor."""
+        fig, ax = plt.subplots(figsize=(10, 6))
+        ax.set_aspect("equal")
+
+        _csys = self.solid_motor._csys
+
+        chamber = self._generate_combustion_chamber(_csys=_csys)
+        nozzle = _generate_nozzle(self.solid_motor)
+        grains = self._generate_grains(_csys=_csys, translate=(0, 0))
+
+        ax.add_patch(chamber)
+        for grain in grains:
+            ax.add_patch(grain)
+        ax.add_patch(nozzle)
+
+        # self._draw_nozzle(ax, nozzle_height, csys)
+        # self._draw_combustion_chamber(ax, csys)
+        # self._draw_grains(ax, csys)
+        self._draw_center_of_mass(ax)
+
+        self._set_plot_properties(ax)
+        plt.show()
+        return None
+
+    def _generate_combustion_chamber(self, translate=(0, 0), csys=1):
+        # csys = self.solid_motor.csys
+        chamber_length = (
+            abs(
+                self.solid_motor.center_of_dry_mass_position
+                - self.solid_motor.nozzle_position
+            )
+            * 2
+        ) * csys
+        x = np.array(
+            [
+                self.solid_motor.nozzle_position,
+                self.solid_motor.nozzle_position,
+                self.solid_motor.nozzle_position + chamber_length,
+                self.solid_motor.nozzle_position + chamber_length,
+            ]
+        )
+        y = np.array(
+            [
+                self.solid_motor.nozzle_radius,
+                self.solid_motor.grain_outer_radius * 1.4,
+                self.solid_motor.grain_outer_radius * 1.4,
+                0,
+            ]
+        )
+        # we need to draw the other half of the nozzle
+        x = np.concatenate([x, x[::-1]])
+        y = np.concatenate([y, -y[::-1]])
+        # now we need to sum the  the translate
+        x = x + translate[0]
+        y = y + translate[1]
+
+        patch = Polygon(
+            np.column_stack([x, y]),
+            label="Combustion Chamber",
+            facecolor="lightslategray",
+            edgecolor="black",
+        )
+        return patch
+
+    def _generate_grains(self, translate=(0, 0), csys=1):
+        patches = []
+        n_total = self.solid_motor.grain_number
+        separation = self.solid_motor.grain_separation
+        height = self.solid_motor.grain_initial_height
+        outer_radius = self.solid_motor.grain_outer_radius
+        inner_radius = self.solid_motor.grain_initial_inner_radius
+
+        cm_teo = (
+            csys * ((n_total / 2) * (height + separation))
+            + self.solid_motor.nozzle_position
+        )
+        cm_real = self.solid_motor.center_of_propellant_mass(0)
+
+        init = abs(cm_teo - cm_real) * csys
+
+        inner_y = np.array([0, inner_radius, inner_radius, 0])
+        outer_y = np.array([inner_radius, outer_radius, outer_radius, inner_radius])
+        inner_y = np.concatenate([inner_y, -inner_y[::-1]])
+        outer_y = np.concatenate([outer_y, -outer_y[::-1]])
+        inner_y = inner_y + translate[1]
+        outer_y = outer_y + translate[1]
+        for n in range(n_total):
+            grain_start = init + csys * (separation / 2 + n * (height + separation))
+            grain_end = grain_start + height * csys
+            x = np.array([grain_start, grain_start, grain_end, grain_end])
+            # draw the other half of the nozzle
+            x = np.concatenate([x, x[::-1]])
+            # sum the translate
+            x = x + translate[0]
+            patch = Polygon(
+                np.column_stack([x, outer_y]),
+                facecolor="olive",
+                edgecolor="khaki",
+            )
+            patches.append(patch)
+
+            patch = Polygon(
+                np.column_stack([x, inner_y]),
+                facecolor="khaki",
+                edgecolor="olive",
+            )
+            if n == 0:
+                patch.set_label("Grains")
+            patches.append(patch)
+        return patches
+
+    def _draw_center_of_mass(self, ax):
+        ax.axhline(0, color="k", linestyle="--", alpha=0.5)  # symmetry line
+        ax.plot(
+            [self.solid_motor.grains_center_of_mass_position],
+            [0],
+            "ro",
+            label="Grains Center of Mass",
+        )
+        ax.plot(
+            [self.solid_motor.center_of_dry_mass_position],
+            [0],
+            "bo",
+            label="Center of Dry Mass",
+        )
+
+    def _set_plot_properties(self, ax):
+        ax.set_title("Solid Motor Representation")
+        ax.set_ylabel("Radius (m)")
+        ax.set_xlabel("Position (m)")
+        # ax.grid(True)
+        plt.ylim(
+            -self.solid_motor.grain_outer_radius * 1.2 * 1.7,
+            self.solid_motor.grain_outer_radius * 1.2 * 1.7,
+        )
+        plt.legend(loc="upper left", bbox_to_anchor=(1, 1))
+
     def all(self):
         """Prints out all graphs available about the SolidMotor. It simply calls
         all the other plotter methods in this class.