Skip to content

Commit

Permalink
ENH: liquid motors draw
Browse files Browse the repository at this point in the history
  • Loading branch information
Gui-FernandesBR committed Oct 11, 2023
1 parent 3f9271e commit 33c3116
Show file tree
Hide file tree
Showing 4 changed files with 140 additions and 0 deletions.
3 changes: 3 additions & 0 deletions rocketpy/motors/liquid_motor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
3 changes: 3 additions & 0 deletions rocketpy/motors/tank.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
76 changes: 76 additions & 0 deletions rocketpy/plots/liquid_motor_plots.py
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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.
Expand Down
58 changes: 58 additions & 0 deletions rocketpy/plots/tank_plots.py
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -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.
Expand Down

0 comments on commit 33c3116

Please sign in to comment.