def plot_3d_traj_metric_trace_init_and_opt(fig, ax, metric, traj_init,
traj_opt):
metric_tensor_init, _, _ = gp_metric_tensor(
traj_init[:, 0:2],
metric.gp.Z,
metric.gp.kernel,
mean_func=metric.gp.mean_func,
f=metric.gp.q_mu,
full_cov=True,
q_sqrt=metric.gp.q_sqrt,
cov_weight=metric.cov_weight,
)
metric_tensor_opt, _, _ = gp_metric_tensor(
traj_opt[:, 0:2],
metric.gp.Z,
metric.gp.kernel,
mean_func=metric.gp.mean_func,
f=metric.gp.q_mu,
full_cov=True,
q_sqrt=metric.gp.q_sqrt,
cov_weight=metric.cov_weight,
)
metric_trace_init = np.trace(metric_tensor_init, axis1=1, axis2=2)
metric_trace_opt = np.trace(metric_tensor_opt, axis1=1, axis2=2)
plot_3d_traj_col(fig, ax, traj_init, zs=metric_trace_init, color="c")
plot_3d_traj_col(fig, ax, traj_opt, zs=metric_trace_opt, color="m")
return fig, ax
def plot_metirc_trace_over_time(metric, solver, traj_init, traj_opt):
metric_tensor_init, _, _ = gp_metric_tensor(
traj_init[:, 0:2],
metric.gp.Z,
metric.gp.kernel,
mean_func=metric.gp.mean_func,
f=metric.gp.q_mu,
full_cov=True,
q_sqrt=metric.gp.q_sqrt,
cov_weight=metric.cov_weight,
)
metric_trace_init = np.trace(metric_tensor_init, axis1=1, axis2=2)
metric_tensor_opt, _, _ = gp_metric_tensor(
traj_opt[:, 0:2],
metric.gp.Z,
metric.gp.kernel,
mean_func=metric.gp.mean_func,
f=metric.gp.q_mu,
full_cov=True,
q_sqrt=metric.gp.q_sqrt,
cov_weight=metric.cov_weight,
)
metric_trace_opt = np.trace(metric_tensor_opt, axis1=1, axis2=2)
fig, ax = plt.subplots(1, 1, figsize=(6.4, 2.8))
ax.set_xlabel("Time $t$")
ax.set_ylabel("Tr$(G(\mathbf{x}_t))$")
ax.plot(
solver.times,
metric_trace_init,
color=color_init,
label="Initial trajectory",
)
ax.plot(
solver.times,
metric_trace_opt,
color=color_opt,
label="Optimised trajectory",
)
ax.legend()
def plot_3d_traj_metric_trace(fig, ax, metric, traj):
metric_tensor, mu_j, cov_j = gp_metric_tensor(
traj[:, 0:2],
metric.gp.Z,
metric.gp.kernel,
mean_func=metric.gp.mean_func,
f=metric.gp.q_mu,
full_cov=True,
q_sqrt=metric.gp.q_sqrt,
cov_weight=metric.cov_weight,
)
metric_trace = np.trace(metric_tensor, axis1=1, axis2=2)
plot_3d_traj(fig, ax, traj, zs=metric_trace)
return fig, ax
def plot_svgp_metric_and_start_end(metric, solver, traj_opt=None):
params = {
"text.usetex": True,
"text.latex.preamble": [
"\\usepackage{amssymb}",
"\\usepackage{amsmath}",
],
}
plt.rcParams.update(params)
input_dim = metric.gp.X.shape[1]
# plot original GP
Xnew, xx, yy = create_grid(metric.gp.X, N=961)
metric_tensor, mu_j, cov_j = gp_metric_tensor(
Xnew,
metric.gp.Z,
metric.gp.kernel,
mean_func=metric.gp.mean_func,
f=metric.gp.q_mu,
full_cov=True,
q_sqrt=metric.gp.q_sqrt,
cov_weight=metric.cov_weight,
)
print("metric yo yo")
print(metric_tensor.shape)
fig, axs = plt.subplots(input_dim, input_dim)
for i in range(input_dim):
for j in range(input_dim):
plot_contourf(
fig,
axs[i, j],
xx,
yy,
metric_tensor[:, i, j],
label="$G(\mathbf{x})$",
)
axs[i, j].set_xlabel("$x$")
axs[i, j].set_ylabel("$y$")
return fig, axs
def plot_3d_metric_trace(metric, solver):
# plot original GP
Xnew, xx, yy = create_grid(metric.gp.X, N=961)
metric_tensor, mu_j, cov_j = gp_metric_tensor(
Xnew,
metric.gp.Z,
metric.gp.kernel,
mean_func=metric.gp.mean_func,
f=metric.gp.q_mu,
full_cov=True,
q_sqrt=metric.gp.q_sqrt,
cov_weight=metric.cov_weight,
)
metric_trace = np.trace(metric_tensor, axis1=1, axis2=2)
# fig = plt.figure(figsize=plt.figaspect(0.5))
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection="3d")
surf = plot_3d_surf(fig, ax, xx, yy, metric_trace)
ax.set_zlabel("Tr$(G(\mathbf{x}_*))$")
return fig, ax
def plot_svgp_metric_trace_and_start_end(metric, solver, traj_opt=None):
# plot original GP
Xnew, xx, yy = create_grid(metric.gp.X, N=961)
metric_tensor, mu_j, cov_j = gp_metric_tensor(
Xnew,
metric.gp.Z,
metric.gp.kernel,
mean_func=metric.gp.mean_func,
f=metric.gp.q_mu,
full_cov=True,
q_sqrt=metric.gp.q_sqrt,
cov_weight=metric.cov_weight,
)
metric_trace = np.trace(metric_tensor, axis1=1, axis2=2)
fig, ax = plt.subplots(1, 1)
# fig, ax = plt.subplots(1, 1, figsize=(12, 4))
plt.subplots_adjust(wspace=0, hspace=0)
surf_traceG = plot_contourf(
fig, ax, xx, yy, metric_trace, label="Tr$(G(\mathbf{x}))$"
)
plot_start_and_end_pos(fig, ax, solver)
ax.set_xlabel("$x$")
ax.set_ylabel("$y$")
plot_traj(
fig,
ax,
solver.state_guesses,
color=color_init,
label="Initial trajectory",
)
if traj_opt is not None:
plot_traj(
fig, ax, traj_opt, color=color_opt, label="Optimised trajectory"
)
ax.legend()
return fig, ax