KinematicChainPathPlot.py
1#!/usr/bin/env python
2
3
36
37# Author: Mark Moll
38
39# This is a helper program to visualize the output of the KinematicChainBenchmark.cpp.
40# Sample usage (when running from the OMPL build directory):
41#
42# ./bin/demo_KinematicChainBenchmark 10 1
43# /path/to/KinematicChainPathPlot.py 10 # produces kinematic_10.pdf
44# /path/to/KinematicChainPathPlot.py 10 movie # produces kinematic_10.mp4
45
46from sys import argv
47import numpy as np
48import matplotlib.pyplot as plt
49from matplotlib.patches import PathPatch
50from matplotlib.path import Path
51from matplotlib import animation, cm, colors
52
53fig = plt.figure()
54plt.axis('equal')
55ax = plt.axes(xlim=(-1, 1), ylim=(-1, 1))
56ln, = plt.plot([], [], animated=True)
57
58def getPositions(pose, linkLength, color='red'):
59 angles = np.cumsum(pose)
60 return PathPatch(Path(
61 np.insert(linkLength * np.cumsum([np.cos(angles), np.sin(angles)], \
62 axis=1), 0, 0., axis=1).T), facecolor='none', edgecolor=color)
63
64def drawEnvironment(env):
65 ax.clear()
66 if env:
67 plt.gca().add_patch(env)
68
69def drawPose(index, env, poses, linkLength):
70 drawEnvironment(env)
71 if index == -1:
72 cMap = cm.ScalarMappable(
73 norm=colors.Normalize(vmin=0, vmax=poses.shape[0] - 1),
74 cmap=plt.get_cmap('viridis'))
75 for (i, pose) in enumerate(poses):
76 plt.gca().add_patch(getPositions(pose, linkLength, \
77 cMap.to_rgba(i)))
78 else:
79 plt.gca().add_patch(getPositions(poses[index, :], linkLength))
80 return ln,
81
82def makeMovie(fname, env, poses, linkLength):
83 ani = animation.FuncAnimation(plt.gcf(), drawPose, \
84 fargs=(env, poses, linkLength), frames=poses.shape[0], \
85 blit=True)
86 ani.save(fname, bitrate=300, fps=20)
87
88if __name__ == '__main__':
89 if len(argv) < 2:
90 print('Usage: %s num_links [movie]' % argv[0])
91 exit(1)
92
93 dims = int(argv[1])
94
95 env = None
96 try:
97 coords = np.loadtxt('environment_%d.dat' % dims)
98 env = PathPatch(Path(coords), facecolor='none', edgecolor='black')
99 except ValueError:
100 pass
101
102 poses = np.loadtxt('kinematic_path_%d.dat' % dims)
103 linkLength = 1. / dims
104
105 if len(argv) > 2:
106 makeMovie('kinematic_%d.mp4' % dims, env, poses, linkLength)
107 else:
108 drawPose(-1, env, poses, linkLength)
109 fig = plt.gcf()
110 fig.savefig('kinematic_%d.pdf' % dims)