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:
113  global sideLength, shipRadius, kouleRadius, propagationStepSize, shipAcceleration, \
114  shipRotVel, shipDelta, shipEps, path
115  sideLength, shipRadius, kouleRadius, propagationStepSize, shipAcceleration, \
116  shipRotVel, shipDelta, shipEps = [float(x) for x in next(f).split()]
117  path = [[float(x) for x in line.split(' ')] for line in f]
118  if len(path) == 0:
119  print('Error: %s contains no solution path' % fname)
120  return
121  step = int(ceil(speedUp / (propagationStepSize * targetFrameRate)))
122  path = path[0:len(path):step]
123  print('Creating a movie with %d frames...' % len(path))
124  print('Printing a \'.\' for every 10th frame:')
125  ani = animation.FuncAnimation(fig, plotSystem, frames = len(path),
126  interval = 1000. / step, blit = True)
127  (base,ext) = splitext(basename(fname))
128  outfname = base + '.mp4'
129  ani.save(outfname, bitrate = 300, fps = targetFrameRate)
130  print('')
131 
132 if __name__ == '__main__':
133  if len(argv) == 1:
134  print('Usage: KoulesPlayback.py <filename> [<filename2> ...]')
135  else:
136  for fname in argv[1:]:
137  makeMovie(fname)