diff --git a/RoboticArm/NLinkArm.py b/RoboticArm/NLinkArm.py new file mode 100644 index 00000000..4ab56054 --- /dev/null +++ b/RoboticArm/NLinkArm.py @@ -0,0 +1,62 @@ +""" +Class for controlling and plotting an arm with an arbitrary number of links. + +Author: Daniel Ingram +""" + +import numpy as np +import matplotlib.pyplot as plt + +class NLinkArm(object): + def __init__(self, link_lengths, joint_angles, goal): + self.n_links = len(link_lengths) + if self.n_links is not len(joint_angles): + raise ValueError() + + self.link_lengths = np.array(link_lengths) + self.joint_angles = np.array(joint_angles) + self.points = [[0, 0] for _ in range(self.n_links+1)] + + self.lim = sum(link_lengths) + self.goal = np.array(goal).T + + self.fig = plt.figure() + self.fig.canvas.mpl_connect('button_press_event', self.click) + + plt.ion() + plt.show() + + self.update_points() + + def update_joints(self, joint_angles): + self.joint_angles = joint_angles + self.update_points() + + def update_points(self): + for i in range(1, self.n_links+1): + self.points[i][0] = self.points[i-1][0] + self.link_lengths[i-1]*np.cos(np.sum(self.joint_angles[:i])) + self.points[i][1] = self.points[i-1][1] + self.link_lengths[i-1]*np.sin(np.sum(self.joint_angles[:i])) + + self.end_effector = np.array(self.points[self.n_links]).T + self.plot() + + def plot(self): + plt.cla() + + for i in range(self.n_links+1): + if i is not self.n_links: + plt.plot([self.points[i][0], self.points[i+1][0]], [self.points[i][1], self.points[i+1][1]], 'r-') + plt.plot(self.points[i][0], self.points[i][1], 'ko') + + plt.plot(self.goal[0], self.goal[1], 'gx') + + plt.plot([self.end_effector[0], self.goal[0]], [self.end_effector[1], self.goal[1]], 'g--') + + plt.xlim([-self.lim, self.lim]) + plt.ylim([-self.lim, self.lim]) + plt.draw() + plt.pause(0.0001) + + def click(self, event): + self.goal = np.array([event.xdata, event.ydata]).T + self.plot()