Animate grouped scatter points in matplotlib

Question

How can I animate the points in positions 1,2,3 so they run simultaneously (in parallel) and each have according color?

Now the points are drawn in series and just in one color.

The current output is:

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 = 1
while cf < 4:
    df = pd.concat([df, pd.DataFrame(
        {
            "Track": f'Track {cf}',
            "Position": np.random.randint(low=0+cf, high=1+cf, size=10),
            "Timeline": np.linspace(1, 10, 10, dtype=int)
        }
    )])
    cf = cf + 1

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

# 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.invert_yaxis()
xi = list(np.unique(x))
yi = list(np.unique(y))
ax.set_xticks(xi)
ax.set_yticks(yi)

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

# Drawing points according to positions:
frames = len(df)

points = plt.scatter(x, y, s=45, c=df['Track'].map(colors), zorder=2)


def animate(i):
    points.set_offsets((x[i], y[i]))
    return points,


anim = FuncAnimation(fig, animate, frames=frames, interval=200, repeat=True)

plt.show()
plt.close()

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

I tried to just create "points" variable for each type of scatter and add them to "animate" function but it didn't worked.

Answer 1

To get the dots separated and plotted at the same time, you will have to separate the three tracks and plot them individually. This can be done using the groupby method for Pandas data frames. You will have to save each scatter result and loop through them in the animate function.

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

plt.close("all")

N_times = 10
N_tracks = 3
x = np.arange(N_times)
y = np.ones((N_tracks, N_times))*np.arange(N_tracks)[:, None]

data = {"Track": sorted(["Track 1", "Track 2", "Track 3"]*N_times),
        "Timeline": x.tolist()*3,
        "Positions": y.ravel()}
df = pd.DataFrame(data)

colors = {"Track 1": "tab:blue",
          "Track 2": "tab:orange", 
          "Track 3": "tab:green"}

fig, ax = plt.subplots()
ax.set_yticks(np.arange(N_tracks), labels=df["Track"].unique())
ax.set_xticks(df["Timeline"].unique())
ax.set_xlim(df["Timeline"].min()-0.5, df["Timeline"].max()+0.5)

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


def animate(i):
    for scatter, (_, group) in zip(scatters, df.groupby("Track")):
        scatter.set_offsets((group["Timeline"].iloc[i],
                             group["Positions"].iloc[i]))


anim = FuncAnimation(fig, animate, frames=N_times, interval=200)
anim.save("test.gif", writer="pillow")