Skip to content

Commit 3d0edf0

Browse files
MateusStanoMateusStano
committed
ENH: reestructure draw to better handle weird positions and diameter changes
1 parent 8bb9fdb commit 3d0edf0

1 file changed

Lines changed: 213 additions & 107 deletions

File tree

‎rocketpy/plots/rocket_plots.py‎

Lines changed: 213 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import numpy as np
12
import matplotlib.pyplot as plt
23

34
from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail
@@ -105,7 +106,7 @@ def thrust_to_weight(self):
105106

106107
return None
107108

108-
def draw(self, vis_args=None):
109+
def draw2(self, vis_args=None):
109110
"""Draws the rocket in a matplotlib figure.
110111
111112
Parameters
@@ -117,7 +118,7 @@ def draw(self, vis_args=None):
117118
"background": "#EEEEEE",
118119
"tail": "black",
119120
"nose": "black",
120-
"body": "dimgrey",
121+
"body": "black",
121122
"fins": "black",
122123
"motor": "black",
123124
"buttons": "black",
@@ -131,144 +132,248 @@ def draw(self, vis_args=None):
131132
"background": "#EEEEEE",
132133
"tail": "black",
133134
"nose": "black",
134-
"body": "dimgrey",
135+
"body": "black",
135136
"fins": "black",
136137
"motor": "black",
137138
"buttons": "black",
138-
"line_width": 2.0,
139+
"line_width": 1.0,
139140
}
140141

141142
# Create the figure and axis
142-
_, ax = plt.subplots(figsize=(10, 5))
143+
_, ax = plt.subplots(figsize=(8, 5))
143144
ax.set_aspect("equal")
144145
ax.set_facecolor(vis_args["background"])
145146
ax.grid(True, linestyle="--", linewidth=0.5)
146147

147148
csys = self.rocket._csys
148-
149-
# Draw rocket body
150149
reverse = csys == 1
151150
self.rocket.aerodynamic_surfaces.sort_by_position(reverse=reverse)
152-
y_tube, x_tube = [], []
153-
for surface, position in self.rocket.aerodynamic_surfaces:
154-
if isinstance(surface, (NoseCone, Tail)):
155-
# Append the x and y coordinates of the surface shape_vec to the respective lists
156-
x_tube.extend((-csys) * surface.shape_vec[0] + position)
157-
y_tube.extend(surface.shape_vec[1])
158-
if isinstance(surface, Fins):
159-
pass
160-
161-
# Negate each element in the y_tube list using a list comprehension
162-
y_tube_negated = [-y for y in y_tube]
163-
plt.plot(
164-
x_tube, y_tube, color=vis_args["body"], linewidth=vis_args["line_width"]
165-
)
166-
plt.plot(
167-
x_tube,
168-
y_tube_negated,
169-
color=vis_args["body"],
170-
linewidth=vis_args["line_width"],
171-
)
172151

173-
# Get nozzle position (kinda not working for EmptyMotor class)
174-
x_nozzle = self.rocket.motor_position + self.rocket.motor.nozzle_position
175-
176-
# Find the last point of the rocket
177-
idx = -1 if csys == 1 else 0
178-
surface, position = self.rocket.aerodynamic_surfaces[idx]
179-
length = surface.shape_vec[0][-1] - surface.shape_vec[0][0]
180-
x_last = position + (-1 * csys) * length
181-
y_last = surface.shape_vec[1][-1]
182-
183-
plt.plot(
184-
[x_nozzle, x_last],
185-
[0, y_last],
186-
color=vis_args["body"],
187-
linewidth=vis_args["line_width"],
188-
)
189-
plt.plot(
190-
[x_nozzle, x_last],
191-
[0, -y_last],
192-
color=vis_args["body"],
193-
linewidth=vis_args["line_width"],
194-
)
195-
196-
# Draw nosecone
197-
nosecones = self.rocket.aerodynamic_surfaces.get_tuple_by_type(NoseCone)
198-
for nose, position in nosecones:
199-
x_nosecone = -csys * nose.shape_vec[0] + position
200-
y_nosecone = nose.shape_vec[1]
201-
202-
plt.plot(
203-
x_nosecone,
204-
y_nosecone,
205-
color=vis_args["nose"],
206-
linewidth=vis_args["line_width"] - 0.05,
207-
)
208-
plt.plot(
209-
x_nosecone,
210-
-y_nosecone,
211-
color=vis_args["nose"],
212-
linewidth=vis_args["line_width"] - 0.05,
213-
)
152+
# List of drawn surfaces with the position of points of interest
153+
# and the radius of the rocket at that point
154+
drawn_surfaces = []
214155

215-
# Draw transitions
216-
tails = self.rocket.aerodynamic_surfaces.get_tuple_by_type(Tail)
217-
for tail, position in tails:
218-
x_tail = -csys * tail.shape_vec[0] + position
219-
y_tail = tail.shape_vec[1]
156+
# Ideia is to get the shape of each aerodynamic surface in their own
157+
# coordinate system and then plot them in the rocket coordinate system
158+
# using the position of each surface
159+
# For the tubes, the surfaces need to be checked in order to check for
160+
# diameter changes. The final point of the last surface is the final
161+
# point of the last tube
220162

221-
plt.plot(
222-
x_tail,
223-
y_tail,
224-
color=vis_args["tail"],
163+
for surface, position in self.rocket.aerodynamic_surfaces:
164+
if isinstance(surface, NoseCone):
165+
x_nosecone = -csys * surface.shape_vec[0] + position
166+
y_nosecone = surface.shape_vec[1]
167+
168+
ax.plot(
169+
x_nosecone,
170+
y_nosecone,
171+
color=vis_args["nose"],
172+
linewidth=vis_args["line_width"],
173+
)
174+
ax.plot(
175+
x_nosecone,
176+
-y_nosecone,
177+
color=vis_args["nose"],
178+
linewidth=vis_args["line_width"],
179+
)
180+
# close the nosecone
181+
ax.plot(
182+
[x_nosecone[-1], x_nosecone[-1]],
183+
[y_nosecone[-1], -y_nosecone[-1]],
184+
color=vis_args["nose"],
185+
linewidth=vis_args["line_width"],
186+
)
187+
188+
# Add the nosecone to the list of drawn surfaces
189+
drawn_surfaces.append(
190+
(surface, x_nosecone[-1], surface.rocket_radius, x_nosecone[-1])
191+
)
192+
193+
elif isinstance(surface, Tail):
194+
x_tail = -csys * surface.shape_vec[0] + position
195+
y_tail = surface.shape_vec[1]
196+
197+
ax.plot(
198+
x_tail,
199+
y_tail,
200+
color=vis_args["tail"],
201+
linewidth=vis_args["line_width"],
202+
)
203+
ax.plot(
204+
x_tail,
205+
-y_tail,
206+
color=vis_args["tail"],
207+
linewidth=vis_args["line_width"],
208+
)
209+
# close above and below the tail
210+
ax.plot(
211+
[x_tail[-1], x_tail[-1]],
212+
[y_tail[-1], -y_tail[-1]],
213+
color=vis_args["tail"],
214+
linewidth=vis_args["line_width"],
215+
)
216+
ax.plot(
217+
[x_tail[0], x_tail[0]],
218+
[y_tail[0], -y_tail[0]],
219+
color=vis_args["tail"],
220+
linewidth=vis_args["line_width"],
221+
)
222+
223+
# Add the tail to the list of drawn surfaces
224+
drawn_surfaces.append(
225+
(surface, position, surface.bottom_radius, x_tail[-1])
226+
)
227+
228+
# Draw fins
229+
elif isinstance(surface, Fins):
230+
num_fins = surface.n
231+
x_fin = -csys * surface.shape_vec[0] + position
232+
y_fin = surface.shape_vec[1] + surface.rocket_radius
233+
234+
# Calculate the rotation angles for the other two fins (symmetrically)
235+
rotation_angles = [2 * np.pi * i / num_fins for i in range(num_fins)]
236+
237+
# Apply rotation transformations to get points for the other fins in 2D space
238+
for angle in rotation_angles:
239+
# Create a rotation matrix for the current angle around the x-axis
240+
rotation_matrix = np.array([[1, 0], [0, np.cos(angle)]])
241+
242+
# Apply the rotation to the original fin points
243+
rotated_points_2d = np.dot(
244+
rotation_matrix, np.vstack((x_fin, y_fin))
245+
)
246+
247+
# Extract x and y coordinates of the rotated points
248+
x_rotated, y_rotated = rotated_points_2d
249+
250+
# Project points above the XY plane back into the XY plane (set z-coordinate to 0)
251+
x_rotated = np.where(
252+
rotated_points_2d[1] > 0, rotated_points_2d[0], x_rotated
253+
)
254+
y_rotated = np.where(
255+
rotated_points_2d[1] > 0, rotated_points_2d[1], y_rotated
256+
)
257+
258+
# Plot the fins
259+
ax.plot(
260+
x_rotated,
261+
y_rotated,
262+
color=vis_args["fins"],
263+
linewidth=vis_args["line_width"],
264+
)
265+
266+
# Add the fin to the list of drawn surfaces
267+
drawn_surfaces.append(
268+
(surface, position, surface.rocket_radius, x_rotated[-1])
269+
)
270+
271+
# Draw tubes
272+
for i, d_surface in enumerate(drawn_surfaces):
273+
# Draw the tubes, from the end of the first surface to the beggining
274+
# of the next surface, with the radius of the rocket at that point
275+
surface, position, radius, last_x = d_surface
276+
277+
if i == len(drawn_surfaces) - 1:
278+
# If the last surface is a tail, do nothing
279+
if isinstance(surface, Tail):
280+
continue
281+
# Else goes to the end of the surface
282+
else:
283+
x_tube = [position, last_x]
284+
y_tube = [radius, radius]
285+
y_tube_negated = [-radius, -radius]
286+
else:
287+
# If it is not the last surface, the tube goes to the beggining
288+
# of the next surface
289+
next_surface, next_position, next_radius, next_last_x = drawn_surfaces[
290+
i + 1
291+
]
292+
x_tube = [last_x, next_position]
293+
y_tube = [radius, radius]
294+
y_tube_negated = [-radius, -radius]
295+
296+
ax.plot(
297+
x_tube,
298+
y_tube,
299+
color=vis_args["body"],
225300
linewidth=vis_args["line_width"],
226301
)
227-
plt.plot(
228-
x_tail,
229-
-y_tail,
230-
color=vis_args["tail"],
302+
ax.plot(
303+
x_tube,
304+
y_tube_negated,
305+
color=vis_args["body"],
231306
linewidth=vis_args["line_width"],
232307
)
233308

234-
# Draw fins
235-
fins = self.rocket.aerodynamic_surfaces.get_tuple_by_type(Fins)
236-
for fin, position in fins:
237-
x_fin = -csys * fin.shape_vec[0] + position
238-
y_fin = fin.shape_vec[1] + self.rocket.radius
239-
240-
plt.plot(
241-
x_fin,
242-
y_fin,
243-
color=vis_args["fins"],
244-
linewidth=vis_args["line_width"],
245-
)
246-
plt.plot(
247-
x_fin,
248-
-y_fin,
249-
color=vis_args["fins"],
250-
linewidth=vis_args["line_width"],
251-
)
309+
# TODO - Draw motor
310+
nozzle_position = (
311+
self.rocket.motor_position
312+
+ self.rocket.motor.nozzle_position
313+
* self.rocket._csys
314+
* self.rocket.motor._csys
315+
)
316+
ax.scatter(
317+
nozzle_position, 0, label="Nozzle Outlet", color="brown", s=10, zorder=10
318+
)
319+
# Check if nozzle is beyond the last surface, if so draw a tube
320+
# to it, with the radius of the last surface
321+
if self.rocket._csys == 1:
322+
if nozzle_position < last_x:
323+
x_tube = [last_x, nozzle_position]
324+
y_tube = [radius, radius]
325+
y_tube_negated = [-radius, -radius]
326+
327+
ax.plot(
328+
x_tube,
329+
y_tube,
330+
color=vis_args["body"],
331+
linewidth=vis_args["line_width"],
332+
)
333+
ax.plot(
334+
x_tube,
335+
y_tube_negated,
336+
color=vis_args["body"],
337+
linewidth=vis_args["line_width"],
338+
)
339+
else: # if self.rocket._csys == -1:
340+
if nozzle_position > last_x:
341+
x_tube = [last_x, nozzle_position]
342+
y_tube = [radius, radius]
343+
y_tube_negated = [-radius, -radius]
344+
345+
ax.plot(
346+
x_tube,
347+
y_tube,
348+
color=vis_args["body"],
349+
linewidth=vis_args["line_width"],
350+
)
351+
ax.plot(
352+
x_tube,
353+
y_tube_negated,
354+
color=vis_args["body"],
355+
linewidth=vis_args["line_width"],
356+
)
252357

253358
# Draw rail buttons
254359
buttons, pos = self.rocket.rail_buttons[0]
255360
lower = pos
256361
upper = pos + buttons.buttons_distance * csys
257-
plt.scatter(
258-
lower, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=10
362+
ax.scatter(
363+
lower, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=15
259364
)
260-
plt.scatter(
261-
upper, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=10
365+
ax.scatter(
366+
upper, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=15
262367
)
263368

264369
# Draw center of mass and center of pressure
265370
cm = self.rocket.center_of_mass(0)
266-
plt.scatter(cm, 0, color="black", label="Center of Mass", s=30)
267-
plt.scatter(cm, 0, facecolors="none", edgecolors="black", s=100)
371+
ax.scatter(cm, 0, color="black", label="Center of Mass", s=30)
372+
ax.scatter(cm, 0, facecolors="none", edgecolors="black", s=100)
268373

269374
cp = self.rocket.cp_position
270-
plt.scatter(cp, 0, label="Center Of Pressure", color="red", s=30, zorder=10)
271-
plt.scatter(cp, 0, facecolors="none", edgecolors="red", s=100, zorder=10)
375+
ax.scatter(cp, 0, label="Center Of Pressure", color="red", s=30, zorder=10)
376+
ax.scatter(cp, 0, facecolors="none", edgecolors="red", s=100, zorder=10)
272377

273378
# Set plot attributes
274379
plt.title(f"Rocket Geometry")
@@ -278,6 +383,7 @@ def draw(self, vis_args=None):
278383
plt.legend(loc="best")
279384
plt.tight_layout()
280385
plt.show()
386+
281387
return None
282388

283389
def all(self):

0 commit comments

Comments
 (0)