From 69aaece6534debdb9493efb57d40f35dee5c8e87 Mon Sep 17 00:00:00 2001 From: Niklas Freund Date: Sun, 10 Nov 2024 01:28:08 +0100 Subject: [PATCH] ENH: Implement optional plot saving (#597) * ENH: Add filename option to most plot functions * MNT: Use f-strings in plot_helpers.py Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com> * MNT: Run isort * DOC: Make docstring fit into 80 columns and reference matplotlib formats * MNT: Run black * MNT: Move get_matplotlib_supported_file_endings into tools * Apply suggestions from code review Co-authored-by: Lucas Prates <57069366+Lucas-Prates@users.noreply.github.com> * Fix function/__call__ check * Change plot_helpers warnings from UserWarning to ValueError * Update docstrings * Fix issues from rebase * MNT: post merge fixes regarding plots. * MNT: post merge fixes and linting. * TST: implement testing of show or save plots. * MNT: Correct CHANGELOG file. * MNT: Improve error handling of unsupported file endings. * DOC: add doc section for plot saving. --------- Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com> Co-authored-by: Lucas Prates <57069366+Lucas-Prates@users.noreply.github.com> Co-authored-by: Pedro Bressan --- CHANGELOG.md | 2 +- docs/user/first_simulation.rst | 39 +++++ rocketpy/mathutils/function.py | 40 ++++- rocketpy/motors/hybrid_motor.py | 29 ++-- rocketpy/motors/liquid_motor.py | 26 ++- rocketpy/motors/motor.py | 17 +- rocketpy/motors/solid_motor.py | 29 ++-- rocketpy/motors/tank.py | 19 ++- rocketpy/plots/aero_surface_plots.py | 67 ++++++-- rocketpy/plots/compare/compare_flights.py | 33 ++-- rocketpy/plots/environment_analysis_plots.py | 167 +++++++++++++++---- rocketpy/plots/environment_plots.py | 50 +++++- rocketpy/plots/flight_plots.py | 142 +++++++++++++--- rocketpy/plots/hybrid_motor_plots.py | 88 +++++++--- rocketpy/plots/liquid_motor_plots.py | 13 +- rocketpy/plots/motor_plots.py | 105 ++++++++---- rocketpy/plots/plot_helpers.py | 67 ++++++++ rocketpy/plots/rocket_plots.py | 35 +++- rocketpy/plots/solid_motor_plots.py | 19 ++- rocketpy/plots/tank_plots.py | 13 +- rocketpy/rocket/aero_surface/fins/fins.py | 12 +- rocketpy/rocket/aero_surface/nose_cone.py | 16 +- rocketpy/rocket/rocket.py | 9 +- rocketpy/simulation/monte_carlo.py | 4 +- rocketpy/tools.py | 18 ++ rocketpy/utilities.py | 24 ++- tests/integration/test_plots.py | 42 ++++- tests/unit/test_plots.py | 53 +++++- 28 files changed, 953 insertions(+), 225 deletions(-) create mode 100644 rocketpy/plots/plot_helpers.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ccef502b..271c3a2bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,7 +32,7 @@ Attention: The newest changes should be on top --> ### Added -- +- ENH: Implement optional plot saving [#597](https://github.com/RocketPy-Team/RocketPy/pull/597) ### Changed diff --git a/docs/user/first_simulation.rst b/docs/user/first_simulation.rst index e2f6b4d99..5624ed926 100644 --- a/docs/user/first_simulation.rst +++ b/docs/user/first_simulation.rst @@ -654,6 +654,45 @@ Finally, the :meth:`rocketpy.Flight.export_data` method also provides a convenie os.remove("calisto_flight_data.csv") +Saving and Storing Plots +------------------------ + +Here we show how to save the plots and drawings generated by RocketPy. +For instance, we can save our rocket drawing as a ``.png`` file: + +.. jupyter-execute:: + + calisto.draw(filename="calisto_drawing.png") + +Also, if you want to save a specific rocketpy plot, every RocketPy +attribute of type :class:`rocketpy.Function` is capable of saving its plot +as an image file. For example, we can save our rocket's speed plot and the +trajectory plot as ``.jpg`` files: + +.. jupyter-execute:: + + test_flight.speed.plot(filename="speed_plot.jpg") + test_flight.plots.trajectory_3d(filename="trajectory_plot.jpg") + +The supported file formats are ``.eps``, ``.jpg``, ``.jpeg``, ``.pdf``, +``.pgf``, ``.png``, ``.ps``, ``.raw``, ``.rgba``, ``.svg``, ``.svgz``, +``.tif``, ``.tiff`` and ``.webp``. More details on manipulating data and +results can be found in the :ref:`Function Class Usage `. + +.. note:: + + The ``filename`` argument is optional. If not provided, the plot will be + shown instead. If the provided filename is in a directory that does not + exist, the directory will be created. + +.. jupyter-execute:: + :hide-code: + :hide-output: + + os.remove("calisto_drawing.png") + os.remove("speed_plot.jpg") + os.remove("trajectory_plot.jpg") + Further Analysis ---------------- diff --git a/rocketpy/mathutils/function.py b/rocketpy/mathutils/function.py index ba05636a5..314e6415b 100644 --- a/rocketpy/mathutils/function.py +++ b/rocketpy/mathutils/function.py @@ -22,6 +22,8 @@ RBFInterpolator, ) +from ..plots.plot_helpers import show_or_save_plot + # Numpy 1.x compatibility, # TODO: remove these lines when all dependencies support numpy>=2.0.0 if np.lib.NumpyVersion(np.__version__) >= "2.0.0b1": @@ -1378,7 +1380,7 @@ def remove_outliers_iqr(self, threshold=1.5): ) # Define all presentation methods - def __call__(self, *args): + def __call__(self, *args, filename=None): """Plot the Function if no argument is given. If an argument is given, return the value of the function at the desired point. @@ -1392,13 +1394,18 @@ def __call__(self, *args): evaluated at all points in the list and a list of floats will be returned. If the function is N-D, N arguments must be given, each one being an scalar or list. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- ans : None, scalar, list """ if len(args) == 0: - return self.plot() + return self.plot(filename=filename) else: return self.get_value(*args) @@ -1459,8 +1466,11 @@ def plot(self, *args, **kwargs): Function.plot_2d if Function is 2-Dimensional and forward arguments and key-word arguments.""" if isinstance(self, list): + # Extract filename from kwargs + filename = kwargs.get("filename", None) + # Compare multiple plots - Function.compare_plots(self) + Function.compare_plots(self, filename) else: if self.__dom_dim__ == 1: self.plot_1d(*args, **kwargs) @@ -1488,6 +1498,7 @@ def plot_1d( # pylint: disable=too-many-statements force_points=False, return_object=False, equal_axis=False, + filename=None, ): """Plot 1-Dimensional Function, from a lower limit to an upper limit, by sampling the Function several times in the interval. The title of @@ -1518,6 +1529,11 @@ def plot_1d( # pylint: disable=too-many-statements Setting force_points to True will plot all points, as a scatter, in which the Function was evaluated in the dataset. Default value is False. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1558,7 +1574,7 @@ def plot_1d( # pylint: disable=too-many-statements plt.title(self.title) plt.xlabel(self.__inputs__[0].title()) plt.ylabel(self.__outputs__[0].title()) - plt.show() + show_or_save_plot(filename) if return_object: return fig, ax @@ -1581,6 +1597,7 @@ def plot_2d( # pylint: disable=too-many-statements disp_type="surface", alpha=0.6, cmap="viridis", + filename=None, ): """Plot 2-Dimensional Function, from a lower limit to an upper limit, by sampling the Function several times in the interval. The title of @@ -1620,6 +1637,11 @@ def plot_2d( # pylint: disable=too-many-statements cmap : string, optional Colormap of plotted graph, which can be any of the color maps available in matplotlib. Default value is viridis. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1692,7 +1714,7 @@ def plot_2d( # pylint: disable=too-many-statements axes.set_xlabel(self.__inputs__[0].title()) axes.set_ylabel(self.__inputs__[1].title()) axes.set_zlabel(self.__outputs__[0].title()) - plt.show() + show_or_save_plot(filename) @staticmethod def compare_plots( # pylint: disable=too-many-statements @@ -1707,6 +1729,7 @@ def compare_plots( # pylint: disable=too-many-statements force_points=False, return_object=False, show=True, + filename=None, ): """Plots N 1-Dimensional Functions in the same plot, from a lower limit to an upper limit, by sampling the Functions several times in @@ -1751,6 +1774,11 @@ def compare_plots( # pylint: disable=too-many-statements False. show : bool, optional If True, shows the plot. Default value is True. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1826,7 +1854,7 @@ def compare_plots( # pylint: disable=too-many-statements plt.ylabel(ylabel) if show: - plt.show() + show_or_save_plot(filename) if return_object: return fig, ax diff --git a/rocketpy/motors/hybrid_motor.py b/rocketpy/motors/hybrid_motor.py index 2916486d0..b3527c604 100644 --- a/rocketpy/motors/hybrid_motor.py +++ b/rocketpy/motors/hybrid_motor.py @@ -601,16 +601,19 @@ def add_tank(self, tank, position): ) reset_funcified_methods(self) - def draw(self): - """Draws a representation of the HybridMotor.""" - self.plots.draw() - - def info(self): - """Prints out basic data about the Motor.""" - self.prints.all() - self.plots.thrust() - - def all_info(self): - """Prints out all data and graphs available about the Motor.""" - self.prints.all() - self.plots.all() + def draw(self, filename=None): + """Draws a representation of the HybridMotor. + + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + + Returns + ------- + None + """ + self.plots.draw(filename) diff --git a/rocketpy/motors/liquid_motor.py b/rocketpy/motors/liquid_motor.py index cde0e9d03..cc7bf5027 100644 --- a/rocketpy/motors/liquid_motor.py +++ b/rocketpy/motors/liquid_motor.py @@ -463,21 +463,19 @@ def add_tank(self, tank, position): self.positioned_tanks.append({"tank": tank, "position": position}) reset_funcified_methods(self) - def draw(self): - """Draw a representation of the LiquidMotor.""" - self.plots.draw() + def draw(self, filename=None): + """Draw a representation of the LiquidMotor. - def info(self): - """Prints out basic data about the Motor.""" - self.prints.all() - self.plots.thrust() - - def all_info(self): - """Prints out all data and graphs available about the Motor. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.prints.all() - self.plots.all() + self.plots.draw(filename) diff --git a/rocketpy/motors/motor.py b/rocketpy/motors/motor.py index e0e6dfc9a..558c9f9a2 100644 --- a/rocketpy/motors/motor.py +++ b/rocketpy/motors/motor.py @@ -1083,15 +1083,26 @@ def get_attr_value(obj, attr_name, multiplier=1): # Write last line file.write(f"{self.thrust.source[-1, 0]:.4f} {0:.3f}\n") - def info(self): + def info(self, filename=None): """Prints out a summary of the data and graphs available about the Motor. + + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + + Returns + ------- + None """ # Print motor details self.prints.all() - self.plots.thrust() + self.plots.thrust(filename=filename) - @abstractmethod def all_info(self): """Prints out all data and graphs available about the Motor.""" self.prints.all() diff --git a/rocketpy/motors/solid_motor.py b/rocketpy/motors/solid_motor.py index f6f09967e..09a930522 100644 --- a/rocketpy/motors/solid_motor.py +++ b/rocketpy/motors/solid_motor.py @@ -727,16 +727,19 @@ def propellant_I_13(self): def propellant_I_23(self): return 0 - def draw(self): - """Draw a representation of the SolidMotor.""" - self.plots.draw() - - def info(self): - """Prints out basic data about the SolidMotor.""" - self.prints.all() - self.plots.thrust() - - def all_info(self): - """Prints out all data and graphs available about the SolidMotor.""" - self.prints.all() - self.plots.all() + def draw(self, filename=None): + """Draw a representation of the SolidMotor. + + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + + Returns + ------- + None + """ + self.plots.draw(filename) diff --git a/rocketpy/motors/tank.py b/rocketpy/motors/tank.py index a3b21f434..418124cf0 100644 --- a/rocketpy/motors/tank.py +++ b/rocketpy/motors/tank.py @@ -476,9 +476,22 @@ def underfill_height_exception(param_name, param): elif (height < bottom_tolerance).any(): underfill_height_exception(name, height) - def draw(self): - """Draws the tank geometry.""" - self.plots.draw() + def draw(self, filename=None): + """Draws the tank geometry. + + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + + Returns + ------- + None + """ + self.plots.draw(filename) class MassFlowRateBasedTank(Tank): diff --git a/rocketpy/plots/aero_surface_plots.py b/rocketpy/plots/aero_surface_plots.py index 2a696b1fc..a0b7d4c4b 100644 --- a/rocketpy/plots/aero_surface_plots.py +++ b/rocketpy/plots/aero_surface_plots.py @@ -4,6 +4,8 @@ import numpy as np from matplotlib.patches import Ellipse +from .plot_helpers import show_or_save_plot + class _AeroSurfacePlots(ABC): """Abstract class that contains all aero surface plots.""" @@ -23,7 +25,7 @@ def __init__(self, aero_surface): self.aero_surface = aero_surface @abstractmethod - def draw(self): + def draw(self, filename=None): pass def lift(self): @@ -52,10 +54,18 @@ class _NoseConePlots(_AeroSurfacePlots): """Class that contains all nosecone plots. This class inherits from the _AeroSurfacePlots class.""" - def draw(self): + def draw(self, filename=None): """Draw the nosecone shape along with some important information, including the center line and the center of pressure position. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -122,8 +132,7 @@ def draw(self): ax.set_ylabel("Radius") ax.set_title(self.aero_surface.kind + " Nose Cone") ax.legend(bbox_to_anchor=(1, -0.2)) - - plt.show() + show_or_save_plot(filename) class _FinsPlots(_AeroSurfacePlots): @@ -131,7 +140,7 @@ class _FinsPlots(_AeroSurfacePlots): _AeroSurfacePlots class.""" @abstractmethod - def draw(self): + def draw(self, filename=None): pass def airfoil(self): @@ -192,10 +201,18 @@ class _TrapezoidalFinsPlots(_FinsPlots): """Class that contains all trapezoidal fin plots.""" # pylint: disable=too-many-statements - def draw(self): + def draw(self, filename=None): """Draw the fin shape along with some important information, including the center line, the quarter line and the center of pressure position. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -306,17 +323,25 @@ def draw(self): ax.legend(bbox_to_anchor=(1.05, 1.0), loc="upper left") plt.tight_layout() - plt.show() + show_or_save_plot(filename) class _EllipticalFinsPlots(_FinsPlots): """Class that contains all elliptical fin plots.""" # pylint: disable=too-many-statements - def draw(self): + def draw(self, filename=None): """Draw the fin shape along with some important information. These being: the center line and the center of pressure position. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -377,16 +402,24 @@ def draw(self): ax.legend(bbox_to_anchor=(1.05, 1.0), loc="upper left") plt.tight_layout() - plt.show() + show_or_save_plot(filename) class _FreeFormFinsPlots(_FinsPlots): """Class that contains all free form fin plots.""" # pylint: disable=too-many-statements - def draw(self): - """Draw the fin shape along with some important information, including - the center line, the quarter line and the center of pressure position. + def draw(self, filename=None): + """Draw the fin shape along with some important information. + These being: the center line and the center of pressure position. + + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -444,13 +477,13 @@ def draw(self): ax.legend(bbox_to_anchor=(1.05, 1.0), loc="upper left") plt.tight_layout() - plt.show() + show_or_save_plot(filename) class _TailPlots(_AeroSurfacePlots): """Class that contains all tail plots.""" - def draw(self): + def draw(self, filename=None): # This will de done in the future pass @@ -465,7 +498,7 @@ def drag_coefficient_curve(self): else: return self.aero_surface.drag_coefficient.plot() - def draw(self): + def draw(self, filename=None): raise NotImplementedError def all(self): @@ -481,12 +514,12 @@ def all(self): class _GenericSurfacePlots(_AeroSurfacePlots): """Class that contains all generic surface plots.""" - def draw(self): + def draw(self, filename=None): pass class _LinearGenericSurfacePlots(_AeroSurfacePlots): """Class that contains all linear generic surface plots.""" - def draw(self): + def draw(self, filename=None): pass diff --git a/rocketpy/plots/compare/compare_flights.py b/rocketpy/plots/compare/compare_flights.py index 3bf2fb8ed..2689ac482 100644 --- a/rocketpy/plots/compare/compare_flights.py +++ b/rocketpy/plots/compare/compare_flights.py @@ -3,6 +3,7 @@ import matplotlib.pyplot as plt import numpy as np +from ..plot_helpers import show_or_save_fig, show_or_save_plot from .compare import Compare @@ -89,8 +90,8 @@ def __process_savefig(self, filename, fig): ------- None """ + show_or_save_fig(fig, filename) if filename: - fig.savefig(filename) print("Plot saved to file: " + filename) else: plt.show() @@ -1153,9 +1154,11 @@ def compare_trajectories_3d( # pylint: disable=too-many-statements Tuple with the size of the figure. The default is (7,7). legend : boolean, optional Whether legend will or will not be plotted. Default is True - filename : string, optional - If a filename is passed, the figure will be saved in the current - directory. The default is None. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1221,10 +1224,7 @@ def compare_trajectories_3d( # pylint: disable=too-many-statements fig1.tight_layout() # Save figure - if filename: - plt.savefig(filename) - else: - plt.show() + show_or_save_plot(filename) def trajectories_3d(self, figsize=(7, 7), legend=None, filename=None): """Creates a trajectory plot that is the combination of the trajectories @@ -1237,8 +1237,11 @@ def trajectories_3d(self, figsize=(7, 7), legend=None, filename=None): must be in the form (width, height). legend : boolean, optional Whether legend will or will not be included. Default is True - savefig : string, optional - If a string is passed, the figure will be saved in the path passed. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1411,6 +1414,11 @@ def __plot_xz( # pylint: disable=too-many-statements .svg, .pgf, .eps figsize : tuple, optional Tuple with the size of the figure. The default is (7, 7). + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1454,10 +1462,7 @@ def __plot_xz( # pylint: disable=too-many-statements fig.tight_layout() # Save figure - if filename: - plt.savefig(filename) - else: - plt.show() + show_or_save_plot(filename) def __plot_yz( # pylint: disable=too-many-statements self, flights, names_list, figsize=(7, 7), legend=None, filename=None diff --git a/rocketpy/plots/environment_analysis_plots.py b/rocketpy/plots/environment_analysis_plots.py index 0b2c28990..b5b63642c 100644 --- a/rocketpy/plots/environment_analysis_plots.py +++ b/rocketpy/plots/environment_analysis_plots.py @@ -8,6 +8,7 @@ from rocketpy.units import convert_units from ..tools import find_two_closest_integers, import_optional_dependency +from .plot_helpers import show_or_save_plot # TODO: `wind_speed_limit` and `clear_range_limits` and should be numbers, not booleans @@ -73,12 +74,20 @@ def __beaufort_wind_scale(self, units, max_wind_speed=None): # Surface level plots - def wind_gust_distribution(self): + def wind_gust_distribution(self, filename=None): """Get all values of wind gust speed (for every date and hour available) and plot a single distribution. Expected result is a Weibull distribution, however, the result is not always a perfect fit, and sometimes it may look like a normal distribution. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -113,9 +122,9 @@ def wind_gust_distribution(self): plt.title("Wind Gust Speed Distribution (at surface)") plt.xlim(0, max(self.env_analysis.wind_gust_list)) plt.legend() - plt.show() + show_or_save_plot(filename) - def surface10m_wind_speed_distribution(self, wind_speed_limit=False): + def surface10m_wind_speed_distribution(self, wind_speed_limit=False, filename=None): """Get all values of sustained surface wind speed (for every date and hour available) and plot a single distribution. Expected result is a Weibull distribution. The wind speed limit is plotted as a vertical line. @@ -125,6 +134,11 @@ def surface10m_wind_speed_distribution(self, wind_speed_limit=False): wind_speed_limit : bool, optional If True, plots the wind speed limit as a vertical line. The default is False. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -172,14 +186,23 @@ def surface10m_wind_speed_distribution(self, wind_speed_limit=False): plt.title("Sustained Wind Speed Distribution (at surface+10m)") plt.xlim(0, max(self.env_analysis.surface_10m_wind_speed_list)) plt.legend() - plt.show() + show_or_save_plot(filename) def average_surface_temperature_evolution( self, + filename=None, ): # pylint: disable=too-many-statements """Plots average temperature progression throughout the day, including sigma contours. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -239,10 +262,10 @@ def average_surface_temperature_evolution( plt.title("Average Temperature Along Day") plt.grid(alpha=0.25) plt.legend() - plt.show() + show_or_save_plot(filename) def average_surface10m_wind_speed_evolution( - self, wind_speed_limit=False + self, wind_speed_limit=False, filename=None ): # pylint: disable=too-many-statements """Plots average surface wind speed progression throughout the day, including sigma contours. @@ -252,6 +275,11 @@ def average_surface10m_wind_speed_evolution( wind_speed_limit : bool, optional If True, plots the wind speed limit as a horizontal line. The default is False. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -334,14 +362,23 @@ def average_surface10m_wind_speed_evolution( plt.title("Average Sustained Surface Wind Speed Along Day") plt.grid(alpha=0.25) plt.legend() - plt.show() + show_or_save_plot(filename) def average_surface100m_wind_speed_evolution( self, + filename=None, ): # pylint: disable=too-many-statements """Plots average surface wind speed progression throughout the day, including sigma contours. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -408,11 +445,11 @@ def average_surface100m_wind_speed_evolution( plt.title("Average 100m Wind Speed Along Day") plt.grid(alpha=0.25) plt.legend() - plt.show() + show_or_save_plot(filename) # Average profiles plots (pressure level data) - def average_wind_speed_profile(self, clear_range_limits=False): + def average_wind_speed_profile(self, clear_range_limits=False, filename=None): """Average wind speed for all datetimes available. The plot also includes sigma contours. @@ -420,6 +457,11 @@ def average_wind_speed_profile(self, clear_range_limits=False): ---------- clear_range_limits : bool, optional If True, clears the range limits. The default is False. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -510,9 +552,9 @@ def average_wind_speed_profile(self, clear_range_limits=False): ) ), ) - plt.show() + show_or_save_plot(filename) - def average_wind_velocity_xy_profile(self, clear_range_limits=False): + def average_wind_velocity_xy_profile(self, clear_range_limits=False, filename=None): """Average wind X and wind Y for all datetimes available. The X component is the wind speed in the direction of East, and the Y component is the wind speed in the direction of North. @@ -521,6 +563,11 @@ def average_wind_velocity_xy_profile(self, clear_range_limits=False): ---------- clear_range_limits : bool, optional If True, clears the range limits. The default is False. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -572,15 +619,20 @@ def average_wind_velocity_xy_profile(self, clear_range_limits=False): plt.title("Average Wind X and Y Profile") plt.legend() plt.grid() - plt.show() + show_or_save_plot(filename) - def average_wind_heading_profile(self, clear_range_limits=False): + def average_wind_heading_profile(self, clear_range_limits=False, filename=None): """Average wind heading for all datetimes available. Parameters ---------- clear_range_limits : bool, optional If True, clears the range limits. The default is False. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -625,9 +677,9 @@ def average_wind_heading_profile(self, clear_range_limits=False): plt.xlim(0, 360) plt.title("Average Wind heading Profile") plt.legend() - plt.show() + show_or_save_plot(filename) - def average_pressure_profile(self, clear_range_limits=False): + def average_pressure_profile(self, clear_range_limits=False, filename=None): """Average pressure profile for all datetimes available. The plot also includes sigma contours. @@ -635,6 +687,11 @@ def average_pressure_profile(self, clear_range_limits=False): ---------- clear_range_limits : bool, optional If True, clears the range limits. The default is False. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -713,9 +770,9 @@ def average_pressure_profile(self, clear_range_limits=False): 0, max(np.percentile(self.env_analysis.pressure_profiles_list, 99.85, axis=0)), ) - plt.show() + show_or_save_plot(filename) - def average_temperature_profile(self, clear_range_limits=False): + def average_temperature_profile(self, clear_range_limits=False, filename=None): """Average temperature profile for all datetimes available. The plot also includes sigma contours. @@ -723,6 +780,11 @@ def average_temperature_profile(self, clear_range_limits=False): ---------- clear_range_limits : bool, optional If True, clears the range limits. The default is False. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -808,7 +870,7 @@ def average_temperature_profile(self, clear_range_limits=False): ) ), ) - plt.show() + show_or_save_plot(filename) # Wind roses (surface level data) @@ -856,7 +918,7 @@ def plot_wind_rose( ax.yaxis.set_major_formatter(mtick.PercentFormatter(decimals=0)) return ax - def average_wind_rose_specific_hour(self, hour, fig=None): + def average_wind_rose_specific_hour(self, hour, fig=None, filename=None): """Plot a specific hour of the average windrose Parameters @@ -865,6 +927,11 @@ def average_wind_rose_specific_hour(self, hour, fig=None): Hour to be plotted fig: matplotlib.pyplot.figure Figure to plot the windrose + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -882,11 +949,21 @@ def average_wind_rose_specific_hour(self, hour, fig=None): ), fig=fig, ) - plt.show() + show_or_save_plot(filename) - def average_wind_rose_grid(self): # pylint: disable=too-many-statements + def average_wind_rose_grid( + self, filename=None + ): # pylint: disable=too-many-statements """Plot wind roses for all hours of a day, in a grid like plot. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -950,7 +1027,7 @@ def average_wind_rose_grid(self): # pylint: disable=too-many-statements y=1, ) plt.bbox_inches = "tight" - plt.show() + show_or_save_plot(filename) def animate_average_wind_rose(self, figsize=(5, 5), filename="wind_rose.gif"): """Animates the wind_rose of an average day. The inputs of a wind_rose @@ -1010,10 +1087,20 @@ def animate_average_wind_rose(self, figsize=(5, 5), filename="wind_rose.gif"): # More plots and animations - def wind_gust_distribution_grid(self): # pylint: disable=too-many-statements + def wind_gust_distribution_grid( + self, filename=None + ): # pylint: disable=too-many-statements """Plots shown in the animation of how the wind gust distribution varies throughout the day. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -1079,7 +1166,7 @@ def wind_gust_distribution_grid(self): # pylint: disable=too-many-statements f"Wind Gust Speed ({self.env_analysis.unit_system['wind_speed']})" ) fig.supylabel("Probability") - plt.show() + show_or_save_plot(filename) def animate_wind_gust_distribution(self): # pylint: disable=too-many-statements """Animation of how the wind gust distribution varies throughout the day. @@ -1177,7 +1264,7 @@ def update(frame): return HTML(animation.to_jshtml()) def surface_wind_speed_distribution_grid( - self, wind_speed_limit=False + self, wind_speed_limit=False, filename=None ): # pylint: disable=too-many-statements """Plots shown in the animation of how the sustained surface wind speed distribution varies throughout the day. The plots are histograms of the @@ -1188,6 +1275,11 @@ def surface_wind_speed_distribution_grid( ---------- wind_speed_limit : bool, optional Whether to plot the wind speed limit as a vertical line + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1275,11 +1367,10 @@ def surface_wind_speed_distribution_grid( f"Sustained Surface Wind Speed ({self.env_analysis.unit_system['wind_speed']})" ) fig.supylabel("Probability") - plt.show() + show_or_save_plot(filename) - def animate_surface_wind_speed_distribution( - self, wind_speed_limit=False - ): # pylint: disable=too-many-statements + # pylint: disable=too-many-statements + def animate_surface_wind_speed_distribution(self, wind_speed_limit=False): """Animation of how the sustained surface wind speed distribution varies throughout the day. Each frame is a histogram of the wind speed distribution for a specific hour. @@ -1402,7 +1493,7 @@ def update(frame): return HTML(animation.to_jshtml()) def wind_speed_profile_grid( - self, clear_range_limits=False + self, clear_range_limits=False, filename=None ): # pylint: disable=too-many-statements """Creates a grid of plots with the wind profile over the average day. Each subplot represents a different hour of the day. @@ -1411,6 +1502,11 @@ def wind_speed_profile_grid( ---------- clear_range_limits : bool, optional Whether to clear the sky range limits or not, by default False + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1493,10 +1589,10 @@ def wind_speed_profile_grid( fig.suptitle("Average Wind Profile") fig.supxlabel(f"Wind speed ({self.env_analysis.unit_system['wind_speed']})") fig.supylabel(f"Altitude AGL ({self.env_analysis.unit_system['length']})") - plt.show() + show_or_save_plot(filename) def wind_heading_profile_grid( - self, clear_range_limits=False + self, clear_range_limits=False, filename=None ): # pylint: disable=too-many-statements """Creates a grid of plots with the wind heading profile over the average day. Each subplot represents a different hour of the day. @@ -1506,6 +1602,11 @@ def wind_heading_profile_grid( clear_range_limits : bool, optional Whether to clear the sky range limits or not, by default False. This is useful when the launch site is constrained in terms or altitude. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- @@ -1582,7 +1683,7 @@ def wind_heading_profile_grid( fig.suptitle("Average Wind Heading Profile") fig.supxlabel(f"Wind heading ({self.env_analysis.unit_system['angle']})") fig.supylabel(f"Altitude AGL ({self.env_analysis.unit_system['length']})") - plt.show() + show_or_save_plot(filename) def animate_wind_speed_profile( self, clear_range_limits=False diff --git a/rocketpy/plots/environment_plots.py b/rocketpy/plots/environment_plots.py index 39fb9548e..f15d066af 100644 --- a/rocketpy/plots/environment_plots.py +++ b/rocketpy/plots/environment_plots.py @@ -1,6 +1,8 @@ import matplotlib.pyplot as plt import numpy as np +from .plot_helpers import show_or_save_plot + class _EnvironmentPlots: """Class that holds plot methods for Environment class. @@ -171,10 +173,18 @@ def __pressure_temperature(self, ax): return ax - def gravity_model(self): + def gravity_model(self, filename=None): """Plots the gravity model graph that represents the gravitational acceleration as a function of height. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -192,13 +202,21 @@ def gravity_model(self): ax.set_ylim(self.grid[0], self.grid[-1]) plt.xticks(rotation=45) - plt.show() + show_or_save_plot(filename) - def atmospheric_model(self): + def atmospheric_model(self, filename=None): """Plots all atmospheric model graphs available. This includes wind speed and wind direction, density and speed of sound, wind u and wind v, and pressure and temperature. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -224,13 +242,21 @@ def atmospheric_model(self): ax4 = self.__pressure_temperature(ax4) plt.subplots_adjust(wspace=0.5, hspace=0.3) - plt.show() + show_or_save_plot(filename) # pylint: disable=too-many-statements - def ensemble_member_comparison(self): + def ensemble_member_comparison(self, filename=None): """Plots ensemble member comparisons. It requires that the environment model has been set as Ensemble. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -321,16 +347,24 @@ def ensemble_member_comparison(self): # Display plot plt.subplots_adjust(wspace=0.5, hspace=0.3) - plt.show() + show_or_save_plot(filename) # Clean up self.environment.select_ensemble_member(current_member) - def info(self): + def info(self, filename=None): """Plots a summary of the atmospheric model, including wind speed and wind direction, density and speed of sound. This is important for the Environment.info() method. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -346,7 +380,7 @@ def info(self): ax2 = self.__density_speed_of_sound(ax2) plt.subplots_adjust(wspace=0.5) - plt.show() + show_or_save_plot(filename) def all(self): """Prints out all graphs available about the Environment. This includes diff --git a/rocketpy/plots/flight_plots.py b/rocketpy/plots/flight_plots.py index 0babaff4c..edc30b956 100644 --- a/rocketpy/plots/flight_plots.py +++ b/rocketpy/plots/flight_plots.py @@ -3,6 +3,8 @@ import matplotlib.pyplot as plt import numpy as np +from .plot_helpers import show_or_save_plot + class _FlightPlots: """Class that holds plot methods for Flight class. @@ -52,9 +54,17 @@ def first_event_time_index(self): else: return -1 - def trajectory_3d(self): # pylint: disable=too-many-statements + def trajectory_3d(self, filename=None): # pylint: disable=too-many-statements """Plot a 3D graph of the trajectory + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -121,11 +131,21 @@ def trajectory_3d(self): # pylint: disable=too-many-statements ax1.set_zlim(min_z, max_z) ax1.view_init(15, 45) ax1.set_box_aspect(None, zoom=0.95) # 95% for label adjustment - plt.show() + show_or_save_plot(filename) - def linear_kinematics_data(self): # pylint: disable=too-many-statements + def linear_kinematics_data( + self, filename=None + ): # pylint: disable=too-many-statements """Prints out all Kinematics graphs available about the Flight + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -193,11 +213,19 @@ def linear_kinematics_data(self): # pylint: disable=too-many-statements ax4up.tick_params("y", colors="#1f77b4") plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) - def attitude_data(self): # pylint: disable=too-many-statements + def attitude_data(self, filename=None): # pylint: disable=too-many-statements """Prints out all Angular position graphs available about the Flight + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -243,12 +271,20 @@ def attitude_data(self): # pylint: disable=too-many-statements ax4.grid(True) plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) - def flight_path_angle_data(self): + def flight_path_angle_data(self, filename=None): """Prints out Flight path and Rocket Attitude angle graphs available about the Flight + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -285,12 +321,22 @@ def flight_path_angle_data(self): ax2.grid(True) plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) - def angular_kinematics_data(self): # pylint: disable=too-many-statements + def angular_kinematics_data( + self, filename=None + ): # pylint: disable=too-many-statements """Prints out all Angular velocity and acceleration graphs available about the Flight + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -351,11 +397,19 @@ def angular_kinematics_data(self): # pylint: disable=too-many-statements ax3up.tick_params("y", colors="#1f77b4") plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) - def rail_buttons_forces(self): # pylint: disable=too-many-statements + def rail_buttons_forces(self, filename=None): # pylint: disable=too-many-statements """Prints out all Rail Buttons Forces graphs available about the Flight. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -434,11 +488,19 @@ def rail_buttons_forces(self): # pylint: disable=too-many-statements ax2.set_title("Rail Buttons Shear Force") plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) - def aerodynamic_forces(self): # pylint: disable=too-many-statements + def aerodynamic_forces(self, filename=None): # pylint: disable=too-many-statements """Prints out all Forces and Moments graphs available about the Flight + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -514,11 +576,19 @@ def aerodynamic_forces(self): # pylint: disable=too-many-statements ax4.grid() plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) - def energy_data(self): # pylint: disable=too-many-statements + def energy_data(self, filename=None): # pylint: disable=too-many-statements """Prints out all Energy components graphs available about the Flight + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -625,12 +695,22 @@ def energy_data(self): # pylint: disable=too-many-statements ax4.grid() plt.subplots_adjust(hspace=1) - plt.show() + show_or_save_plot(filename) - def fluid_mechanics_data(self): # pylint: disable=too-many-statements + def fluid_mechanics_data( + self, filename=None + ): # pylint: disable=too-many-statements """Prints out a summary of the Fluid Mechanics graphs available about the Flight + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -715,12 +795,22 @@ def fluid_mechanics_data(self): # pylint: disable=too-many-statements ax6.grid() plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) - def stability_and_control_data(self): # pylint: disable=too-many-statements + def stability_and_control_data( + self, filename=None + ): # pylint: disable=too-many-statements """Prints out Rocket Stability and Control parameters graphs available about the Flight + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -795,11 +885,19 @@ def stability_and_control_data(self): # pylint: disable=too-many-statements ax2.grid() plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) - def pressure_rocket_altitude(self): + def pressure_rocket_altitude(self, filename=None): """Plots out pressure at rocket's altitude. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -816,7 +914,7 @@ def pressure_rocket_altitude(self): ax1.set_xlim(0, self.flight.t_final) ax1.grid() - plt.show() + show_or_save_plot(filename) def pressure_signals(self): """Plots out all Parachute Trigger Pressure Signals. diff --git a/rocketpy/plots/hybrid_motor_plots.py b/rocketpy/plots/hybrid_motor_plots.py index 21a415986..eb13c4d69 100644 --- a/rocketpy/plots/hybrid_motor_plots.py +++ b/rocketpy/plots/hybrid_motor_plots.py @@ -1,6 +1,7 @@ from matplotlib import pyplot as plt from .motor_plots import _MotorPlots +from .plot_helpers import show_or_save_plot class _HybridMotorPlots(_MotorPlots): @@ -13,7 +14,7 @@ class _HybridMotorPlots(_MotorPlots): """ - def grain_inner_radius(self, lower_limit=None, upper_limit=None): + def grain_inner_radius(self, lower_limit=None, upper_limit=None, filename=None): """Plots grain_inner_radius of the hybrid_motor as a function of time. Parameters @@ -24,15 +25,22 @@ def grain_inner_radius(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.solid.grain_inner_radius.plot(lower=lower_limit, upper=upper_limit) + self.motor.solid.grain_inner_radius.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def grain_height(self, lower_limit=None, upper_limit=None): + def grain_height(self, lower_limit=None, upper_limit=None, filename=None): """Plots grain_height of the hybrid_motor as a function of time. Parameters @@ -43,15 +51,22 @@ def grain_height(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.solid.grain_height.plot(lower=lower_limit, upper=upper_limit) + self.motor.solid.grain_height.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def burn_rate(self, lower_limit=None, upper_limit=None): + def burn_rate(self, lower_limit=None, upper_limit=None, filename=None): """Plots burn_rate of the hybrid_motor as a function of time. Parameters @@ -62,15 +77,22 @@ def burn_rate(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.solid.burn_rate.plot(lower=lower_limit, upper=upper_limit) + self.motor.solid.burn_rate.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def burn_area(self, lower_limit=None, upper_limit=None): + def burn_area(self, lower_limit=None, upper_limit=None, filename=None): """Plots burn_area of the hybrid_motor as a function of time. Parameters @@ -81,15 +103,22 @@ def burn_area(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.solid.burn_area.plot(lower=lower_limit, upper=upper_limit) + self.motor.solid.burn_area.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def Kn(self, lower_limit=None, upper_limit=None): + def Kn(self, lower_limit=None, upper_limit=None, filename=None): """Plots Kn of the hybrid_motor as a function of time. Parameters @@ -100,17 +129,32 @@ def Kn(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.solid.Kn.plot(lower=lower_limit, upper=upper_limit) + self.motor.solid.Kn.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def draw(self): + def draw(self, filename=None): """Draw a representation of the HybridMotor. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -144,7 +188,7 @@ def draw(self): ax.set_title("Hybrid Motor Representation") self._draw_center_of_mass(ax) self._set_plot_properties(ax) - plt.show() + show_or_save_plot(filename) def all(self): """Prints out all graphs available about the HybridMotor. It simply calls diff --git a/rocketpy/plots/liquid_motor_plots.py b/rocketpy/plots/liquid_motor_plots.py index 3e5e7703b..822364fc2 100644 --- a/rocketpy/plots/liquid_motor_plots.py +++ b/rocketpy/plots/liquid_motor_plots.py @@ -1,6 +1,7 @@ import matplotlib.pyplot as plt from .motor_plots import _MotorPlots +from .plot_helpers import show_or_save_plot class _LiquidMotorPlots(_MotorPlots): @@ -13,9 +14,17 @@ class _LiquidMotorPlots(_MotorPlots): """ - def draw(self): + def draw(self, filename=None): """Draw a representation of the LiquidMotor. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -41,7 +50,7 @@ def draw(self): ax.set_title("Liquid Motor Representation") self._draw_center_of_mass(ax) self._set_plot_properties(ax) - plt.show() + show_or_save_plot(filename) def all(self): """Prints out all graphs available about the LiquidMotor. It simply calls diff --git a/rocketpy/plots/motor_plots.py b/rocketpy/plots/motor_plots.py index 39cdcaeb6..e7ddb7efa 100644 --- a/rocketpy/plots/motor_plots.py +++ b/rocketpy/plots/motor_plots.py @@ -2,6 +2,8 @@ import numpy as np from matplotlib.patches import Polygon +from ..plots.plot_helpers import show_or_save_plot + class _MotorPlots: """Class that holds plot methods for Motor class. @@ -27,7 +29,7 @@ def __init__(self, motor): """ self.motor = motor - def thrust(self, lower_limit=None, upper_limit=None): + def thrust(self, lower_limit=None, upper_limit=None, filename=None): """Plots thrust of the motor as a function of time. Parameters @@ -38,14 +40,19 @@ def thrust(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.thrust.plot(lower=lower_limit, upper=upper_limit) + self.motor.thrust.plot(lower=lower_limit, upper=upper_limit, filename=filename) - def total_mass(self, lower_limit=None, upper_limit=None): + def total_mass(self, lower_limit=None, upper_limit=None, filename=None): """Plots total_mass of the motor as a function of time. Parameters @@ -56,14 +63,21 @@ def total_mass(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.total_mass.plot(lower=lower_limit, upper=upper_limit) + self.motor.total_mass.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def propellant_mass(self, lower_limit=None, upper_limit=None): + def propellant_mass(self, lower_limit=None, upper_limit=None, filename=None): """Plots propellant_mass of the motor as a function of time. Parameters @@ -74,14 +88,21 @@ def propellant_mass(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is None, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.propellant_mass.plot(lower=lower_limit, upper=upper_limit) + self.motor.propellant_mass.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def center_of_mass(self, lower_limit=None, upper_limit=None): + def center_of_mass(self, lower_limit=None, upper_limit=None, filename=None): """Plots center_of_mass of the motor as a function of time. Parameters @@ -92,14 +113,21 @@ def center_of_mass(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.center_of_mass.plot(lower=lower_limit, upper=upper_limit) + self.motor.center_of_mass.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def mass_flow_rate(self, lower_limit=None, upper_limit=None): + def mass_flow_rate(self, lower_limit=None, upper_limit=None, filename=None): """Plots mass_flow_rate of the motor as a function of time. Parameters @@ -110,14 +138,21 @@ def mass_flow_rate(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.mass_flow_rate.plot(lower=lower_limit, upper=upper_limit) + self.motor.mass_flow_rate.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def exhaust_velocity(self, lower_limit=None, upper_limit=None): + def exhaust_velocity(self, lower_limit=None, upper_limit=None, filename=None): """Plots exhaust_velocity of the motor as a function of time. Parameters @@ -128,14 +163,23 @@ def exhaust_velocity(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ - self.motor.exhaust_velocity.plot(lower=lower_limit, upper=upper_limit) + self.motor.exhaust_velocity.plot( + lower=lower_limit, upper=upper_limit, filename=filename + ) - def inertia_tensor(self, lower_limit=None, upper_limit=None, show_products=False): + def inertia_tensor( + self, lower_limit=None, upper_limit=None, show_products=False, filename=None + ): """Plots all inertia tensors (I_11, I_22, I_33, I_12, I_13, I_23) of the motor as a function of time in a single chart. @@ -152,9 +196,14 @@ def inertia_tensor(self, lower_limit=None, upper_limit=None, show_products=False in the plot. Default is False. These are kept as hidden by default because they are usually very small compared to the main inertia components. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). - Return - ------ + Returns + ------- None """ lower = lower_limit or self.motor.burn_start_time @@ -181,7 +230,7 @@ def inertia_tensor(self, lower_limit=None, upper_limit=None, show_products=False ax.legend() ax.grid(True) - plt.show() + show_or_save_plot(filename) def _generate_nozzle(self, translate=(0, 0), csys=1): """Generates a patch that represents the nozzle of the motor. It is diff --git a/rocketpy/plots/plot_helpers.py b/rocketpy/plots/plot_helpers.py new file mode 100644 index 000000000..f795d46e9 --- /dev/null +++ b/rocketpy/plots/plot_helpers.py @@ -0,0 +1,67 @@ +from pathlib import Path + +import matplotlib.pyplot as plt +from matplotlib.figure import Figure + +from ..tools import get_matplotlib_supported_file_endings + +SAVEFIG_DPI = 300 + + +def show_or_save_plot(filename=None): + """Shows or saves the current matplotlib plot. If a filename is given, the + plot will be saved, otherwise it will be shown. + + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to, by default None. Supported file + endings are: eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, + tif, tiff and webp (these are the formats supported by matplotlib). + """ + if filename is None: + plt.show() + else: + file_ending = Path(filename).suffix + supported_endings = get_matplotlib_supported_file_endings() + if file_ending not in supported_endings: + raise ValueError( + f"Unsupported file ending '{file_ending}'." + f"Supported file endings are: {supported_endings}." + ) + + # Before export, ensure the folder the file should go into exists + Path(filename).parent.mkdir(parents=True, exist_ok=True) + + plt.savefig(filename, dpi=SAVEFIG_DPI) + plt.close() + + +def show_or_save_fig(fig: Figure, filename=None): + """Shows or saves the given matplotlib Figure. If a filename is given, the + figure will be saved, otherwise it will be shown. + + Parameters + ---------- + fig : matplotlib.figure.Figure + The figure to be saved or shown. + filename : str | None, optional + The path the figure should be saved to, by default None. Supported file + endings are: eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, + tif, tiff and webp (these are the formats supported by matplotlib). + """ + if filename is None: + fig.show() + else: + file_ending = Path(filename).suffix + supported_endings = get_matplotlib_supported_file_endings() + if file_ending not in supported_endings: + raise ValueError( + f"Unsupported file ending '{file_ending}'." + f"Supported file endings are: {supported_endings}." + ) + + # Before export, ensure the folder the file should go into exists + Path(filename).parent.mkdir(parents=True, exist_ok=True) + + fig.savefig(filename, dpi=SAVEFIG_DPI) diff --git a/rocketpy/plots/rocket_plots.py b/rocketpy/plots/rocket_plots.py index 4ae47cdc4..9689a257b 100644 --- a/rocketpy/plots/rocket_plots.py +++ b/rocketpy/plots/rocket_plots.py @@ -7,6 +7,8 @@ from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail from rocketpy.rocket.aero_surface.generic_surface import GenericSurface +from .plot_helpers import show_or_save_plot + class _RocketPlots: """Class that holds plot methods for Rocket class. @@ -52,15 +54,23 @@ def reduced_mass(self): self.rocket.reduced_mass() - def static_margin(self): + def static_margin(self, filename=None): """Plots static margin of the rocket as a function of time. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None """ - self.rocket.static_margin() + self.rocket.static_margin(filename=filename) def stability_margin(self): """Plots static margin of the rocket as a function of time. @@ -113,9 +123,17 @@ def power_off_drag(self): self.rocket.power_off_drag() # pylint: disable=too-many-statements - def drag_curves(self): + def drag_curves(self, filename=None): """Plots power off and on drag curves of the rocket as a function of time. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -150,7 +168,7 @@ def drag_curves(self): ax.axvspan(0.8, 1.2, alpha=0.3, color="gray", label="Transonic Region") ax.legend(loc="best", shadow=True) plt.grid(True) - plt.show() + show_or_save_plot(filename) def thrust_to_weight(self): """ @@ -161,7 +179,7 @@ def thrust_to_weight(self): lower=0, upper=self.rocket.motor.burn_out_time ) - def draw(self, vis_args=None, plane="xz"): + def draw(self, vis_args=None, plane="xz", filename=None): """Draws the rocket in a matplotlib figure. Parameters @@ -188,6 +206,11 @@ def draw(self, vis_args=None, plane="xz"): plane : str, optional Plane in which the rocket will be drawn. Default is 'xz'. Other options is 'yz'. Used only for sensors representation. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). """ if vis_args is None: vis_args = { @@ -223,7 +246,7 @@ def draw(self, vis_args=None, plane="xz"): plt.ylabel("Radius (m)") plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left") plt.tight_layout() - plt.show() + show_or_save_plot(filename) def _draw_aerodynamic_surfaces(self, ax, vis_args, plane): """Draws the aerodynamic surfaces and saves the position of the points diff --git a/rocketpy/plots/solid_motor_plots.py b/rocketpy/plots/solid_motor_plots.py index 57af737b8..202220102 100644 --- a/rocketpy/plots/solid_motor_plots.py +++ b/rocketpy/plots/solid_motor_plots.py @@ -1,6 +1,7 @@ import matplotlib.pyplot as plt from .motor_plots import _MotorPlots +from .plot_helpers import show_or_save_plot class _SolidMotorPlots(_MotorPlots): @@ -100,7 +101,11 @@ def Kn(self, lower_limit=None, upper_limit=None): upper_limit : float Upper limit of the plot. Default is none, which means that the plot limits will be automatically calculated. - + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Return ------ None @@ -108,9 +113,17 @@ def Kn(self, lower_limit=None, upper_limit=None): self.motor.Kn.plot(lower=lower_limit, upper=upper_limit) - def draw(self): + def draw(self, filename=None): """Draw a representation of the SolidMotor. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -139,7 +152,7 @@ def draw(self): ax.set_title("Solid Motor Representation") self._draw_center_of_mass(ax) self._set_plot_properties(ax) - plt.show() + show_or_save_plot(filename) def all(self): """Prints out all graphs available about the SolidMotor. It simply calls diff --git a/rocketpy/plots/tank_plots.py b/rocketpy/plots/tank_plots.py index 7c0541eb2..a4bc0dd2a 100644 --- a/rocketpy/plots/tank_plots.py +++ b/rocketpy/plots/tank_plots.py @@ -2,6 +2,8 @@ import numpy as np from matplotlib.patches import Polygon +from .plot_helpers import show_or_save_plot + class _TankPlots: """Class that holds plot methods for Tank class. @@ -67,9 +69,17 @@ def _generate_tank(self, translate=(0, 0), csys=1): # Don't set any plot config here. Use the draw methods for that return tank - def draw(self): + def draw(self, filename=None): """Draws the tank geometry. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None @@ -89,6 +99,7 @@ def draw(self): 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) + show_or_save_plot(filename) def all(self): """Prints out all graphs available about the Tank. It simply calls diff --git a/rocketpy/rocket/aero_surface/fins/fins.py b/rocketpy/rocket/aero_surface/fins/fins.py index 29f4e8cff..c7973bc69 100644 --- a/rocketpy/rocket/aero_surface/fins/fins.py +++ b/rocketpy/rocket/aero_surface/fins/fins.py @@ -426,12 +426,20 @@ def compute_forces_and_moments( M3 = M3_forcing - M3_damping return R1, R2, R3, M1, M2, M3 - def draw(self): + def draw(self, filename=None): """Draw the fin shape along with some important information, including the center line, the quarter line and the center of pressure position. + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). + Returns ------- None """ - self.plots.draw() + self.plots.draw(filename) diff --git a/rocketpy/rocket/aero_surface/nose_cone.py b/rocketpy/rocket/aero_surface/nose_cone.py index 7d59473e3..e9fbe4324 100644 --- a/rocketpy/rocket/aero_surface/nose_cone.py +++ b/rocketpy/rocket/aero_surface/nose_cone.py @@ -486,15 +486,23 @@ def evaluate_center_of_pressure(self): self.cp = (self.cpx, self.cpy, self.cpz) return self.cp - def draw(self): - """Draw the fin shape along with some important information, including - the center line, the quarter line and the center of pressure position. + def draw(self, filename=None): + """Draw the nosecone shape along with some important information, + including the center of pressure position. + + Parameters + ---------- + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). Returns ------- None """ - self.plots.draw() + self.plots.draw(filename) def info(self): """Prints and plots summarized information of the nose cone. diff --git a/rocketpy/rocket/rocket.py b/rocketpy/rocket/rocket.py index ed376f582..58067a03f 100644 --- a/rocketpy/rocket/rocket.py +++ b/rocketpy/rocket/rocket.py @@ -1839,7 +1839,7 @@ def add_thrust_eccentricity(self, x, y): self.thrust_eccentricity_x = y return self - def draw(self, vis_args=None, plane="xz"): + def draw(self, vis_args=None, plane="xz", filename=None): """Draws the rocket in a matplotlib figure. Parameters @@ -1866,8 +1866,13 @@ def draw(self, vis_args=None, plane="xz"): plane : str, optional Plane in which the rocket will be drawn. Default is 'xz'. Other options is 'yz'. Used only for sensors representation. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case + the plot will be shown instead of saved. Supported file endings are: + eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff + and webp (these are the formats supported by matplotlib). """ - self.plots.draw(vis_args, plane) + self.plots.draw(vis_args, plane, filename) def info(self): """Prints out a summary of the data and graphs available about diff --git a/rocketpy/simulation/monte_carlo.py b/rocketpy/simulation/monte_carlo.py index c6cd1bed4..4ed051a4f 100644 --- a/rocketpy/simulation/monte_carlo.py +++ b/rocketpy/simulation/monte_carlo.py @@ -870,10 +870,10 @@ def export_ellipses_to_kml( # pylint: disable=too-many-statements for i, points in enumerate(kml_data): if i < len(impact_ellipses): - name = f"Impact Ellipse {i+1}" + name = f"Impact Ellipse {i + 1}" ellipse_color = colors[0] # default is blue else: - name = f"Apogee Ellipse {i +1- len(impact_ellipses)}" + name = f"Apogee Ellipse {i + 1 - len(impact_ellipses)}" ellipse_color = colors[1] # default is green mult_ell = kml.newmultigeometry(name=name) diff --git a/rocketpy/tools.py b/rocketpy/tools.py index 7668ecbc8..46660755a 100644 --- a/rocketpy/tools.py +++ b/rocketpy/tools.py @@ -15,6 +15,7 @@ import time from bisect import bisect_left +import matplotlib.pyplot as plt import numpy as np import pytz from cftime import num2pydate @@ -1149,6 +1150,23 @@ def euler313_to_quaternions(phi, theta, psi): return e0, e1, e2, e3 +def get_matplotlib_supported_file_endings(): + """Gets the file endings supported by matplotlib. + + Returns + ------- + list[str] + List of file endings prepended with a dot + """ + # Get matplotlib's supported file ending and return them (without descriptions, hence only keys) + filetypes = plt.gcf().canvas.get_supported_filetypes().keys() + + # Ensure the dot is included in the filetype endings + filetypes = ["." + filetype for filetype in filetypes] + + return filetypes + + if __name__ == "__main__": import doctest diff --git a/rocketpy/utilities.py b/rocketpy/utilities.py index adb925eee..fb728dbb6 100644 --- a/rocketpy/utilities.py +++ b/rocketpy/utilities.py @@ -9,6 +9,7 @@ from .environment.environment import Environment from .mathutils.function import Function +from .plots.plot_helpers import show_or_save_plot from .rocket.aero_surface import TrapezoidalFins from .simulation.flight import Flight @@ -202,7 +203,12 @@ def du(z, u): # pylint: disable=too-many-statements def fin_flutter_analysis( - fin_thickness, shear_modulus, flight, see_prints=True, see_graphs=True + fin_thickness, + shear_modulus, + flight, + see_prints=True, + see_graphs=True, + filename=None, ): """Calculate and plot the Fin Flutter velocity using the pressure profile provided by the selected atmospheric model. It considers the Flutter @@ -225,6 +231,11 @@ def fin_flutter_analysis( see_graphs : boolean, optional True if you want to see the graphs, False otherwise. If False, the function will return the vectors containing the data for the graphs. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case the + plot will be shown instead of saved. Supported file endings are: eps, + jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and webp + (these are the formats supported by matplotlib). Return ------ @@ -270,7 +281,7 @@ def fin_flutter_analysis( flight, ) if see_graphs: - _flutter_plots(flight, flutter_mach, safety_factor) + _flutter_plots(flight, flutter_mach, safety_factor, filename) else: return flutter_mach, safety_factor @@ -309,7 +320,7 @@ def _flutter_safety_factor(flight, flutter_mach): return safety_factor -def _flutter_plots(flight, flutter_mach, safety_factor): +def _flutter_plots(flight, flutter_mach, safety_factor, filename=None): """Plot the Fin Flutter Mach Number and the Safety Factor for the flutter. Parameters @@ -322,6 +333,11 @@ def _flutter_plots(flight, flutter_mach, safety_factor): safety_factor : rocketpy.Function Function containing the Safety Factor for the fin flutter. See fin_flutter_analysis for more details. + filename : str | None, optional + The path the plot should be saved to. By default None, in which case the + plot will be shown instead of saved. Supported file endings are: eps, + jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and webp + (these are the formats supported by matplotlib). Returns ------- @@ -357,7 +373,7 @@ def _flutter_plots(flight, flutter_mach, safety_factor): ax2.grid() plt.subplots_adjust(hspace=0.5) - plt.show() + show_or_save_plot(filename) def _flutter_prints( diff --git a/tests/integration/test_plots.py b/tests/integration/test_plots.py index 232ef71c6..7595855e4 100644 --- a/tests/integration/test_plots.py +++ b/tests/integration/test_plots.py @@ -2,12 +2,50 @@ import os from unittest.mock import patch +import matplotlib.pyplot as plt + from rocketpy import Flight from rocketpy.plots.compare import CompareFlights @patch("matplotlib.pyplot.show") -def test_compare_flights(mock_show, calisto, example_plain_env): +def test_compare(mock_show, flight_calisto): + """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. + flight_calisto : rocketpy.Flight + Flight object to be used in the tests. See conftest.py for more details. + """ + flight = flight_calisto + + objects = [flight, flight, flight] + + comparison = CompareFlights(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) + + +@patch("matplotlib.pyplot.show") +@patch("matplotlib.figure.Figure.show") +def test_compare_flights(mock_show, mock_figure_show, calisto, example_plain_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. @@ -16,6 +54,8 @@ def test_compare_flights(mock_show, calisto, example_plain_env): ---------- mock_show : Mocks the matplotlib.pyplot.show() function to avoid showing the plots. + mock_figure_show : + Mocks the matplotlib.figure.Figure.show() function to avoid showing the plots. calisto : rocketpy.Rocket Rocket object to be used in the tests. See conftest.py for more details. example_plain_env : rocketpy.Environment diff --git a/tests/unit/test_plots.py b/tests/unit/test_plots.py index cd35f8d11..6d7c0cce2 100644 --- a/tests/unit/test_plots.py +++ b/tests/unit/test_plots.py @@ -1,8 +1,11 @@ -from unittest.mock import patch +import os +from unittest.mock import MagicMock, patch import matplotlib.pyplot as plt +import pytest from rocketpy.plots.compare import Compare +from rocketpy.plots.plot_helpers import show_or_save_fig, show_or_save_plot @patch("matplotlib.pyplot.show") @@ -38,3 +41,51 @@ def test_compare(mock_show, flight_calisto): # pylint: disable=unused-argument ) assert isinstance(fig, plt.Figure) + + +@patch("matplotlib.pyplot.show") +@pytest.mark.parametrize("filename", [None, "test.png"]) +def test_show_or_save_plot(mock_show, filename): + """This test is to check if the show_or_save_plot function is + working properly. + + Parameters + ---------- + mock_show : + Mocks the matplotlib.pyplot.show() function to avoid showing + the plots. + filename : str + Name of the file to save the plot. If None, the plot will be + shown instead. + """ + plt.subplots() + show_or_save_plot(filename) + + if filename is None: + mock_show.assert_called_once() + else: + assert os.path.exists(filename) + os.remove(filename) + + +@pytest.mark.parametrize("filename", [None, "test.png"]) +def test_show_or_save_fig(filename): + """This test is to check if the show_or_save_fig function is + working properly. + + Parameters + ---------- + filename : str + Name of the file to save the plot. If None, the plot will be + shown instead. + """ + fig, _ = plt.subplots() + + fig.show = MagicMock() + show_or_save_fig(fig, filename) + + if filename is None: + fig.show.assert_called_once() + else: + assert os.path.exists(filename) + os.remove(filename)