Skip to content

Commit 0c8790e

Browse files
ENH: Add save functionality to _MonteCarloPlots.all method (#884)
* Initial plan * Add save functionality to _MonteCarloPlots.all method Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com> * Improve test cleanup for monte_carlo_plots_all_save test Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com> * DOC: Add PR #848 to CHANGELOG.md Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com>
1 parent d63add7 commit 0c8790e

File tree

3 files changed

+60
-10
lines changed

3 files changed

+60
-10
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ Attention: The newest changes should be on top -->
3232

3333
### Added
3434

35+
- ENH: Add save functionality to `_MonteCarloPlots.all` method [#848](https://github.com/RocketPy-Team/RocketPy/pull/848)
3536
- ENH: Add persistent caching for ThrustCurve API [#881](https://github.com/RocketPy-Team/RocketPy/pull/881)
3637
- ENH: Compatibility with MERRA-2 atmosphere reanalysis files [#825](https://github.com/RocketPy-Team/RocketPy/pull/825)
3738
- ENH: Enable only radial burning [#815](https://github.com/RocketPy-Team/RocketPy/pull/815)

rocketpy/plots/monte_carlo_plots.py

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
1+
from pathlib import Path
2+
13
import matplotlib.pyplot as plt
24
import numpy as np
35
from matplotlib.transforms import offset_copy
46

57
from ..tools import generate_monte_carlo_ellipses, import_optional_dependency
8+
from .plot_helpers import show_or_save_plot
69

710

811
class _MonteCarloPlots:
@@ -159,7 +162,7 @@ def ellipses(
159162
else:
160163
plt.show()
161164

162-
def all(self, keys=None):
165+
def all(self, keys=None, *, filename=None):
163166
"""
164167
Plot the histograms of the Monte Carlo simulation results.
165168
@@ -168,6 +171,14 @@ def all(self, keys=None):
168171
keys : str, list or tuple, optional
169172
The keys of the results to be plotted. If None, all results will be
170173
plotted. Default is None.
174+
filename : str | None, optional
175+
The path the plot should be saved to. By default None, in which case
176+
the plot will be shown instead of saved. If a filename is provided,
177+
each histogram will be saved with the key name appended to the
178+
filename (e.g., "plots/histogram_apogee.png" for key "apogee" with
179+
filename "plots/histogram.png"). Supported file endings are: eps,
180+
jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff and
181+
webp (these are the formats supported by matplotlib).
171182
172183
Returns
173184
-------
@@ -207,7 +218,16 @@ def all(self, keys=None):
207218
ax1.set_xticks([])
208219

209220
plt.tight_layout()
210-
plt.show()
221+
222+
# Generate the filename for this specific key if saving
223+
if filename is not None:
224+
file_path = Path(filename)
225+
key_filename = str(
226+
file_path.parent / f"{file_path.stem}_{key}{file_path.suffix}"
227+
)
228+
show_or_save_plot(key_filename)
229+
else:
230+
show_or_save_plot(None)
211231

212232
def plot_comparison(self, other_monte_carlo):
213233
"""

tests/integration/simulation/test_monte_carlo.py

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,18 @@
1111

1212
def _post_test_file_cleanup():
1313
"""Clean monte carlo files after test session if they exist."""
14-
if os.path.exists("monte_carlo_class_example.kml"):
15-
os.remove("monte_carlo_class_example.kml")
16-
if os.path.exists("monte_carlo_test.errors.txt"):
17-
os.remove("monte_carlo_test.errors.txt")
18-
if os.path.exists("monte_carlo_test.inputs.txt"):
19-
os.remove("monte_carlo_test.inputs.txt")
20-
if os.path.exists("monte_carlo_test.outputs.txt"):
21-
os.remove("monte_carlo_test.outputs.txt")
14+
files_to_cleanup = [
15+
"monte_carlo_class_example.kml",
16+
"monte_carlo_test.errors.txt",
17+
"monte_carlo_test.inputs.txt",
18+
"monte_carlo_test.outputs.txt",
19+
"test_histogram_apogee.png",
20+
"test_multi_apogee.png",
21+
"test_multi_x_impact.png",
22+
]
23+
for filepath in files_to_cleanup:
24+
if os.path.exists(filepath):
25+
os.remove(filepath)
2226

2327

2428
@pytest.mark.slow
@@ -134,6 +138,31 @@ def test_monte_carlo_plots(mock_show, monte_carlo_calisto_pre_loaded):
134138
_post_test_file_cleanup()
135139

136140

141+
def test_monte_carlo_plots_all_save(monte_carlo_calisto_pre_loaded):
142+
"""Tests the plots.all method with save functionality.
143+
144+
Parameters
145+
----------
146+
monte_carlo_calisto_pre_loaded : MonteCarlo
147+
The MonteCarlo object, this is a pytest fixture.
148+
"""
149+
try:
150+
# Test saving with a single key
151+
monte_carlo_calisto_pre_loaded.plots.all(
152+
keys="apogee", filename="test_histogram.png"
153+
)
154+
assert os.path.exists("test_histogram_apogee.png")
155+
156+
# Test saving with multiple keys
157+
monte_carlo_calisto_pre_loaded.plots.all(
158+
keys=["apogee", "x_impact"], filename="test_multi.png"
159+
)
160+
assert os.path.exists("test_multi_apogee.png")
161+
assert os.path.exists("test_multi_x_impact.png")
162+
finally:
163+
_post_test_file_cleanup()
164+
165+
137166
def test_monte_carlo_export_ellipses_to_kml(monte_carlo_calisto_pre_loaded):
138167
"""Tests the export_ellipses_to_kml method of the MonteCarlo class.
139168

0 commit comments

Comments
 (0)