mirror of
https://github.com/AtsushiSakai/PythonRobotics.git
synced 2026-02-10 23:45:03 -05:00
77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
"""
|
|
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, show_animation):
|
|
self.show_animation = show_animation
|
|
self.n_links = len(link_lengths)
|
|
if self.n_links != 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
|
|
|
|
if show_animation: # pragma: no cover
|
|
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
|
|
if self.show_animation: # pragma: no cover
|
|
self.plot()
|
|
|
|
def plot(self): # pragma: no cover
|
|
plt.cla()
|
|
# for stopping simulation with the esc key.
|
|
plt.gcf().canvas.mpl_connect('key_release_event',
|
|
lambda event: [exit(0) if event.key == 'escape' else None])
|
|
|
|
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()
|