diff --git a/rocketpy/plots/__init__.py b/rocketpy/plots/__init__.py index 130d50aff..e69de29bb 100644 --- a/rocketpy/plots/__init__.py +++ b/rocketpy/plots/__init__.py @@ -1,18 +0,0 @@ -from .aero_surface_plots import ( - _EllipticalFinsPlots, - _FinsPlots, - _NoseConePlots, - _TailPlots, - _TrapezoidalFinsPlots, -) -from .compare import Compare, CompareFlights -from .environment_analysis_plots import _EnvironmentAnalysisPlots -from .environment_plots import _EnvironmentPlots -from .flight_plots import _FlightPlots -from .hybrid_motor_plots import _HybridMotorPlots -from .liquid_motor_plots import _LiquidMotorPlots -from .motor_plots import _MotorPlots -from .rocket_plots import _RocketPlots -from .solid_motor_plots import _SolidMotorPlots -from .tank_geometry_plots import _TankGeometryPlots -from .tank_plots import _TankPlots diff --git a/rocketpy/plots/rocket_plots.py b/rocketpy/plots/rocket_plots.py index 9764eb467..27eee40c7 100644 --- a/rocketpy/plots/rocket_plots.py +++ b/rocketpy/plots/rocket_plots.py @@ -1,5 +1,7 @@ -import matplotlib.pyplot as plt import numpy as np +import matplotlib.pyplot as plt + +from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail class _RocketPlots: @@ -104,6 +106,290 @@ def thrust_to_weight(self): return None + def draw(self, vis_args=None): + """Draws the rocket in a matplotlib figure. + + Parameters + ---------- + vis_args : dict, optional + Determines the visual aspects when drawing the rocket. If None, + default values are used. Default values are: + { + "background": "#EEEEEE", + "tail": "black", + "nose": "black", + "body": "black", + "fins": "black", + "motor": "black", + "buttons": "black", + "line_width": 2.0, + } + A full list of color names can be found at: + https://matplotlib.org/stable/gallery/color/named_colors + """ + # TODO: we need to modularize this function, it is too big + if vis_args is None: + vis_args = { + "background": "#EEEEEE", + "tail": "black", + "nose": "black", + "body": "black", + "fins": "black", + "motor": "black", + "buttons": "black", + "line_width": 1.0, + } + + # Create the figure and axis + _, ax = plt.subplots(figsize=(8, 5)) + ax.set_aspect("equal") + ax.set_facecolor(vis_args["background"]) + ax.grid(True, linestyle="--", linewidth=0.5) + + csys = self.rocket._csys + reverse = csys == 1 + self.rocket.aerodynamic_surfaces.sort_by_position(reverse=reverse) + + # List of drawn surfaces with the position of points of interest + # and the radius of the rocket at that point + drawn_surfaces = [] + + # Ideia is to get the shape of each aerodynamic surface in their own + # coordinate system and then plot them in the rocket coordinate system + # using the position of each surface + # For the tubes, the surfaces need to be checked in order to check for + # diameter changes. The final point of the last surface is the final + # point of the last tube + + for surface, position in self.rocket.aerodynamic_surfaces: + if isinstance(surface, NoseCone): + x_nosecone = -csys * surface.shape_vec[0] + position + y_nosecone = surface.shape_vec[1] + + ax.plot( + x_nosecone, + y_nosecone, + color=vis_args["nose"], + linewidth=vis_args["line_width"], + ) + ax.plot( + x_nosecone, + -y_nosecone, + color=vis_args["nose"], + linewidth=vis_args["line_width"], + ) + # close the nosecone + ax.plot( + [x_nosecone[-1], x_nosecone[-1]], + [y_nosecone[-1], -y_nosecone[-1]], + color=vis_args["nose"], + linewidth=vis_args["line_width"], + ) + + # Add the nosecone to the list of drawn surfaces + drawn_surfaces.append( + (surface, x_nosecone[-1], surface.rocket_radius, x_nosecone[-1]) + ) + + elif isinstance(surface, Tail): + x_tail = -csys * surface.shape_vec[0] + position + y_tail = surface.shape_vec[1] + + ax.plot( + x_tail, + y_tail, + color=vis_args["tail"], + linewidth=vis_args["line_width"], + ) + ax.plot( + x_tail, + -y_tail, + color=vis_args["tail"], + linewidth=vis_args["line_width"], + ) + # close above and below the tail + ax.plot( + [x_tail[-1], x_tail[-1]], + [y_tail[-1], -y_tail[-1]], + color=vis_args["tail"], + linewidth=vis_args["line_width"], + ) + ax.plot( + [x_tail[0], x_tail[0]], + [y_tail[0], -y_tail[0]], + color=vis_args["tail"], + linewidth=vis_args["line_width"], + ) + + # Add the tail to the list of drawn surfaces + drawn_surfaces.append( + (surface, position, surface.bottom_radius, x_tail[-1]) + ) + + # Draw fins + elif isinstance(surface, Fins): + num_fins = surface.n + x_fin = -csys * surface.shape_vec[0] + position + y_fin = surface.shape_vec[1] + surface.rocket_radius + + # Calculate the rotation angles for the other two fins (symmetrically) + rotation_angles = [2 * np.pi * i / num_fins for i in range(num_fins)] + + # Apply rotation transformations to get points for the other fins in 2D space + for angle in rotation_angles: + # Create a rotation matrix for the current angle around the x-axis + rotation_matrix = np.array([[1, 0], [0, np.cos(angle)]]) + + # Apply the rotation to the original fin points + rotated_points_2d = np.dot( + rotation_matrix, np.vstack((x_fin, y_fin)) + ) + + # Extract x and y coordinates of the rotated points + x_rotated, y_rotated = rotated_points_2d + + # Project points above the XY plane back into the XY plane (set z-coordinate to 0) + x_rotated = np.where( + rotated_points_2d[1] > 0, rotated_points_2d[0], x_rotated + ) + y_rotated = np.where( + rotated_points_2d[1] > 0, rotated_points_2d[1], y_rotated + ) + + # Plot the fins + ax.plot( + x_rotated, + y_rotated, + color=vis_args["fins"], + linewidth=vis_args["line_width"], + ) + + # Add the fin to the list of drawn surfaces + drawn_surfaces.append( + (surface, position, surface.rocket_radius, x_rotated[-1]) + ) + + # Draw tubes + for i, d_surface in enumerate(drawn_surfaces): + # Draw the tubes, from the end of the first surface to the beginning + # of the next surface, with the radius of the rocket at that point + surface, position, radius, last_x = d_surface + + if i == len(drawn_surfaces) - 1: + # If the last surface is a tail, do nothing + if isinstance(surface, Tail): + continue + # Else goes to the end of the surface + else: + x_tube = [position, last_x] + y_tube = [radius, radius] + y_tube_negated = [-radius, -radius] + else: + # If it is not the last surface, the tube goes to the beginning + # of the next surface + next_surface, next_position, next_radius, next_last_x = drawn_surfaces[ + i + 1 + ] + x_tube = [last_x, next_position] + y_tube = [radius, radius] + y_tube_negated = [-radius, -radius] + + ax.plot( + x_tube, + y_tube, + color=vis_args["body"], + linewidth=vis_args["line_width"], + ) + ax.plot( + x_tube, + y_tube_negated, + color=vis_args["body"], + linewidth=vis_args["line_width"], + ) + + # TODO - Draw motor + nozzle_position = ( + self.rocket.motor_position + + self.rocket.motor.nozzle_position + * self.rocket._csys + * self.rocket.motor._csys + ) + ax.scatter( + nozzle_position, 0, label="Nozzle Outlet", color="brown", s=10, zorder=10 + ) + # Check if nozzle is beyond the last surface, if so draw a tube + # to it, with the radius of the last surface + if self.rocket._csys == 1: + if nozzle_position < last_x: + x_tube = [last_x, nozzle_position] + y_tube = [radius, radius] + y_tube_negated = [-radius, -radius] + + ax.plot( + x_tube, + y_tube, + color=vis_args["body"], + linewidth=vis_args["line_width"], + ) + ax.plot( + x_tube, + y_tube_negated, + color=vis_args["body"], + linewidth=vis_args["line_width"], + ) + else: # if self.rocket._csys == -1: + if nozzle_position > last_x: + x_tube = [last_x, nozzle_position] + y_tube = [radius, radius] + y_tube_negated = [-radius, -radius] + + ax.plot( + x_tube, + y_tube, + color=vis_args["body"], + linewidth=vis_args["line_width"], + ) + ax.plot( + x_tube, + y_tube_negated, + color=vis_args["body"], + linewidth=vis_args["line_width"], + ) + + # Draw rail buttons + try: + buttons, pos = self.rocket.rail_buttons[0] + lower = pos + upper = pos + buttons.buttons_distance * csys + ax.scatter( + lower, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=15 + ) + ax.scatter( + upper, -self.rocket.radius, marker="s", color=vis_args["buttons"], s=15 + ) + except IndexError: + pass + + # Draw center of mass and center of pressure + cm = self.rocket.center_of_mass(0) + ax.scatter(cm, 0, color="black", label="Center of Mass", s=30) + ax.scatter(cm, 0, facecolors="none", edgecolors="black", s=100) + + cp = self.rocket.cp_position + ax.scatter(cp, 0, label="Center Of Pressure", color="red", s=30, zorder=10) + ax.scatter(cp, 0, facecolors="none", edgecolors="red", s=100, zorder=10) + + # Set plot attributes + plt.title(f"Rocket Geometry") + plt.ylim([-self.rocket.radius * 4, self.rocket.radius * 6]) + plt.xlabel("Position (m)") + plt.ylabel("Radius (m)") + plt.legend(loc="best") + plt.tight_layout() + plt.show() + + return None + def all(self): """Prints out all graphs available about the Rocket. It simply calls all the other plotter methods in this class. diff --git a/rocketpy/rocket/__init__.py b/rocketpy/rocket/__init__.py index 337fc9b7b..0fd8d27f0 100644 --- a/rocketpy/rocket/__init__.py +++ b/rocketpy/rocket/__init__.py @@ -1,4 +1,4 @@ -from .aero_surface import ( +from rocketpy.rocket.aero_surface import ( AeroSurface, EllipticalFins, Fins, @@ -7,6 +7,6 @@ Tail, TrapezoidalFins, ) -from .components import Components -from .parachute import Parachute -from .rocket import Rocket +from rocketpy.rocket.components import Components +from rocketpy.rocket.parachute import Parachute +from rocketpy.rocket.rocket import Rocket diff --git a/rocketpy/rocket/aero_surface.py b/rocketpy/rocket/aero_surface.py index f285356b2..026ef26ea 100644 --- a/rocketpy/rocket/aero_surface.py +++ b/rocketpy/rocket/aero_surface.py @@ -1,7 +1,6 @@ from abc import ABC, abstractmethod import warnings -import matplotlib.pyplot as plt import numpy as np from scipy.optimize import fsolve @@ -20,6 +19,8 @@ _TrapezoidalFinsPrints, ) +# TODO: all the evaluate_shape() methods need tests and documentation + class AeroSurface(ABC): """Abstract class used to define aerodynamic surfaces.""" @@ -1215,6 +1216,30 @@ def evaluate_geometrical_parameters(self): self.roll_damping_interference_factor = roll_damping_interference_factor self.roll_forcing_interference_factor = roll_forcing_interference_factor + self.evaluate_shape() + return None + + def evaluate_shape(self): + if self.sweep_length: + points = [ + (0, 0), + (self.sweep_length, self.span), + (self.sweep_length + self.tip_chord, self.span), + (self.root_chord, 0), + ] + else: + points = [ + (0, 0), + (self.root_chord - self.tip_chord, self.span), + (self.root_chord, self.span), + (self.root_chord, 0), + ] + + x_array, y_array = zip(*points) + self.shape_vec = [np.array(x_array), np.array(y_array)] + + return None + def info(self): self.prints.geometry() self.prints.lift() @@ -1521,6 +1546,16 @@ def evaluate_geometrical_parameters(self): self.roll_damping_interference_factor = roll_damping_interference_factor self.roll_forcing_interference_factor = roll_forcing_interference_factor + self.evaluate_shape() + return None + + def evaluate_shape(self): + angles = np.arange(0, 360, 5) + x_array = self.root_chord / 2 + self.root_chord / 2 * np.cos(np.radians(angles)) + y_array = self.span * np.sin(np.radians(angles)) + self.shape_vec = [x_array, y_array] + return None + def info(self): self.prints.geometry() self.prints.lift() @@ -1675,6 +1710,16 @@ def evaluate_geometrical_parameters(self): self.surface_area = ( np.pi * self.slant_length * (self.top_radius + self.bottom_radius) ) + self.evaluate_shape() + return None + + def evaluate_shape(self): + # Assuming the tail is a cone, calculate the shape vector + self.shape_vec = [ + np.array([0, self.length]), + np.array([self.top_radius, self.bottom_radius]), + ] + return None def evaluate_lift_coefficient(self): """Calculates and returns tail's lift coefficient. diff --git a/rocketpy/rocket/components.py b/rocketpy/rocket/components.py index fe033dafe..1d1d33e56 100644 --- a/rocketpy/rocket/components.py +++ b/rocketpy/rocket/components.py @@ -25,7 +25,13 @@ def __init__(self): def __repr__(self): """Return a string representation of the Components instance.""" - return repr(self._components) + components_str = "\n".join( + [ + f"\tComponent: {str(c.component):80} Position: {c.position:>6.3f}" + for c in self._components + ] + ) + return f"Components:\n{components_str}" def __len__(self): """Return the number of components in the list of components.""" @@ -155,3 +161,19 @@ def clear(self): None """ self._components.clear() + + def sort_by_position(self, reverse=False): + """Sort the list of components by position. + + Parameters + ---------- + reverse : bool + If True, sort in descending order. If False, sort in ascending + order. + + Returns + ------- + None + """ + self._components.sort(key=lambda x: x.position, reverse=reverse) + return None diff --git a/rocketpy/rocket/rocket.py b/rocketpy/rocket/rocket.py index 0ebee9ae1..d350e51c1 100644 --- a/rocketpy/rocket/rocket.py +++ b/rocketpy/rocket/rocket.py @@ -2,11 +2,11 @@ import numpy as np -from ..mathutils.function import Function -from ..motors.motor import EmptyMotor -from ..plots.rocket_plots import _RocketPlots -from ..prints.rocket_prints import _RocketPrints -from .aero_surface import ( +from rocketpy.mathutils.function import Function +from rocketpy.motors.motor import EmptyMotor +from rocketpy.plots.rocket_plots import _RocketPlots +from rocketpy.prints.rocket_prints import _RocketPrints +from rocketpy.rocket.aero_surface import ( EllipticalFins, Fins, NoseCone, @@ -14,8 +14,8 @@ Tail, TrapezoidalFins, ) -from .components import Components -from .parachute import Parachute +from rocketpy.rocket.components import Components +from rocketpy.rocket.parachute import Parachute class Rocket: @@ -1183,6 +1183,30 @@ def add_thrust_eccentricity(self, x, y): # Return self return self + def draw(self, vis_args=None): + """Draws the rocket in a matplotlib figure. + + Parameters + ---------- + vis_args : dict, optional + Determines the visual aspects when drawing the rocket. If None, + default values are used. Default values are: + { + "background": "#EEEEEE", + "tail": "black", + "nose": "black", + "body": "dimgrey", + "fins": "black", + "motor": "black", + "buttons": "black", + "line_width": 2.0, + } + A full list of color names can be found at: + https://matplotlib.org/stable/gallery/color/named_colors + """ + self.plots.draw(vis_args) + return None + def info(self): """Prints out a summary of the data and graphs available about the Rocket.