diff --git a/rocketpy/Environment.py b/rocketpy/Environment.py index 0292b0e1c..96f05bdb9 100644 --- a/rocketpy/Environment.py +++ b/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") \ No newline at end of file diff --git a/rocketpy/EnvironmentAnalysis.py b/rocketpy/EnvironmentAnalysis.py index 29ad0565c..eb812d293 100644 --- a/rocketpy/EnvironmentAnalysis.py +++ b/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") \ No newline at end of file diff --git a/rocketpy/Flight.py b/rocketpy/Flight.py index 50ac5d97d..22f516fa0 100644 --- a/rocketpy/Flight.py +++ b/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[:]