Drawing grouped animated plt.step using matplotlib

Question

I'm just learning python and want to draw the plt.step lines that will follow the scatter points. The scatter points are drawn perfectly, so I used pretty much the same approach to draw the plt.steps lines, but it doesn't working which is unclear for me why, since ax.set_data can take the same 2D type arrays as ax.set_offsets do.

Unfortunately I am stuck with the following issues:

a) there is only one plt.step line (should be the same amount as 'Tracks' which are 3)

b) the line is drawn one cell behind the point

c) the color of the line is not corresponding to the scatter point

I appreciate any advise!

Current output:

The example of the code:

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation

df = pd.DataFrame()
cf = 0
while cf < 3:
    df = pd.concat([df, pd.DataFrame(
        {
            "Track": f'Track {cf + 1}',
            "Timeline": np.linspace(0, 9, 10, dtype=int),
            "Position": np.random.randint(low=0+cf, high=3+cf, size=10)
        }
    )])
    cf = cf + 1

df = df.reset_index(drop=True)
print(df)
df.info()

# plot:
fig, ax = plt.subplots()

# Point coordinates:
y = df['Position']
x = df['Timeline']

# Labels with axes:
ax.set_xlabel('Timeline')
ax.set_ylabel('Position')
ax.set_ylim(-0.2, 4.2)
ax.invert_yaxis()
ax.set_xticks(list(np.unique(x)))
ax.set_yticks(list(np.unique(y)))
ax.set_xlim(df["Timeline"].min()-0.5, df["Timeline"].max()+0.5)

# Colors:
colors = {'Track 1': 'tab:red', 'Track 2': 'tab:blue', 'Track 3': 'blue'}

# Drawing points and lines according to positions:
frames = (len(df.groupby(['Timeline'])))

steps = []
scatters = []
for track, group in df.groupby("Track"):
    scatters.append(ax.scatter(group["Timeline"].iloc[0],
                               group["Position"].iloc[0],
                               s=45, c=colors[track]))
    steps = plt.step(group["Timeline"].iloc[0],
                     group["Position"].iloc[0],
                     color=colors[track])


def animate(i):
    for scatter, (_, group) in zip(scatters, df.groupby("Track")):
        scatter.set_offsets((group["Timeline"].iloc[i],
                             group["Position"].iloc[i]))
    for step, (_, group) in zip(steps, df.groupby('Track')):
        step.set_data(group['Timeline'].iloc[:i],
                       group['Position'].iloc[:i])
    print('step', i) #for some reason there are three 0 steps in the beginning


anim = FuncAnimation(fig, animate, frames=frames, interval=400)
ax.grid()

anim.save('test.gif', writer='pillow')

Answer 1

Let me address your 3 issues.

1. There is only one line because although you intended to add the lines to the list steps, you instead just overwrote the list every loop. Because ax.step returns a list of one item, you need to index the first item otherwise you'll have a list rather than the Line2D object you wanted.
2. When using slicing in python, list[start:stop], the slice is inclusive of the start index and and exclusive of the stop index. So, when you do .iloc[:i], you are not including index i, so you are leaving out the point you want. Instead, you need to do .iloc[:i+1].
3. The colors don't match because of the first issue, which is that it is only taking the result of the last item in the loop (because you keep overwriting steps). Once you add them to a list, they will be colored properly.

Here is the corrected code:

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation

df = pd.DataFrame()
cf = 0
while cf < 3:
    df = pd.concat([df, pd.DataFrame(
        {
            "Track": f'Track {cf + 1}',
            "Timeline": np.linspace(0, 9, 10, dtype=int),
            "Position": np.random.randint(low=0+cf, high=3+cf, size=10)
        }
    )])
    cf = cf + 1

df = df.reset_index(drop=True)
print(df)
df.info()

# plot:
fig, ax = plt.subplots()

# Point coordinates:
y = df['Position']
x = df['Timeline']

# Labels with axes:
ax.set_xlabel('Timeline')
ax.set_ylabel('Position')
ax.set_ylim(-0.2, 4.2)
ax.invert_yaxis()
ax.set_xticks(list(np.unique(x)))
ax.set_yticks(list(np.unique(y)))
ax.set_xlim(df["Timeline"].min()-0.5, df["Timeline"].max()+0.5)
ax.grid()

# Colors:
colors = {'Track 1': 'tab:red', 'Track 2': 'tab:blue', 'Track 3': 'blue'}

# Drawing points and lines according to positions:
frames = (len(df.groupby(['Timeline'])))

steps = []
scatters = []
for track, group in df.groupby("Track"):
    scatters.append(ax.scatter(group["Timeline"].iloc[0],
                               group["Position"].iloc[0],
                               s=45, c=colors[track]))
    steps.append(ax.step(group["Timeline"].iloc[0],
                     group["Position"].iloc[0],
                     color=colors[track], alpha=0.5, linewidth=3)[0])


def animate(i):
    for scatter, (_, group) in zip(scatters, df.groupby("Track")):
        scatter.set_offsets((group["Timeline"].iloc[i],
                             group["Position"].iloc[i]))
    for step, (_, group) in zip(steps, df.groupby('Track')):
        step.set_data(group['Timeline'].iloc[:i+1],
                       group['Position'].iloc[:i+1])


anim = FuncAnimation(fig, animate, frames=frames, interval=400)
anim.save('test.gif', writer='pillow')

I also thickened the lines and added some opacity to make it clearer.