Matplotlib中实现分组散点并行动画并设置对应颜色

问题需求 #

• 让位置1、2、3的点**同时（并行）**进行动画
• 每个点使用对应的专属颜色

当前问题 #

• 点是串行依次绘制的，同一时间仅显示一个点
• 所有点共用单一颜色，无法区分不同轨迹

当前输出效果 #

动画中只有单个点按时间线逐个移动，无法同时展示三条轨迹的实时位置

示例代码 #

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

尝试过的方案 #

• 尝试为每条轨迹创建单独的points变量并加入动画函数，但未成功

修改后的解决方案代码 #

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)

# 按轨迹拆分数据，便于单独处理
track_groups = df.groupby('Track')
track_names = list(track_groups.groups.keys())

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

# Labels with axes:
ax.set_xlabel('Timeline')
ax.set_ylabel('Position')
ax.invert_yaxis()
xi = list(np.unique(df['Timeline']))
yi = list(np.unique(df['Position']))
ax.set_xticks(xi)
ax.set_yticks(yi)

# 颜色映射（调整第三轨颜色便于区分）
colors = {'Track 1': 'tab:red', 'Track 2': 'tab:blue', 'Track 3': 'tab:green'}

# 为每条轨迹创建散点对象，初始位置设为第一个时间点
scatter_list = []
for track in track_names:
    track_data = track_groups.get_group(track)
    scatter = ax.scatter(track_data['Timeline'].iloc[0], 
                         track_data['Position'].iloc[0], 
                         s=45, 
                         c=colors[track], 
                         zorder=2,
                         label=track)
    scatter_list.append((scatter, track_data))

# 动画帧数改为时间点数量（10帧，对应1-10的时间线）
frames = len(xi)

def animate(frame_idx):
    # 遍历所有轨迹，更新当前时间点的位置
    for scatter, track_data in scatter_list:
        current_data = track_data[track_data['Timeline'] == xi[frame_idx]]
        scatter.set_offsets((current_data['Timeline'].values[0], current_data['Position'].values[0]))
    # 返回所有散点对象用于动画更新
    return [scatter for scatter, _ in scatter_list]

# 创建动画
anim = FuncAnimation(fig, animate, frames=frames, interval=200, repeat=True)
ax.legend()  # 添加图例区分轨迹

plt.show()
plt.close()

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

关键修改说明 #

1. 数据拆分：用groupby按Track拆分数据，每条轨迹单独处理
2. 多散点对象：为每条轨迹创建独立的scatter对象，绑定对应颜色
3. 动画逻辑调整：将帧数改为时间点数量，每一帧同时更新所有轨迹的点位置，实现并行动画
4. 图例添加：加入图例方便区分不同轨迹

内容的提问来源于stack exchange，提问作者muted_buddy