issue-RocketPy-Team#90-team-05

rocketpy/Environment.py

@@ -11,6 +11,7 @@
 import warnings
 from datetime import datetime, timedelta
 
+import os
 import matplotlib.pyplot as plt
 import numpy as np
 import pytz
@@ -300,6 +301,8 @@ def __init__(
         elevation=0,
         datum="SIRGAS2000",
         timeZone="UTC",
+        saveImagesPng=False,
+        saveImagesPdf=False,
     ):
         """Initialize Environment class, saving launch rail length,
         launch date, location coordinates and elevation. Note that
@@ -365,6 +368,9 @@ def __init__(
             self.localDate = None
             self.timeZone = None
 
+        self.saveImagesPng = saveImagesPng
+        self.saveImagesPdf = saveImagesPdf
+
         # Initialize constants
         self.earthRadius = 6.3781 * (10**6)
         self.airGasConstant = 287.05287  # in J/K/Kg
@@ -2972,7 +2978,7 @@ def info(self):
 
         plt.subplots_adjust(wspace=0.5)
         plt.show()
-
+        self.saveImages("AtmosphericModelPlots")
     def allInfo(self):
         """Prints out all data and graphs available about the Environment.
 
@@ -3208,6 +3214,7 @@ def allInfo(self):
         # Display plot
         plt.subplots_adjust(wspace=0.5, hspace=0.3)
         plt.show()
+        self.saveImages("EnsembleMembersComparison")
 
         # Clean up
         self.selectEnsembleMember(currentMember)
@@ -3675,3 +3682,13 @@ def printEarthDetails(self):
         print("Gravity acceleration at launch site: Still not implemented :(")
 
         return None
+
+    def saveImages(self, name):
+        if isinstance(self.saveImagesPng, str):
+            plt.savefig(os.path.join(self.saveImagesPng, name+".png"))
+        elif self.saveImagesPng is True:
+            plt.savefig(name+".png")
+        elif isinstance(self.saveImagesPdf, str):
+            plt.savefig(os.path.join(self.saveImagesPdf, name+".pdf"))
+        elif self.saveImagesPdf is True:
+            plt.savefig(name+".pdf")

rocketpy/EnvironmentAnalysis.py

@@ -8,6 +8,7 @@
 import warnings
 from collections import defaultdict
 
+import os
 import ipywidgets as widgets
 import matplotlib.ticker as mtick
 import netCDF4
