Merge pull request #434 from RocketPy-Team/enh/liquid-motors-draw

ENH: liquid motors draw

Author: Gui-FernandesBR

rocketpy/motors/liquid_motor.py

@@ -444,6 +444,9 @@ def add_tank(self, tank, position):
         self.positioned_tanks.append({"tank": tank, "position": position})
         reset_funcified_methods(self)
 
+    def draw(self):
+        return self.plots.draw()
+
     def info(self):
         """Prints out basic data about the Motor."""
         self.prints.all()

rocketpy/motors/tank.py

@@ -410,6 +410,9 @@ def inertia(self):
         """
         return self.liquid_inertia + self.gas_inertia
 
+    def draw(self):
+        """Draws the tank geometry."""
+        return self.plots.draw()
 
 class MassFlowRateBasedTank(Tank):
     """Class to define a tank based on mass flow rates inputs. This class

rocketpy/plots/liquid_motor_plots.py

@@ -1,3 +1,12 @@
+import copy
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.patches import Polygon
+
+from rocketpy.plots import _generate_nozzle
+
+
 class _LiquidMotorPlots:
     """Class that holds plot methods for LiquidMotor class.
 
@@ -256,6 +265,73 @@ def I_23(self, lower_limit=None, upper_limit=None):
 
         return None
 
+    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
+
+        Returns
+        -------
+        patches : list
+            List of patches that represent the tanks of the liquid_motor.
+        """
+        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.add_patch(patch)
+
+        # 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")
+        plt.show()
+
     def all(self):
         """Prints out all graphs available about the LiquidMotor. It simply calls
         all the other plotter methods in this class.

rocketpy/plots/tank_plots.py

@@ -1,3 +1,9 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.animation import FuncAnimation
+from matplotlib.patches import Polygon, Rectangle
+
+
 class _TankPlots:
     """Class that holds plot methods for Tank class.
 
@@ -22,9 +28,61 @@ def __init__(self, tank):
         """
 
         self.tank = tank
+        self.name = tank.name
+        self.geometry = tank.geometry
 
         return None
 
+    def _generate_tank(self, translate=(0, 0), csys=1):
+        """Generates a matplotlib patch object that represents the tank.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes, optional
+            Axes object to plot the tank on. If None, a new figure and axes
+            will be created.
+        translate : tuple, optional
+            Tuple of floats that represents the translation of the tank
+            geometry.
+
+        Returns
+        -------
+        tank : matplotlib.patches.Polygon
+            Polygon object that represents the tank.
+        """
+        # get positions of all points
+        x = csys * self.geometry.radius.x_array + translate[0]
+        y = csys * self.geometry.radius.y_array + translate[1]
+        x = np.concatenate([x, x[::-1]])
+        y = np.concatenate([y, -y[::-1]])
+        xy = np.column_stack([x, y])
+
+        tank = Polygon(
+            xy,
+            label=self.name,
+            facecolor="dimgray",
+            edgecolor="black",
+        )
+        # Don't set any plot config here. Use the draw methods for that
+        return tank
+
+    def draw(self):
+        fig, ax = plt.subplots(facecolor="#EEEEEE")
+
+        ax.add_patch(self._generate_tank())
+
+        ax.set_aspect("equal")
+        ax.grid(True, linestyle="--", linewidth=0.5)
+
+        ax.set_xlabel("Length (m)")
+        ax.set_ylabel("Radius (m)")
+        ax.set_title("Tank Geometry")
+
+        x_max = self.geometry.radius.x_array.max()
+        y_max = self.geometry.radius.y_array.max()
+        ax.set_xlim(-1.2 * x_max, 1.2 * x_max)
+        ax.set_ylim(-1.5 * y_max, 1.5 * y_max)
+
     def all(self):
         """Prints out all graphs available about the Tank. It simply calls
         all the other plotter methods in this class.