Как мне получить легенду на моем категориальном цветном графике разброса?

Question

Я создал диаграмму рассеяния с помощью mathplotlib и раскрасил ее, используя столбец в моем фрейме данных. Теперь я хочу добавить легенду, чтобы прояснить, какой цвет представляет какие данные. Однако простое добавление plt.legend() без меток не помогает, а добавление меток в моей команде plt.scatter тоже не помогает.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

%matplotlib notebook

color = {
    "Africa" : "red",
    "Americas" : "green",
    "Eastern Mediterranean": "blue",
    "Europe" : "yellow",
    "South-East Asia": "black",
    "Western Pacific" : "orange"
}

data.columns = ['Country', 'GDP', 'Region', 'Air pollution (ug/m3)']
data['Color'] = data['Region'].map(color)

plt.scatter(data['GDP'], data['Air pollution (ug/m3)'], picker= 0, c = data['Color'], label = data['Region'])
plt.legend()
    
def onpick(event):
    origin = data.iloc[event.ind[0]]['Country']
    plt.gca().set_title('Selected item came from {}'.format(origin))

plt.gcf().canvas.mpl_connect('pick_event', onpick)

В настоящее время это выглядит так:

1

Однако я бы хотел, чтобы легенда выглядела как словарь color, с цветом в виде маркера и Region за ним. Как мне лучше всего это сделать?

Answer 1

Если вы не заинтересованы в использовании именно этих цветов, вы можете просто использовать sns.scatterplot, как в этом коде, без необходимости отображать каждый цвет:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from random import sample
import seaborn as sns

N = 100
data = pd.DataFrame({'GDP': np.random.random(N),
                     'Air pollution (ug/m3)': np.random.random(N),
                     'Region': sample(['Africa',
                                       'Americas',
                                       'Eastern Mediterranean',
                                       'Europe',
                                       'South-East Asia',
                                       'Western Pacific']*N, N)})

sns.scatterplot(data = data,
                x = 'GDP',
                y = 'Air pollution (ug/m3)',
                hue = 'Region')
plt.legend(bbox_to_anchor = (1.05, 0.98), loc = 'upper left')

plt.show()

Otherwise, if you want to mantain your colors, you can re-define the cycler:

import matplotlib.pyplot as plt
from random import sample
import seaborn as sns
from cycler import cycler

N = 100
data = pd.DataFrame({'GDP': np.random.random(N),
                     'Air pollution (ug/m3)': np.random.random(N),
                     'Region': sample(['Africa',
                                       'Americas',
                                       'Eastern Mediterranean',
                                       'Europe',
                                       'South-East Asia',
                                       'Western Pacific']*N, N)})

default_cycler = cycler(color=['red', 'green', 'blue', 'yellow', 'black', 'orange'])
plt.rc('axes', prop_cycle=default_cycler)

sns.scatterplot(data = data,
                x = 'GDP',
                y = 'Air pollution (ug/m3)',
                hue = 'Region')
plt.legend(bbox_to_anchor = (1.05, 0.98), loc = 'upper left')

plt.show()

---

Regarding interactivity, as exposed здесь :

Как и в любом другом случае, вы определяете аргумент picker и подключаете функцию обратного вызова

В вашем случае:

sns.scatterplot(data = data,
                x = 'GDP',
                y = 'Air pollution (ug/m3)',
                hue = 'Region',
                picker = 4)
plt.legend(bbox_to_anchor = (1.05, 0.98), loc = 'upper left')

def onpick(event):
    origin = data.iloc[event.ind[0]]['Country']
    plt.gca().set_title('Selected item came from {}'.format(origin))

plt.gcf().canvas.mpl_connect('pick_event', onpick)