добавление легенды к анимированному графику рассеяния - PullRequest
Новая
       19

добавление легенды к анимированному графику рассеяния

0 голосов
/ 01 июля 2018

Я строю анимацию с точечной диаграммой, показывающей данные для нескольких групп с течением времени.
Когда я хочу добавить легенду, лучшее, что я могу получить, показывает только одну группу.
Примерный набор данных: 

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

Я получаю только группу, появляющуюся в легенде.

Если я только наберу ax.legend(loc=1), он показывает что-то похожее на это: 

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

Числа меняются в каждом кадре.

Я уже проверил эти ответы:
Как мне получить это, чтобы показать легенду на сюжете? : доставил меня туда, где я сейчас.
Как добавить легенду / метку в анимации Python : я получаю UnboundLocalError: local variable 'legend' referenced before assignment на legend.remove()
Добавить легенду для анимации (Исполнителей) в matplotlib : показывает только группу a.

1 Ответ

0 голосов
/ 02 июля 2018

Я нашел решение.
Мне нужно создать один точечный график на группу. Затем я обновляю каждый график рассеяния в моем методе 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()
...