-
Notifications
You must be signed in to change notification settings - Fork 4
/
animate.py
118 lines (101 loc) · 4.98 KB
/
animate.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
from __future__ import print_function
import click
import csv
import numpy as np
import matplotlib.animation as animation
import matplotlib.pyplot as plt
from collections import defaultdict
class Animate:
def __init__(self, data_points, rod_length, delta_t, blit=True, **fig_kwargs):
self.animation_length = len(data_points)
self.rod_length = rod_length
self.delta_t = delta_t
self.data = defaultdict(list)
for point in data_points:
for k, v in point.items():
self.data[k].append(float(v))
self.fig = plt.figure(**fig_kwargs)
self._create_pendulum()
self._create_graph()
self.animation = animation.FuncAnimation(self.fig, self.animation_step,
frames=self.animation_length,
init_func=self.animation_init,
interval=self.delta_t,
blit=blit)
def _create_pendulum(self):
""" creates the plot for the animated pendulum (left) """
xmin, xmax = (min(self.data['x']) - self.rod_length,
max(self.data['x']) + self.rod_length)
ymin, ymax = (-self.rod_length * 2, self.rod_length * 2)
self.pendulum = self.fig.add_subplot(1, 2, 1,
aspect='equal', autoscale_on=False,
xlim=(xmin, xmax), ylim=(ymin, ymax))
self.pendulum.grid()
self.pendulum.title.set_text('Animated Pendulum')
self.pendulum.plot([0., 0.], [ymin, ymax], 'b--')
self.pendulum.xaxis.label.set_text('Cart Position (m)')
self.rod, = self.pendulum.plot([], [], 'o-', lw=2)
def _create_graph(self):
""" creates the plot for the actual data (right) """
ymin = min(self.data['x'] + self.data['theta'] + self.data['force'])
ymax = max(self.data['x'] + self.data['theta'] + self.data['force'])
self.graph = self.fig.add_subplot(1, 2, 2,
xlim=(0, max(self.data['t']) + 1),
ylim=(ymin - 1, ymax + 1))
self.theta_plot, = self.graph.plot([], [], 'b-', label='Theta (deg)')
self.force_plot, = self.graph.plot([], [], 'g-', label='Force (norm)')
self.pos_plot, = self.graph.plot([], [], 'r-', label='Position (m)')
self.graph.grid(); self.graph.legend()
self.graph.xaxis.label.set_text('Time (s)')
def _plot_pendulum(self, t):
""" plots the state of the pendulum """
# theta=0 is upward in our data but righward everywhere else
x0, y0 = self.data['x'][t], 0
x1, y1 = (x0 + self.rod_length * np.cos(self.data['theta'][t] + np.pi / 2),
y0 + self.rod_length * np.sin(self.data['theta'][t] + np.pi / 2))
self.rod.set_data([x0, x1], [y0, y1])
def _plot_graph(self, t):
""" plots the simulation data up until now """
time = self.data['t'][:t]
self.theta_plot.set_data(time, self.data['theta'][:t])
self.force_plot.set_data(time, self.data['force'][:t])
self.pos_plot.set_data(time, self.data['x'][:t])
def animation_init(self):
self.rod.set_data([], [])
self.theta_plot.set_data([], [])
self.force_plot.set_data([], [])
self.pos_plot.set_data([], [])
return self.theta_plot, self.force_plot, self.pos_plot, self.rod
def animation_step(self, i):
t = i % self.animation_length
self._plot_pendulum(t)
self._plot_graph(t)
return self.theta_plot, self.force_plot, self.pos_plot, self.rod
def show(self):
plt.show()
@click.command()
@click.argument('input-run', type=click.File('r'))
@click.option('--rod-length', default=0.5, help='Length of the rod (m)')
@click.option('--delta-t', default=20, help='Simulation delta t (ms)')
@click.option('--gif-name', type=click.STRING,
help='If specified, save a GIF with this filename')
@click.option('--gif-fps', type=click.INT, default=24)
@click.option('--width', type=click.FLOAT, default=8.,
help='Width of the animation (inches)')
@click.option('--height', type=click.FLOAT, default=8.,
help='Height of the animation (inches)')
@click.option('--dpi', type=click.FLOAT, default=96.)
@click.option('--blit/--no-blit', default=True)
@click.option('--csv-delimiter', default=',')
def main(input_run, rod_length, delta_t, gif_name, gif_fps, width, height, dpi, blit, csv_delimiter):
data = [p for p in csv.DictReader(input_run, delimiter=csv_delimiter)]
anim = Animate(data, rod_length, delta_t, blit, dpi=dpi, tight_layout=True,
figsize=(width, height))
if gif_name and gif_fps:
print('Saving as animated gif')
anim.animation.save(gif_name, writer='imagemagick', fps=gif_fps)
else:
print('Showing animation')
anim.show()
if __name__ == '__main__':
main()