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.
@@ -1088,13 +1090,15 @@ 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.
 
         Returns
         -------
         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 +1159,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 +1191,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 +1222,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: eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and webp.
 
         Returns
         -------
@@ -1254,7 +1264,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 +1287,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 +1327,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: eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and webp.
 
         Returns
         -------
@@ -1389,7 +1402,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 +1417,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 +1462,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: eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and webp.
 
         Returns
         -------
@@ -1522,7 +1538,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,34 @@ def add_tank(self, tank, position):
         )
         reset_funcified_methods(self)
 
-    def draw(self):
-        """Draws a representation of the HybridMotor."""
-        self.plots.draw()
+    def draw(self, filename=None):
+        """Draws a representation of the HybridMotor.
 
-    def info(self):
-        """Prints out basic data 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.
+
+        Returns
+        -------
+        None
+        """
+        self.plots.draw(filename)
+
+    def info(self, filename=None):
+        """Prints out basic data 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.
+
+        Returns
+        -------
+        None
+        """
         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,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.
+
+        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.
+
+        Returns
+        -------
+        None
+        """
+        self.plots.draw(filename)
 
     def info(self):
         """Prints out basic data about the Motor."""

rocketpy/motors/motor.py

@@ -1035,13 +1035,22 @@ 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.
+
+        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.
+
+        Returns
+        -------
+        None
         """
         # 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,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 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.
+
+        Returns
+        -------
+        None
+        """
+        self.plots.draw(filename)
 
     def info(self):
         """Prints out basic data about the SolidMotor."""

rocketpy/motors/tank.py

@@ -485,9 +485,19 @@ 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.
+
+        Returns
+        -------
+        None
+        """
+        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: eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and webp.
+
         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: eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and webp.
+
         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: eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and webp.
+
         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):