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
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]
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
80  elif deltaTheta > 0:
81  # rotate counterclockwise, draw thruster on right side
83  else:
84  # rotate clockwise, draw thruster on left side
86  # draw ship
88  # draw two blue "eyes"
93
94 def plotKoules(state):
95  numKoules = int(len(state)/4)
96  for i in range(numKoules):
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
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