ENH: draw motors

rocketpy/motors/hybrid_motor.py

@@ -41,6 +41,26 @@ class HybridMotor(Motor):
         Solid motor object that composes the hybrid motor.
     HybridMotor.liquid : LiquidMotor
         Liquid motor object that composes the hybrid motor.
+    HybridMotor.positioned_tanks : list
+        List containing the motor's added tanks and their respective
+        positions.
+    HybridMotor.grains_center_of_mass_position : float
+        Position of the center of mass of the grains in meters, specified in
+        the motor's coordinate system.
+        See :doc:`Positions and Coordinate Systems </user/positions>`
+        for more information.
+    HybridMotor.grain_number : int
+        Number of solid grains.
+    HybridMotor.grain_density : float
+        Density of each grain in kg/meters cubed.
+    HybridMotor.grain_outer_radius : float
+        Outer radius of each grain in meters.
+    HybridMotor.grain_initial_inner_radius : float
+        Initial inner radius of each grain in meters.
+    HybridMotor.grain_initial_height : float
+        Initial height of each grain in meters.
+    HybridMotor.grain_separation : float
+        Distance between two grains in meters.
     HybridMotor.dry_mass : float
         The total mass of the motor structure, including chambers
         and tanks, when it is empty and does not contain any propellant.
@@ -325,6 +345,17 @@ def __init__(
             interpolation_method,
             coordinate_system_orientation,
         )
+
+        self.positioned_tanks = self.liquid.positioned_tanks
+        self.grain_number = grain_number
+        self.grain_density = grain_density
+        self.grain_outer_radius = grain_outer_radius
+        self.grain_initial_inner_radius = grain_initial_inner_radius
+        self.grain_initial_height = grain_initial_height
+        self.grain_separation = grain_separation
+        self.grains_center_of_mass_position = grains_center_of_mass_position
+        self.throat_radius = throat_radius
+
         # Initialize plots and prints object
         self.prints = _HybridMotorPrints(self)
         self.plots = _HybridMotorPlots(self)
@@ -522,6 +553,11 @@ def add_tank(self, tank, position):
         )
         reset_funcified_methods(self)
 
+    def draw(self):
+        """Draws a representation of the HybridMotor."""
+        self.plots.draw()
+        return None
+
     def info(self):
         """Prints out basic data about the Motor."""
         self.prints.all()

rocketpy/motors/liquid_motor.py

@@ -445,6 +445,7 @@ def add_tank(self, tank, position):
         reset_funcified_methods(self)
 
     def draw(self):
+        """Draw a representation of the LiquidMotor."""
         return self.plots.draw()
 
     def info(self):

rocketpy/motors/solid_motor.py

@@ -702,6 +702,7 @@ def propellant_I_23(self):
         return 0
 
     def draw(self):
+        """Draw a representation of the SolidMotor."""
         self.plots.draw()
         return None
 

rocketpy/plots/__init__.py

