KoulesPlayback.py
1 #!/usr/bin/env python
2 
3 
36 
37 # Author: Beck Chen, Mark Moll
38 
39 from sys import argv, stdout
40 from os.path import basename, splitext
41 from math import cos, sin, atan2, pi, ceil
42 import matplotlib.pyplot as plt
43 import matplotlib.animation as animation
44 
45 targetFrameRate = 30 # desired number of frames per second
46 speedUp = 1.
47 # the parameters will be read from file
48 sideLength = 0
49 shipRadius = 0
50 kouleRadius = 0
51 propagationStepSize = 0
52 shipAcceleration = 0
53 shipRotVel = 0
54 shipDelta = 0
55 shipEps = 0
56 
57 fig = plt.figure(figsize=(6, 6))
58 ax = plt.axes(xlim=(0, 1), ylim=(0, 1))
59 fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
60 handle, = ax.plot([], [])
61 path = None
62 
63 def normalizeAngle(theta):
64  if theta < -pi:
65  return theta + 2. * pi
66  if theta > pi:
67  return theta - 2. * pi
68  return theta
69 
70 def 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 
94 def 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 
99 def 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 
109 def makeMovie(fname):
110  with open(fname, 'r') as f: global sideLength, shipRadius, kouleRadius, propagationStepSize, shipAcceleration, \
111  shipRotVel, shipDelta, shipEps, path
112  sideLength, shipRadius, kouleRadius, propagationStepSize, shipAcceleration, \
113  shipRotVel, shipDelta, shipEps = [float(x) for x in next(f).split()]
114  path = [[float(x) for x in line.split(' ')] for line in f]
115  if not path:
116  print('Error: %s contains no solution path' % fname)
117  return
118  step = int(ceil(speedUp / (propagationStepSize * targetFrameRate)))
119  path = path[0:len(path):step]
120  print('Creating a movie with %d frames...' % len(path))
121  print('Printing a \'.\' for every 10th frame:')
122  ani = animation.FuncAnimation(fig, plotSystem, frames=len(path), \
123  interval=1000. / step, blit=True)
124  (base, _) = splitext(basename(fname))
125  outfname = base + '.mp4'
126  ani.save(outfname, bitrate=300, fps=targetFrameRate)
127  print('')
128 
129 if __name__ == '__main__':
130  if len(argv) == 1:
131  print('Usage: KoulesPlayback.py <filename> [<filename2> ...]')
132  else:
133  for trajectory_file in argv[1:]:
134  makeMovie(trajectory_file)
135