@@ -62,6 +63,8 @@ def __init__(
         longitude,
         start_hour=0,
         end_hour=24,
+        saveImagesPng=False,
+        saveImagesPdf=False,
         surfaceDataFile=None,
         pressureLevelDataFile=None,
         timezone=None,
@@ -116,6 +119,8 @@ def __init__(
         self.surfaceDataFile = surfaceDataFile
         self.pressureLevelDataFile = pressureLevelDataFile
         self.preferred_timezone = timezone
+        self.saveImagesPng = saveImagesPng
+        self.saveImagesPdf = saveImagesPdf
 
         # Manage units and timezones
         self.__init_data_parsing_units()
@@ -1185,6 +1190,7 @@ def plot_wind_gust_distribution(self):
         plt.title("Wind Gust Speed Distribution")
         plt.legend()
         plt.show()
+        self.saveImages("WindGustSpeedDistribution")
 
         return None
 
@@ -1239,6 +1245,7 @@ def plot_surface10m_wind_speed_distribution(self, SAcup_wind_constraints=False):
         plt.title("Sustained Surface Wind Speed Distribution")
         plt.legend()
         plt.show()
+        self.saveImages("SustainedSurfaceWindSpeedDistribution")
 
         return None
 
@@ -1321,6 +1328,7 @@ def plot_average_temperature_along_day(self):
         plt.grid(alpha=0.25)
         plt.legend()
         plt.show()
+        self.saveImages("AverageTemperatureAlongDay")
 
     def calculate_average_sustained_surface10m_wind_along_day(self):
         """Computes average sustained wind speed progression throughout the
@@ -1422,6 +1430,7 @@ def plot_average_surface10m_wind_speed_along_day(
         plt.grid(alpha=0.25)
         plt.legend()
         plt.show()
+        self.saveImages("AverageSustainedSurfaceWindSpeedAlongDay")
 
     def calculate_average_sustained_surface100m_wind_along_day(self):
         """Computes average sustained wind speed progression throughout the
@@ -1512,6 +1521,7 @@ def plot_average_sustained_surface100m_wind_speed_along_day(self):
         plt.grid(alpha=0.25)
         plt.legend()
         plt.show()
+        self.saveImages("Average100mWindSpeedAlongDay")
 
     def plot_average_wind_speed_profile(self, SAcup_altitude_constraints=False):
         """Average wind speed for all datetimes available."""
@@ -1585,6 +1595,7 @@ def plot_average_wind_speed_profile(self, SAcup_altitude_constraints=False):
         plt.title("Average Wind Speed Profile")
         plt.legend()
         plt.show()
+        self.saveImages("AverageWindSpeedProfile")
 
     def process_wind_speed_and_direction_data_for_average_day(self):
         """Process the wind_speed and wind_direction data to generate lists of all the wind_speeds recorded
@@ -1714,6 +1725,7 @@ def plot_average_pressure_profile(self, SAcup_altitude_constraints=False):
         plt.title("Average Pressure Profile")
         plt.legend()
         plt.show()
+        self.saveImages("AveragePressureProfile")
 
     @staticmethod
     def plot_wind_rose(
@@ -1779,6 +1791,10 @@ def plot_average_day_wind_rose_specific_hour(self, hour, fig=None):
             fig=fig,
         )
         plt.show()
+        if self.saveImagesPng is True:
+            plt.savefig("WindRoseofanAverageDaySpecificHour", format="png")
+        if self.saveImagesPdf is True:
+            plt.savefig("WindRoseofanAverageDaySpecificHour", format="pdf")
 
     def plot_average_day_wind_rose_all_hours(self):
         """Plot windroses for all hours of a day, in a grid like plot."""
@@ -2545,3 +2561,13 @@ def allInfo(self):
         print(
             f"Percentage of Days Without Clouds: {100*self.percentage_of_days_with_no_cloud_coverage:.1f} %"
         )
+
+    def saveImages(self, name):
+        if isinstance(self.saveImagesPng, str):
+            plt.savefig(os.path.join(self.saveImagesPng, name+".png"))
+        elif self.saveImagesPng is True:
+            plt.savefig(name+".png")
+        elif isinstance(self.saveImagesPdf, str):
+            plt.savefig(os.path.join(self.saveImagesPdf, name+".pdf"))
+        elif self.saveImagesPdf is True:
+            plt.savefig(name+".pdf")

rocketpy/Flight.py

@@ -6,8 +6,10 @@
 __copyright__ = "Copyright 20XX, RocketPy Team"
 __license__ = "MIT"
 
+import os
 import math
 import time
+from tokenize import String
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -516,6 +518,8 @@ def __init__(
         heading=90,
         initialSolution=None,
         terminateOnApogee=False,
+        saveImagesPng=False,
+        saveImagesPdf=False,
         maxTime=600,
         maxTimeStep=np.inf,
         minTimeStep=0,
@@ -603,7 +607,8 @@ def __init__(
         self.initialSolution = initialSolution
         self.timeOvershoot = timeOvershoot
         self.terminateOnApogee = terminateOnApogee
-
+        self.saveImagesPng = saveImagesPng
+        self.saveImagesPdf = saveImagesPdf
         # Modifying Rail Length for a better out of rail condition
         upperRButton = max(self.rocket.railButtons[0])
         lowerRButton = min(self.rocket.railButtons[0])
@@ -2578,6 +2583,7 @@ def plot3dTrajectory(self):
         ax1.set_xlim3d([minXY, maxXY])
         ax1.view_init(15, 45)
         plt.show()
+        self.saveImages("Flight Trajectory")
 
         return None
 
@@ -2657,6 +2663,7 @@ def plotLinearKinematicsData(self):
 
         plt.subplots_adjust(hspace=0.5)
         plt.show()
+        self.saveImages("LinearKinematicsData")
         return None
 
     def plotAttitudeData(self):
@@ -2723,6 +2730,8 @@ def plotAttitudeData(self):
 
         plt.subplots_adjust(hspace=0.5)
         plt.show()
+        self.saveImages("AttitudeData")
+
 
         return None
 
@@ -2778,7 +2787,7 @@ def plotFlightPathAngleData(self):
 
         plt.subplots_adjust(hspace=0.5)
         plt.show()
-
+        self.saveImages("FlightPathAngleData")
         return None
 
     def plotAngularKinematicsData(self):
@@ -2863,6 +2872,7 @@ def plotAngularKinematicsData(self):
 
         plt.subplots_adjust(hspace=0.5)
         plt.show()
+        self.saveImages("AngularKinematicsData")
 
         return None
 
@@ -2936,6 +2946,10 @@ def plotTrajectoryForceData(self):
 
         plt.subplots_adjust(hspace=0.5)
         plt.show()
+        if self.saveImagesPng is True:
+            plt.savefig("TrajectoryForceData.png")
+        if self.saveImagesPdf is True:
+            plt.savefig("TrajectoryForceData.pdf")
 
         # Aerodynamic force and moment plots
         fig7 = plt.figure(figsize=(9, 12))
@@ -2994,6 +3008,7 @@ def plotTrajectoryForceData(self):
 
         plt.subplots_adjust(hspace=0.5)
         plt.show()
+        self.saveImages("TrajectoryForceData")
 
         return None
 
@@ -3087,6 +3102,7 @@ def plotEnergyData(self):
 
         plt.subplots_adjust(hspace=1)
         plt.show()
+        self.saveImages("EnergyData")
 
         return None
 
@@ -3163,7 +3179,7 @@ def plotFluidMechanicsData(self):
 
         plt.subplots_adjust(hspace=0.5)
         plt.show()
-
+        self.saveImages("FluidMechanicsData")
         return None
 
     def calculateFinFlutterAnalysis(self, finThickness, shearModulus):
@@ -3387,6 +3403,7 @@ def plotStabilityAndControlData(self):
 
         plt.subplots_adjust(hspace=0.5)
         plt.show()
+        self.saveImages("StabilityAndControlData")
 
         return None
 
@@ -3424,6 +3441,7 @@ def plotPressureSignals(self):
             ax1.grid()
 
             plt.show()
+            self.saveImages("PressureSignals")
 
         else:
             for parachute in self.rocket.parachutes:
@@ -3787,6 +3805,16 @@ def timeIterator(self, nodeList):
             yield i, nodeList[i]
             i += 1
 
+    def saveImages(self, name):
+        if isinstance(self.saveImagesPng, str):
+            plt.savefig(os.path.join(self.saveImagesPng, name+".png"))
+        elif self.saveImagesPng is True:
+            plt.savefig(name+".png")
+        elif isinstance(self.saveImagesPdf, str):
+            plt.savefig(os.path.join(self.saveImagesPdf, name+".pdf"))
+        elif self.saveImagesPdf is True:
+            plt.savefig(name+".pdf")
+
     class FlightPhases:
         def __init__(self, init_list=[]):
             self.list = init_list[:]