def compute_measure_postep(data, trajecs, S_M=None):
if S_M is None:
S_M = data['S_M']
# apex state after step
XN = [traj.y[:, -1] for traj in trajecs]
# state in low-dim state-space
SN = [sys.xp2s_y_xdot(xn, data['p']) for xn in XN]
# digitalize, to bin index
SN_dig = [vibly.digitize_s(sn, data['grids']['states']) for sn in SN]
# measure of each point
SNM = [interp_measure(sbin, S_M, data['grids']) for sbin in SN_dig]
return SNM
def plot_ground_perturbations(ax,
trajectories,
S_M,
grids,
p,
v_threshold=0.1,
col_offset=0.65,
col_norm=1.0,
draw_ground=True,
draw_LC=False,
Q_M=None,
colormap=None,
norm=1):
'''
Plot a series of trajectories, centered around level-ground as nominal
inputs:
trajectories: list of traj objects (sol to scipy.integrate.solve_ivp)
v_threshold: minimum safety, otherwise don't plot it
'''
# TODO redo this with trajectories colored by measure
print(" ")
# * plot step-ups
up_cmap = plt.get_cmap("Blues")
down_cmap = plt.get_cmap("Reds")
idx0 = -1
pert_min = 0
pert_max = 0
lc_viab = 0
for idx in range(len(trajectories)):
traj = trajectories[idx]
x = traj.y[:, -1] # ground
s = sys.xp2s_y_xdot(x, p)
sbin = vibly.digitize_s(s, grids['states'])
s_m = interp_measure(sbin, S_M, grids)
if s_m >= v_threshold:
if np.isclose(x[-1], 0): # limit cycle
# print afterwards, so it's on top of other circles
idx0 = idx
# keep track of max viability at LC
lc_viab = s_m
continue
elif x[-1] < 0:
col = down_cmap(col_offset - np.abs(x[-1]) / col_norm)
# col = down_cmap(np.abs(x[-1])/col_norm)
zod = 2
# keep track of min and max perturbations
if x[-1] < pert_min:
pert_min = x[-1]
elif x[-1] > 0:
col = up_cmap(col_offset - np.abs(x[-1]) / col_norm)
# col = up_cmap(np.abs(x[-1])/col_norm)
zod = 1
if x[-1] > pert_max:
pert_max = x[-1]
ax.plot(traj.y[0], traj.y[1], color=col, linewidth=2, zorder=zod)
# and plot ground
td_index = np.abs(traj.t - traj.t_events[1]).argmin()
to_index = np.abs(traj.t - traj.t_events[3]).argmin()
ax.plot(traj.y[0, td_index:to_index],
traj.y[-1, td_index:to_index],
color=col,
linewidth=2)
if idx0 > 0:
zod = 3
traj = trajectories[idx0]
ax.plot(traj.y[0],
traj.y[1],
linewidth=2.5,
color='black',
zorder=zod)
ax.plot(traj.y[0], traj.y[-1], color='black', zorder=zod)
print("Min perturbation: " + str(pert_min))
print("Max perturbation: " + str(pert_max))
print("Viability Measure at Limit Cycle: " + str(lc_viab))
if draw_LC:
if idx0 <= 0:
print("WARNING: no LC fond.")
x_next = trajectories[idx0].y[:, -1]
# plot pointmass
ax.scatter(x_next[0], x_next[1], s=300, color='black', zorder=3)
# plot leg
ax.plot([x_next[0], x_next[4]], [x_next[1], x_next[5]],
linewidth=3,
color='black',
zorder=3)
if Q_M is not None:
# pick out the correct slice
s_next = sys.xp2s_y_xdot(x_next, p)
sbin = vibly.digitize_s(s_next, grids['states'])
A_slice = np.copy(Q_M[tuple(sbin) + (slice(None), )])
viable_actions = grids['actions'][0][np.nonzero(A_slice)]
viable_action_M = A_slice[np.nonzero(A_slice)]
# color = plt.cm.hsv(viable_action_M)
# create an array of foot-points
foot_x = np.zeros_like(viable_actions)
foot_y = np.zeros_like(viable_actions)
for i, a in np.ndenumerate(viable_actions):
p['angle_of_attack'] = a
xtemp = sys.reset_leg(x_next, p)
foot_x[i] = xtemp[4]
foot_y[i] = xtemp[5]
cmap = plt.get_cmap("viridis")
color = cmap(norm(interp_measure(sbin, S_M, grids)))
ax.plot(foot_x, foot_y, zorder=2, color=color)
ax.plot([x_next[0], foot_x[0]], [x_next[1], foot_y[0]],
zorder=2,
color=color)
ax.plot([x_next[0], foot_x[-1]], [x_next[1], foot_y[-1]],
zorder=2,
color=color)
ax.grid = True
add_title(p)
def poincare_plot(fig,
ax,
data,
vmax=1,
trajectories=None,
min_M=0.0,
col_offset=0.65,
col_norm=1.0):
grids = data['grids']
extent = [
grids['states'][1][0], grids['states'][1][-1], grids['states'][0][0],
grids['states'][0][-1]
]
# vmax = get_max_measure((data['S_M'] for data in data_list))
ax.imshow(data['S_M'],
origin='lower',
extent=extent,
aspect='auto',
interpolation='none',
vmin=0,
vmax=vmax,
cmap='viridis')
add_title(data['p'])
if trajectories is not None:
# X0 = [traj.y[:, 0] for traj in trajectories]
# XN = [traj.y[:, -1] for traj in trajectories]
# SN = [sys.xp2s_y_xdot(xn, data['p']) for xn in XN]
# s_grid_shape = list(map(np.size, grids['states']))
# # s_bin_shape = tuple(dim+1 for dim in s_grid_shape)
# SN_dig = [vibly.digitize_s(sn, grids['states']) for sn in SN]
# SNM = np.array([interp_measure(sbin, data['S_M'], grids)
# for sbin in SN_dig])
# ground_heights = [traj.y[-1, 0] for traj in trajectories]
# indices = np.arange(SNM.size)
# SNp = np.array(SN).T
# * get index of nominal trajectory, max step up/down
index0, max_up, max_down = get_perturbation_indices(trajectories)
# index0 = -1
# up_indices = list()
# down_indices = list()
# for idx, traj in enumerate(trajectories):
# * plot each step at next step
up_cmap = plt.get_cmap("Blues")
down_cmap = plt.get_cmap("Reds")
idx0 = -1
for idx in range(len(trajectories)):
traj = trajectories[idx]
xn = traj.y[:, -1] # state at next step
sn = sys.xp2s_y_xdot(xn, data['p'])
sbin = vibly.digitize_s(sn, grids['states'])
s_m = interp_measure(sbin, data['S_M'], grids)
if s_m > min_M:
if np.isclose(xn[-1], 0):
# print afterwards, so it's on top of other circles
sn0 = sn.copy()
idx0 = 1
continue
elif xn[-1] < 0:
col = down_cmap(col_offset - np.abs(xn[-1]) / col_norm)
elif xn[-1] > 0:
col = up_cmap(col_offset - np.abs(xn[-1]) / col_norm)
ax.scatter(sn[1],
sn[0],
facecolors='none',
edgecolors=col,
s=20)
if idx0 > 0:
ax.scatter(sn0[1],
sn0[0],
facecolors='none',
edgecolors='black',
s=20)
data[idx]['grids']['states'][0][0],
data[idx]['grids']['states'][0][-1]]
plt.figure(idx)
plt.imshow(data[idx]['S_M'], origin='lower', extent=extent,
interpolation='bessel', vmin=0, vmax=vmax, cmap='viridis')
# sns.heatmap(data[idx]['S_M'],
# vmin=0, vmax=vmax, cmap='viridis')
plt.title(str(np.round(data[idx]['p']['linear_normalized_damping_coefficient'], decimals=2)))
X0 = [traj.y[:, 0] for traj in data[idx]['trajectories']]
XN = [traj.y[:, -1] for traj in data[idx]['trajectories']]
SN = [sys.xp2s_y_xdot(xn, data[idx]['p']) for xn in XN]
s_grid_shape = list(map(np.size, data[idx]['grids']['states']))
s_bin_shape = tuple(dim+1 for dim in s_grid_shape)
SN_dig = [vibly.digitize_s(sn, data[idx]['grids']['states']) for sn in SN]
SNM = np.array([interp_measure(sbin, data[idx]['S_M'], data[idx]['grids']) for sbin in SN_dig])
# pick out only the ones inside the grid
ground_heights = [x[-1] for x in X0]
indices = np.arange(SNM.size)
SNp = np.array(SN).T
plt.scatter(SNp[1, indices[SNM>0.1]], SNp[0, indices[SNM>0.1]],
facecolors='none', edgecolors=[0.8, 0.3, 0.3], s=20)
plt.show()
# * make a collage of 4
# todo: color-code with same coloring as trajectories plot
# todo: make trajectories plot
if FLAG_SM4:
S_M = data_list[0]["S_M"]
# * Ground Truth XV
XV = S_M > 0.0
mynorm = colors.TwoSlopeNorm(vmin=-100, vcenter=0.0, vmax=1)
# mynorm = colors.CenteredNorm()
mymap = vplot.get_vmap(0)
extent = [
grids['states'][1][0], grids['states'][1][-1], grids['states'][0][0],
grids['states'][0][-1]
]
# pick out value of x[0] to slice through
x0 = vibly.digitize_s([12., 0], grids['states'], to_bin=False)
x1_slice = grids['states'][1]
XV_slice = XV[x0[0], :]
these_values = [0, 3, 5] # to plot entire value functions
num_plots = len(these_values) + 1
fig, axs = plt.subplots(num_plots, 1)
# * color settings for slice
p_cmap = plt.get_cmap("Dark2")
col_norm = 20.
col_offset = 1.
v_min = 0.
v_max = 0.
found = False