Merge pull request #296 from RocketPy-Team/enh/environment_plots

ENH: Moved Environment plots to environment_plots.py

Author: MateusStano

rocketpy/Environment.py

@@ -18,8 +18,10 @@
 import pytz
 import requests
 
+from .plots.environment_plots import _EnvironmentPlots
 from .prints.environment_prints import _EnvironmentPrints
 
+
 try:
     import netCDF4
 except ImportError:
@@ -401,6 +403,9 @@ def __init__(
         # Recalculate Earth Radius
         self.earthRadius = self.calculateEarthRadius(self.lat, self.datum)  # in m
 
+        # Initialize plots and prints object
+        self.plots = _EnvironmentPlots(self)
+
         return None
 
     def setDate(self, date, timeZone="UTC"):
@@ -2865,53 +2870,7 @@ def info(self):
 
         # Plot graphs
         print("\n\nAtmospheric Model Plots")
-        # Create height grid
-        grid = np.linspace(self.elevation, self.maxExpectedHeight)
-
-        # Create figure
-        plt.figure(figsize=(9, 4.5))
-
-        # Create wind speed and wind direction subplot
-        ax1 = plt.subplot(121)
-        ax1.plot(
-            [self.windSpeed(i) for i in grid], grid, "#ff7f0e", label="Speed of Sound"
-        )
-        ax1.set_xlabel("Wind Speed (m/s)", color="#ff7f0e")
-        ax1.tick_params("x", colors="#ff7f0e")
-        ax1up = ax1.twiny()
-        ax1up.plot(
-            [self.windDirection(i) for i in grid],
-            grid,
-            color="#1f77b4",
-            label="Density",
-        )
-        ax1up.set_xlabel("Wind Direction (°)", color="#1f77b4")
-        ax1up.tick_params("x", colors="#1f77b4")
-        ax1up.set_xlim(0, 360)
-        ax1.set_ylabel("Height Above Sea Level (m)")
-        ax1.grid(True)
-
-        # Create density and speed of sound subplot
-        ax2 = plt.subplot(122)
-        ax2.plot(
-            [self.speedOfSound(i) for i in grid],
-            grid,
-            "#ff7f0e",
-            label="Speed of Sound",
-        )
-        ax2.set_xlabel("Speed of Sound (m/s)", color="#ff7f0e")
-        ax2.tick_params("x", colors="#ff7f0e")
-        ax2up = ax2.twiny()
-        ax2up.plot(
-            [self.density(i) for i in grid], grid, color="#1f77b4", label="Density"
-        )
-        ax2up.set_xlabel("Density (kg/m³)", color="#1f77b4")
-        ax2up.tick_params("x", colors="#1f77b4")
-        ax2.set_ylabel("Height Above Sea Level (m)")
-        ax2.grid(True)
-
-        plt.subplots_adjust(wspace=0.5)
-        plt.show()
+        self.plots.atmospheric_model()
 
     def allInfo(self):
         """Prints out all data and graphs available about the Environment.
@@ -2928,165 +2887,7 @@ def allInfo(self):
         self.prints.all()
 
         # Plot graphs
-        print("\n\nAtmospheric Model Plots")
-        # Create height grid
-        grid = np.linspace(self.elevation, self.maxExpectedHeight)
-
-        # Create figure
-        plt.figure(figsize=(9, 9))
-
-        # Create wind speed and wind direction subplot
-        ax1 = plt.subplot(221)
-        ax1.plot(
-            [self.windSpeed(i) for i in grid], grid, "#ff7f0e", label="Speed of Sound"
-        )
-        ax1.set_xlabel("Wind Speed (m/s)", color="#ff7f0e")
-        ax1.tick_params("x", colors="#ff7f0e")
-        ax1up = ax1.twiny()
-        ax1up.plot(
-            [self.windDirection(i) for i in grid],
-            grid,
-            color="#1f77b4",
-            label="Density",
-        )
-        ax1up.set_xlabel("Wind Direction (°)", color="#1f77b4")
-        ax1up.tick_params("x", colors="#1f77b4")
-        ax1up.set_xlim(0, 360)
-        ax1.set_ylabel("Height Above Sea Level (m)")
-        ax1.grid(True)
-
-        # Create density and speed of sound subplot
-        ax2 = plt.subplot(222)
-        ax2.plot(
-            [self.speedOfSound(i) for i in grid],
-            grid,
-            "#ff7f0e",
-            label="Speed of Sound",
-        )
-        ax2.set_xlabel("Speed of Sound (m/s)", color="#ff7f0e")
-        ax2.tick_params("x", colors="#ff7f0e")
-        ax2up = ax2.twiny()
-        ax2up.plot(
-            [self.density(i) for i in grid], grid, color="#1f77b4", label="Density"
-        )
-        ax2up.set_xlabel("Density (kg/m³)", color="#1f77b4")
-        ax2up.tick_params("x", colors="#1f77b4")
-        ax2.set_ylabel("Height Above Sea Level (m)")
-        ax2.grid(True)
-
-        # Create wind u and wind v subplot
-        ax3 = plt.subplot(223)
-        ax3.plot([self.windVelocityX(i) for i in grid], grid, label="Wind U")
-        ax3.plot([self.windVelocityY(i) for i in grid], grid, label="Wind V")
-        ax3.legend(loc="best").set_draggable(True)
-        ax3.set_ylabel("Height Above Sea Level (m)")
-        ax3.set_xlabel("Wind Speed (m/s)")
-        ax3.grid(True)
-
-        # Create pressure and temperature subplot
-        ax4 = plt.subplot(224)
-        ax4.plot(
-            [self.pressure(i) / 100 for i in grid], grid, "#ff7f0e", label="Pressure"
-        )
-        ax4.set_xlabel("Pressure (hPa)", color="#ff7f0e")
-        ax4.tick_params("x", colors="#ff7f0e")
-        ax4up = ax4.twiny()
-        ax4up.plot(
-            [self.temperature(i) for i in grid],
-            grid,
-            color="#1f77b4",
-            label="Temperature",
-        )
-        ax4up.set_xlabel("Temperature (K)", color="#1f77b4")
-        ax4up.tick_params("x", colors="#1f77b4")
-        ax4.set_ylabel("Height Above Sea Level (m)")
-        ax4.grid(True)
-
-        plt.subplots_adjust(wspace=0.5, hspace=0.3)
-        plt.show()
-
-        # Plot ensemble member comparison
-        if self.atmosphericModelType != "Ensemble":
-            return None
-
-        print("\n\nEnsemble Members Comparison")
-        currentMember = self.ensembleMember
-
-        # Create figure
-        plt.figure(figsize=(9, 13.5))
-
-        # Create wind u subplot
-        ax5 = plt.subplot(321)
-        for i in range(self.numEnsembleMembers):
-            self.selectEnsembleMember(i)
-            ax5.plot([self.windVelocityX(i) for i in grid], grid, label=i)
-        # ax5.legend(loc='best').set_draggable(True)
-        ax5.set_ylabel("Height Above Sea Level (m)")
-        ax5.set_xlabel("Wind Speed (m/s)")
-        ax5.set_title("Wind U - Ensemble Members")
-        ax5.grid(True)
-
-        # Create wind v subplot
-        ax6 = plt.subplot(322)
-        for i in range(self.numEnsembleMembers):
-            self.selectEnsembleMember(i)
-            ax6.plot([self.windVelocityY(i) for i in grid], grid, label=i)
-        # ax6.legend(loc='best').set_draggable(True)
-        ax6.set_ylabel("Height Above Sea Level (m)")
-        ax6.set_xlabel("Wind Speed (m/s)")
-        ax6.set_title("Wind V - Ensemble Members")
-        ax6.grid(True)
-
-        # Create wind speed subplot
-        ax7 = plt.subplot(323)
-        for i in range(self.numEnsembleMembers):
-            self.selectEnsembleMember(i)
-            ax7.plot([self.windSpeed(i) for i in grid], grid, label=i)
-        # ax7.legend(loc='best').set_draggable(True)
-        ax7.set_ylabel("Height Above Sea Level (m)")
-        ax7.set_xlabel("Wind Speed (m/s)")
-        ax7.set_title("Wind Speed Magnitude - Ensemble Members")
-        ax7.grid(True)
-
-        # Create wind direction subplot
-        ax8 = plt.subplot(324)
-        for i in range(self.numEnsembleMembers):
-            self.selectEnsembleMember(i)
-            ax8.plot([self.windDirection(i) for i in grid], grid, label=i)
-        # ax8.legend(loc='best').set_draggable(True)
-        ax8.set_ylabel("Height Above Sea Level (m)")
-        ax8.set_xlabel("Degrees True (°)")
-        ax8.set_title("Wind Direction - Ensemble Members")
-        ax8.grid(True)
-
-        # Create pressure subplot
-        ax9 = plt.subplot(325)
-        for i in range(self.numEnsembleMembers):
-            self.selectEnsembleMember(i)
-            ax9.plot([self.pressure(i) for i in grid], grid, label=i)
-        # ax9.legend(loc='best').set_draggable(True)
-        ax9.set_ylabel("Height Above Sea Level (m)")
-        ax9.set_xlabel("Pressure (P)")
-        ax9.set_title("Pressure - Ensemble Members")
-        ax9.grid(True)
-
-        # Create temperature subplot
-        ax10 = plt.subplot(326)
-        for i in range(self.numEnsembleMembers):
-            self.selectEnsembleMember(i)
-            ax10.plot([self.temperature(i) for i in grid], grid, label=i)
-        # ax10.legend(loc='best').set_draggable(True)
-        ax10.set_ylabel("Height Above Sea Level (m)")
-        ax10.set_xlabel("Temperature (K)")
-        ax10.set_title("Temperature - Ensemble Members")
-        ax10.grid(True)
-
-        # Display plot
-        plt.subplots_adjust(wspace=0.5, hspace=0.3)
-        plt.show()
-
-        # Clean up
-        self.selectEnsembleMember(currentMember)
+        self.plots.all()
 
         return None