From 36d0786d10ea9c8854d746909437dada441825c2 Mon Sep 17 00:00:00 2001 From: Maxim Vochten Date: Thu, 14 Nov 2024 15:06:58 +0100 Subject: [PATCH] plot screw axis and objet trajectory together --- examples/calculate_invariants_pose.py | 198 ++++++++++++++++++++++++++ 1 file changed, 198 insertions(+) create mode 100644 examples/calculate_invariants_pose.py diff --git a/examples/calculate_invariants_pose.py b/examples/calculate_invariants_pose.py new file mode 100644 index 0000000..1d758fd --- /dev/null +++ b/examples/calculate_invariants_pose.py @@ -0,0 +1,198 @@ + +# In progress +# TODO: implement calculate twist from pose with discrete formulas + +from invariants_py.reparameterization import interpT +import invariants_py.kinematics.rigidbody_kinematics as SE3 +import invariants_py.kinematics.orientation_kinematics as S03 +import numpy as np +import matplotlib.pyplot as plt +from invariants_py.calculate_invariants.opti_calculate_screw_invariants_pose_fatrop import OCP_calc_pose + +def plot_pose_frames(T_obj, length=0.1, skip_frames=5): + """ + Plots the trajectory and pose frames in a 3D plot. + + Parameters: + - T_obj (numpy.ndarray): The pose data as a 3D array of shape (N, 4, 4). + - length (float, optional): The length of the quiver arrows representing the pose frames. Default is 0.075. + - skip_frames (int, optional): The number of frames to skip when plotting the pose frames. Default is 5. + + Returns: + None + """ + # Extract rotation matrices and positions + R = T_obj[:, :3, :3] + p = np.squeeze(T_obj[:, :3, 3]) + + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + ax.plot(p[:, 0], p[:, 1], p[:, 2], '-', color='blue') + for i, (R_i, p_i) in enumerate(zip(R[::skip_frames], p[::skip_frames])): + ax.quiver(p_i[0], p_i[1], p_i[2], R_i[0, 0], R_i[1, 0], R_i[2, 0], color='r', length=length) + ax.quiver(p_i[0], p_i[1], p_i[2], R_i[0, 1], R_i[1, 1], R_i[2, 1], color='g', length=length) + ax.quiver(p_i[0], p_i[1], p_i[2], R_i[0, 2], R_i[1, 2], R_i[2, 2], color='b', length=length) + ax.set_xlabel('X [m]') + ax.set_ylabel('Y [m]') + ax.set_zlabel('Z [m]') + plt.title('Trajectory and pose frames') + if plt.get_backend() != 'agg': + plt.show() + +def plot_instantaneous_screw_axis(T_isa): + """ + Plots the instantaneous screw axis in a 3D plot. + + Parameters: + - T_isa (numpy.ndarray): The pose data as a 3D array of shape (N, 4, 4). + + Returns: + None + """ + N = T_isa.shape[0] + + # Extract points and directions + points = T_isa[:, :3, 3] # First three elements of the fourth column + directions = T_isa[:, :3, 0] # First three elements of the first column + + # Create a 3D plot + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + + # Plot the points + ax.scatter(points[:, 0], points[:, 1], points[:, 2], color='r', label='Points on the line') + + # Plot the directions as lines + for i in range(N): + start_point = points[i] - directions[i] * 1.0 + end_point = points[i] + directions[i] * 1.0 # Scale the direction for better visualization + ax.plot([start_point[0], end_point[0]], [start_point[1], end_point[1]], [start_point[2], end_point[2]], color='b') + + # Set axis limits + # ax.set_xlim([-0.1, +0.1]) + # ax.set_ylim([-0.1, +0.1]) + # ax.set_zlim([-0.1, +0.1]) + + # Set labels + ax.set_xlabel('X') + ax.set_ylabel('Y') + ax.set_zlabel('Z') + ax.set_title('Plotting the instantaneous screw axis') + + # Add legend + ax.legend() + + # Show plot + plt.show() + +# Plot both the pose frames and the instantaneous screw axis +def plot_combined(T_obj, T_isa, length=0.1, skip_frames=5): + """ + Plots the trajectory, pose frames, and instantaneous screw axis in a 3D plot. + + Parameters: + - T_obj (numpy.ndarray): The pose data as a 3D array of shape (N, 4, 4). + - T_isa (numpy.ndarray): The instantaneous screw axis data as a 3D array of shape (N, 4, 4). + - length (float, optional): The length of the quiver arrows representing the pose frames. Default is 0.075. + - skip_frames (int, optional): The number of frames to skip when plotting the pose frames. Default is 5. + + Returns: + None + """ + # Extract rotation matrices and positions from T_obj + R = T_obj[:, :3, :3] + p = np.squeeze(T_obj[:, :3, 3]) + + # Extract points and directions from T_isa + points = T_isa[:, :3, 3] # First three elements of the fourth column + directions = T_isa[:, :3, 0] # First three elements of the first column + + # Create a 3D plot + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + + # Plot the trajectory + ax.plot(p[:, 0], p[:, 1], p[:, 2], '-', color='blue', label='Trajectory') + + # Plot the pose frames + for i, (R_i, p_i) in enumerate(zip(R[::skip_frames], p[::skip_frames])): + ax.quiver(p_i[0], p_i[1], p_i[2], R_i[0, 0], R_i[1, 0], R_i[2, 0], color='r', length=length) + ax.quiver(p_i[0], p_i[1], p_i[2], R_i[0, 1], R_i[1, 1], R_i[2, 1], color='g', length=length) + ax.quiver(p_i[0], p_i[1], p_i[2], R_i[0, 2], R_i[1, 2], R_i[2, 2], color='b', length=length) + + # Plot the points on the instantaneous screw axis + ax.scatter(points[:, 0], points[:, 1], points[:, 2], color='r', label='Points on the line') + + # Plot the directions as lines + for i in range(points.shape[0]): + start_point = points[i] - directions[i] * 1.0 + end_point = points[i] + directions[i] * 1.0 # Scale the direction for better visualization + ax.plot([start_point[0], end_point[0]], [start_point[1], end_point[1]], [start_point[2], end_point[2]], color='b') + + # Set labels + ax.set_xlabel('X [m]') + ax.set_ylabel('Y [m]') + ax.set_zlabel('Z [m]') + ax.set_title('Trajectory, Pose Frames, and Instantaneous Screw Axis') + + # Add legend + ax.legend() + + # Show plot + plt.show() + +# Synthetic pose data +N = 100 +# T_start = np.eye(4) # pose matrix 1 +# T_mid = np.eye(4) +# T_mid[:3, :3] = S03.rotate_z(np.pi) # pose matrix 3 +# T_mid[:3, 3] = np.array([0.5, 0.5, 0.5]) +# T_end = np.eye(4) +# T_end[:3, :3] = S03.RPY(np.pi/2, 0, np.pi/2) # pose matrix 2 +# T_end[:3, 3] = np.array([1, 1, 1]) + +# # Interpolate between R_start and R_end +# T_obj_m = interpT(np.linspace(0,1,N), np.array([0,0.5,1]), np.stack([T_start, T_mid, T_end],0)) + +T_start = np.eye(4) # pose matrix 1 +#T_mid = np.eye(4) +#T_mid[:3, :3] = S03.rotate_z(np.pi) # pose matrix 3 +#T_mid[:3, 3] = np.array([0.5, 0.5, 0.5]) +T_end = np.eye(4) +T_end[:3, :3] = S03.rotate_z(np.pi) # pose matrix 2 +T_end[:3, 3] = np.array([1, 1, 1]) + +# Interpolate between R_start and R_end +T_obj_m = interpT(np.linspace(0,1,N), np.array([0,1]), np.stack([T_start, T_end],0)) + + +# Call the function with the synthetic data +plot_pose_frames(T_obj_m) + +# # calculate screw twist from the pose data + +# # Initialize an array to store the twist vectors +# twist_vectors = np.zeros((N, 6)) + +# # Loop through each sample and calculate the twist vector +# for i in range(N): +# twist_matrix = SE3.logm_T(T_obj_m[i]) +# twist_vector = SE3.crossvec(twist_matrix) +# twist_vectors[i, :] = twist_vector +# print(twist_vectors) + +OCP = OCP_calc_pose(N, rms_error_traj_pos = 10e-3, rms_error_traj_rot = 10e-3, bool_unsigned_invariants=True, solver='fatrop') + +# Example: calculate invariants for a given trajectory +h = 0.01 # step size for integration of dynamic equations +U, T_obj, T_isa = OCP.calculate_invariants(T_obj_m, h) +print("T_obj: ", T_obj) +print("T_isa: ", T_isa) + +# Plot the instantaneous screw axis +plot_instantaneous_screw_axis(T_isa) + +# Assuming T_isa is already defined and is an Nx4x4 matrix +# Call the combined plotting function with the synthetic data +plot_combined(T_obj_m, T_isa) +plot_combined(T_obj, T_isa) \ No newline at end of file