@@ -1,38 +0,0 @@
-import numpy as np
-from matplotlib.patches import Polygon
-
-
-def _generate_nozzle(motor, translate=(0, 0), csys=1):
-    nozzle_radius = motor.nozzle_radius
-    nozzle_position = motor.nozzle_position
-    try:
-        throat_radius = motor.throat_radius
-    except AttributeError:
-        # Liquid motors don't have throat radius, let's estimate it
-        throat_radius = 0.01
-
-    # calculate length between throat and nozzle outlet using 15º angle
-    major_axis = (nozzle_radius - throat_radius) / np.tan(np.deg2rad(15))
-    # calculate minor axis considering a 45º angle
-    minor_axis = (nozzle_radius - throat_radius) / np.tan(np.deg2rad(45))
-
-    # calculate x and y coordinates of the nozzle
-    x = csys * np.array(
-        [0, 0, major_axis, major_axis + minor_axis, major_axis + minor_axis]
-    )
-    y = csys * np.array([0, nozzle_radius, throat_radius, nozzle_radius, 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 position and the translate
-    x = x + nozzle_position + translate[0]
-    y = y + translate[1]
-
-    patch = Polygon(
-        np.column_stack([x, y]),
-        label="Nozzle",
-        facecolor="black",
-        edgecolor="black",
-    )
-    motor.nozzle_length = major_axis + minor_axis
-    return patch

rocketpy/plots/hybrid_motor_plots.py

@@ -1,3 +1,5 @@
+from matplotlib import pyplot as plt
+
 from .motor_plots import _MotorPlots
 
 
@@ -120,6 +122,47 @@ def Kn(self, lower_limit=None, upper_limit=None):
 
         self.motor.solid.Kn.plot(lower=lower_limit, upper=upper_limit)
 
+    def draw(self):
+        """Draw a representation of the HybridMotor.
+
+        Returns
+        -------
+        None
+        """
+        _, ax = plt.subplots(figsize=(8, 6), facecolor="#EEEEEE")
+
+        tanks_and_centers = self._generate_positioned_tanks(csys=self.motor._csys)
+        nozzle = self._generate_nozzle(
+            translate=(self.motor.nozzle_position, 0), csys=self.motor._csys
+        )
+        chamber = self._generate_combustion_chamber(
+            translate=(self.motor.grains_center_of_mass_position, 0)
+        )
+        grains = self._generate_grains(
+            translate=(self.motor.grains_center_of_mass_position, 0)
+        )
+        outline = self._generate_motor_region(
+            list_of_patches=[nozzle, chamber, *grains]
+            + [tank for tank, _ in tanks_and_centers]
+        )
+
+        ax.add_patch(outline)
+        ax.add_patch(chamber)
+        for grain in grains:
+            ax.add_patch(grain)
+        for patch, center in tanks_and_centers:
+            ax.add_patch(patch)
+            ax.plot(
+                center[0], center[1], marker="o", color="red", markersize=2
+            )  # Adjust markersize for better visibility
+        ax.add_patch(nozzle)
+
+        ax.set_title("Hybrid Motor Representation")
+        self._draw_center_of_mass(ax)
+        self._set_plot_properties(ax)
+        plt.show()
+        return None
+
     def all(self):
         """Prints out all graphs available about the HybridMotor. It simply calls
         all the other plotter methods in this class.

rocketpy/plots/liquid_motor_plots.py

@@ -1,13 +1,9 @@
-import copy
-
 import matplotlib.pyplot as plt
-import numpy as np
-from matplotlib.patches import Polygon
 
-from rocketpy.plots import _generate_nozzle
+from .motor_plots import _MotorPlots
 
 
-class _LiquidMotorPlots:
+class _LiquidMotorPlots(_MotorPlots):
     """Class that holds plot methods for LiquidMotor class.
 
     Attributes
@@ -31,71 +27,34 @@ def __init__(self, liquid_motor):
         """
         super().__init__(liquid_motor)
 
-    def _generate_positioned_tanks(self, translate=(0, 0), csys=1):
-        """Generates a list of patches that represent the tanks of the
-        liquid_motor.
-
-        Parameters
-        ----------
-        None
+    def draw(self):
+        """Draw a representation of the LiquidMotor.
 
         Returns
         -------
-        patches : list
-            List of patches that represent the tanks of the liquid_motor.
+        None
         """
-        colors = {
-            0: ("black", "dimgray"),
-            1: ("darkblue", "cornflowerblue"),
-            2: ("darkgreen", "limegreen"),
-            3: ("darkorange", "gold"),
-            4: ("darkred", "tomato"),
-            5: ("darkviolet", "violet"),
-        }
-        patches = []
-        for idx, pos_tank in enumerate(self.liquid_motor.positioned_tanks):
-            tank = pos_tank["tank"]
-            position = pos_tank["position"]
-            trans = (position + translate[0], translate[1])
-            patch = tank.plots._generate_tank(trans, csys)
-            patch.set_facecolor(colors[idx][1])
-            patch.set_edgecolor(colors[idx][0])
-            patches.append(patch)
-        return patches
-
-    def _draw_center_of_interests(self, ax, translate=(0, 0)):
-        # center of dry mass position
-        # center of wet mass time = 0
-        # center of wet mass time = end
-        return None
-
-    def draw(self):
-        fig, ax = plt.subplots(facecolor="#EEEEEE")
-
-        patches = self._generate_positioned_tanks()
-        for patch in patches:
+        _, ax = plt.subplots(figsize=(8, 6), facecolor="#EEEEEE")
+
+        tanks_and_centers = self._generate_positioned_tanks(csys=self.motor._csys)
+        nozzle = self._generate_nozzle(
+            translate=(self.motor.nozzle_position, 0), csys=self.motor._csys
+        )
+        outline = self._generate_motor_region(
+            list_of_patches=[nozzle] + [tank for tank, _ in tanks_and_centers]
+        )
+
+        ax.add_patch(outline)
+        for patch, center in tanks_and_centers:
             ax.add_patch(patch)
+            ax.plot(center[0], center[1], marker="o", color="red", markersize=2)
 
         # add the nozzle
-        ax.add_patch(_generate_nozzle(self.liquid_motor, translate=(0, 0)))
-
-        # find the maximum and minimum x and y values of the tanks
-        x_min = y_min = np.inf
-        x_max = y_max = -np.inf
-        for patch in patches:
-            x_min = min(x_min, patch.xy[:, 0].min())
-            x_max = max(x_max, patch.xy[:, 0].max())
-            y_min = min(y_min, patch.xy[:, 1].min())
-            y_max = max(y_max, patch.xy[:, 1].max())
-
-        ax.set_aspect("equal")
-        ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
-        ax.grid(True, linestyle="--", linewidth=0.5)
-        ax.set_ylim(y_min - 0.25, y_max + 0.25)
-        ax.set_xlim(x_min - 0.10, x_max + 0.10)
-        ax.set_xlabel("Position (m)")
-        ax.set_ylabel("Radius (m)")
-        ax.set_title("Liquid Motor Geometry")
+        ax.add_patch(nozzle)
+
+        ax.set_title("Liquid Motor Representation")
+        self._draw_center_of_mass(ax)
+        self._set_plot_properties(ax)
         plt.show()
 
     def all(self):