def __init__(self, N, ts = 0.1, to = 60.0, save=False):
"""
simulate N particles connected by a teather at a central point
and update the simulation every ts seconds for up to to seconds
"""
self.ts = ts
self.to = to
self.particles = [particle(q=10e-6,m=1e-19) for n in range(N)]
self.fig = figure()
self.ax = self.fig.add_axes([0, 0, 1, 1], projection='3d')
#self.ax.axis('off')
self.colors = cm.Dark2(linspace(0, 1, len(self.particles)))
self.pts = [self.ax.plot([], [], [], 'o', c=c)[0] for c in self.colors]
self.lines = [self.ax.plot([], [], [], '-', c=c)[0] for c in self.colors]
self.force_labels = [self.ax.text3D(p.pos[0]/2, p.pos[1]/2, p.pos[2]/2, "", color=c, size=10) for p,c in zip(self.particles, self.colors)]
self.force_vectors = [self.ax.add_artist(Arrow3D([p.pos[0], p.Fdir[0]], [p.pos[1], p.Fdir[1]], [p.pos[2], p.Fdir[2]], axes=self.ax, mutation_scale=20, lw=0.5, arrowstyle="<-", color=c)) for p,c in zip(self.particles, self.colors)]
self.ax.set_xlim(-2, 2)
self.ax.set_ylim(-2, 2)
self.ax.set_zlim(-2, 2)
self.anim = animation.FuncAnimation(self.fig, self.animate, init_func=self.animate_init, frames=int(self.to/self.ts), interval=self.ts*1000, repeat=False, blit=True)
self.writer = animation.AVConvWriter(fps=25)
if save:
self.anim.save(filename='tps.mp4', writer=self.writer)
return init
def create_animate(self):
def animate(i):
y = np.cos(2.0 * np.pi * self._x) * np.exp(-0.005 * self._x**2 * i)
self._line.set_data(self._x, y)
return [self._line]
return animate
arg_parser = ArgumentParser(description='create an MP4 animation of a '
'function')
arg_parser.add_argument('file', help='MP4 output file name')
options = arg_parser.parse_args()
figure = plt.figure()
cosDecay = CosDecay(figure, -6.0 * np.pi, 6.0 * np.pi, 2000)
init_f = cosDecay.create_init()
anim_f = cosDecay.create_animate()
anim = animation.FuncAnimation(figure,
anim_f,
init_func=init_f,
frames=300,
interval=20,
blit=True)
anim.save(options.file, fps=2, writer=animation.AVConvWriter())
def single_animation(system,
ex,
fig_size=(8, 8),
hide_axes=True,
filename=None):
"""Creates and saves an animation of a single example/system.
The animation is saved in `./animations` subdirectory.
Animations are done in 50 fps.
Bob size is proportional to bob mass.
Trailing dots show previous pendulum positions.
Args:
system (`ERK_DAE1` class instance): DAE system with the results in
`ys` and `zs` attributes.
ex (`DoublePendulum` class instance)
fig_size (tuple (float, float)): size of the figure
hide_axes (bool): should axis label be hidden of visible. Hidden is
the default, as it makes animation creation much faster
filename (string): name of the file (without an extension) where
animation should be saved
Returns:
filename (string)
"""
fig = plt.figure(figsize=fig_size)
ax = plt.axes(xlim=(system.ys[:, 1].min() * 1.2,
system.ys[:, 1].max() * 1.2),
ylim=(system.ys[:, 3].min() * 1.2,
max(0.2, system.ys[:, (2, 3)].max()) * 1.2))
if hide_axes:
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.set_aspect('equal')
ax.plot(0, 0, 'o', ms=8, c='C0')
line, = ax.plot([], [], '-', lw=3, c='C0')
m1, = ax.plot([], [], 'o', ms=4 * ex.M[0, 0], c='C1')
m2, = ax.plot([], [], 'o', ms=4 * ex.M[1, 1], c='C2')
p1, = ax.plot([], [], 'o', ms=1, c='C1')
p2, = ax.plot([], [], 'o', ms=1, c='C2')
fig.tight_layout()
skip = int(0.02 / system.h) # 50fps --> 0.02 sec per frame
def animate(i):
i *= skip
line.set_data([0, system.ys[i, 0], system.ys[i, 1]],
[0, system.ys[i, 2], system.ys[i, 3]])
p1.set_data(system.ys[0:i:skip, 0], system.ys[0:i:skip, 2])
p2.set_data(system.ys[0:i:skip, 1], system.ys[0:i:skip, 3])
m1.set_data(system.ys[i, 0], system.ys[i, 2])
m2.set_data(system.ys[i, 1], system.ys[i, 3])
return line, p1, p2, m1, m2
anim = animation.FuncAnimation(
fig=fig,
func=animate,
frames=len(system.ys) // skip,
interval=20,
blit=True,
)
filename = datetime.now().strftime(
'%Y-%m-%d_%H-%M') if filename is None else filename
metadata = create_metadata(filename, comment=create_comment(ex))
if not os.path.isdir(os.path.join(os.getcwd(), 'animations')):
os.mkdir('animations')
try:
writer = animation.AVConvWriter(fps=50, bitrate=-1, metadata=metadata)
anim.save('./animations/{}.mp4'.format(filename), writer=writer)
except FileNotFoundError:
writer = animation.FFMpegWriter(fps=50, bitrate=-1, metadata=metadata)
anim.save('./animations/{}.mp4'.format(filename), writer=writer)
return filename
def multi_animation(systems,
ex,
fig_size=(8, 8),
hide_axes=True,
filename=None):
"""Creates and saves an animation of a multiple systems.
The animation is saved in `./animations` subdirectory.
beforehand.
Animations are done in 50 fps.
Bob size is proportional to bob mass.
Trailing dots show previous pendulum positions.
Args:
systems (list of `ERK_DAE1` class instance): list of DAE systems with
the results in `ys` and `zs` attributes.
ex (`DoublePendulum` class instance)
fig_size (tuple (float, float)): size of the figure
hide_axes (bool): should axis label be hidden of visible. Hidden is
the default, as it makes animation creation much faster
filename (string): name of the file (without an extension) where
animation should be saved
Returns:
filename (string)
"""
ys = np.array([s.ys for s in systems])
fig = plt.figure(figsize=fig_size)
ax = plt.axes(xlim=(ys[:, :, 1].min() * 1.2, ys[:, :, 1].max() * 1.2),
ylim=(ys[:, :, 3].min() * 1.2,
max(0.2, ys[:, :, (2, 3)].max()) * 1.2))
if hide_axes:
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.set_aspect('equal')
ax.plot(0, 0, 'o', ms=8, c='k')
lines, m1s, m2s, trails = [], [], [], []
for k in range(len(systems)):
lines.append(ax.plot([], [], '-', lw=3, c='C{}'.format(k))[0])
m1s.append(
ax.plot([], [], 'o', ms=4 * ex.M[0, 0], c='C{}'.format(k))[0])
m2s.append(
ax.plot([], [], 'o', ms=4 * ex.M[1, 1], c='C{}'.format(k))[0])
trails.append(ax.plot([], [], 'o', ms=1, c='C{}'.format(k))[0])
fig.tight_layout()
skip = int(0.02 / systems[0].h) # 50fps --> 0.02 sec per frame
def animate(i):
i *= skip
for k in range(len(systems)):
lines[k].set_data([0, ys[k, i, 0], ys[k, i, 1]],
[0, ys[k, i, 2], ys[k, i, 3]])
m1s[k].set_data(ys[k, i, 0], ys[k, i, 2])
m2s[k].set_data(ys[k, i, 1], ys[k, i, 3])
trails[k].set_data(ys[k, max(0, i - 100 * skip):i:skip, 1],
ys[k, max(0, i - 100 * skip):i:skip, 3])
return lines, m1s, m2s, trails
anim = animation.FuncAnimation(
fig=fig,
func=animate,
frames=ys.shape[1] // skip,
interval=20,
blit=False,
)
filename = datetime.now().strftime(
'%Y-%m-%d_%H-%M') if filename is None else filename
metadata = create_metadata(filename, comment=create_comment(ex))
if not os.path.isdir(os.path.join(os.getcwd(), 'animations')):
os.mkdir('animations')
try:
writer = animation.AVConvWriter(fps=50, bitrate=-1, metadata=metadata)
anim.save('./animations/{}.mp4'.format(filename), writer=writer)
except FileNotFoundError:
writer = animation.FFMpegWriter(fps=50, bitrate=-1, metadata=metadata)
anim.save('./animations/{}.mp4'.format(filename), writer=writer)
return filename
def save(self, filename):
writer = anim.AVConvWriter(fps=10, bitrate=10000) # 400 kbits/s is 240p video
BaseQuickAnimation.save(self, filename, writer=writer)
def build_neuron_video(algorithm='snider', dim=3, alpha=0.5, **kwargs):
fig = plt.figure()
ax = p3.Axes3D(fig)
num_points = kwargs['npoints']
xmin = kwargs['xmin']
xmax = kwargs['xmax']
ymin = kwargs['ymin']
ymax = kwargs['ymax']
zmin = kwargs['zmin']
zmax = kwargs['zmax']
init_points = random_points_uniform(num_points=400, xmin=xmin, xmax=xmax,\
ymin=ymin, ymax=ymax,\
zmin=zmin, zmax=zmax)
for i in xrange(len(init_points)):
init_points[i] = init_points[i][:dim]
points = init_points[:]
G = init_graph(dim=dim)
graphs = [G.copy()]
unmarked_points = [init_points]
done = False
max_steps = kwargs['max_steps']
for i in xrange(max_steps):
if len(points) == 0:
break
can_extend = update_graph(G, algorithm, points, alpha, dim=dim, **kwargs)
graphs.append(G.copy())
unmarked_points.append(points[:])
if not can_extend:
break
retract_graph(G)
graphs.append(G.copy())
unmarked_points.append(points[:])
def init():
x, y, z = zip(*init_points)
ax.scatter(x, y, zs=z, s=25)
def init2d():
x, y = zip(*init_points)
plt.scatter(x, y, s=250)
def redraw(frame):
plt.clf()
G = graphs[frame]
for u in G.nodes():
if G.node[u]['label'] == 'synapse':
points.remove(G.node[u]['coord'])
init()
draw_graph_3d(graphs[frame], ax)
def redraw2d(frame):
print frame
plt.clf()
G = graphs[frame]
if len(unmarked_points[frame]) > 0:
x, y, z = zip(*unmarked_points[frame])
plt.scatter(x, y, s=250)
pos = {}
for u in G.nodes():
coord = G.node[u]['coord']
pos[u] = (coord[0], coord[1])
viz_tree(G, save=False)
ani = animation.FuncAnimation(fig, redraw2d, init_func=init, frames=len(graphs), \
interval = 200)
mywriter = animation.AVConvWriter()
ani.save('neuron_builder/neuron_builder.mp4', writer=mywriter)
plt.close()
def create_movie(self,
filename,
use_frames=True,
integral_x=False,
integral_y=False,
fps=1.0,
fading=False,
subframes=10,
dpi=100,
writer_name="libav"):
"""
Creates and saves a movie from the frames.
Args:
filename (string): name of the movie
use_frames (bool): put the frame image in the movie
integral_x (bool): put the frame horizontal integral in the movie
integral_y (bool): put the frame vertical integral in the movie
fps (float): frame per second speed
fading (bool): set the fading on/off between frames
subframes (int): if the fading is active, each frame is divided
into this number of subframes, and the transition
will use 3 frames (min 3)
dpi (float): resolution of the movie (max 200)
writer_name (string): name of a valid movie writer (see the code
forlibav, ffmpeg, imagemagick)
"""
#check if plot can be done
if not (use_frames or integral_x or integral_y):
print "Error: nothing to plot in the video"
return
if not use_frames and integral_x and integral_y:
print "Error: wrong movie elements selections"
return
#frame figure width in inches
inch_size = 5.0
#number of subrames (that must be > 3 for the fading)
subframes = int(max(subframes, 4))
#avoid too big resolution that would crash the cpu
dpi = float(min(dpi, 200))
#get an initial frame, and the dimension from the main slices
frame = self.frames[self.valid_frames_id[0]]
frame_h, frame_w = frame[self.main_slices].shape
frame_ratio = float(frame_h) / float(frame_w)
#calculate the movie size according to what must be plotted
if use_frames:
fig_size = [inch_size, inch_size * frame_ratio]
if integral_x:
fig_size[1] = fig_size[1] + inch_size / 2.0
if integral_y:
fig_size[0] = fig_size[0] + inch_size / 2.0
elif integral_x:
fig_size = [inch_size, inch_size / 2.0]
elif integral_y:
fig_size = [inch_size / 2.0, inch_size]
plt.rcParams['figure.figsize'] = fig_size
metadata = dict(title='Atom Movie',
artist='Matplotlib',
comment='Video of ultracold gases')
if fading:
#if fading is active the fps must be multiplied for the subframes
fps_real = float(fps) * subframes
else:
subframes = 1
fps_real = float(fps)
#initialize the writer
if writer_name == "ffmpeg":
writer = animation.FFMpegWriter(fps=float(fps_real),
metadata=metadata)
elif writer_name == "libav":
writer = animation.AVConvWriter(fps=float(fps_real),
metadata=metadata)
elif writer_name == "imagemagick":
writer = animation.ImageMagickWriter(fps=float(fps_real),
metadata=metadata)
else:
print "Error, movie writer not supported."
print "Supported writers: libav, ffmpeg, imagemagick"
return
#TODO: check GIF first frame colors
#initialize figure and axes
fig = plt.figure()
ax1 = None #frame axe
axx = None #horizontal integral axe
axy = None #vertical integral axe
if integral_x and not integral_y and not use_frames:
axx = plt.subplot2grid((1, 1), (0, 0))
elif not integral_x and integral_y and not use_frames:
axy = plt.subplot2grid((1, 1), (0, 0))
elif integral_x and not integral_y and use_frames:
ax1 = plt.subplot2grid((3, 1), (0, 0), rowspan=2)
axx = plt.subplot2grid((3, 1), (2, 0), sharex=ax1)
elif not integral_x and integral_y and use_frames:
ax1 = plt.subplot2grid((1, 3), (0, 1), colspan=2)
axy = plt.subplot2grid((1, 3), (0, 0), sharey=ax1)
elif integral_x and integral_y and use_frames:
ax1 = plt.subplot2grid((3, 3), (0, 1), rowspan=2, colspan=2)
axx = plt.subplot2grid((3, 3), (2, 1), colspan=2, sharex=ax1)
axy = plt.subplot2grid((3, 3), (0, 0), rowspan=2, sharey=ax1)
elif not integral_x and not integral_y and use_frames:
ax1 = plt.subplot2grid((1, 1), (0, 0))
#TODO control better space between subplots
fig.subplots_adjust(left=0,
bottom=0,
right=1,
top=1,
wspace=None,
hspace=None)
fig.tight_layout(pad=0.1)
def get_frame(frames, n_frame, n_subfr):
"""
Calculates and returns the frame given the current frame id.
"""
if n_subfr < subframes - 3 or not fading:
#if the subframes are at the beginning or the fading is off
frame = frames[n_frame]
else:
#id of the next frame (ciclical)
if n_frame == self.valid_frames_id[-1]:
n_frame_next = self.valid_frames_id[0]
else:
n_frame_next = n_frame + 1
#basic fading algoritm between frames
if n_subfr == subframes - 3:
frame = 0.75 * frames[n_frame] + 0.25 * frames[n_frame_next]
elif n_subfr == subframes - 2:
frame = 0.5 * frames[n_frame] + 0.5 * frames[n_frame_next]
elif n_subfr == subframes - 1:
frame = 0.25 * frames[n_frame] + 0.75 * frames[n_frame_next]
return frame
def calc_x_int(frame):
"""
Calculate the horizontal integral of a frame. Returns the
coordinates and the normalized integral.
Args:
frame (np.array): frame to integrate
"""
x_int = np.sum(frame[self.main_slices], axis=0)
#indices of the main and max selection
ind_main = self.main_slices[1].indices(frame.shape[1])
ind_max = self.maximum_slices[1].indices(frame.shape[1])
#slices of the frame that will be used to calculate the
#normalization of the integral
norm_slice = [
self.main_slices[0],
slice(max(ind_main[0], ind_max[0]),
min(ind_main[1], ind_max[1]))
]
norm = np.mean(np.sum(frame[norm_slice], axis=0))
return np.arange(frame_w), x_int / norm
def calc_y_int(frame):
"""
Calculate the vertical integral of a frame. Returns the
normalized integral and the coordinates.
Args:
frame (np.array): frame to integrate
"""
y_int = np.sum(frame[self.main_slices], axis=1)
ind_main = self.main_slices[0].indices(frame.shape[0])
ind_max = self.maximum_slices[0].indices(frame.shape[0])
norm_slice = [
slice(max(ind_main[0], ind_max[0]),
min(ind_main[1], ind_max[1])), self.main_slices[1]
]
norm = np.mean(np.sum(frame[norm_slice], axis=1))
return y_int / norm, np.arange(frame_h)
#create initial plots
if use_frames:
ax1_plot = ax1.imshow(frame[self.main_slices],
vmin=self.get_vlim_img()[0],
vmax=self.get_vlim_img()[1],
aspect='auto')
ax1.axis("off")
ax1_plot.set_cmap(self.cmap_img)
if integral_x:
axx_plot, = axx.plot(*calc_x_int(frame), lw=4)
axx.axis("off")
axx.set_xlim(0, frame_w)
axx.set_ylim(-0.05, 1.1)
if integral_y:
axy_plot, = axy.plot(*calc_y_int(frame), lw=4)
axy.axis("off")
axy.set_ylim(0, frame_h)
axy.set_xlim(-0.05, 1.1)
#reconstruct frames from the total image (this is done in order
#to include in the viedeo all effect applied to the final image)
frames = []
for n_frame in range(self.frame_num):
row = int(n_frame / self.frame_cols)
col = int(n_frame % self.frame_cols)
frm = self.all_frames_img[row * self.frame_h:(row + 1) *
self.frame_h, col *
self.frame_w:(col + 1) * self.frame_w]
frames.append(frm)
frames = np.array(frames)
#write the movie from the figure
with writer.saving(fig, str(filename), dpi):
#iterate over frames and subframes
for n_frame in self.valid_frames_id:
for n_subfr in range(subframes):
frame = get_frame(frames, n_frame, n_subfr)
#update plot data and grab instead of pltting again
#for efficiency
if use_frames:
ax1_plot.set_data(frame[self.main_slices])
if integral_x:
axx_plot.set_data(*calc_x_int(frame))
if integral_y:
axy_plot.set_data(*calc_y_int(frame))
writer.grab_frame()
x = np.linspace(0, 2, 1000)
y = (1 - (.01 * i)) * np.sin(2 * np.pi * (x - 0.01 * i))
line.set_data(x, y)
return line,
#call animator. blit only redraws changed elements
#great example of functional programming language!
anim = animation.FuncAnimation(fig,
animate,
init_func=init,
frames=400,
interval=25,
blit=False)
#save the animation as mp4
#Requires avconv
#forces x264 to support HTML5 embedding.
mywriter = animation.AVConvWriter(fps=40)
anim.save('basicanimation.mp4',
writer=mywriter,
extra_args=['-vcodev', 'libx264', '-verbose-debug'])
print("Done rendering animation!")
#TODO: embed this in an ipython notebook
#plt.show()
def create_animation(simulation, save=False, hide_axes=True, filename=None):
fig = plot.figure(figsize=(15, 10))
ax = plot.axes()
# ax.plot(0,0,'o', ms=8, c='C0')
if simulation.type == 'Moon':
color_1 = 'b'
color_2 = 'gray'
else:
color_1 = 'yellow'
color_2 = 'b'
size1 = 20
for key, point in simulation.eq_points.items():
plot.plot(point[0], point[1], 'rx')
# line, = ax.plot([],[],'-',lw=3,c='C0', animated=True)
m3, = ax.plot([], [], 'o', ms=3, c='g')
path3, = ax.plot([], [], '-', c='k')
m1 = ax.plot(-1 * simulation.mu, 0, 'o', c=color_1, ms=size1)
m2 = ax.plot(1 - simulation.mu, 0, 'o', c=color_2, ms=size1 / 2)
lims = 1.5
ax.set_xlim(-1 * lims, lims)
ax.set_ylim(-1 * lims, lims)
if hide_axes:
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.set_aspect('equal')
# xlims, ylims = ax.get_xlim(), ax.get_ylim()
x, y = np.arange(-2, 2, 0.01), np.arange(-2, 2, 0.01)
X, Y = np.meshgrid(x, y)
Z = V(X, Y, simulation.mu)
print(simulation.contour_levels)
contours = plot.contour(X,
Y,
Z,
levels=simulation.contour_levels,
linestyles='dashed')
fig.tight_layout()
skip = int(0.02 / (simulation.time[-1] /
len(simulation.time))) # Should be the time step?
# skip = int()
def animate(i):
# i *= skip
# line.set_data([0,simulation.trajectory[i,0]])
# m3.set_data(simulation.trajectory[0:i:skip,0], simulation.trajectory[0:i:skip,1])
path3.set_data(simulation.trajectory[0:i, 0],
simulation.trajectory[0:i, 1])
m3.set_data(simulation.trajectory[i, 0], simulation.trajectory[i, 1])
return m3, path3
anim = animation.FuncAnimation(fig=fig,
func=animate,
frames=len(simulation.time) // skip,
interval=0.2,
blit=True)
plot.title("Orbit Starting at {0} in the Earth-{1} System".format(
simulation.initial_point_str, simulation.type))
if save:
filename = datetime.now().strftime(
'%Y-%m-%d_%H-%M') if filename is None else filename
if not os.path.isdir(os.path.join(os.getcwd(), 'animations')):
os.mkdir('animations')
try:
writer = animation.AVConvWriter(fps=50, bitrate=-1)
anim.save('./animations/{}.mp4'.format(filename), writer=writer)
except FileNotFoundError:
writer = animation.FFMpegWriter(fps=50, bitrate=-1)
anim.save('./animations/{}.mp4'.format(filename), writer=writer)
else:
plot.show()