Why is the linewidth inconsistent for this animated LineCollection?

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.

Fake data summary values

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)

Inflate data to get positions of all particles in each frame

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)

Run the animation

# 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()

Solution

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])