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:
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 
130 if __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)