ENH: Add filename option to most plot functions

Author: nalquas

rocketpy/mathutils/function.py

@@ -16,6 +16,8 @@
 import numpy as np
 from scipy import integrate, linalg, optimize
 
+from ..plots.plot_helpers import show_or_save_plot
+
 NUMERICAL_TYPES = (float, int, complex, np.ndarray, np.integer, np.floating)
 INTERPOLATION_TYPES = {
     "linear": 0,
@@ -1074,7 +1076,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.
@@ -1093,8 +1095,8 @@ def __call__(self, *args):
         -------
         ans : None, scalar, list
         """
-        if len(args) == 0:
-            return self.plot()
+        if len(args) == 0 or (len(args) == 1 and filename != None):
+            return self.plot(filename=filename)
         else:
             return self.get_value(*args)
 
@@ -1155,8 +1157,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)
@@ -1184,6 +1189,7 @@ def plot_1d(
         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
@@ -1214,6 +1220,8 @@ def plot_1d(
             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: png, pdf, ps, eps and svg.
 
         Returns
         -------
@@ -1254,7 +1262,7 @@ def plot_1d(
         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
 
@@ -1277,6 +1285,7 @@ def plot_2d(
         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
@@ -1316,6 +1325,8 @@ def plot_2d(
         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: png, pdf, ps, eps and svg.
 
         Returns
         -------
@@ -1389,7 +1400,7 @@ def plot_2d(
         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(
@@ -1404,6 +1415,7 @@ def compare_plots(
         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
@@ -1448,6 +1460,8 @@ def compare_plots(
             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: png, pdf, ps, eps and svg.
 
         Returns
         -------
@@ -1522,7 +1536,7 @@ def compare_plots(
         plt.ylabel(ylabel)
 
         if show:
-            plt.show()
+            show_or_save_plot(filename)
 
         if return_object:
             return fig, ax

rocketpy/motors/hybrid_motor.py

@@ -600,14 +600,14 @@ def add_tank(self, tank, position):
         )
         reset_funcified_methods(self)
 
-    def draw(self):
+    def draw(self, filename=None):
         """Draws a representation of the HybridMotor."""
-        self.plots.draw()
+        self.plots.draw(filename)
 
-    def info(self):
+    def info(self, filename=None):
         """Prints out basic data about the Motor."""
         self.prints.all()
-        self.plots.thrust()
+        self.plots.thrust(filename=filename)
         return None
 
     def all_info(self):

rocketpy/motors/liquid_motor.py

@@ -466,9 +466,9 @@ def add_tank(self, tank, position):
         self.positioned_tanks.append({"tank": tank, "position": position})
         reset_funcified_methods(self)
 
-    def draw(self):
+    def draw(self, filename=None):
         """Draw a representation of the LiquidMotor."""
-        self.plots.draw()
+        self.plots.draw(filename)
 
     def info(self):
         """Prints out basic data about the Motor."""

rocketpy/motors/motor.py

@@ -1035,13 +1035,13 @@ def get_attr_value(obj, attr_name, multiplier=1):
 
         return None
 
-    def info(self):
+    def info(self, filename=None):
         """Prints out a summary of the data and graphs available about the
         Motor.
         """
         # Print motor details
         self.prints.all()
-        self.plots.thrust()
+        self.plots.thrust(filename=filename)
         return None
 
     @abstractmethod

rocketpy/motors/solid_motor.py

@@ -699,9 +699,9 @@ def propellant_I_13(self):
     def propellant_I_23(self):
         return 0
 
-    def draw(self):
+    def draw(self, filename=None):
         """Draw a representation of the SolidMotor."""
-        self.plots.draw()
+        self.plots.draw(filename)
 
     def info(self):
         """Prints out basic data about the SolidMotor."""

rocketpy/motors/tank.py

@@ -485,9 +485,9 @@ def underfill_height_exception(param_name, param):
             elif (height < bottom_tolerance).any():
                 underfill_height_exception(name, height)
 
-    def draw(self):
+    def draw(self, filename=None):
         """Draws the tank geometry."""
-        self.plots.draw()
+        self.plots.draw(filename)
 
 
 class MassFlowRateBasedTank(Tank):

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."""
@@ -24,7 +26,7 @@ def __init__(self, aero_surface):
         return None
 
     @abstractmethod
-    def draw(self):
+    def draw(self, filename=None):
         pass
 
     def lift(self):
@@ -70,10 +72,15 @@ def __init__(self, nosecone):
         super().__init__(nosecone)
         return None
 
-    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: png, pdf, ps, eps and svg.
+
         Returns
         -------
         None
@@ -141,7 +148,7 @@ def draw(self):
         ax.set_title(self.aero_surface.kind + " Nose Cone")
         ax.legend(bbox_to_anchor=(1, -0.2))
         # Show Plot
-        plt.show()
+        show_or_save_plot(filename)
         return None
 
 
@@ -165,7 +172,7 @@ def __init__(self, fin_set):
         return None
 
     @abstractmethod
-    def draw(self):
+    def draw(self, filename=None):
         pass
 
     def airfoil(self):
@@ -233,10 +240,15 @@ def __init__(self, fin_set):
         super().__init__(fin_set)
         return None
 
-    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: png, pdf, ps, eps and svg.
+
         Returns
         -------
         None
@@ -347,7 +359,7 @@ 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)
         return None
 
 
@@ -358,10 +370,15 @@ def __init__(self, fin_set):
         super().__init__(fin_set)
         return None
 
-    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: png, pdf, ps, eps and svg.
+
         Returns
         -------
         None
@@ -422,7 +439,7 @@ 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)
 
         return None
 
@@ -445,7 +462,7 @@ def __init__(self, tail):
         super().__init__(tail)
         return None
 
-    def draw(self):
+    def draw(self, filename=None):
         # This will de done in the future
         return None
 
@@ -474,7 +491,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):

rocketpy/plots/compare/compare_flights.py

@@ -3,6 +3,8 @@
 
 from .compare import Compare
 
+from ..plot_helpers import show_or_save_plot, show_or_save_fig
+
 
 class CompareFlights(Compare):
     """A class to compare the results of multiple flights.
@@ -90,11 +92,9 @@ 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()
         return None
 
     def __process_legend(self, legend, fig):
@@ -1276,10 +1276,7 @@ def compare_trajectories_3d(
         fig1.tight_layout()
 
         # Save figure
-        if filename:
-            plt.savefig(filename)
-        else:
-            plt.show()
+        show_or_save_plot(filename)
 
         return None
 
@@ -1474,6 +1471,8 @@ def __plot_xz(
             .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: png, pdf, ps, eps and svg.
 
         Returns
         -------
@@ -1517,10 +1516,7 @@ def __plot_xz(
         fig.tight_layout()
 
         # Save figure
-        if filename:
-            plt.savefig(filename)
-        else:
-            plt.show()
+        show_or_save_plot(filename)
 
         return None