This animation is supposed to particles in motion with the leading component having alpha=1 and lw=4, with a tail of reducing lw and alpha. The alpha works, the lw is inconsistent even for the same particle - sometimes the leading edge is fat and the tail thin, sometimes the obverse is true. Why?!
Note - I'm, actually modelling traffic around a network, but I thought this multiple particle random walk would be a good example. For my use case the tail will be fairly crucial for the clarity and visibility of the final graphic.
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.collections import LineCollection
from matplotlib import cm
from functools import partial
np.random.seed(42) # to match my example
n_particles = 50 # number of particles to simulate
n_frames = 500 # number of frames to simulate
fade_frames = 5 # number of frames to fade out the animated plots for
# set up particle stating points & ids
particle_ids = [f"p{i}" for i in range(1, n_particles+1)]
starting_x = np.random.uniform(40, 60, n_particles)
starting_y = np.random.uniform(40, 60, n_particles)
# randomly assign start & end lives of points
start_frames = np.random.randint(0, n_frames/2, n_particles)
end_frames = start_frames + np.random.randint(n_frames/4, n_frames/2, n_particles)
colours = np.random.choice(range(10), n_particles) # Range 10 to match colours in the "tab10" colourmap
particles = pd.DataFrame({"id": particle_ids, "starting_x": starting_x, "starting_y": starting_y, "start_frames": start_frames, "end_frames": end_frames, "colour":colours})#.sort_values("start_frames").reset_index(drop=True)
My true data looks like this:
particle_ids = []
x_loc = []
y_loc = []
frames = []
colours = []
for id in particles.id:
start_frame = particles.loc[particles.id == id, "start_frames"].values[0]
end_frame = particles.loc[particles.id == id, "end_frames"].values[0]
colour = particles.loc[particles.id == id, "colour"].values[0]
particle_ids.append(id)
colours.append(colour)
x_loc.append(particles.loc[particles.id == id, "starting_x"].values[0])
y_loc.append(particles.loc[particles.id == id, "starting_y"].values[0])
frames.append(start_frame)
for frame in range(start_frame+1, end_frame+1):
particle_ids.append(id)
colours.append(colour)
x_loc.append(x_loc[-1] + np.random.uniform(-2, 2))
y_loc.append(y_loc[-1] + np.random.uniform(-2, 2))
frames.append(frame)
movements = (pd.DataFrame({
"id": particle_ids,
"x_loc": x_loc,
"y_loc": y_loc,
"frame": frames,
"colour":colours})
.sort_values("frame")
.reset_index(drop=True))
# replace colour code with rgb values. Set alpha during animation
cmap = dict(zip(range(10), cm.tab10.colors))
movements['colour'] = movements.colour.map(cmap)
# Set up plot
xmin = movements.x_loc.min()
xmax = movements.x_loc.max()
ymin = movements.y_loc.min()
ymax = movements.y_loc.max()
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim(0,100)
ax.set_ylim(0, 100)
lines = LineCollection([], ) # empty container to update later
ax.add_collection(lines)
def animate(frame, lines):
# Get data for the current frame
start_frame = frame - fade_frames
if frame >= movements.frame.max():
end_frame = movements.frame.max()
else:
end_frame = frame
if start_frame <= 0:
start_frame = 0
end_frame = 1
all_segments = []
all_colours = []
all_linewidths = []
frame_slice = movements[(start_frame<=movements.frame) & (movements.frame<=end_frame)]
# plot LineCollection for each particle separately
for id, grp in frame_slice.groupby('id'):
coords = list(zip(grp.x_loc.tolist(), grp.y_loc.tolist()))
all_segments.extend([(coords[i-1], coords[i]) for i in range(1, len(coords))])
colours = grp.colour.tolist()[:-1]
fade_values = [(frame-start_frame)/fade_frames for frame in grp.frame]
rgba = [(r,g,b,a) for (r,g,b), a in zip(colours, fade_values)]
all_colours.extend(rgba)
all_linewidths.extend([4*val for val in fade_values])
ax.draw_artist(lines)
lines.set_segments(all_segments)
lines.set_color(all_colours)
lines.set_linewidth(all_linewidths)
return lines,
# Run animation
ani = animation.FuncAnimation(
fig,
partial(animate, lines=lines),
frames=n_frames,
interval=100,
repeat_delay=200,
blit=True)
ani.save('test.gif')
plt.show()
You have an off-by-1 indexing problem with creating fade_values
.
frame_slice = movements[(start_frame<=movements.frame) & (movements.frame<=end_frame)]
#this generates 6 frames because you include both the start and end frames (which you need because N line segments requires N+1 vertexes.
for id, grp in frame_slice.groupby('id'):
coords = list(zip(grp.x_loc.tolist(), grp.y_loc.tolist()))
all_segments.extend([(coords[i-1], coords[i]) for i in range(1, len(coords))])
#iterating from 1-len here reduces the frame count (correctly?) to 5 because you need
# 6 vertices to make 5 line segments
colours = grp.colour.tolist()[1:]
#all points from groupby should have the same colour anyway
fade_values = [(frame-start_frame)/fade_frames for frame in grp.frame.tolist()[1:]]
#Here's the fix! you need to throw out the first frame as you have 1 more frame than line segments.
rgba = [(r,g,b,a) for (r,g,b), a in zip(colours, fade_values)]
all_colours.extend(rgba)
all_linewidths.extend([4*val for val in fade_values])