Merge pull request #326 from RocketPy-Team/maint/extract-logic-prints-env

MAINT: move last plots from Environment class

Author: Gui-FernandesBR

rocketpy/Environment.py

@@ -12,7 +12,6 @@
 import warnings
 from datetime import datetime, timedelta
 
-import matplotlib.pyplot as plt
 import numpy as np
 import numpy.ma as ma
 import pytz
@@ -3022,58 +3021,8 @@ def info(self):
         None
         """
 
-        # All prints
         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, 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.info()
         return None
 
     def allInfo(self):

rocketpy/Parachute.py

@@ -141,13 +141,11 @@ def __str__(self):
         )
 
     def info(self):
-
         self.prints.all()
 
         return None
 
     def allInfo(self):
-
         self.info()
         # self.plots.all() # Parachutes still doesn't have plots
 

rocketpy/plots/environment_plots.py

@@ -1,4 +1,4 @@
-__author__ = "Mateus Stano Junqueira"
+__author__ = "Mateus Stano Junqueira, Guilherme Fernandes Alves"
 __copyright__ = "Copyright 20XX, RocketPy Team"
 __license__ = "MIT"
 
@@ -39,112 +39,184 @@ def __init__(self, environment):
 
         return None
 
-    def atmospheric_model(self):
-        """Plots all atmospheric model graphs available
+    def __wind(self, ax):
+        """Adds wind speed and wind direction graphs to the same axis.
 
         Parameters
         ----------
-        None
+        ax : matplotlib.pyplot.axis
+            Axis to add the graphs.
 
-        Return
-        ------
-        None
+        Returns
+        -------
+        ax : matplotlib.pyplot.axis
+            Axis with the graphs.
         """
