SR-ARPOD/stellar/benchmark/plot_all_state.py
2026-04-01 22:48:53 +08:00

112 lines
2.9 KiB
Python

import os
import matplotlib.pyplot as plt
#read from run folder
def graph100(runs_dir, graph_all=True):
print(runs_dir)
if graph_all != True:
opt = str(input('Which component to plot in isolation [x] [y] [z]'))
for filename in os.listdir(runs_dir):
if filename.endswith('.txt'):
f = open(f'{runs_dir}/{filename}', "r")
pos_data = f.readlines()
xs = []
ys = []
zs = []
for idx, pos in enumerate(pos_data):
pos_split = pos.split(',')
if idx > 0:
x = float(pos_split[0])
y = float(pos_split[1])
z = float(pos_split[2])
xs.append(x)
ys.append(y)
zs.append(z)
f.close()
# plt.plot(xs, label='x')
# plt.plot(ys, label='y')
# plt.plot(zs, label='z')
if graph_all:
l1, = plt.plot(xs, color = 'red')
l2, = plt.plot(ys, color = 'green')
l3, = plt.plot(zs, color = 'blue')
else:
if opt == 'x':
l1, = plt.plot(xs, color = 'red')
elif opt =='y':
l2, = plt.plot(ys, color = 'green')
elif opt == 'z':
l3, = plt.plot(zs, color = 'blue')
else:
print(f'OPTION {opt} is invalid')
raise Exception
plt.xlim(0, 450)
if graph_all:
lhandle = [l1, l2, l3]
llabel = ['x', 'y', 'z']
else:
if opt == 'x':
lhandle = [l1]
if runs_dir == 'velocities/vbar0':
llabel = ['xdot']
plt.ylim(-5, 5)
else:
llabel = ['x']
elif opt =='y':
lhandle = [l2]
if runs_dir == 'velocities/vbar0':
llabel = ['ydot']
plt.ylim(-5, 5)
else:
llabel = ['y']
elif opt == 'z':
lhandle = [l3]
if runs_dir == 'velocities/vbar0':
llabel = ['zdot']
plt.ylim(-5, 5)
else:
llabel = ['z']
else:
print(f'OPTION {opt} is invalid')
raise Exception
plt.grid(visible=True)
plt.legend(handles=lhandle, labels=llabel)
#leg = plt.legend(loc='upper right')
plt.xlabel("Time [s]", fontsize=15)
plt.ylabel("Position [m]", fontsize=15)
plt.savefig('100pos.eps', format='eps')
plt.show()
usr_vel = str(input('Graph velocity? [Y] [N]'))
if usr_vel == 'Y':
data_dir = 'velocities/vbar0'
else:
data_dir = 'runs/vbar0'
usr_input = str(input('Graph all components? [Y] [N]'))
if usr_input == 'Y':
graph100(data_dir)
elif usr_input == 'N':
graph100(data_dir, False)
else:
print(f'INVALID OPTION {usr_input}')
raise Exception