diff --git a/CHANGELOG.md b/CHANGELOG.md index 8e3d514c0..c65e36614 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,7 @@ Attention: The newest changes should be on top --> ### Added +- ENH: Add save functionality to `_MonteCarloPlots.all` method [#848](https://github.com/RocketPy-Team/RocketPy/pull/848) - ENH: Add persistent caching for ThrustCurve API [#881](https://github.com/RocketPy-Team/RocketPy/pull/881) - ENH: Compatibility with MERRA-2 atmosphere reanalysis files [#825](https://github.com/RocketPy-Team/RocketPy/pull/825) - ENH: Enable only radial burning [#815](https://github.com/RocketPy-Team/RocketPy/pull/815) diff --git a/rocketpy/plots/monte_carlo_plots.py b/rocketpy/plots/monte_carlo_plots.py index 22d320b6b..933bf9301 100644 --- a/rocketpy/plots/monte_carlo_plots.py +++ b/rocketpy/plots/monte_carlo_plots.py @@ -1,8 +1,11 @@ +from pathlib import Path + import matplotlib.pyplot as plt import numpy as np from matplotlib.transforms import offset_copy from ..tools import generate_monte_carlo_ellipses, import_optional_dependency +from .plot_helpers import show_or_save_plot class _MonteCarloPlots: @@ -159,7 +162,7 @@ def ellipses( else: plt.show() - def all(self, keys=None): + def all(self, keys=None, *, filename=None): """ Plot the histograms of the Monte Carlo simulation results. @@ -168,6 +171,14 @@ def all(self, keys=None): keys : str, list or tuple, optional The keys of the results to be plotted. If None, all results will be plotted. 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. If a filename is provided, + each histogram will be saved with the key name appended to the + filename (e.g., "plots/histogram_apogee.png" for key "apogee" with + filename "plots/histogram.png"). 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 ------- @@ -207,7 +218,16 @@ def all(self, keys=None): ax1.set_xticks([]) plt.tight_layout() - plt.show() + + # Generate the filename for this specific key if saving + if filename is not None: + file_path = Path(filename) + key_filename = str( + file_path.parent / f"{file_path.stem}_{key}{file_path.suffix}" + ) + show_or_save_plot(key_filename) + else: + show_or_save_plot(None) def plot_comparison(self, other_monte_carlo): """ diff --git a/tests/integration/simulation/test_monte_carlo.py b/tests/integration/simulation/test_monte_carlo.py index 50ebdaa5e..4b1b82392 100644 --- a/tests/integration/simulation/test_monte_carlo.py +++ b/tests/integration/simulation/test_monte_carlo.py @@ -11,14 +11,18 @@ def _post_test_file_cleanup(): """Clean monte carlo files after test session if they exist.""" - if os.path.exists("monte_carlo_class_example.kml"): - os.remove("monte_carlo_class_example.kml") - if os.path.exists("monte_carlo_test.errors.txt"): - os.remove("monte_carlo_test.errors.txt") - if os.path.exists("monte_carlo_test.inputs.txt"): - os.remove("monte_carlo_test.inputs.txt") - if os.path.exists("monte_carlo_test.outputs.txt"): - os.remove("monte_carlo_test.outputs.txt") + files_to_cleanup = [ + "monte_carlo_class_example.kml", + "monte_carlo_test.errors.txt", + "monte_carlo_test.inputs.txt", + "monte_carlo_test.outputs.txt", + "test_histogram_apogee.png", + "test_multi_apogee.png", + "test_multi_x_impact.png", + ] + for filepath in files_to_cleanup: + if os.path.exists(filepath): + os.remove(filepath) @pytest.mark.slow @@ -134,6 +138,31 @@ def test_monte_carlo_plots(mock_show, monte_carlo_calisto_pre_loaded): _post_test_file_cleanup() +def test_monte_carlo_plots_all_save(monte_carlo_calisto_pre_loaded): + """Tests the plots.all method with save functionality. + + Parameters + ---------- + monte_carlo_calisto_pre_loaded : MonteCarlo + The MonteCarlo object, this is a pytest fixture. + """ + try: + # Test saving with a single key + monte_carlo_calisto_pre_loaded.plots.all( + keys="apogee", filename="test_histogram.png" + ) + assert os.path.exists("test_histogram_apogee.png") + + # Test saving with multiple keys + monte_carlo_calisto_pre_loaded.plots.all( + keys=["apogee", "x_impact"], filename="test_multi.png" + ) + assert os.path.exists("test_multi_apogee.png") + assert os.path.exists("test_multi_x_impact.png") + finally: + _post_test_file_cleanup() + + def test_monte_carlo_export_ellipses_to_kml(monte_carlo_calisto_pre_loaded): """Tests the export_ellipses_to_kml method of the MonteCarlo class.