-
-        # Create figure
-        plt.figure(figsize=(9, 9))
-
-        # Create wind speed and wind direction subplot
-        ax1 = plt.subplot(221)
-        ax1.plot(
+        ax.plot(
             [self.environment.windSpeed(i) for i in self.grid],
             self.grid,
             "#ff7f0e",
-            label="Speed of Sound",
+            label="Wind Speed",
         )
-        ax1.set_xlabel("Wind Speed (m/s)", color="#ff7f0e")
-        ax1.tick_params("x", colors="#ff7f0e")
-        ax1up = ax1.twiny()
-        ax1up.plot(
+        ax.set_xlabel("Wind Speed (m/s)", color="#ff7f0e")
+        ax.tick_params("x", colors="#ff7f0e")
+        axup = ax.twiny()
+        axup.plot(
             [self.environment.windDirection(i) for i in self.grid],
             self.grid,
             color="#1f77b4",
-            label="Density",
+            label="Wind Direction",
         )
-        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)
+        axup.set_xlabel("Wind Direction (°)", color="#1f77b4")
+        axup.tick_params("x", colors="#1f77b4")
+        axup.set_xlim(0, 360)
+        ax.set_ylabel("Height Above Sea Level (m)")
+        ax.grid(True)
 
-        # Create density and speed of sound subplot
-        ax2 = plt.subplot(222)
-        ax2.plot(
+        return ax
+
+    def __density_speed_of_sound(self, ax):
+        """Adds density and speed of sound graphs to the same axis.
+
+        Parameters
+        ----------
+        ax : matplotlib.pyplot.axis
+            Axis to add the graphs.
+
+        Returns
+        -------
+        ax : matplotlib.pyplot.axis
+            Axis with the graphs.
+        """
+        ax.plot(
             [self.environment.speedOfSound(i) for i in self.grid],
             self.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(
+        ax.set_xlabel("Speed of Sound (m/s)", color="#ff7f0e")
+        ax.tick_params("x", colors="#ff7f0e")
+        axup = ax.twiny()
+        axup.plot(
             [self.environment.density(i) for i in self.grid],
             self.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)
+        axup.set_xlabel("Density (kg/m³)", color="#1f77b4")
+        axup.tick_params("x", colors="#1f77b4")
+        ax.set_ylabel("Height Above Sea Level (m)")
+        ax.grid(True)
 
-        # Create wind u and wind v subplot
-        ax3 = plt.subplot(223)
-        ax3.plot(
+        return ax
+
+    def __wind_components(self, ax):
+        """Adds wind u and wind v graphs to the same axis.
+
+        Parameters
+        ----------
+        ax : matplotlib.pyplot.axis
+            Axis to add the graphs.
+
+        Returns
+        -------
+        ax : matplotlib.pyplot.axis
+            Axis with the graphs.
+        """
+        ax.plot(
             [self.environment.windVelocityX(i) for i in self.grid],
             self.grid,
             label="Wind U",
         )
-        ax3.plot(
+        ax.plot(
             [self.environment.windVelocityY(i) for i in self.grid],
             self.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)
+        ax.legend(loc="best").set_draggable(True)
+        ax.set_ylabel("Height Above Sea Level (m)")
+        ax.set_xlabel("Wind Speed (m/s)")
+        ax.grid(True)
 
-        # Create pressure and temperature subplot
-        ax4 = plt.subplot(224)
-        ax4.plot(
+        return ax
+
+    def __pressure_temperature(self, ax):
+        """Adds pressure and temperature graphs to the same axis.
+
+        Parameters
+        ----------
+        ax : matplotlib.pyplot.axis
+            Axis to add the graphs.
+
+        Returns
+        -------
+        ax : matplotlib.pyplot.axis
+            Axis with the graphs.
+        """
+        ax.plot(
             [self.environment.pressure(i) / 100 for i in self.grid],
             self.grid,
             "#ff7f0e",
             label="Pressure",
         )
-        ax4.set_xlabel("Pressure (hPa)", color="#ff7f0e")
-        ax4.tick_params("x", colors="#ff7f0e")
-        ax4up = ax4.twiny()
-        ax4up.plot(
+        ax.set_xlabel("Pressure (hPa)", color="#ff7f0e")
+        ax.tick_params("x", colors="#ff7f0e")
+        axup = ax.twiny()
+        axup.plot(
             [self.environment.temperature(i) for i in self.grid],
             self.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)
+        axup.set_xlabel("Temperature (K)", color="#1f77b4")
+        axup.tick_params("x", colors="#1f77b4")
+        ax.set_ylabel("Height Above Sea Level (m)")
+        ax.grid(True)
+
+        return ax
+
+    def atmospheric_model(self):
+        """Plots all atmospheric model graphs available. This includes wind speed
+        and wind direction, density and speed of sound, wind u and wind v, and
+        pressure and temperature.
+
+        Parameters
+        ----------
+        None
+
+        Return
+        ------
+        None
+        """
+
+        # Create figure
+        plt.figure(figsize=(9, 9))
+
+        # Create wind speed and wind direction subplot
+        ax1 = plt.subplot(221)
+        ax1 = self.__wind(ax1)
+
+        # Create density and speed of sound subplot
+        ax2 = plt.subplot(222)
+        ax2 = self.__density_speed_of_sound(ax2)
+
+        # Create wind u and wind v subplot
+        ax3 = plt.subplot(223)
+        ax3 = self.__wind_components(ax3)
+        ax3.legend(loc="best").set_draggable(True)
+
+        # Create pressure and temperature subplot
+        ax4 = plt.subplot(224)
+        ax4 = self.__pressure_temperature(ax4)
 
         plt.subplots_adjust(wspace=0.5, hspace=0.3)
         plt.show()
 
         return None
 
     def ensemble_member_comparison(self):
-        """Plots ensemble member comparisons.
+        """Plots ensemble member comparisons. It requires that the environment
+        model has been set as Ensemble.
 
         Parameters
         ----------
@@ -253,8 +325,33 @@ def ensemble_member_comparison(self):
 
         return None
 
+    def info(self):
+        """Plots a summary of the atmospheric model, including wind speed and
+        wind direction, density and speed of sound. This is important for the
+        Environment.info() method.
+
+        Returns
+        -------
+        None
+        """
+        print("\nAtmospheric Model Plots\n")
+        plt.figure(figsize=(9, 4.5))
+        # Create wind speed and wind direction subplot
+        ax1 = plt.subplot(121)
+        ax1 = self.__wind(ax1)
+
+        # Create density and speed of sound subplot
+        ax2 = plt.subplot(122)
+        ax2 = self.__density_speed_of_sound(ax2)
+
+        plt.subplots_adjust(wspace=0.5)
+        plt.show()
+        return None
+
     def all(self):
-        """Prints out all graphs available about the Environment.
+        """Prints out all graphs available about the Environment. This includes
+        a complete description of the atmospheric model and the ensemble members
+        comparison if the atmospheric model is an ensemble.
 
         Parameters
         ----------