Skip to content

Commit 64eb8a1

Browse files
committed
enh: moved plots from Environment.py to environment_plots.py
1 parent bb5d7c5 commit 64eb8a1

File tree

2 files changed

+279
-209
lines changed

2 files changed

+279
-209
lines changed

rocketpy/Environment.py

Lines changed: 7 additions & 206 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
import pytz
1717
import requests
1818

19+
from .plots.environment_plots import _EnvironmentPlots
20+
1921
try:
2022
import netCDF4
2123
except ImportError:
@@ -394,6 +396,9 @@ def __init__(
394396
# Recalculate Earth Radius
395397
self.earthRadius = self.calculateEarthRadius(self.lat, self.datum) # in m
396398

399+
# Initialize plots and prints object
400+
self.plots = _EnvironmentPlots(self)
401+
397402
return None
398403

399404
def setDate(self, date, timeZone="UTC"):
@@ -2925,53 +2930,7 @@ def info(self):
29252930

29262931
# Plot graphs
29272932
print("\n\nAtmospheric Model Plots")
2928-
# Create height grid
2929-
grid = np.linspace(self.elevation, self.maxExpectedHeight)
2930-
2931-
# Create figure
2932-
plt.figure(figsize=(9, 4.5))
2933-
2934-
# Create wind speed and wind direction subplot
2935-
ax1 = plt.subplot(121)
2936-
ax1.plot(
2937-
[self.windSpeed(i) for i in grid], grid, "#ff7f0e", label="Speed of Sound"
2938-
)
2939-
ax1.set_xlabel("Wind Speed (m/s)", color="#ff7f0e")
2940-
ax1.tick_params("x", colors="#ff7f0e")
2941-
ax1up = ax1.twiny()
2942-
ax1up.plot(
2943-
[self.windDirection(i) for i in grid],
2944-
grid,
2945-
color="#1f77b4",
2946-
label="Density",
2947-
)
2948-
ax1up.set_xlabel("Wind Direction (°)", color="#1f77b4")
2949-
ax1up.tick_params("x", colors="#1f77b4")
2950-
ax1up.set_xlim(0, 360)
2951-
ax1.set_ylabel("Height Above Sea Level (m)")
2952-
ax1.grid(True)
2953-
2954-
# Create density and speed of sound subplot
2955-
ax2 = plt.subplot(122)
2956-
ax2.plot(
2957-
[self.speedOfSound(i) for i in grid],
2958-
grid,
2959-
"#ff7f0e",
2960-
label="Speed of Sound",
2961-
)
2962-
ax2.set_xlabel("Speed of Sound (m/s)", color="#ff7f0e")
2963-
ax2.tick_params("x", colors="#ff7f0e")
2964-
ax2up = ax2.twiny()
2965-
ax2up.plot(
2966-
[self.density(i) for i in grid], grid, color="#1f77b4", label="Density"
2967-
)
2968-
ax2up.set_xlabel("Density (kg/m³)", color="#1f77b4")
2969-
ax2up.tick_params("x", colors="#1f77b4")
2970-
ax2.set_ylabel("Height Above Sea Level (m)")
2971-
ax2.grid(True)
2972-
2973-
plt.subplots_adjust(wspace=0.5)
2974-
plt.show()
2933+
self.plots.atmospheric_model()
29752934

29762935
def allInfo(self):
29772936
"""Prints out all data and graphs available about the Environment.
@@ -3052,165 +3011,7 @@ def allInfo(self):
30523011
)
30533012

30543013
# Plot graphs
3055-
print("\n\nAtmospheric Model Plots")
3056-
# Create height grid
3057-
grid = np.linspace(self.elevation, self.maxExpectedHeight)
3058-
3059-
# Create figure
3060-
plt.figure(figsize=(9, 9))
3061-
3062-
# Create wind speed and wind direction subplot
3063-
ax1 = plt.subplot(221)
3064-
ax1.plot(
3065-
[self.windSpeed(i) for i in grid], grid, "#ff7f0e", label="Speed of Sound"
3066-
)
3067-
ax1.set_xlabel("Wind Speed (m/s)", color="#ff7f0e")
3068-
ax1.tick_params("x", colors="#ff7f0e")
3069-
ax1up = ax1.twiny()
3070-
ax1up.plot(
3071-
[self.windDirection(i) for i in grid],
3072-
grid,
3073-
color="#1f77b4",
3074-
label="Density",
3075-
)
3076-
ax1up.set_xlabel("Wind Direction (°)", color="#1f77b4")
3077-
ax1up.tick_params("x", colors="#1f77b4")
3078-
ax1up.set_xlim(0, 360)
3079-
ax1.set_ylabel("Height Above Sea Level (m)")
3080-
ax1.grid(True)
3081-
3082-
# Create density and speed of sound subplot
3083-
ax2 = plt.subplot(222)
3084-
ax2.plot(
3085-
[self.speedOfSound(i) for i in grid],
3086-
grid,
3087-
"#ff7f0e",
3088-
label="Speed of Sound",
3089-
)
3090-
ax2.set_xlabel("Speed of Sound (m/s)", color="#ff7f0e")
3091-
ax2.tick_params("x", colors="#ff7f0e")
3092-
ax2up = ax2.twiny()
3093-
ax2up.plot(
3094-
[self.density(i) for i in grid], grid, color="#1f77b4", label="Density"
3095-
)
3096-
ax2up.set_xlabel("Density (kg/m³)", color="#1f77b4")
3097-
ax2up.tick_params("x", colors="#1f77b4")
3098-
ax2.set_ylabel("Height Above Sea Level (m)")
3099-
ax2.grid(True)
3100-
3101-
# Create wind u and wind v subplot
3102-
ax3 = plt.subplot(223)
3103-
ax3.plot([self.windVelocityX(i) for i in grid], grid, label="Wind U")
3104-
ax3.plot([self.windVelocityY(i) for i in grid], grid, label="Wind V")
3105-
ax3.legend(loc="best").set_draggable(True)
3106-
ax3.set_ylabel("Height Above Sea Level (m)")
3107-
ax3.set_xlabel("Wind Speed (m/s)")
3108-
ax3.grid(True)
3109-
3110-
# Create pressure and temperature subplot
3111-
ax4 = plt.subplot(224)
3112-
ax4.plot(
3113-
[self.pressure(i) / 100 for i in grid], grid, "#ff7f0e", label="Pressure"
3114-
)
3115-
ax4.set_xlabel("Pressure (hPa)", color="#ff7f0e")
3116-
ax4.tick_params("x", colors="#ff7f0e")
3117-
ax4up = ax4.twiny()
3118-
ax4up.plot(
3119-
[self.temperature(i) for i in grid],
3120-
grid,
3121-
color="#1f77b4",
3122-
label="Temperature",
3123-
)
3124-
ax4up.set_xlabel("Temperature (K)", color="#1f77b4")
3125-
ax4up.tick_params("x", colors="#1f77b4")
3126-
ax4.set_ylabel("Height Above Sea Level (m)")
3127-
ax4.grid(True)
3128-
3129-
plt.subplots_adjust(wspace=0.5, hspace=0.3)
3130-
plt.show()
3131-
3132-
# Plot ensemble member comparison
3133-
if self.atmosphericModelType != "Ensemble":
3134-
return None
3135-
3136-
print("\n\nEnsemble Members Comparison")
3137-
currentMember = self.ensembleMember
3138-
3139-
# Create figure
3140-
plt.figure(figsize=(9, 13.5))
3141-
3142-
# Create wind u subplot
3143-
ax5 = plt.subplot(321)
3144-
for i in range(self.numEnsembleMembers):
3145-
self.selectEnsembleMember(i)
3146-
ax5.plot([self.windVelocityX(i) for i in grid], grid, label=i)
3147-
# ax5.legend(loc='best').set_draggable(True)
3148-
ax5.set_ylabel("Height Above Sea Level (m)")
3149-
ax5.set_xlabel("Wind Speed (m/s)")
3150-
ax5.set_title("Wind U - Ensemble Members")
3151-
ax5.grid(True)
3152-
3153-
# Create wind v subplot
3154-
ax6 = plt.subplot(322)
3155-
for i in range(self.numEnsembleMembers):
3156-
self.selectEnsembleMember(i)
3157-
ax6.plot([self.windVelocityY(i) for i in grid], grid, label=i)
3158-
# ax6.legend(loc='best').set_draggable(True)
3159-
ax6.set_ylabel("Height Above Sea Level (m)")
3160-
ax6.set_xlabel("Wind Speed (m/s)")
3161-
ax6.set_title("Wind V - Ensemble Members")
3162-
ax6.grid(True)
3163-
3164-
# Create wind speed subplot
3165-
ax7 = plt.subplot(323)
3166-
for i in range(self.numEnsembleMembers):
3167-
self.selectEnsembleMember(i)
3168-
ax7.plot([self.windSpeed(i) for i in grid], grid, label=i)
3169-
# ax7.legend(loc='best').set_draggable(True)
3170-
ax7.set_ylabel("Height Above Sea Level (m)")
3171-
ax7.set_xlabel("Wind Speed (m/s)")
3172-
ax7.set_title("Wind Speed Magnitude - Ensemble Members")
3173-
ax7.grid(True)
3174-
3175-
# Create wind direction subplot
3176-
ax8 = plt.subplot(324)
3177-
for i in range(self.numEnsembleMembers):
3178-
self.selectEnsembleMember(i)
3179-
ax8.plot([self.windDirection(i) for i in grid], grid, label=i)
3180-
# ax8.legend(loc='best').set_draggable(True)
3181-
ax8.set_ylabel("Height Above Sea Level (m)")
3182-
ax8.set_xlabel("Degrees True (°)")
3183-
ax8.set_title("Wind Direction - Ensemble Members")
3184-
ax8.grid(True)
3185-
3186-
# Create pressure subplot
3187-
ax9 = plt.subplot(325)
3188-
for i in range(self.numEnsembleMembers):
3189-
self.selectEnsembleMember(i)
3190-
ax9.plot([self.pressure(i) for i in grid], grid, label=i)
3191-
# ax9.legend(loc='best').set_draggable(True)
3192-
ax9.set_ylabel("Height Above Sea Level (m)")
3193-
ax9.set_xlabel("Pressure (P)")
3194-
ax9.set_title("Pressure - Ensemble Members")
3195-
ax9.grid(True)
3196-
3197-
# Create temperature subplot
3198-
ax10 = plt.subplot(326)
3199-
for i in range(self.numEnsembleMembers):
3200-
self.selectEnsembleMember(i)
3201-
ax10.plot([self.temperature(i) for i in grid], grid, label=i)
3202-
# ax10.legend(loc='best').set_draggable(True)
3203-
ax10.set_ylabel("Height Above Sea Level (m)")
3204-
ax10.set_xlabel("Temperature (K)")
3205-
ax10.set_title("Temperature - Ensemble Members")
3206-
ax10.grid(True)
3207-
3208-
# Display plot
3209-
plt.subplots_adjust(wspace=0.5, hspace=0.3)
3210-
plt.show()
3211-
3212-
# Clean up
3213-
self.selectEnsembleMember(currentMember)
3014+
self.plots.all()
32143015

32153016
return None
32163017

0 commit comments

Comments
 (0)