!apt update -y !apt install ffmpeg -y import numpy as np import sympy import matplotlib.pyplot as plt from matplotlib.animation import FFMpegWriter from matplotlib.patches import Circle from numpy.linalg import solve from numpy.linalg import inv # Defines our kinematic tree class KinematicTreeNode: def __init__(self): #Constructor self.translation = None self.rotation = None self.child = None # Define our rotation matrix def RotZ(theta): c = np.cos(theta) s = np.sin(theta) return np.array([[c, -s], [s, c]]) # FORWARD KINEMATICS # Function that returns ALL joint positions relative to the input: def GetAllJoints(X): if X.child == None: # This is the last node in this branch of the tree return X.translation else: returned_joints = GetAllJoints(X.child) # Calls itself to go to the next child temp, translate = np.meshgrid(returned_joints[0,:], X.translation) transformed_joints = np.matmul(RotZ(X.rotation),returned_joints)+translate list_of_joints = np.hstack((X.translation, transformed_joints)) return list_of_joints # Build the tree base = KinematicTreeNode() L1 = KinematicTreeNode() L2 = KinematicTreeNode() L3 = KinematicTreeNode() base.translation = np.array([[0],[0]], dtype=float) # translation to the joint base.rotation = sympy.pi # rotation of the joint L1.translation = np.array([[1],[0]], dtype=float) L1.rotation = -sympy.pi/2 L2.translation = np.array([[1],[0]], dtype=float) L2.rotation = -sympy.pi/2 L3.translation = np.array([[1],[0]], dtype=float) L3.rotation = 0 # Loop for Inverse Kinematics with Newton-Raphson base.child = L1 L1.child = L2 L2.child = L3 L3.child = None # Step 1: Start with an initial guess theta_current = np.array([base.rotation, L1.rotation, L2.rotation], dtype=float) theta_temp= theta_current for x_iteration in range(1, 201): # Inverse Kinematics: # Define a desired position x_des = x_iteration/100 y_des = 1 theta_des = 0 done = False # Newton-Raphson Search # Define our desired END EFFECTOR POSITION desired_position = np.array([x_des,y_des,theta_des], dtype=float) # Small Epsilon for Finite Difference: epsilon = 0.001 while not done: # Step 2: Find f(theta) base.rotation = theta_current[0] L1.rotation = theta_current[1] L2.rotation = theta_current[2] mainTheta= base.rotation + L1.rotation + L2.rotation joint_locations = GetAllJoints(base) end_effector = np.array([ joint_locations[0,-1], joint_locations[1,-1], mainTheta ]) # Step 3: Find the Jacobian of f w.r.t. theta jac = np.array([[0., 0., 0.],[0., 0., 0.],[0., 0., 0.]], dtype=float) for j in range(0, 3): for k in range(0, 3): delta_theta = np.array([[0.],[0.],[0.]]) delta_theta[k,0] = epsilon base.rotation = theta_current[0] + delta_theta[0,0] L1.rotation = theta_current[1] + delta_theta[1,0] L2.rotation = theta_current[2] + delta_theta[2,0] mainTheta= base.rotation + L1.rotation + L2.rotation joint_locations = GetAllJoints(base) delta_end_effector = np.array([ joint_locations[0,-1], joint_locations[1,-1], mainTheta ]) partial_der = (delta_end_effector - end_effector) / epsilon jac[j,k] = partial_der[j] # Step 4: Find Step Length delta_theta_step = np.matmul(inv(jac), (desired_position - end_effector)) for x in range(1, 101): hold = theta_current + ((delta_theta_step/100) * x) theta_temp= np.vstack(( theta_temp , hold) ) # Theta current: theta_current= theta_current + delta_theta_step # Terminating Condition: if np.linalg.norm(delta_theta_step) < epsilon: done = True theta= theta_temp #Define Joint Locations: x_Base = np.zeros(len(theta)) y_Base = np.zeros(len(theta)) x_L1 = np.zeros(len(theta)) y_L1 = np.zeros(len(theta)) x_L2 = np.zeros(len(theta)) y_L2 = np.zeros(len(theta)) x_L3 = np.zeros(len(theta)) y_L3 = np.zeros(len(theta)) for p in range(0, len(theta)): base.rotation = theta[p, 0] L1.rotation = theta[p, 1] L2.rotation = theta[p, 2] joint_locations = GetAllJoints(base) x_Base[p] = joint_locations[0,0] y_Base[p] = joint_locations[1,0] x_L1[p] = joint_locations[0,1] y_L1[p] = joint_locations[1,1] x_L2[p] = joint_locations[0,2] y_L2[p] = joint_locations[1,2] x_L3[p] = joint_locations[0,3] y_L3[p] = joint_locations[1,3] #Define Time and Theta dt = 0.001 sim_time = 5 t_vec = np.arange(0, sim_time, dt) sim_length = len(t_vec) #Setting Up Animation fig, ax = plt.subplots() p1, = ax.plot([], [], color='black') ax.axis('equal') ax.set_xlim([-4, 4]) ax.set_xlabel('X') # X Label ax.set_ylabel('Y') # Y Label ax.set_title('Question 2b Animation') video_title = "Q2bAnimation" # Initialize joints: base_joint = Circle((0, 0), radius=0.1, color='black', zorder=10) joint_1 = Circle((0, 0), radius=0.1, color='red', zorder=10) joint_2 = Circle((0, 0), radius=0.1, color='green', zorder=10) joint_3 = Circle((0, 0), radius=0.1, color='blue', zorder=10) ax.add_patch(base_joint) ax.add_patch(joint_1) ax.add_patch(joint_2) ax.add_patch(joint_3) # Setup Animation Writer: FPS = 30 sample_rate = int(1 / (dt * FPS)) dpi = 300 writerObj = FFMpegWriter(fps=FPS) theta1 = np.zeros(int (sim_length/sample_rate)+1) theta2 = np.zeros(int (sim_length/sample_rate)+1) theta3 = np.zeros(int (sim_length/sample_rate)+1) count=0 # Plot and Create Animation: with writerObj.saving(fig, video_title+".mp4", dpi): for i in range(0, sim_length, sample_rate): # Draw Arms: align= int( len(theta)/sim_length*i ) x_data_points = [x_Base[align], x_L1[align], x_L2[align], x_L3[align]] y_data_points = [y_Base[align], y_L1[align], y_L2[align], y_L3[align]] p1.set_data(x_data_points, y_data_points) # Draw Joints: joint_1.center = x_L1[align], y_L1[align] joint_2.center = x_L2[align], y_L2[align] joint_3.center = x_L3[align], y_L3[align] # Update Drawing: fig.canvas.draw() # Grab and Save Frame: writerObj.grab_frame() # Store Angles theta1[count] = theta[align, 0] theta2[count] = theta[align, 1] theta3[count] = theta[align, 2] count= count + 1 plt.clf() count1 = np.arange(0, 5, 5/152) count2 = np.arange(0, 5, 5/152) count3 = np.arange(0, 5, 5/152) # Plot angles over time fig, ax = plt.subplots() ax.plot(count1, theta1, color='blue') ax.plot(count2, theta2, color='red') ax.plot(count3, theta3, color='yellow') plt.legend(['Theta 1','Theta 2','Theta 3']) plt.show() from IPython.display import Video Video("/work/Q2bAnimation.mp4", embed=True, width=640, height=480)