Add more plotting to the planar pcs simulation example

Author: mstoelzle

examples/simulate_planar_pcs.py

@@ -43,7 +43,7 @@
 strain_selector = jnp.array([True, False, False])
 
 # define initial configuration
-q0 = jnp.array([10 * jnp.pi])
+q0 = jnp.array([5 * jnp.pi])
 # number of generalized coordinates
 n_q = q0.shape[0]
 
@@ -142,6 +142,27 @@ def draw_robot(
     # the evolution of the generalized velocities
     q_d_ts = sol.ys[:, n_q:]
 
+    # evaluate the forward kinematics along the trajectory
+    chi_ee_ts = vmap(forward_kinematics_fn, in_axes=(None, 0, None))(params, q_ts, jnp.array([jnp.sum(params["l"])]))
+    # plot the end-effector position along the trajectory
+    plt.figure()
+    plt.plot(chi_ee_ts[0, :], chi_ee_ts[1, :])
+    plt.axis("equal")
+    plt.grid(True)
+    plt.xlabel("x [m]")
+    plt.ylabel("y [m]")
+    plt.show()
+    # plot end-effector position vs time
+    plt.figure()
+    plt.plot(video_ts, chi_ee_ts[:, 0], label="x")
+    plt.plot(video_ts, chi_ee_ts[:, 1], label="y")
+    plt.xlabel("Time [s]")
+    plt.ylabel("Position [m]")
+    plt.legend()
+    plt.grid(True)
+    plt.box(True)
+    plt.show()
+
     # plot the energy along the trajectory
     kinetic_energy_fn_vmapped = vmap(partial(auxiliary_fns["kinetic_energy_fn"], params))
     potential_energy_fn_vmapped = vmap(partial(auxiliary_fns["potential_energy_fn"], params))