Question

我正在构建一个带有散点图的动画，该散点图显示一段时间内多个组的数据。
当我想添加图例时，我能得到的最好的只是显示一组。
样本数据集：

import pandas as pd 

df = pd.DataFrame([
    [1, 'a', 0.39, 0.73],
    [1, 'b', 0.87, 0.94],
    [1, 'c', 0.87, 0.23],
    [2, 'a', 0.17, 0.37],
    [2, 'b', 0.03, 0.12],
    [2, 'c', 0.86, 0.22],
    [3, 'a', 0.01, 0.15],
    [3, 'b', 0.03, 0.1],
    [3, 'c', 0.29, 0.19],
    columns=['period', 'group', 'x', 'y']
)

我这样制作动画：

import matplotlib.pyplot as plt
import matplotlib.animation as animation

fig, ax = plt.subplots()

ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
colors = {
'a': 'r',
'b': 'b',
'c': 'g'
        }
scat = ax.scatter([], [],
                c=df['group'].map(colors),
                )

def init():
    scat.set_offsets([])
    return scat,

def update(period):
    scat.set_offsets(df[df['period'] == period][['x', 'y']])
    scat.set_label(df[df['period']  == period]['group'])
    ax.legend([scat], df['group'].unique().tolist(), loc=1)
    ax.set_title(period)
    return scat,

ani = animation.FuncAnimation(fig, update, init_func=init,
                            frames=[1,2,3,4,5],
                            interval=500,
                            repeat=True)

plt.show()

我只在图例中出现了A组。

如果我只键入ax.legend(loc=1)，它将显示如下内容：

6  a
7  b
8  c
Name: group, dtype:object

数字在每帧中改变。

我已经检查了这些答案：
How do I get this to show the legend on the plot?：带我到现在的位置。
How to add legend/label in python animation：我在UnboundLocalError: local variable 'legend' referenced before assignment上得到legend.remove()
Add a legend for an animation (of Artists) in matplotlib：仅显示组a。

Answer 1

我找到了解决方法。
我需要为每个组创建一个散点图。然后，我使用update()方法更新每个散点图。
这是我的最终代码：

fig, ax = plt.subplots()
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
colors = {
    'a': 'r',
    'b': 'b',
    'c': 'g'
}

scats = []
groups = df.groupby('group')
for name, grp in groups:
    scat = ax.scatter([], [],
                      color=colors[name],
                      label=name)
    scats.append(scat)
ax.legend(loc=4)

def init():
    for scat in scats:
        scat.set_offsets([])
    return scats,

def update(period):
    for scat, (name, data) in zip(scats, groups):
        sample = data[data['period'] == period][['x', 'y']]
        scat.set_offsets(sample)
return scats,

ani = animation.FuncAnimation(fig, update, init_func=init
                              frames=[1, 2, 3, 4, 5],
                              interval=500,
                              repeat=True)

plt.show()

向动画散点图添加图例

1 个答案: