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.

score 1 · Accepted Answer · answered Jun 27 '23 at 15:57

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

Thank you! Btw, can it be applied to lines as well? I tried to apply exact same approach with ax.plot but it gave me the error AttributeError: 'list' object has no attribute 'set_data' — muted_buddy, Jun 28 '23 at 09:43
@muted_buddy Lines can be animated as well. Without seeing your code I cannot say exactly what you're doing wrong. Look at [this answer](https://stackoverflow.com/questions/28074461/animating-growing-line-plot) for how to animate a line. — jared, Jun 28 '23 at 14:25

Animate grouped scatter points in matplotlib

1 Answers1

Linked