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