Loading...
Searching...
No Matches
KoulesPlayback.py
1#!/usr/bin/env python
2
3
36
37# Author: Beck Chen, Mark Moll
38
39from sys import argv, stdout
40from os.path import basename, splitext
41from math import cos, sin, atan2, pi, ceil
42import matplotlib.pyplot as plt
43import matplotlib.animation as animation
44
45targetFrameRate = 30 # desired number of frames per second
46speedUp = 1.
47# the parameters will be read from file
48sideLength = 0
49shipRadius = 0
50kouleRadius = 0
51propagationStepSize = 0
52shipAcceleration = 0
53shipRotVel = 0
54shipDelta = 0
55shipEps = 0
56
57fig = plt.figure(figsize=(6, 6))
58ax = plt.axes(xlim=(0, 1), ylim=(0, 1))
59fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
60handle, = ax.plot([], [])
61path = None
62
63def normalizeAngle(theta):
64 if theta < -pi:
65 return theta + 2. * pi
66 if theta > pi:
67 return theta - 2. * pi
68 return theta
69
70def plotShip(x, u):
71 pos = (x[0], x[1])
72 theta = x[4]
73 (cs, ss) = (shipRadius*cos(theta), shipRadius * sin(theta))
74 v = [u[0] - x[2], u[1] - x[3]]
75 deltaTheta = normalizeAngle(atan2(v[1], v[0]) - theta)
76 if v[0]*v[0] + v[1]*v[1] >= shipDelta * shipDelta:
77 if abs(deltaTheta) < shipEps:
78 # accelerate forward, draw thruster on the back
79 ax.add_patch(plt.Circle((pos[0] - cs, pos[1] - ss), .3 * shipRadius, color="red"))
80 elif deltaTheta > 0:
81 # rotate counterclockwise, draw thruster on right side
82 ax.add_patch(plt.Circle((pos[0] + ss, pos[1] - cs), .3 * shipRadius, color="red"))
83 else:
84 # rotate clockwise, draw thruster on left side
85 ax.add_patch(plt.Circle((pos[0] - ss, pos[1] + cs), .3 * shipRadius, color="red"))
86 # draw ship
87 ax.add_patch(plt.Circle(x[:2], shipRadius, color="yellow"))
88 # draw two blue "eyes"
89 ax.add_patch(plt.Circle((pos[0] + .7*shipRadius*cos(theta + .75), \
90 pos[1] + .7*shipRadius*sin(theta + .75)), .2 * shipRadius, color="blue"))
91 ax.add_patch(plt.Circle((pos[0] + .7*shipRadius*cos(theta - .75), \
92 pos[1] + .7*shipRadius*sin(theta - .75)), .2 * shipRadius, color="blue"))
93
94def plotKoules(state):
95 numKoules = int(len(state)/4)
96 for i in range(numKoules):
97 ax.add_patch(plt.Circle((state[4 * i], state[4 * i + 1]), kouleRadius, color="red"))
98
99def plotSystem(index):
100 ax.clear()
101 ax.add_patch(plt.Rectangle((0, 0), 1, 1, color='black'))
102 plotKoules(path[index][5:-3])
103 plotShip(path[index][0:5], path[index][-3:])
104 if index % 10 == 0:
105 stdout.write('.')
106 stdout.flush()
107 return handle,
108
109def makeMovie(fname):
110 with open(fname, 'r') as f:
111 global sideLength, shipRadius, kouleRadius, propagationStepSize, shipAcceleration, \
112 shipRotVel, shipDelta, shipEps, path
113 sideLength, shipRadius, kouleRadius, propagationStepSize, shipAcceleration, \
114 shipRotVel, shipDelta, shipEps = [float(x) for x in next(f).split()]
115 path = [[float(x) for x in line.split(' ')] for line in f]
116 if not path:
117 print('Error: %s contains no solution path' % fname)
118 return
119 step = int(ceil(speedUp / (propagationStepSize * targetFrameRate)))
120 path = path[0:len(path):step]
121 print('Creating a movie with %d frames...' % len(path))
122 print('Printing a \'.\' for every 10th frame:')
123 ani = animation.FuncAnimation(fig, plotSystem, frames=len(path), \
124 interval=1000. / step, blit=True)
125 (base, _) = splitext(basename(fname))
126 outfname = base + '.mp4'
127 ani.save(outfname, bitrate=300, fps=targetFrameRate)
128 print('')
129
130if __name__ == '__main__':
131 if len(argv) == 1:
132 print('Usage: KoulesPlayback.py <filename> [<filename2> ...]')
133 else:
134 for trajectory_file in argv[1:]:
135 makeMovie(trajectory_file)