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 <[email protected]>

* 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 <[email protected]>

* 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 <[email protected]>
Co-authored-by: Lucas Prates <[email protected]>
Co-authored-by: Pedro Bressan <[email protected]>

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
 

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 <functionusage>`.
+
+.. 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
 ----------------
 

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

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)

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)

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()

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)

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):