diff --git a/rocketpy/motors/liquid_motor.py b/rocketpy/motors/liquid_motor.py index d89157e72..064267c11 100644 --- a/rocketpy/motors/liquid_motor.py +++ b/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() diff --git a/rocketpy/motors/tank.py b/rocketpy/motors/tank.py index 0b7069141..ff575afb4 100644 --- a/rocketpy/motors/tank.py +++ b/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 diff --git a/rocketpy/plots/liquid_motor_plots.py b/rocketpy/plots/liquid_motor_plots.py index 0d53792ea..0f228074e 100644 --- a/rocketpy/plots/liquid_motor_plots.py +++ b/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. diff --git a/rocketpy/plots/tank_plots.py b/rocketpy/plots/tank_plots.py index 0205a3414..b9239defb 100644 --- a/rocketpy/plots/tank_plots.py +++ b/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.