Skip to content

Commit 3b065f7

Browse files
Gui-FernandesBRGui-FernandesBR
authored
Merge pull request #419 from RocketPy-Team/enh/rocket-drawing
ENH: rocket drawing
2 parents 82767ea + de4e9ca commit 3b065f7

6 files changed

Lines changed: 391 additions & 32 deletions

File tree

‎rocketpy/plots/__init__.py‎

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +0,0 @@
1-
from .aero_surface_plots import (
2-
_EllipticalFinsPlots,
3-
_FinsPlots,
4-
_NoseConePlots,
5-
_TailPlots,
6-
_TrapezoidalFinsPlots,
7-
)
8-
from .compare import Compare, CompareFlights
9-
from .environment_analysis_plots import _EnvironmentAnalysisPlots
10-
from .environment_plots import _EnvironmentPlots
11-
from .flight_plots import _FlightPlots
12-
from .hybrid_motor_plots import _HybridMotorPlots
13-
from .liquid_motor_plots import _LiquidMotorPlots
14-
from .motor_plots import _MotorPlots
15-
from .rocket_plots import _RocketPlots
16-
from .solid_motor_plots import _SolidMotorPlots
17-
from .tank_geometry_plots import _TankGeometryPlots
18-
from .tank_plots import _TankPlots

‎rocketpy/plots/rocket_plots.py‎

Lines changed: 287 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
1-
import matplotlib.pyplot as plt
21
import numpy as np
2+
import matplotlib.pyplot as plt
3+
4+
from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail
35

46

57
class _RocketPlots:
@@ -122,6 +124,290 @@ def thrust_to_weight(self):
122124

123125
return None
124126

