From 2110f857cafdb9a9d8529b943fac876c2449ddc2 Mon Sep 17 00:00:00 2001 From: Guilherme Date: Thu, 5 Jan 2023 06:47:46 +0100 Subject: [PATCH 1/4] MAINT: avoid DRY - duplicate calls at rocket_plots --- rocketpy/plots/rocket_plots.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/rocketpy/plots/rocket_plots.py b/rocketpy/plots/rocket_plots.py index dc2ca4a9c..1826be771 100644 --- a/rocketpy/plots/rocket_plots.py +++ b/rocketpy/plots/rocket_plots.py @@ -132,7 +132,9 @@ def thrustToWeight(self): return None def all(self): - """Prints out all graphs available about the Rocket. + """Prints out all graphs available about the Rocket. It simply calls + all the other plotter methods in this class. + Parameters ---------- None @@ -143,12 +145,12 @@ def all(self): # Show plots print("\nMass Plots") - self.rocket.totalMass() - self.rocket.reducedMass() + self.totalMass() + self.reducedMass() print("\nAerodynamics Plots") - self.rocket.staticMargin() - self.rocket.powerOnDrag() - self.rocket.powerOffDrag() - self.rocket.thrustToWeight.plot(lower=0, upper=self.rocket.motor.burnOutTime) + self.staticMargin() + self.powerOnDrag() + self.powerOffDrag() + self.thrustToWeight() return None From 9354ec87a8766e24d657278852cba37a7d5d7839 Mon Sep 17 00:00:00 2001 From: Guilherme Date: Thu, 5 Jan 2023 06:48:28 +0100 Subject: [PATCH 2/4] TST: moving fixtures to conftest.py --- tests/conftest.py | 30 +++++++++++++++++++++++++++--- tests/test_environment.py | 22 ---------------------- 2 files changed, 27 insertions(+), 25 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index b08ad8752..135b216da 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,9 @@ -import pytest -from rocketpy import SolidMotor -from rocketpy import Rocket +import datetime + import numericalunits +import pytest + +from rocketpy import Environment, Rocket, SolidMotor def pytest_addoption(parser): @@ -100,3 +102,25 @@ def pytest_collection_modifyitems(config, items): for item in items: if "slow" in item.keywords: item.add_marker(skip_slow) + + +@pytest.fixture +def example_env(): + Env = Environment(railLength=5, datum="WGS84") + return Env + + +@pytest.fixture +def example_env_robust(): + Env = Environment( + railLength=5, + latitude=32.990254, + longitude=-106.974998, + elevation=1400, + datum="WGS84", + ) + tomorrow = datetime.date.today() + datetime.timedelta(days=1) + Env.setDate( + (tomorrow.year, tomorrow.month, tomorrow.day, 12) + ) # Hour given in UT/C time + return Env diff --git a/tests/test_environment.py b/tests/test_environment.py index 078a23cd9..20f826275 100644 --- a/tests/test_environment.py +++ b/tests/test_environment.py @@ -6,28 +6,6 @@ from rocketpy import Environment, Flight, Rocket, SolidMotor -@pytest.fixture -def example_env(): - Env = Environment(railLength=5, datum="WGS84") - return Env - - -@pytest.fixture -def example_env_robust(): - Env = Environment( - railLength=5, - latitude=32.990254, - longitude=-106.974998, - elevation=1400, - datum="WGS84", - ) - tomorrow = datetime.date.today() + datetime.timedelta(days=1) - Env.setDate( - (tomorrow.year, tomorrow.month, tomorrow.day, 12) - ) # Hour given in UTC time - return Env - - def test_env_set_date(example_env): tomorrow = datetime.date.today() + datetime.timedelta(days=1) example_env.setDate((tomorrow.year, tomorrow.month, tomorrow.day, 12)) From c910f8994c4663196439394c88e262a2e75947d3 Mon Sep 17 00:00:00 2001 From: Guilherme Date: Thu, 5 Jan 2023 06:49:04 +0100 Subject: [PATCH 3/4] MAINT: extract logic methods, avoid repeating --- rocketpy/plots/compare/compare_flights.py | 236 ++++++++++------------ 1 file changed, 107 insertions(+), 129 deletions(-) diff --git a/rocketpy/plots/compare/compare_flights.py b/rocketpy/plots/compare/compare_flights.py index bc04b61d5..36cebf213 100644 --- a/rocketpy/plots/compare/compare_flights.py +++ b/rocketpy/plots/compare/compare_flights.py @@ -52,6 +52,74 @@ def __init__(self, flights): return None + def __process_xlim(self, x_lim): + """Function to process the x_lim key word argument. It is simply a + logic to check if the string "apogee" is used as an item for the tuple, + and if so, replace it with the maximum apogee time of all flights. + This garantes that we do not repeat the same code for each plot. + + Parameters + ---------- + x_lim : tuple + A list of two items, where the first item represents the x axis lower limit + and second item, the x axis upper limit. If set to None, will be calculated + automatically by matplotlib. If the string "apogee" is used as a item for the + tuple, the maximum apogee time of all flights will be used instead. + + Returns + ------- + x_lim + The processed x_lim keyword argument. + """ + if x_lim: + x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] + x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + return x_lim + + def __process_savefig(self, filename, fig): + """Function to either save the plot or show it, depending on the + filename key word argument. This way we do not repeat the same code + for each plot. + + Parameters + ---------- + filename : str, optional + If a filename is provided, the plot will be saved to a file, by default None. + Image options are: png, pdf, ps, eps and svg. + fig : matplotlib.figure.Figure + The figure to be saved or shown. + + Returns + ------- + None + """ + if filename: + fig.savefig(filename) + print("Plot saved to file: " + filename) + else: + plt.show() + return None + + def __process_legend(self, legend, fig): + """Function to add a legend to the plot, if the legend key word + argument is set to True. This way we do not repeat the same code for + each plot. + + Parameters + ---------- + legend : bool + If set to True, a legend will be added to the plot, by default True. + fig : matplotlib.figure.Figure + The figure to which the legend will be added. + + Returns + ------- + None + """ + if legend: + fig.legend() + return None + def positions( self, figsize=(7, 10), x_lim=None, y_lim=None, legend=True, filename=None ): @@ -83,9 +151,7 @@ def positions( None """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -102,10 +168,7 @@ def positions( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -146,9 +209,7 @@ def velocities( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -165,10 +226,7 @@ def velocities( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -210,9 +268,7 @@ def stream_velocities( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -239,10 +295,7 @@ def stream_velocities( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -283,9 +336,7 @@ def accelerations( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -307,10 +358,7 @@ def accelerations( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -346,9 +394,7 @@ def euler_angles( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -369,10 +415,7 @@ def euler_angles( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -413,9 +456,7 @@ def quaternions( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -437,10 +478,7 @@ def quaternions( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -476,9 +514,7 @@ def attitude_angles( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -499,10 +535,7 @@ def attitude_angles( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -538,14 +571,7 @@ def angular_velocities( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] - - # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -566,8 +592,7 @@ def angular_velocities( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) + self.__process_savefig(filename, fig) return None @@ -603,9 +628,7 @@ def angular_accelerations( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -626,10 +649,7 @@ def angular_accelerations( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -670,9 +690,7 @@ def aerodynamic_forces( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -692,10 +710,7 @@ def aerodynamic_forces( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -735,9 +750,7 @@ def aerodynamic_moments( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -757,10 +770,7 @@ def aerodynamic_moments( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -794,9 +804,7 @@ def energies( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -817,10 +825,7 @@ def energies( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -859,9 +864,7 @@ def powers( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -878,10 +881,7 @@ def powers( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -921,9 +921,7 @@ def rail_buttons_forces( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -950,10 +948,7 @@ def rail_buttons_forces( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -993,9 +988,7 @@ def angles_of_attack( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -1012,10 +1005,7 @@ def angles_of_attack( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -1055,9 +1045,7 @@ def fluid_mechanics( """ # Check if key word is used for x_limit - if x_lim: - x_lim[0] = self.apogee_time if x_lim[0] == "apogee" else x_lim[0] - x_lim[1] = self.apogee_time if x_lim[1] == "apogee" else x_lim[1] + x_lim = self.__process_xlim(x_lim) # Create the figure fig, _ = super().create_comparison_figure( @@ -1084,10 +1072,7 @@ def fluid_mechanics( ) # Saving the plot to a file if a filename is provided, showing the plot otherwise - if filename: - fig.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -1420,16 +1405,12 @@ def __plot_xy( ax.set_xlim([minXY, maxXY]) # Add legend - if legend: - fig.legend() + self.__process_legend(legend, fig) fig.tight_layout() # Save figure - if filename: - plt.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -1488,8 +1469,7 @@ def __plot_xz( ax.set_xlim([minXY, maxXY]) # Add legend - if legend: - fig.legend() + self.__process_legend(legend, fig) fig.tight_layout() @@ -1556,16 +1536,12 @@ def __plot_yz( ax.set_xlim([minXY, maxXY]) # Add legend - if legend: - fig.legend() + self.__process_legend(legend, fig) fig.tight_layout() # Save figure - if filename: - plt.savefig(filename) - else: - plt.show() + self.__process_savefig(filename, fig) return None @@ -1583,6 +1559,8 @@ def all(self): self.trajectories_3d() + self.trajectories_2d() + self.positions() self.velocities() From f5a4d6d67380324a0a7dfb4c4c691681d2a9139f Mon Sep 17 00:00:00 2001 From: Guilherme Date: Thu, 5 Jan 2023 06:49:31 +0100 Subject: [PATCH 4/4] TST: new plots for compare plots --- tests/test_plots.py | 105 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 tests/test_plots.py diff --git a/tests/test_plots.py b/tests/test_plots.py new file mode 100644 index 000000000..9e6a27974 --- /dev/null +++ b/tests/test_plots.py @@ -0,0 +1,105 @@ +import os +from unittest.mock import patch + +import matplotlib.pyplot as plt + +from rocketpy import Flight +from rocketpy.plots.compare import Compare, CompareFlights + + +@patch("matplotlib.pyplot.show") +def test_compare(mock_show, rocket, example_env): + """Here we want to test the 'x_attributes' argument, which is the only one + that is not tested in the other tests. + + Parameters + ---------- + mock_show : + Mocks the matplotlib.pyplot.show() function to avoid showing the plots. + rocket : rocketpy.Rocket + Rocket object to be used in the tests. See conftest.py for more details. + example_env : rocketpy.Environment + Environment object to be used in the tests. See conftest.py for more details. + """ + rocket.setRailButtons([-0.5, 0.2]) + flight = Flight(environment=example_env, rocket=rocket, inclination=85, heading=0) + + objects = [flight, flight, flight] + + comparison = Compare(object_list=objects) + + fig, _ = comparison.create_comparison_figure( + y_attributes=["z"], + n_rows=1, + n_cols=1, + figsize=(10, 10), + legend=False, + title="Test", + x_labels=["Time (s)"], + y_labels=["Altitude (m)"], + x_lim=(0, 3), + y_lim=(0, 1000), + x_attributes=["time"], + ) + + assert isinstance(fig, plt.Figure) == True + + +@patch("matplotlib.pyplot.show") +def test_compare_flights(mock_show, rocket, example_env): + """Tests the CompareFlights class. It simply ensures that all the methods + are being called without errors. It does not test the actual plots, which + would be very difficult to do. + + Parameters + ---------- + mock_show : + Mocks the matplotlib.pyplot.show() function to avoid showing the plots. + rocket : rocketpy.Rocket + Rocket object to be used in the tests. See conftest.py for more details. + example_env : rocketpy.Environment + Environment object to be used in the tests. See conftest.py for more details. + + Returns + ------- + None + """ + example_env.setAtmosphericModel( + type="CustomAtmosphere", + pressure=None, + temperature=300, + wind_u=[(0, 5), (1000, 10)], + wind_v=[(0, -2), (500, 3), (1600, 2)], + ) + + rocket.setRailButtons([-0.5, 0.2]) + + inclinations = [60, 70, 80, 90] + headings = [0, 45, 90, 180] + flights = [] + + # Create (4 * 4) = 16 different flights to be compared + for heading in headings: + for inclination in inclinations: + flight = Flight( + environment=example_env, + rocket=rocket, + inclination=inclination, + heading=heading, + name=f"Incl {inclination} Head {heading}", + ) + flights.append(flight) + + comparison = CompareFlights(flights) + + assert comparison.all() == None + assert comparison.trajectories_2d(plane="xz", legend=False) == None + assert comparison.trajectories_2d(plane="yz", legend=True) == None + + # Test save fig and then remove file + assert comparison.positions(filename="test.png") == None + os.remove("test.png") + + # Test xlim and ylim arguments + assert comparison.positions(x_lim=[0, 100], y_lim=[0, 1000]) == None + assert comparison.positions(x_lim=[0, "apogee"]) == None