!apt update -y
!apt install ffmpeg -y
import numpy as np
from numpy.linalg import inv
from IPython.display import Video
import matplotlib.pyplot as plt
import sympy
from matplotlib.animation import FFMpegWriter
from matplotlib.patches import Circle
from matplotlib.patches import Rectangle
from mpl_toolkits.mplot3d import Axes3D
# Defines kinematic tree
class KinematicTreeNode:
def __init__(self): # Constructor
self.translation = None
self.rotation = None
self.child = None
# Defines rotation matrix
def RotZ(theta):
c = np.cos(theta)
s = np.sin(theta)
return np.array([[c, -s],[s, c]])
# Returns a matrix with translations
def NodePositions(X):
allpositions = []
if X.child == None:
return X.translation
else:
returnjoints = NodePositions(X.child)
temp, translate = np.meshgrid(returnjoints[0,:], X.translation)
transformed_joints = np.matmul(RotZ(X.rotation),returnjoints)+translate
list_of_joints = np.hstack((X.translation, transformed_joints))
return list_of_joints
# Returns end effector translation and rotation
def EndEffectorPosRot(X):
if X.child == None:
return X.translation, X.rotation
else:
trans, rot = EndEffectorPosRot(X.child)
new_trans = X.translation + np.matmul(RotZ(X.rotation),trans)
new_rot = X.rotation+rot
return new_trans, new_rot
# Creating a function to find the Ideal initial velocity based on max range
def FindInitialVelocities(x_initial,x_final,theta):
v_initial = np.zeros(4)
a_y = -9.81 # acceleration of gravity
if ( (x_final-x_initial) > 0):
x_hold = -(x_final-x_initial)
elif ( (x_final-x_initial) < 0):
x_hold = (x_final-x_initial)
v_initial = (((x_hold)*a_y)/np.sin(2*theta))**(1/2)
return v_initial
# Creating the function for ball trajectory
# n is number of balls
# m is mass
# x_initial,y_initial are the coordinates from where the balls are thrown
# x_final,y_final are the coordinates from where the balls will be caught
# v_initial will be an array of the initial magnitude velocities needed
# theta will be an array of initial angles to throw
def makeBallPaths(m,x_initial,y_initial,x_final,y_final,v_initial,theta):
# Parameters
g = 9.81 # m/s^2, gravity
Cd = 0.00 # drag coefficient, no units
check = 0
# Creating time step and time vector
dt = 0.001 # s, time step
t_final = 10 # seconds, time to simulate
t_vec = np.arange(0,t_final,dt) # 10001 values of time in 'dt' increments
vec_len = len(t_vec)
# Creating the position vectors (x,y) and the time vector
x_vec1 = np.zeros(vec_len) # holds same amount of values as time vec
y_vec1 = np.zeros(vec_len)
dx_vec1 = np.zeros(vec_len)
dy_vec1 = np.zeros(vec_len)
v_mag1 = np.zeros(vec_len)
drag_vec1 = np.zeros(vec_len)
# Solves for initial vx,vy, and drag coeff
dx_initial1 = v_initial*np.cos(theta)
dy_initial1 = v_initial*np.sin(theta)
drag_initial1 = Cd*(v_initial**2)
# Put the initial conditions in arrays
x_vec1[0] = x_initial
dx_vec1[0] = dx_initial1
y_vec1[0] = y_initial
dy_vec1[0] = dy_initial1
v_mag1[0] = v_initial
drag_vec1[0] = drag_initial1
for i in range(1,vec_len):
# Euler Integration
# Solve for ddx,ddy of previous iteration through matrices
A = np.array([[m,0],[0,m]]) # A doesnt change because our balls are same mass
B1 = np.array([[-drag_vec1[i-1]*np.cos(theta)],[-m*g-drag_vec1[i-1]*np.sin(theta)]])
[ddx1, ddy1] = np.linalg.solve(A, B1)
# Solve x,y,dx,dy for all balls
x_vec1[i] = x_vec1[i - 1] + dx_vec1[i - 1] * dt
dx_vec1[i] = dx_vec1[i-1] + ddx1 * dt
y_vec1[i] = y_vec1[i - 1] + dy_vec1[i - 1] * dt
dy_vec1[i] = dy_vec1[i-1] + ddy1 * dt
v_mag1[i] = ((dx_vec1[i]**2)+(dy_vec1[i]**2))**(1/2)
drag_vec1[i] = Cd*(v_mag1[i]**2)
t_ball_motion = 2*(-v_initial*np.sin(theta))/-g
if (y_vec1[i] <= (y_final-5)):
x_vec1[i] = x_vec1[i-1]
y_vec1[i] = (y_final-5)
return t_ball_motion,x_vec1,y_vec1
# Build the tree for the arms and set up initial conditions
base = KinematicTreeNode()
L1 = KinematicTreeNode()
L2 = KinematicTreeNode()
L3 = KinematicTreeNode()
# Set up initial conditions for arms
base.translation = np.array([[0],[0]])
base.rotation = 0.0 #radians
base.child = L1
L1.translation = np.array([[1],[0]])
L1.rotation = np.pi/2 #radians
L1.child = L2
L2.translation = np.array([[1],[0]])
L2.rotation = -np.pi/2 #radians
L2.child = L3
L3.translation = np.array([[1],[0]])
L3.rotation = 0.0
L3.child = None
m = 1.0 # mass of all balls
n = 1.0 # the number of balls
x_int = 1.0 # Initial x for ball arc
y_int = 1.0 # Initial y for ball arc
x_fin = -1.0 # Final x for ball arc
y_fin = 1.0 # Final y for ball arc
# Set up time response and delta theta for throwing
dt = 0.001
thetaThrow = (np.pi/180.0)*45
throwVelocity = FindInitialVelocities(x_int,x_fin,thetaThrow)
dtheta_throw = dt*throwVelocity
# Inverse Kinematics: Throwing
# Define a desired path and initial end effector position
if x_int-x_fin > 0: # Check whether arm is throwing left
zeroangle = 0.0
dtheta_throw = -dtheta_throw
throwingangle = np.pi-thetaThrow
elif x_int-x_fin < 0: # Check whether arm is throwing right
zeroangle = np.pi
throwingangle = thetaThrow
else:
zeroangle = -np.pi/4
dtheta_throw = -dtheta_throw
throwingangle = np.pi/2
theta_des_throw1 = np.arange(thetaThrow,zeroangle,dtheta_throw)
theta_des_throw2 = np.arange(zeroangle,thetaThrow,-dtheta_throw)
theta_des_throw = np.hstack((theta_des_throw1,theta_des_throw2))
# Set time so it is moving at dt but makes the same length of vector
t = np.linspace(0,dt*len(theta_des_throw),len(theta_des_throw))
x_des_throw = np.zeros(len(theta_des_throw))
y_des_throw = np.zeros(len(theta_des_throw))
for i in range(0,len(t)):
x_des_throw[i] = np.add((1-np.cos(thetaThrow)),np.cos(theta_des_throw[i]))
y_des_throw[i] = np.add((1-np.sin(thetaThrow)),np.sin(theta_des_throw[i]))
desired_position = np.array([x_des_throw[0],y_des_throw[0],theta_des_throw[0]], dtype=float)
# Newton-Raphson Search
# Step 1: Start with an initial guess
# Initial guess for our joint angles
theta_current = np.array([0.0,np.pi,-np.pi], dtype=float)
done = False
# Small Epsilon for Finite Difference:
epsilon = 0.001
# Initialize vectors for animation and plot
L1_x = np.zeros(len(t))
L1_y = np.zeros(len(t))
L2_x = np.zeros(len(t))
L2_y = np.zeros(len(t))
L3_x = np.zeros(len(t))
L3_y = np.zeros(len(t))
theta1traveled = np.zeros(len(t))
theta2traveled = np.zeros(len(t))
theta3traveled = np.zeros(len(t))
thetastep_plot = np.zeros(len(t))
xstep_plot = np.zeros(len(t))
ystep_plot = np.zeros(len(t))
for i in range(0,len(t)):
done = False
while done == False:
theta_current[0] = (( -theta_current[0] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
theta_current[1] = (( -theta_current[1] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
theta_current[2] = (( -theta_current[2] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
desired_position = np.array([x_des_throw[i], y_des_throw[i], theta_des_throw[i]])
# Step 2: Find f(theta)
base.rotation = theta_current[0]
L1.rotation = theta_current[1]
L2.rotation = theta_current[2]
joint_locations = NodePositions(base)
endposition, rotation = EndEffectorPosRot(base)
end_effector = np.array([0., 0., 0.])
end_effector[0] = endposition[0]
end_effector[1] = endposition[1]
end_effector[2] = rotation
end_effector[2] = (( -end_effector[2] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
# Step 3: Find the Jacobian of f w.r.t. theta
jac = np.array([
[-np.sin(base.rotation) -np.sin(base.rotation + L1.rotation) -np.sin(base.rotation + L1.rotation + L2.rotation) , -np.sin(base.rotation + L1.rotation) -np.sin(base.rotation + L1.rotation + L2.rotation) , -np.sin(base.rotation + L1.rotation + L2.rotation)],
[np.cos(base.rotation) +np.cos(base.rotation + L1.rotation) + np.cos(base.rotation + L1.rotation + L2.rotation) , np.cos(base.rotation + L1.rotation) + np.cos(base.rotation + L1.rotation + L2.rotation) , np.cos(base.rotation + L1.rotation + L2.rotation)],
[1.0, 1.0, 1.0]
])
# Step 4: Find Step Length
delta_theta = desired_position - end_effector
delta_theta_step = np.matmul(np.linalg.pinv(jac), delta_theta)
# Theta current:
theta_current = theta_current + delta_theta_step
# Terminating Condition:
if np.linalg.norm(delta_theta_step) < epsilon:
done = True
L1_x[i] = joint_locations[0,1]
L1_y[i] = joint_locations[1,1]
L2_x[i] = joint_locations[0,2]
L2_y[i] = joint_locations[1,2]
L3_x[i] = joint_locations[0,3]
L3_y[i] = joint_locations[1,3]
theta1traveled[i] = base.rotation
theta2traveled[i] = L1.rotation
theta3traveled[i] = L2.rotation
# Set up ball pathing and trajectory
t_ball,x_ball,y_ball = makeBallPaths(m,x_int,y_int,x_fin,y_fin,throwVelocity,throwingangle)
# Second Step: move to catching position using inverse kinematics
# Define a desired path and initial end effector position
# Set time so it is moving at dt but makes the same length of vector
t_catch = np.arange(0.0,t_ball,dt)
x_des_catch = np.linspace(x_int,x_fin,len(t_catch))
y_des_catch = np.linspace(y_int,y_fin,len(t_catch))
theta_des_catch = np.linspace(thetaThrow,zeroangle+np.pi,len(t_catch))
desired_position = np.array([x_des_catch[0],y_des_catch[0],theta_des_catch[0]], dtype=float)
# Newton-Raphson Search
# Step 1: Start with an initial guess
# Initial guess for our joint angles
done = False
# Small Epsilon for Finite Difference:
epsilon = 0.001
# Initialize vectors for animation and plot
L1_xcatch = np.zeros(len(t_catch))
L1_ycatch = np.zeros(len(t_catch))
L2_xcatch = np.zeros(len(t_catch))
L2_ycatch = np.zeros(len(t_catch))
L3_xcatch = np.zeros(len(t_catch))
L3_ycatch = np.zeros(len(t_catch))
theta1traveled_catch = np.zeros(len(t_catch))
theta2traveled_catch = np.zeros(len(t_catch))
theta3traveled_catch = np.zeros(len(t_catch))
thetastep_plot_catch = np.zeros(len(t_catch))
xstep_plot_catch = np.zeros(len(t_catch))
ystep_plot_catch = np.zeros(len(t_catch))
for i in range(0,len(t_catch)):
done = False
while done == False:
theta_current[0] = (( -theta_current[0] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
theta_current[1] = (( -theta_current[1] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
theta_current[2] = (( -theta_current[2] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
desired_position = np.array([x_des_catch[i], y_des_catch[i], theta_des_catch[i]])
# Step 2: Find f(theta)
base.rotation = theta_current[0]
L1.rotation = theta_current[1]
L2.rotation = theta_current[2]
joint_locations = NodePositions(base)
endposition, rotation = EndEffectorPosRot(base)
end_effector = np.array([0., 0., 0.])
end_effector[0] = endposition[0]
end_effector[1] = endposition[1]
end_effector[2] = rotation
end_effector[2] = (( -end_effector[2] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
# Step 3: Find the Jacobian of f w.r.t. theta
jac = np.array([
[-np.sin(base.rotation) -np.sin(base.rotation + L1.rotation) -np.sin(base.rotation + L1.rotation + L2.rotation) , -np.sin(base.rotation + L1.rotation) -np.sin(base.rotation + L1.rotation + L2.rotation) , -np.sin(base.rotation + L1.rotation + L2.rotation)],
[np.cos(base.rotation) +np.cos(base.rotation + L1.rotation) + np.cos(base.rotation + L1.rotation + L2.rotation) , np.cos(base.rotation + L1.rotation) + np.cos(base.rotation + L1.rotation + L2.rotation) , np.cos(base.rotation + L1.rotation + L2.rotation)],
[1.0, 1.0, 1.0]
])
# Step 4: Find Step Length
delta_theta = desired_position - end_effector
delta_theta_step = np.matmul(np.linalg.pinv(jac), delta_theta)
# Theta current:
theta_current = theta_current + delta_theta_step
# Terminating Condition:
if np.linalg.norm(delta_theta_step) < epsilon:
done = True
L1_xcatch[i] = joint_locations[0,1]
L1_ycatch[i] = joint_locations[1,1]
L2_xcatch[i] = joint_locations[0,2]
L2_ycatch[i] = joint_locations[1,2]
L3_xcatch[i] = joint_locations[0,3]
L3_ycatch[i] = joint_locations[1,3]
theta1traveled_catch[i] = base.rotation
theta2traveled_catch[i] = L1.rotation
theta3traveled_catch[i] = L2.rotation
# Return to starting position for throwing again
# Define a desired path and initial end effector position
t_return = t_catch
x_des_return = np.linspace(x_fin,x_int,len(t_return))
y_des_return = np.linspace(y_fin,y_int,len(t_return))
theta_des_return = np.linspace(zeroangle+np.pi,zeroangle,len(t_return))
desired_position = np.array([x_des_return[0],y_des_return[0],theta_des_return[0]], dtype=float)
# Newton-Raphson Search
# Step 1: Start with an initial guess
# Initial guess for our joint angles
done = False
# Small Epsilon for Finite Difference:
epsilon = 0.001
# Initialize vectors for animation and plot
L1_xreturn = np.zeros(len(t_return))
L1_yreturn = np.zeros(len(t_return))
L2_xreturn = np.zeros(len(t_return))
L2_yreturn = np.zeros(len(t_return))
L3_xreturn = np.zeros(len(t_return))
L3_yreturn = np.zeros(len(t_return))
theta1traveled_return = np.zeros(len(t_return))
theta2traveled_return = np.zeros(len(t_return))
theta3traveled_return = np.zeros(len(t_return))
thetastep_plot_return = np.zeros(len(t_return))
xstep_plot_return = np.zeros(len(t_return))
ystep_plot_return = np.zeros(len(t_return))
for i in range(0,len(t_return)):
done = False
while done == False:
theta_current[0] = (( -theta_current[0] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
theta_current[1] = (( -theta_current[1] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
theta_current[2] = (( -theta_current[2] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
desired_position = np.array([x_des_return[i], y_des_return[i], theta_des_return[i]])
# Step 2: Find f(theta)
base.rotation = theta_current[0]
L1.rotation = theta_current[1]
L2.rotation = theta_current[2]
joint_locations = NodePositions(base)
endposition, rotation = EndEffectorPosRot(base)
end_effector = np.array([0., 0., 0.])
end_effector[0] = endposition[0]
end_effector[1] = endposition[1]
end_effector[2] = rotation
end_effector[2] = (( -end_effector[2] + np.pi) % (2.0 * np.pi ) - np.pi) * -1.0
# Step 3: Find the Jacobian of f w.r.t. theta
jac = np.array([
[-np.sin(base.rotation) -np.sin(base.rotation + L1.rotation) -np.sin(base.rotation + L1.rotation + L2.rotation) , -np.sin(base.rotation + L1.rotation) -np.sin(base.rotation + L1.rotation + L2.rotation) , -np.sin(base.rotation + L1.rotation + L2.rotation)],
[np.cos(base.rotation) +np.cos(base.rotation + L1.rotation) + np.cos(base.rotation + L1.rotation + L2.rotation) , np.cos(base.rotation + L1.rotation) + np.cos(base.rotation + L1.rotation + L2.rotation) , np.cos(base.rotation + L1.rotation + L2.rotation)],
[1.0, 1.0, 1.0]
])
# Step 4: Find Step Length
delta_theta = desired_position - end_effector
delta_theta_step = np.matmul(np.linalg.pinv(jac), delta_theta)
# Theta current:
theta_current = theta_current + delta_theta_step
# Terminating Condition:
if np.linalg.norm(delta_theta_step) < epsilon:
done = True
L1_xreturn[i] = joint_locations[0,1]
L1_yreturn[i] = joint_locations[1,1]
L2_xreturn[i] = joint_locations[0,2]
L2_yreturn[i] = joint_locations[1,2]
L3_xreturn[i] = joint_locations[0,3]
L3_yreturn[i] = joint_locations[1,3]
theta1traveled_return[i] = base.rotation
theta2traveled_return[i] = L1.rotation
theta3traveled_return[i] = L2.rotation
# Create Animation:
# Setup Figure: Initialize Figure / Axe Handles
fig, ax = plt.subplots()
p1, = ax.plot([], [], color='black', linewidth=2)
p2, = ax.plot([], [], color='black', linewidth=2)
p3, = ax.plot([], [], color='black', linewidth=2)
lb, ub = -5, 5
ax.axis('equal')
ax.set_xlim([lb, ub])
ax.set_xlabel('X') # X Label
ax.set_ylabel('Y') # Y Label
ax.set_title('Juggling Simulation:')
video_title = "05 Juggling Sim"
# Setup Animation Writer:
FPS = 20
sample_rate = int(1 / (dt * FPS))
dpi = 300
writerObj = FFMpegWriter(fps=FPS)
# Initialize Patch: Pendulum 1 and 2
joint_1 = Circle((0, 0), radius=0.05, color='blue', zorder=10)
joint_2 = Circle((0, 0), radius=0.05, color='blue', zorder=10)
joint_3 = Circle((0, 0), radius=0.05, color='blue', zorder=10)
ball_1 = Circle((0, 0), radius=0.1, color='red', zorder=10)
ax.add_patch(joint_1)
ax.add_patch(joint_2)
ax.add_patch(joint_3)
ax.add_patch(ball_1)
# Draw Static Objects:
pin_joint = Circle((0, 0), radius=0.05, color='black', zorder=10)
ax.add_patch(pin_joint)
# Plot and Create Animation:
with writerObj.saving(fig, video_title+".mp4", dpi):
for i in range(0, len(t), sample_rate): # Throwing part of animation
# Draw Arm:
L1_arm_x = [0, L1_x[i]]
L1_arm_y = [0, L1_y[i]]
L2_arm_x = [L1_x[i], L2_x[i]]
L2_arm_y = [L1_y[i], L2_y[i]]
L3_arm_x = [L2_x[i], L3_x[i]]
L3_arm_y = [L2_y[i], L3_y[i]]
p1.set_data(L1_arm_x, L1_arm_y)
p2.set_data(L2_arm_x, L2_arm_y)
p3.set_data(L3_arm_x, L3_arm_y)
# Update Patches:
joint_1_center = L1_x[i], L1_y[i]
joint_2_center = L2_x[i], L2_y[i]
joint_3_center = L3_x[i], L3_y[i]
ball_1_center = joint_3_center
joint_1.center = joint_1_center
joint_2.center = joint_2_center
joint_3.center = joint_3_center
ball_1.center = ball_1_center
# Update Drawing:
fig.canvas.draw()
# Grab and Save Frame:
writerObj.grab_frame()
for i in range(0, len(t_catch), sample_rate):
# Draw Arm:
L1_arm_xcatch = [0, L1_xcatch[i]]
L1_arm_ycatch = [0, L1_ycatch[i]]
L2_arm_xcatch = [L1_xcatch[i], L2_xcatch[i]]
L2_arm_ycatch = [L1_ycatch[i], L2_ycatch[i]]
L3_arm_xcatch = [L2_xcatch[i], L3_xcatch[i]]
L3_arm_ycatch = [L2_ycatch[i], L3_ycatch[i]]
p1.set_data(L1_arm_xcatch, L1_arm_ycatch)
p2.set_data(L2_arm_xcatch, L2_arm_ycatch)
p3.set_data(L3_arm_xcatch, L3_arm_ycatch)
# Update Patches:
joint_1_center = L1_xcatch[i], L1_ycatch[i]
joint_2_center = L2_xcatch[i], L2_ycatch[i]
joint_3_center = L3_xcatch[i], L3_ycatch[i]
ball_1_center = joint_3_center
joint_1.center = joint_1_center
joint_2.center = joint_2_center
joint_3.center = joint_3_center
ball_projectile = x_ball[i], y_ball[i]
ball_1.center = ball_projectile
# Update Drawing:
fig.canvas.draw()
# Grab and Save Frame:
writerObj.grab_frame()
for i in range(0, len(t_return), sample_rate):
# Draw Arm:
L1_arm_xreturn = [0, L1_xreturn[i]]
L1_arm_yreturn = [0, L1_yreturn[i]]
L2_arm_xreturn = [L1_xreturn[i], L2_xreturn[i]]
L2_arm_yreturn = [L1_yreturn[i], L2_yreturn[i]]
L3_arm_xreturn = [L2_xreturn[i], L3_xreturn[i]]
L3_arm_yreturn = [L2_yreturn[i], L3_yreturn[i]]
p1.set_data(L1_arm_xreturn, L1_arm_yreturn)
p2.set_data(L2_arm_xreturn, L2_arm_yreturn)
p3.set_data(L3_arm_xreturn, L3_arm_yreturn)
# Update Patches:
joint_1_center = L1_xreturn[i], L1_yreturn[i]
joint_2_center = L2_xreturn[i], L2_yreturn[i]
joint_3_center = L3_xreturn[i], L3_yreturn[i]
ball_1_center = joint_3_center
joint_1.center = joint_1_center
joint_2.center = joint_2_center
joint_3.center = joint_3_center
ball_1.center = ball_1_center
# Update Drawing:
fig.canvas.draw()
# Grab and Save Frame:
writerObj.grab_frame()
Video("/work/"+video_title+".mp4", embed=True, width=640, height=480)