127+
def draw(self, vis_args=None):
128+
"""Draws the rocket in a matplotlib figure.
129+
130+
Parameters
131+
----------
132+
vis_args : dict, optional
133+
Determines the visual aspects when drawing the rocket. If None,
134+
default values are used. Default values are:
135+
{
136+
"background": "#EEEEEE",
137+
"tail": "black",
138+
"nose": "black",
139+
"body": "black",
140+
"fins": "black",
141+
"motor": "black",
142+
"buttons": "black",
143+
"line_width": 2.0,
144+
}
145+
A full list of color names can be found at:
146+
https://matplotlib.org/stable/gallery/color/named_colors
147+
"""
148+
# TODO: we need to modularize this function, it is too big
149+
if vis_args is None:
150+
vis_args = {
151+
"background": "#EEEEEE",
152+
"tail": "black",
153+
"nose": "black",
154+
"body": "black",
155+
"fins": "black",
156+
"motor": "black",
157+
"buttons": "black",
158+
"line_width": 1.0,
159+
}
160+
161+
# Create the figure and axis
162+
_, ax = plt.subplots(figsize=(8, 5))
163+
ax.set_aspect("equal")
164+
ax.set_facecolor(vis_args["background"])
165+
ax.grid(True, linestyle="--", linewidth=0.5)
166+
167+
csys = self.rocket._csys
168+
reverse = csys == 1
169+
self.rocket.aerodynamic_surfaces.sort_by_position(reverse=reverse)
170+
171+
# List of drawn surfaces with the position of points of interest
172+
# and the radius of the rocket at that point
173+
drawn_surfaces = []
174+
175+
# Ideia is to get the shape of each aerodynamic surface in their own
176+
# coordinate system and then plot them in the rocket coordinate system
177+
# using the position of each surface
178+
# For the tubes, the surfaces need to be checked in order to check for
179+
# diameter changes. The final point of the last surface is the final
180+
# point of the last tube
181+
182+
for surface, position in self.rocket.aerodynamic_surfaces:
183+
if isinstance(surface, NoseCone):
184+
x_nosecone = -csys * surface.shape_vec[0] + position
185+
y_nosecone = surface.shape_vec[1]
186+
187+
ax.plot(
188+
x_nosecone,
189+
y_nosecone,
190+
color=vis_args["nose"],
191+
linewidth=vis_args["line_width"],
192+
)
193+
ax.plot(
194+
x_nosecone,
195+
-y_nosecone,
196+
color=vis_args["nose"],
197+
linewidth=vis_args["line_width"],
198+
)
199+
# close the nosecone
200+
ax.plot(
201+
[x_nosecone[-1], x_nosecone[-1]],
202+
[y_nosecone[-1], -y_nosecone[-1]],
203+
color=vis_args["nose"],
204+
linewidth=vis_args["line_width"],
205+
)
206+
207+
# Add the nosecone to the list of drawn surfaces
208+
drawn_surfaces.append(
209+
(surface, x_nosecone[-1], surface.rocket_radius, x_nosecone[-1])
210+
)
211+
212+
elif isinstance(surface, Tail):
213+
x_tail = -csys * surface.shape_vec[0] + position
214+
y_tail = surface.shape_vec[1]
215+
216+
ax.plot(
217+
x_tail,
218+
y_tail,
219+
color=vis_args["tail"],
220+
linewidth=vis_args["line_width"],
221+
)
222+
ax.plot(
223+
x_tail,
224+
-y_tail,
225+
color=vis_args["tail"],
226+
linewidth=vis_args["line_width"],
227+
)
228+
# close above and below the tail
229+
ax.plot(
230+
[x_tail[-1], x_tail[-1]],
231+
[y_tail[-1], -y_tail[-1]],
232+
color=vis_args["tail"],
233+
linewidth=vis_args["line_width"],
234+
)
235+
ax.plot(
236+
[x_tail[0], x_tail[0]],
237+
[y_tail[0], -y_tail[0]],
238+
color=vis_args["tail"],
239+
linewidth=vis_args["line_width"],
240+
)
241+
242+
# Add the tail to the list of drawn surfaces
243+
drawn_surfaces.append(
244+
(surface, position, surface.bottom_radius, x_tail[-1])
245+
)
246+
247+
# Draw fins
248+
elif isinstance(surface, Fins):
249+
num_fins = surface.n
250+
x_fin = -csys * surface.shape_vec[0] + position
251+
y_fin = surface.shape_vec[1] + surface.rocket_radius
252+
253+
# Calculate the rotation angles for the other two fins (symmetrically)
254+
rotation_angles = [2 * np.pi * i / num_fins for i in range(num_fins)]
255+
256+
# Apply rotation transformations to get points for the other fins in 2D space
257+
for angle in rotation_angles:
258+
# Create a rotation matrix for the current angle around the x-axis
259+
rotation_matrix = np.array([[1, 0], [0, np.cos(angle)]])
260+
261+
# Apply the rotation to the original fin points
262+
rotated_points_2d = np.dot(
263+
rotation_matrix, np.vstack((x_fin, y_fin))
264+
)
265+
266+
# Extract x and y coordinates of the rotated points
267+
x_rotated, y_rotated = rotated_points_2d
268+
269+
# Project points above the XY plane back into the XY plane (set z-coordinate to 0)
270+
x_rotated = np.where(
271+
rotated_points_2d[1] > 0, rotated_points_2d[0], x_rotated
272+
)
273+
y_rotated = np.where(
274+
rotated_points_2d[1] > 0, rotated_points_2d[1], y_rotated
275+
)
276+
277+
# Plot the fins
278+
ax.plot(
279+
x_rotated,
280+
y_rotated,
281+
color=vis_args["fins"],
282+
linewidth=vis_args["line_width"],
283+
)
284+
285+
# Add the fin to the list of drawn surfaces
286+
drawn_surfaces.append(
287+
(surface, position, surface.rocket_radius, x_rotated[-1])
288+
)
289+
290+
# Draw tubes
291+
for i, d_surface in enumerate(drawn_surfaces):
292+
# Draw the tubes, from the end of the first surface to the beginning
293+
# of the next surface, with the radius of the rocket at that point
294+
surface, position, radius, last_x = d_surface
295+
296+
if i == len(drawn_surfaces) - 1:
297+
# If the last surface is a tail, do nothing
298+
if isinstance(surface, Tail):
299+
continue
300+
# Else goes to the end of the surface
301+
else:
302+
x_tube = [position, last_x]
303+
y_tube = [radius, radius]
304+
y_tube_negated = [-radius, -radius]
305+
else:
306+
# If it is not the last surface, the tube goes to the beginning
307+
# of the next surface
308+
next_surface, next_position, next_radius, next_last_x = drawn_surfaces[
309+
i + 1
310+
]
311+
x_tube = [last_x, next_position]
312+
y_tube = [radius, radius]
313+
y_tube_negated = [-radius, -radius]
314+
315+
ax.plot(
316+
x_tube,
317+
y_tube,
318+
color=vis_args["body"],
319+
linewidth=vis_args["line_width"],
320+
)
321+
ax.plot(
322+
x_tube,
323+
y_tube_negated,
324+
color=vis_args["body"],
325+
linewidth=vis_args["line_width"],
326+
)
327+
328+
# TODO - Draw motor
329+
nozzle_position = (
330+
self.rocket.motor_position
331+
+ self.rocket.motor.nozzle_position
332+
* self.rocket._csys
333+
* self.rocket.motor._csys
334+
)
335+
ax.scatter(
336+
nozzle_position, 0, label="Nozzle Outlet", color="brown", s=10, zorder=10
337+
)
338+
# Check if nozzle is beyond the last surface, if so draw a tube
339+
# to it, with the radius of the last surface
340+
if self.rocket._csys == 1:
341+
if nozzle_position < last_x:
342+
x_tube = [last_x, nozzle_position]
343+
y_tube = [radius, radius]
344+
y_tube_negated = [-radius, -radius]
345+
346+
ax.plot(
347+
x_tube,
348+
y_tube,
349+
color=vis_args["body"],
350+
linewidth=vis_args["line_width"],
351+
)
352+
ax.plot(
353+
x_tube,
354+
y_tube_negated,
355+
color=vis_args["body"],
356+
linewidth=vis_args["line_width"],
357+
)
358+
else: # if self.rocket._csys == -1:
359+
if nozzle_position > last_x:
360+
x_tube = [last_x, nozzle_position]
361+
y_tube = [radius, radius]
362+
y_tube_negated = [-radius, -radius]
363+
364+
ax.plot(
365+
x_tube,
366+
y_tube,
367+
color=vis_args["body"],
368+
linewidth=vis_args["line_width"],
369+
)
370+
ax.plot(
371+
x_tube,
372+
y_tube_negated,
373+
color=vis_args["body"],
374+
linewidth=vis_args["line_width"],
375+
)
376+
377+
# Draw rail buttons
378+
try:
379+
buttons, pos = self.rocket.rail_buttons[0]
380+
lower = pos
381+
upper = pos + buttons.buttons_distance * csys
382+
ax.scatter(
383+
lower, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=15
384+
)
385+
ax.scatter(
386+
upper, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=15
387+
)
388+
except IndexError:
389+
pass
390+
391+
# Draw center of mass and center of pressure
392+
cm = self.rocket.center_of_mass(0)
393+
ax.scatter(cm, 0, color="black", label="Center of Mass", s=30)
394+
ax.scatter(cm, 0, facecolors="none", edgecolors="black", s=100)
395+
396+
cp = self.rocket.cp_position
397+
ax.scatter(cp, 0, label="Center Of Pressure", color="red", s=30, zorder=10)
398+
ax.scatter(cp, 0, facecolors="none", edgecolors="red", s=100, zorder=10)
399+
400+
# Set plot attributes
401+
plt.title(f"Rocket Geometry")
402+
plt.ylim([-self.rocket.radius * 4, self.rocket.radius * 6])
403+
plt.xlabel("Position (m)")
404+
plt.ylabel("Radius (m)")
405+
plt.legend(loc="best")
406+
plt.tight_layout()
407+
plt.show()
408+
409+
return None
410+
125411
def all(self):
126412
"""Prints out all graphs available about the Rocket. It simply calls
127413
all the other plotter methods in this class.

‎rocketpy/rocket/__init__.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from .aero_surface import (
1+
from rocketpy.rocket.aero_surface import (
22
AeroSurface,
33
EllipticalFins,
44
Fins,
@@ -7,6 +7,6 @@
77
Tail,
88
TrapezoidalFins,
99
)
10-
from .components import Components
11-
from .parachute import Parachute
12-
from .rocket import Rocket
10+
from rocketpy.rocket.components import Components
11+
from rocketpy.rocket.parachute import Parachute
12+
from rocketpy.rocket.rocket import Rocket

0 commit comments

Comments
 (0)