def plot_step_by_step(X_all, U_all, Z_all, B_all, EB_all, model):
fig, axes = plt.subplots(1, 2, figsize=(20, 8))
fig.tight_layout()
xlim = (-5, 35)
ylim = (-5, 30)
Plotter.plot_mpc(
fig, axes, xlim, ylim, model, X_all, Z_all, B_all, EB_all
)
# Run MPC
X_all, U_all, Z_all, B_all, EB_all = Simulator.mpc(model, model_p)
# Cast simulation results for ease of use
x_all = model.x.repeated(X_all)
u_all = model.u.repeated(U_all)
z_all = model.z.repeated(Z_all)
b_all = model.b.repeated(B_all)
# ---------------------- Step-by-step plotting ----------------------------- #
fig, axes = plt.subplots(1, 2, figsize=(20, 8))
fig.tight_layout()
xlim = (-5, 35)
ylim = (-5, 20)
Plotter.plot_mpc(fig, axes, xlim, ylim,
model, X_all, Z_all, B_all, EB_all)
# -------------------------- Plot full simulation -------------------------- #
# Plot 2D
reload(plotter)
from plotter import Plotter
fig, ax = plt.subplots()
handles = Plotter.plot_trajectory(ax, x_all)
handles.extend(Plotter.plot_observed_ball_trajectory(ax, z_all))
handles.extend(Plotter.plot_filtered_trajectory(ax, b_all))
ax.legend(handles=handles, fontsize=9,
bbox_to_anchor=(0.92, 0.59), labelspacing=0.2,
handler_map={mpatches.Patch: AnyObjectHandler()})
ax.set_aspect('equal')